Skip to content

The complete training code of the open-source high-performance Llama model, including the full process from pre-training to RLHF.

License

Notifications You must be signed in to change notification settings

s-JoL/Open-Llama

Repository files navigation

中文 | English

camel

Open-Llama

GitHub GitHub release (latest by date) GitHub top language GitHub last commit

Open-Llama is an open-source project that offers a complete training pipeline for building large language models, ranging from dataset preparation to tokenization, pre-training, prompt tuning, lora, and the reinforcement learning technique RLHF.

You can try this model directly from the Demo.

Join discord to discuss the development of large language models.

Main contents

  • Support Transformers/HuggingFace. The CheckPoint after Instruct-tuning is open-source on Hugging Face: s-JoL/Open-Llama-V2.

  • By adopting the same evaluation method as the FastChat project, Open-Llama's performance is compared to GPT3.5’s. After testing, it can reach 89% of GPT3.5's performance on Chinese questions.

  • The training speed reaches 3620 tokens/s, faster than the 3370 tokens/s reported in the original Llama paper, reaching the current state-of-the-art level.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("s-JoL/Open-Llama-V2", use_fast=False)
model = AutoModelForCausalLM.from_pretrained("s-JoL/Open-Llama-V2", device_map="auto")

inputs = tokenizer('user:implement quick sort in python\nsystem:', return_tensors='pt', return_attention_mask=False, add_special_tokens=False)
for k, v in inputs.items():
   inputs[k] = v.cuda()
pred = model.generate(**inputs, max_new_tokens=512, do_sample=True)
print(tokenizer.decode(pred.cpu()[0], skip_special_tokens=True))

The CheckPoint after pre-training only is also uploaded to s-JoL/Open-Llama-V2-pretrain.

We have completed 330B token pre-training, training a total of 80 K steps. The Global Batch Size is consistent with Llama at 4M. Using a total of 7 parts of data to constitute the Instruction-tuning data, the model has certain programming abilities, mathematical abilities, and multi-turn dialogue abilities. Specific data can be found in the Instruction-Tuning section.

Below is a display of the model's multi-turn dialogue ability regarding code:

image4

Updates

[2023.5.8] Release v2.1

  • This update adds support for larger model training. Using DeepSpeed stage3 + offload + activation checkpoint, you can train a 65B model with A100-80G.

  • The peft library is introduced to support training such as lora.

  • The following table compares the training speed of Open-Llama and the original Llama, and the performance data of Llama is quoted from the original Llama paper.

DeepSpeed Stage Offload Activation Checkpoint Total Token GPU hours Speed token/s/gpu Batch Size
Open-Llama 7B 1 False False 173.7B 13412 3620 2
Open-Llama 13B 3 False True - - 1856 24
Open-Llama 33B 3 False True - - 708 12
Open-Llama 65B 3 True True - - 369 12
Llama 7B - - - 1T 82432 3370 -
Llama 13B - - - 1T 135168 2055 -
Llama 33B - - - 1.4T 530432 733 -
Llama 65B - - - 1.4T 1022362 380 -

[2023.4.28] Release v2.0

This update mainly includes the following aspects, increasing the effective training speed by 50% compared to the v1 version, reducing padding from 30% to 5%, and improving training speed from 3200 tokens/s to 3620 tokens/s. 0.95 * 3620 / (0.7 * 3200) = 1.521

  1. Use Hugging Face's datasets library for data reading, with the process as follows:
    1. Use the transform function to unify data formats from different datasets to {'text': 'xxx'}
    2. Tokenize using Tokenizer
    3. Sample long sequences; currently, three modes are provided: truncation, sampling (refer to the Gopher paper), and splitting
    4. Optional: concatenate texts from different docs, reducing padding in the data and accelerating training. In the v1 version, padding accounted for 30%; after concatenation, padding is reduced to 5%.
  2. Add Trainer, which can be reused for both pre-training and instruction fine-tuning, see solver/trainer.py
  3. Unify the pre-training and instruction fine-tuning training entry to train_lm.py
  4. Provide more convenient configuration, see configs/pretrain_config.yaml
  5. Provide functionality to continue pre-training based on other pre-trained models and supplementing vocabulary
  6. Resuming training from a checkpoint is supported, including loading optimizer parameters/learning rate and skipping duplicate data

[2023.4.16] Release v1.0

Basic pre-training and instruction fine-tuning codes are provided, with a training speed comparable to that of the original Llama. The pre-trained and fine-tuned models are already open-sourced on Hugging Face.

v1 version code can be seen at https://github.com/s-JoL/Open-Llama/tree/v1.0

Features

Easy to use

We believe that ease of use is one of the most important features when building large language models. To make Open-LLAMA more accessible, we have focused on the following aspects:

  • Minimal implementation: We have adopted the simplest implementation methods, lowering the entry threshold and allowing beginners to get started with ease.
  • Complete pipeline: We have published the complete code from dataset construction to training, making every step in the process of building a large language model clear and visible.

High performance

Due to the high cost of training large language models, high performance is also crucial when building them. To achieve high-performance training, we have employed the following techniques:

  • Fused CUDA kernel: Using the fused CUDA kernel provided in xformers can fuse multiple operations, reducing data transfer between the GPU and CPU, thereby improving training efficiency.
  • Parallelized training: We employ the Accelerate library to support parallelized training on multiple GPUs to speed up the training process.

For a 7B model, the training speed with the native PyTorch Llama model in Transformers is 1378 tokens/s/GPU. Using this codebase, the training speed reaches 3626 tokens/s/GPU, exceeding 3370 tokens/s/GPU reported in the original Llama paper.

If pre-training with 500B tokens, 38300 GPU hours are required. According to the hourly price for 8 A100-80G Spot GPUs on Google Cloud, which is 12.6 US dollars, the total cost is 60,300 US dollars. When using the unaccelerated version for training, the cost is 158,744 US dollars. The final training cost is reduced by 98,000 US dollars.

For more testing, see performance comparison with other open-source models.

Versatility

When training language models, our goal is to build a versatile model that can handle different languages and domains. To achieve this, we have employed the following strategies:

  • Multi-language support: We support multiple language corpora, including English, Chinese, Japanese, and many other languages, allowing users to choose according to their requirements.
  • Domain versatility: We hope that the model can not only help with everyday questions but also assist in professional domains such as science, law, etc.
  • Interaction with the world: By incorporating reinforcement learning (RL), we hope to give the model the ability to interact with the world.

Requirements

  • Python 3.7 or higher
  • PyTorch 1.13
  • Transformers library
  • Accelerate library
  • CUDA 11.6 or higher (for GPU acceleration)
  • Hardware configuration: currently using (64 CPU, 1000G Memory, 8xA100-80G) x N. There is a rather curious phenomenon that when more CPUs are used, the system runs slightly slower. I speculate this may have something to do with the multi-processing of dataloader.

Getting Started

Installation

Use the following command to install related dependencies:

pip install -r requirements.txt

Dataset Preparation

Currently provided are the Wudao dataset open-sourced by Zhiyuan and the Pile dataset open-sourced by EleutherAI. Dataset download and processing scripts are located in the data directory. Due to the required agreement for downloading the Wudao dataset, you may need to modify the link in download_wudao. Wudao.

Thanks to @skepsun's suggestion, using scidb to download the wudao dataset does not require login, and the download is more stable. #42.

Note that data download may fail. It is recommended to divide the download and processing in the script into two parts for multiple attempts, which will automatically resume downloads from breakpoints.

Run the following commands to download the data and perform partitioning:

bash data/download_the_pile.sh
bash data/download_wudao.sh

The data will be stored as small files, with a maximum of 16384 lines per file, for easy reading during multi-process training. The storage format is jsonl.zst, compressed using zstd, with a final data size of 519.5 GB, consisting of 16,466 files in total.

The Pile dataset contains 210,607,728 JSON lines, while the Wudao dataset contains 59,132,213 JSON lines.

The specific data format is as follows:

WuDao
{'id': 1, 'dataType': '百科', 'title': 'some title', 'content': 'some content'}

The Pile
{'text': 'some text', 'meta': {'pile_set_name': 'Github'}}

Check the data integrity in issue.

Related Tools

In the utils directory, training tokenizer/supplementing existing tokenizer models and conversion checkpoint code are provided.

Use SentencePiece to train a tokenizer with the following command:

python3 utils/train_tokenizer.py

In configs, a tokenizer model with a 40k vocabulary, trained only using the Wudao dataset (4w_cn_vocab_wudao15.model), is provided.

To supplement the vocabulary based on an existing tokenizer model, refer to:

python3 utils/merge_tokenizer.py

A bilingual English and Chinese tokenizer model (llama_tokenizer_extended.model) is created by merging the META official tokenizer model with the 40k Chinese tokenizer mentioned above.

To convert existing Llama model checkpoints, refer to:

python3 utils/convert_ckpt.py

Data Loading

Data loading-related code can be found in dataset/dataset.py, which includes pre-training and instruction fine-tuning data processing. To add other datasets, only the transform function needs to be modified.

The data loading process is as follows:

  1. Use the transform function to unify data formats from different datasets to {'text': 'xxx'}
  2. Tokenize using Tokenizer
  3. Sample long sequences; currently, three modes are provided: truncation, sampling (refer to the Gopher paper), and splitting
  4. Optional: concatenate texts from different docs, reducing padding in the data and accelerating training. In the v1 version, padding accounted for 30%; after concatenation, padding is reduced to 5%.

Use the following command to view the output of DataLoader and check the correctness of tokenization:

python3 dataset/dataset.py

Model Structure

We modified according to the section 2.4 Efficient implementation of the Llama paper in the Transformers library, and also referenced other papers to introduce some optimizations. Specifically, we used the memory_efficient_attention operation from the xformers library open-sourced by META for Self Attention computation, which has a significant performance improvement of approximately 30%. Further details can be found in modeling_llama.py.

Additionally, we referred to Bloom and introduced Stable Embedding for Token Embedding to better stabilize training.

Finally, we referenced PALM and employed Shared Input-Output Embeddings.

Pre-training

We use multi-GPU parallel training based on the Accelerate library, with the following start command:

accelerate launch --config_file configs/accelerate_configs/ds_stage1.yaml train_lm.py --train_config configs/pretrain_config.yaml --model_config configs/model_configs/7B.json 

In some cases, you may need to specify the following parameters:

--main_process_ip
--main_process_port
--num_processes
--num_machines
--machine_rank

We use Wandb for visualizing training. You need to modify the WANDB_API_KEY environment variable yourself.

Among them, we use DeepSpeed stage1 to reduce memory usage. For Accelerate-related configurations, see configs/accelerate_configs.

Training related hyperparameters can be found in configs/pretrain_config.yaml.

The default parameters use LlamaTokenizer with a supplemented 40k Chinese vocabulary tokenizer model, and the model size is 7B. The specific configuration is as follows:

max_length batch_size learning_rate weight_decay params dimension n heads n layer vocab_size
2048 2 2e-4 1e-1 7.03B 4096 32 32 68762
==============================================================================================================
Layer (type:depth-idx)                                       Output Shape              Param #
==============================================================================================================
OpenLlamaForCausalLM                                         [1, 32, 64, 128]          --
├─OpenLlamaModel: 1-1                                        [1, 32, 64, 128]          --
│    └─Embedding: 2-1                                        [1, 64, 4096]             281,649,152
│    └─ModuleList: 2-2                                       --                        --
│    │    └─OpenLlamaDecoderLayer: 3x32                      [1, 64, 4096]             202,383,360
│    └─OpenLlamaRMSNorm: 2-3                                 [1, 64, 4096]             4,096
├─Linear: 1-2                                                [1, 64, 68762]            281,649,152
==============================================================================================================
Total params: 7,039,569,920
Trainable params: 7,039,569,920
Non-trainable params: 0
Total mult-adds (G): 7.04

Pre-training loss from scratch is shown below:

loss

Instruction-Tuning

We use the currently available seven datasets for Instruction-tuning, and more tasks and our own datasets will be added later.

The ShareGPT_Vicuna_unfiltered dataset has some issues in the datastes processing, so we directly downloaded the original data and reprocessed it. We performed some preprocessing on the original data, with the format as follows:

user: {prompt}\nsystem: {completion}</s>

The startup command is basically the same as pre-training:

accelerate launch --config_file configs/accelerate_configs/ds_stage1.yaml train_lm.py --train_config configs/instruct_config.yaml --model_config configs/model_configs/7B.json 

In some cases, you may need to specify the following parameters:

--main_process_ip
--main_process_port
--num_processes
--num_machines
--machine_rank

The loss during the process is shown below, with a total of 3 epochs:

loss

RLHF

Not available yet.

Server

For multi-turn dialogue, use chat_server.py.

Developed based on Gradio.

Performance Comparison

Training Framework

In terms of training frameworks, we tested Hugging Face's open-source Accelerate library, PyTorch Lightning, and HPC-AI's open-source ColossalAI. We found that their performance differences are relatively small when fully utilizing GPUs. Therefore, we chose the relatively simple-to-implement Accelerate library as the training framework.

The test code can be found in utils/speed_test.py.

The model structure used during the testing process is:

Model n gpu n layer n heads hidden size vocab size seq length
GPT2 2 6 heads 4096 250100 1024

The test results are shown below, indicating that when the GPUs are fully utilized, the differences in speed and memory consumption are not significant.

Hugging Face Hugging Face ColossalAI ColossalAI ColossalAI
config without activation ckpt, bs2 without activation ckpt, max_bs=12 with activation ckpt, bs2 without activation ckpt, bs2 without activation ckpt, max_bs=10
second pre step 0.336, fw=0.033, bw=0.3, opt=5e-6 1.25 0.347 0.308, fw=0.067, bw=0.152, opt=0.088 1.055
gpu memory nvidia-smi 45445 fw+bw+opt=21053.63+22064.12+17987.52, nvidia-smi 40961 fw+bw+opt=24684.74+21087.13+17987.52, nvidia-smi 46821 oom after 10 steps

Performance Optimization

In the earliest version, we used the native Llama implementation from DeepSpeed stage2 + Transformers for training. However, the speed was significantly different from what was claimed in the paper. Therefore, we carried out a series of optimizations afterwards, and we list each step of the performance improvement below for reference.

The paper mentioned that for the 6.7B model, 1T token was used for training and the final GPU time was 82432, from which the training speed was roughly calculated as 3370 token/s/gpu. After using the following optimizations, the speed is now basically consistent with what was claimed in the paper when tested on 20x8 A100-80G. It is expected that more fusion operators will be added in the future to achieve better performance.

V1 V2
Dataset self implemented datasets
Model Transformers Transformers+xformers
Optimizer Pytorch Adam Fused Adam
DeepSpeed stage2 stage1
Grad Accumulation 4 12
Return Padding Mask yes no
Speed token/s/gpu 1378 3637

Comparison with Other Open-source Models

The following table summarizes the performance of currently available open-source models. In all cases, the GPU device used is A100. Due to differences in the size and structure of the models, it is difficult to make accurate performance comparisons. As a rough estimate, it can be assumed that the speed is generally inversely proportional to the size of the model parameters, which is confirmed by the performance of Llama with models of different sizes. Based on this rough estimate, it can be seen that the performance using our project is significantly better than that of other projects.

Model Open-Llama LLAMA LLAMA LLAMA OPT Bloom GLM GPT-NEOX CPM-ANT CodeGeeX
Model size 7.0B 6.7B 13B 65B 175B 175B 130B 20B 10B 13B
Token 1T 1T 1.4T 180B 366B 400B 402B 200B 13.9B
GPU Hour 82,432 135,168 1,022,362 809,472 1,082,990 43776 175680 47040 3072
speed token/s/gpu 3637 3370 2055 380 61.8 93.9 105.7 635.6 1181 1257
相关依赖 xformers xformers measeq Megatron-DeepSpeed BMtrain MindSpore
speed token*params B/s/gpu 25728 22579 26715 24700 10815 16432 13741 12712 11810 16341

Future Plans

  1. Integrate RLHF code.
  2. Use Triton to add more high-performance operators to further improve performance.
  3. Add code for building pre-training datasets based on Common Crawl and open related datasets.
  4. Add code for multimodal training.

Citation

@misc{openllama,
  title={Open-Llama},
  author={s-JoL},
  year={2023},
  howpublished={\url{https://github.com/s-JoL/Open-Llama}},
}

Star History Chart

About

The complete training code of the open-source high-performance Llama model, including the full process from pre-training to RLHF.

Resources

License

Stars

Watchers

Forks

Packages

No packages published