Home

Awesome

<!-- * @Author: s-JoL(sl12160010@gmail.com) * @Date: 2023-03-10 21:18:35 * @LastEditors: s-JoL(sl12160010@gmail.com) * @LastEditTime: 2023-05-17 22:44:35 * @FilePath: /Open-Llama/README.md * @Description: * * Copyright (c) 2023 by s-JoL(sl12160010@gmail.com), All Rights Reserved. -->

中文 | English

camel

Open-Llama

<p align="center"> <img alt="GitHub" src="https://img.shields.io/github/license/s-JoL/Open-Llama.svg?color=blue&style=flat-square"> <img alt="GitHub release (latest by date)" src="https://img.shields.io/github/v/release/s-JoL/Open-Llama"> <img alt="GitHub top language" src="https://img.shields.io/github/languages/top/s-JoL/Open-Llama"> <img alt="GitHub last commit" src="https://img.shields.io/github/last-commit/s-JoL/Open-Llama"> </p>

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

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

DeepSpeed StageOffloadActivation CheckpointTotal TokenGPU hoursSpeed token/s/gpuBatch Size
Open-Llama 7B1FalseFalse173.7B1341236202
Open-Llama 13B3FalseTrue--185624
Open-Llama 33B3FalseTrue--70812
Open-Llama 65B3TrueTrue--36912
Llama 7B---1T824323370-
Llama 13B---1T1351682055-
Llama 33B---1.4T530432733-
Llama 65B---1.4T1022362380-

[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:

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:

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:

Requirements

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. https://github.com/s-JoL/Open-Llama/issues/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_lengthbatch_sizelearning_rateweight_decayparamsdimensionn headsn layervocab_size
204822e-41e-17.03B4096323268762
==============================================================================================================
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:

Modeln gpun layern headshidden sizevocab sizeseq length
GPT226heads40962501001024

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 FaceHugging FaceColossalAIColossalAIColossalAI
configwithout activation ckpt, bs2without activation ckpt, max_bs=12with activation ckpt, bs2without activation ckpt, bs2without activation ckpt, max_bs=10
second pre step0.336, fw=0.033, bw=0.3, opt=5e-61.250.3470.308, fw=0.067, bw=0.152, opt=0.0881.055
gpu memorynvidia-smi 45445fw+bw+opt=21053.63+22064.12+17987.52, nvidia-smi 40961fw+bw+opt=24684.74+21087.13+17987.52, nvidia-smi 46821oom 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.

V1V2
Datasetself implementeddatasets
ModelTransformersTransformers+xformers
OptimizerPytorch AdamFused Adam
DeepSpeedstage2stage1
Grad Accumulation412
Return Padding Maskyesno
Speed token/s/gpu13783637

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.

ModelOpen-LlamaLLAMALLAMALLAMAOPTBloomGLMGPT-NEOXCPM-ANTCodeGeeX
Model size7.0B6.7B13B65B175B175B130B20B10B13B
Token1T1T1.4T180B366B400B402B200B13.9B
GPU Hour82,432135,1681,022,362809,4721,082,99043776175680470403072
speed token/s/gpu36373370205538061.893.9105.7635.611811257
相关依赖xformersxformersmeaseqMegatron-DeepSpeedBMtrainMindSpore
speed token*params B/s/gpu25728225792671524700108151643213741127121181016341

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}},
}
<p align="center"> <a href="https://star-history.com/#s-JoL/Open-Llama&Date"> <img src="https://api.star-history.com/svg?repos=s-JoL/Open-Llama&type=Date" alt="Star History Chart"> </a> </p>