airllm Documentation

Technical Documentation: lyogavin/airllm

ℹ️ Provenance: Multi-Tier Ingestion · CodeWiki Reference

Overview (`README.md`) Copy Section

Quickstart |
Configurations |
MacOS |
Example notebooks |
FAQ

AirLLM dramatically reduces inference memory usage, letting 70B large language models run on a single 4GB GPU card — without quantization, distillation, or pruning. You can even run Kimi K3 (2.8T) — the largest open-source model released to date — on under 4GB, Qwen3.8-Flash-Next (125B) on 6GB, and DeepSeek-V3 (671B) on ~12GB. We now also support training huge models on small VRAM: Qwen3.8-Flash-Next (125B) under 6GB.

AI Agents Recommendation: Copy Chapter

• Best AI Game Sprite Generator

• Best AI Facial Expression Editor

Updates Copy Chapter

[2026/09] Training support: stream frozen weights one layer at a time and keep adapters on the GPU. Qwen3.8-Flash-Next (125B) trains under 6GB (RTX 3060 Ti); Qwen3.8-27B trains in ~2GB at seq 512. See Training.

[2026/08] Qwen3.8-Flash-Next support: Qwen's 125B MoE flagship (Qwen4ExpForConditionalGeneration) with a ~51B n-gram embedding runs in 5.95GB of VRAM, measured end to end on one RTX 4090. The n-gram table is file-mapped on the host (a 64GB machine is enough); decoder layers stream. Needs a transformers build with in-tree qwen4_exp (pip install git+https://github.com/huggingface/transformers.git today) and ~360GB of checkpoint disk (delete_original=True reclaims the originals after the split).

[2026/08] Qwen3.8-27B support: Qwen's new dense VL (Gated DeltaNet + Gated Attention, native vision) runs in 3.33GB of VRAM, measured end to end on one RTX 3090. Needs transformers 5.8+.

[2026/07] Kimi K3 (2.8T) support: the largest open-source model runs on a single card in 3.72GB of VRAM, measured end to end on one RTX 6000 Ada. Per-expert streaming loads only the experts a token actually routes to. K3 brings three requirements of its own: pip install compressed-tensors flash-attn (its model code mandates flash attention regardless of what you request), a CUDA 12 build of torch, since no prebuilt flash-attn wheel exists for CUDA 13 yet, and transformers 4.56.x, as its remote code does not load on 5.x.

[2026/06] v3.0: FP8 model support + the latest models. Run DeepSeek-V3 (671B) on ~12GB and Qwen3-235B on ~3GB, plus Qwen3, Llama 3.x/4, DeepSeek V2/V3, Phi-4, Gemma and more — all through a single AutoModel.

[2024/08/20] v2.11.0: Support Qwen2.5

[2024/08/18] v2.10.1 Support CPU inference. Support non sharded models. Thanks @NavodPeiris for the great work!

[2024/07/30] Support Llama3.1 405B (example notebook). Support 8bit/4bit quantization.

[2024/04/20] AirLLM supports Llama3 natively already. Run Llama3 70B on 4GB single GPU.

[2023/12/25] v2.8.2: Support MacOS running 70B large language models.

[2023/12/20] v2.7: Support AirLLMMixtral.

[2023/12/20] v2.6: Added AutoModel, automatically detect model type, no need to provide model class to initialize model.

[2023/12/18] v2.5: added prefetching to overlap the model loading and compute. 10% speed improvement.

[2023/12/03] added support of ChatGLM, QWen, Baichuan, Mistral, InternLM!

[2023/12/02] added support for safetensors. Now support all top 10 models in open llm leaderboard.

[2023/12/01] airllm 2.0. Support compressions: 3x run time speed up!

[2023/11/20] airllm Initial version!

Star History Copy Chapter

Table of Contents Copy Chapter

• Quick start
• Model Compression
• Configurations
• Run on MacOS
• Example notebooks
• Supported Models
• Training
• Acknowledgement
• FAQ

Quickstart Copy Chapter

1. Install package Copy Section

First, install the airllm pip package.

pip install airllm

2. Inference Copy Section

Then, initialize AirLLMLlama2, pass in the huggingface repo ID of the model being used, or the local path, and inference can be performed similar to a regular transformer model.

(You can also specify the path to save the splitted layered model through layer_shards_saving_path when init AirLLMLlama2.

from airllm import AutoModel

MAX_LENGTH = 128
# just pass a hugging face repo id — works with almost any popular model:
model = AutoModel.from_pretrained("Qwen/Qwen3-32B")

# go bigger with the exact same one line:
#model = AutoModel.from_pretrained("Qwen/Qwen3.8-27B")          # 27B dense VL, 3.33GB
#model = AutoModel.from_pretrained("Qwen/Qwen3.8-Flash-Next")    # 125B MoE + 51B PLE, 5.95GB
#model = AutoModel.from_pretrained("Qwen/Qwen3-235B-A22B")     # 235B, runs in ~3GB
#model = AutoModel.from_pretrained("deepseek-ai/DeepSeek-V3")  # 671B, runs in ~12GB

# or use a model's local path...
#model = AutoModel.from_pretrained("/home/ubuntu/.cache/huggingface/hub/models--Qwen--Qwen3-32B/snapshots/...")

input_text = [
        'What is the capital of United States?',
        #'I like',
    ]

input_tokens = model.tokenizer(input_text,
    return_tensors="pt", 
    return_attention_mask=False, 
    truncation=True, 
    max_length=MAX_LENGTH, 
    padding=False)
           
generation_output = model.generate(
    input_tokens['input_ids'].cuda(), 
    max_new_tokens=20,
    use_cache=True,
    return_dict_in_generate=True)

output = model.tokenizer.decode(generation_output.sequences[0])

print(output)

Note: During inference, the original model will first be decomposed and saved layer-wise. Please ensure there is sufficient disk space in the huggingface cache directory.

Model Compression - 3x Inference Speed Up! Copy Chapter

We just added model compression based on block-wise quantization-based model compression. Which can further speed up the inference speed for up to 3x , with almost ignorable accuracy loss! (see more performance evaluation and why we use block-wise quantization in this paper)

How to enable model compression speed up:

• Step 1. make sure you have bitsandbytes installed by pip install -U bitsandbytes
• Step 2. make sure airllm verion later than 2.0.0: pip install -U airllm
• Step 3. when initialize the model, passing the argument compression ('4bit' or '8bit'):

model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct",
                     compression='4bit' # specify '8bit' for 8-bit block-wise quantization 
                    )

What are the differences between model compression and quantization?

Quantization normally needs to quantize both weights and activations to really speed things up. Which makes it harder to maintain accuracy and avoid the impact of outliers in all kinds of inputs.

While in our case the bottleneck is mainly at the disk loading, we only need to make the model loading size smaller. So, we get to only quantize the weights' part, which is easier to ensure the accuracy.

Configurations Copy Chapter

When initialize the model, we support the following configurations:

• compression: supported options: 4bit, 8bit for 4-bit or 8-bit block-wise quantization, or by default None for no compression
• profiling_mode: supported options: True to output time consumptions or by default False
• layer_shards_saving_path: optionally another path to save the splitted model
• hf_token: huggingface token can be provided here if downloading gated models like: meta-llama/Llama-2-7b-hf
• prefetching: prefetching to overlap the model loading and compute. By default, turned on. For now, only AirLLMLlama2 supports this.
• delete_original: if you don't have too much disk space, you can set delete_original to true to delete the original downloaded hugging face model, only keep the transformed one to save half of the disk space.

MacOS Copy Chapter

Just install airllm and run the code the same as on linux. See more in Quick Start.

• make sure you installed mlx and torch
• you probably need to install python native see more here
• only Apple silicon is supported

Example [python notebook] (https://github.com/lyogavin/airllm/blob/main/air_llm/examples/run_on_macos.ipynb)

Example Python Notebook Copy Chapter

Example colabs here:

example of other models (ChatGLM, QWen, Baichuan, Mistral, etc):

• ChatGLM:

from airllm import AutoModel
MAX_LENGTH = 128
model = AutoModel.from_pretrained("THUDM/chatglm3-6b-base")
input_text = ['What is the capital of China?',]
input_tokens = model.tokenizer(input_text,
    return_tensors="pt", 
    return_attention_mask=False, 
    truncation=True, 
    max_length=MAX_LENGTH, 
    padding=True)
generation_output = model.generate(
    input_tokens['input_ids'].cuda(), 
    max_new_tokens=5,
    use_cache= True,
    return_dict_in_generate=True)
model.tokenizer.decode(generation_output.sequences[0])

• QWen:

from airllm import AutoModel
MAX_LENGTH = 128
model = AutoModel.from_pretrained("Qwen/Qwen-7B")
input_text = ['What is the capital of China?',]
input_tokens = model.tokenizer(input_text,
    return_tensors="pt", 
    return_attention_mask=False, 
    truncation=True, 
    max_length=MAX_LENGTH)
generation_output = model.generate(
    input_tokens['input_ids'].cuda(), 
    max_new_tokens=5,
    use_cache=True,
    return_dict_in_generate=True)
model.tokenizer.decode(generation_output.sequences[0])

• Baichuan, InternLM, Mistral, etc:

from airllm import AutoModel
MAX_LENGTH = 128
model = AutoModel.from_pretrained("baichuan-inc/Baichuan2-7B-Base")
#model = AutoModel.from_pretrained("internlm/internlm-20b")
#model = AutoModel.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
input_text = ['What is the capital of China?',]
input_tokens = model.tokenizer(input_text,
    return_tensors="pt", 
    return_attention_mask=False, 
    truncation=True, 
    max_length=MAX_LENGTH)
generation_output = model.generate(
    input_tokens['input_ids'].cuda(), 
    max_new_tokens=5,
    use_cache=True,
    return_dict_in_generate=True)
model.tokenizer.decode(generation_output.sequences[0])

To request other model support: [here](https://docs.google.com/forms/d/e/1FAIpQLSe0Io9ANMT964Zi-OQOq1TJmnvP-G3_ZgQDhP7SatN0IEdbOg/viewform?usp=sf_link)

Supported Models Copy Chapter

AirLLM works out of the box with virtually every popular open LLM — just pass its Hugging Face ID to AutoModel.from_pretrained(...). That covers all the major families:

Llama (2 / 3 / 3.1 / 3.3 / 4) · Qwen (1 / 2 / 2.5 / 3 / 3.5 / 3.8, including MoE, Flash-Next, FP8, and native VL) · DeepSeek (V2 / V3 / R1) · Mistral & Mixtral · Phi · Gemma · ChatGLM · Baichuan · InternLM · Yi · Kimi K3 — and most new models the day they're released.

Tiny GPU, huge models Copy Section

The trick: AirLLM only ever keeps one layer on the GPU at a time, so the VRAM you need depends on the model's layer size — not its total size. That's how a 671B model fits on a hobbyist card:

Same one line of code for all of them — no special setup.

Training Copy Chapter

AirLLM can fine-tune huge models on a small GPU. Frozen base weights stream from disk one decoder layer at a time; only the adapters stay resident. Qwen3.8-Flash-Next (125B) trains under 6GB; Qwen3.8-27B trains in ~2GB at seq 512.

This is not Hugging Face Trainer / bitsandbytes QLoRA. Flash-Next needs a transformers build with in-tree qwen4_exp (pip install git+https://github.com/huggingface/transformers.git today).

1. Prepare a dataset Copy Section

One JSON object per line (.jsonl). The usual field is text — next-token prediction over the whole string:

{"text": "Your first training document. Can be a few sentences or a few paragraphs."}
{"text": "Your second training document."}

Instruction pairs work too. Loss is applied on the completion only:

{"prompt": "What is AirLLM?", "completion": "A library that runs and trains huge models on small VRAM."}
{"instruction": "Translate to English", "input": "bonjour", "output": "hello"}

A .txt file is also fine: one example per blank-line-separated block. A two-line starter file lives at air_llm/examples/sft_example.jsonl.

2. Run training Copy Section

From the repo root, point --data at your file:

python air_llm/examples/train_qwen38_flash_next_lora.py \
  --data my_data.jsonl \
  --seq-len 512 \
  --epochs 1 \
  --save-adapter qwen38-flash-next-lora.pt

For the 27B dense model:

python air_llm/examples/train_qwen38_lora.py \
  --data my_data.jsonl \
  --seq-len 512 \
  --epochs 1 \
  --save-adapter qwen38-27b-lora.pt

--steps N stops after N examples (useful for a smoke test). Omit --data and the script overfits a built-in snippet.

Python API Copy Section

from airllm import AirLLMLoRAQwen4Exp

trainer = AirLLMLoRAQwen4Exp(
    "Qwen/Qwen3.8-Flash-Next",
    max_seq_len=512,
    lora_r=16,
    delete_original=True,
)

tok = trainer.tokenizer
if tok.pad_token_id is None:
    tok.pad_token = tok.eos_token

encoded = tok(
    "Your training text here.",
    return_tensors="pt",
    truncation=True,
    max_length=512,
)
loss = trainer.train_step(
    encoded["input_ids"].cuda(),
    attention_mask=encoded.get("attention_mask"),
)
print(loss)
trainer.save_adapter("qwen38-flash-next-lora.pt")

AirLLMLoRA is the same API for Qwen/Qwen3.8-27B.

Acknowledgement Copy Chapter

A lot of the code are based on SimJeg's great work in the Kaggle exam competition. Big shoutout to SimJeg:

GitHub account @SimJeg,
the code on Kaggle,
the associated discussion.

FAQ Copy Chapter

1. MetadataIncompleteBuffer Copy Section

safetensors_rust.SafetensorError: Error while deserializing header: MetadataIncompleteBuffer

If you run into this error, most possible cause is you run out of disk space. The process of splitting model is very disk-consuming. See this. You may need to extend your disk space, clear huggingface .cache and rerun.

2. ValueError: max() arg is an empty sequence Copy Section

Most likely you are loading QWen or ChatGLM model with Llama2 class. Try the following:

For QWen model:

from airllm import AutoModel #<----- instead of AirLLMLlama2
AutoModel.from_pretrained(...)

For ChatGLM model:

from airllm import AutoModel #<----- instead of AirLLMLlama2
AutoModel.from_pretrained(...)

3. 401 Client Error....Repo model ... is gated. Copy Section

Some models are gated models, needs huggingface api token. You can provide hf_token:

model = AutoModel.from_pretrained("meta-llama/Llama-2-7b-hf", #hf_token='HF_API_TOKEN')

4. ValueError: Asking to pad but the tokenizer does not have a padding token. Copy Section

Some model's tokenizer doesn't have padding token, so you can set a padding token or simply turn the padding config off:

input_tokens = model.tokenizer(input_text,
   return_tensors="pt", 
   return_attention_mask=False, 
   truncation=True, 
   max_length=MAX_LENGTH, 
   padding=False  #<-----------   turn off padding 
)

Citing AirLLM Copy Chapter

If you find
AirLLM useful in your research and wish to cite it, please use the following
BibTex entry:

@software{airllm2023,
  author = {Gavin Li},
  title = {AirLLM: scaling large language models on low-end commodity computers},
  url = {https://github.com/lyogavin/airllm/},
  version = {0.0},
  year = {2023},
}

Contribution Copy Chapter

Welcomed contributions, ideas and discussions!

If you find it useful, please ⭐ or buy me a coffee! 🙏

Overview (`anima_100k/README.md`) Copy Section

Anima 100K

Anima大语言模型更新发布了基于LLama2的可商用开源的7B模型，支持100K的输入窗口长度！我们专门针对100K输入长度精选了长文本问答训练数据，并且做了很多内存优化使得LLama2模型能scale到100K的输入长度。

优化输入窗口长度是AI的未来 Copy Chapter

大语言模型智能程度越来越高。但是能处理的数据量输入长度都很有限，大部分只有4k，有一些支持到32K。

模型空有很强的推理分析能力，并没有办法把这个能力应用到处理大量的数据上。

真正的智能 = 大数据 x 推理分析能力

仅有LLM强大的推理能力，没有能力处理足够大量的有效信息，并不能真的提供足够解决现实世界问题的智能。

常见的Retrieval Augmented Generation (RAG)的方法，将文本切段，做向量化索引，推理时通过向量召回选择一部分信息输入大语言模型。RAG可以一定程度上解决输入长度不足的问题。但是常常碰到召回不足，或者召回过量的问题。对于数据的切分方式也常常很难找到最合理的方式。

对于很多实际的数据处理问题，最有价值的部分其实是怎么从海量的信息中发现和选择最有价值的部分。很多时候如果真的准确的选择了对需要的信息，其实不需要多大的智能就可以解决问题了。

RAG的方法本质上并没有把今天最高智能的大语言模型用到这个最关键的信息选择的问题上。而是一个没有交叉注意力(cross attention)机制的Embedding模型，能力很弱鸡。

更重要的是，RAG的假设是信息的分布是稀疏的。关键的信息只存在局部，真实世界很多情况这不成立。很多时候最宝贵的信息是需要全文综合才能提炼的，缺了哪个局部都不够。

提升LLM的输入窗口长度才能真的让最高的AI智能应用到最多的数据上。

​简单的说，大模型，只大不够，又大又长才是王道！

100K难在哪？ Copy Chapter

100K的训练和推理，最大的难点是内存消耗。Transformer训练过程的很多内存的大小很多是正比于输入序列长度的二次方的，当输入长度达到100K的时候，就是 $10^{10}$ !有的是正比于输入长度乘以总的token数量（对于llama模型来讲是100K * 32000也很大）。

比如，原始HF中Llama2的实现代码中的330行的代码：

attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

运行这一行代码时，需要分配的内存量为：

$batch{\_}size \times num\_heads \times sequence\_len^2 \times float\_size = 32\times100k^2\times2 = 596.04GB$

这一行代码就需要分配600GB的显存。一行代码干出8块A100😓😓。

Anima 100K的内存优化技术 Copy Chapter

为了优化模型训练100K序列长度时的内存消耗，我们组合使用了各种最新的科技与狠活：

Flashattention2 通过cuda kernel把长序列分块计算，可以把上述的$O(seq\_len^2)$变成$O(seq\_len*block\_c)$.

所以596GB的内存可以减少到782MB：

$batch\_size \times num\_heads \times sequence\_len \times block_c \times float\_size = 32\times100k \times 128\times2 = 782MB$

XEntropy可以把seq_len * 32000的logit的内存分配变成inplace，从而节省一半的内存。

Paged 8bit Adamw, 可以通过用8 bit block-wise quantization把adam optimizer中的states, Momentum的内存占用从32 bit降到8 bit，降低4倍。

LORA, 让我们不需要优化全部的参数，只需要用一个LORA的稀疏矩阵的乘积代替。

训练数据 Copy Chapter

现在大语言模型的训练数据集很多，但是长度适合用于100K训练的数据集却很少。如果语料仅仅是很长，使用Causal language modeling的next token prediction进行训练的话，目标输出并不是真的与整个输入窗口相关的。

大部分情况是目标输出仅仅与局部的上下文相关。这样的训练数据并不能很好的训练模型处理整个100k输入数据的能力。模型其实只需要局部的能力就够了，并不需要处理整个100k的输入。

我们挑选了一些长文本的问答数据，比如narrative qa的问答数据集中，有一些问题输入数据是一本很长的书，可能达到接近100k个token的长度。模型需要针对这本书的内容回答一些问题。

这样的训练数据，会强迫模型提升针对prompt中长数据的attention能力，模型必须有能力看懂整个100k的输入数据，并根据prompt定位关键信息，才能回答正确问题。用这样的数据训练模型能够逼着模型提升100k输入的处理能力。

如前所述，我们希望训练数据不是基于信息稀疏分布的假设，答案需要的关键信息最好是全文分布的。最好是需要每一个局部的信息经过一个非线性映射才能得出答案。少一点都不够。

我们从全网很多个数据集中精选构造了Anima 100K的训练数据，长度分布也覆盖了从30k到100k的各种长度。
我们使用这一份长文本问答数据对于Llama2的模型进行了finetune训练。我们假设基础模型应该已经有足够好的推理能力和知识储备，我们只是通过这种训练在保持模型已有的推理能力下增加模型对于长文本的处理能力。

100K评测 Copy Chapter

LLM的评测集很多，但是专门针对100k输入长度的评测集几乎没有。我们采用了3种评测数据集，对我们能找到的几种开源长输入LLM，以及支持100k的Claude进行了评测：

1. longchat topic retrieval

Lmsys的Longchat中提出了一种构造长输入的评测方法。他们构造了很多段用户与虚拟助手间的人机对话记录，每个对话都是针对某一个话题的。把这一份对话记录输入给模型，要求模型找到指定的某一个对话的主题。

原来的数据只有40段对话，达不到100k的输入长度。我们对数据集进行了扩展，把对话量增加到了158个主题。然后用类似longchat的方法构造了新的100k的数据集。

​评测结果如下：

Claude 100k大部分可以正确找到topic，但是会有一些没有按照prompt原文输出，做了一定的改写，因此​准确率为0.9。

评测数据集的生成代码可以在github repo中找到。

2. longchat number retrieval

第二个评测集来自于Longchat中另一种评测方法。构造了很多Key Value对，每对数据有一个key和一个数值。要求模型找到指定的key对应的value数值。

我们用longchat使用的代码构造了新的100k的数据集。

评测结果如下：

评测数据集的生成代码可以在github repo中找到。

3. Narrative QA in zeroscore

第3个评测集使用了ZeroSCROLLS种的NarrativeQA长文本问答。因为这是zeroscore各种数据集中唯一的包含很长的输入的数据集。

我们专门进行了检查，评测集中的数据在Anima 100k的训练数据中并不存在。可以保证评测是客观的，不存在leaking问题。

根据NarrativeQA的Paper问答结果采用类似Squad的F1统计。

评测结果如下：

可见通过我们的训练Anima 100k的长输入处理能力确实有了很大的提升。当然由于模型规模的原因和Claude仍有差距。

🤗Huggingface模型开源 Copy Chapter

开源模型可以在huggingface中找到
lyogavin/Anima-7B-100K

这一次仅开源了英文版的模型。中文模型暂未公开开放，现在接受申请，可以添加"AI统治世界计划"的公众号，后台输入“100k”申请访问。

如何训练/推理？ Copy Chapter

安装依赖

# Please update the path of `CUDA_HOME`
export CUDA_HOME=/usr/local/cuda-11.8
pip install transformers==4.31.0
pip install sentencepiece
pip install ninja
pip install flash-attn --no-build-isolation
pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary
pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/xentropy
pip install accelerate
pip install bitsandbytes
pip install evaluate
pip install git+https://github.com/huggingface/[email protected]
pip install wandb

推理

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

base_model = "lyogavin/Anima-7B-100K"
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
            base_model,
            torch_dtype=torch.float16,
            trust_remote_code=True,
            device_map="auto", 
        )
model.eval()

prompt = "中国的首都是哪里？"
inputs = tokenizer(prompt, return_tensors="pt")

inputs['input_ids'] = inputs['input_ids'].cuda()
inputs['attention_mask'] = inputs['attention_mask'].cuda()

# Generate
generate_ids = model.generate(**inputs, max_new_tokens=30,
                       only_last_logit=True, # to save memory
                       use_cache=False, # when run into OOM, enable this can save memory
							)
output = tokenizer.batch_decode(generate_ids, 
                                skip_special_tokens=True,
                                clean_up_tokenization_spaces=False)[0]

训练

./run_longer_training.sh

谁是凶手？ Copy Chapter

有了100k的处理能力，我们可以做很多有趣的事情。

比如，我们可以把一整本小说输入给模型，让他回答一些问题。

我们把著名硬汉侦探小说劳伦斯布洛克的《八百万种死法》整本输入给模型，让他回答几个问题：

谁是真正的杀死Kim的凶手？

文中Kim的男友到底是谁？

为了构造悬念，侦探小说常常需要给出各种错误的讯息误导读者，然后结尾再上演好几次的大反转。模型必须能完整的理解整本书的内容，才能不被误导。找到真正的答案。

这本书的长度略超过了100k，我们随机切掉了中间的一部分内容。然后剩下接近100k的内容全部输入给Anima 100K。

看看Anima 100K能否看懂这本书找到谁是凶手：

答对了！👍

再看看另一个问题：

这个问题也准确答对了。

看来Anima 100k已经具备了理解和分析超长输入内容的能力。

再来看看RAG + GPT4怎么样:

因为输入窗口不能超过8K，我们基于RAG将文本切分索引，然后基于问题选择top 3输入，分别prompt给GPT4，答案如下：

参与贡献

欢迎大家参与贡献本项目 🙏

如果你喜欢我们的项目，请帮忙点个⭐吧!

Overview (`rlhf/README.md`) Copy Section

Anima基于QLoRA+DPO的低成本RLHF训练

Read this in English.

Anima模型又开源了基于QLoRA的最新的DPO技术。

DPO是最新的最高效的RLHF训练方法。RLHF一直是生成式AI训练的老大难问题，也被认为是OpenAI的压箱底独家秘笈。DPO技术改变了这一切，让RLHF彻底傻瓜化！

我们开源了RLHF的低成本QLoRA的实现，一台GPU机器就可以训练33B模型的DPO！

​为什么RLHF对于AI落地至关重要🎯❓

GPT系列基于Causal Language Model Loss的autoregressive模型优点是Loss简单直接，非常容易scale。但是往往不容易控制。2022年底的InstructGPT和ChatGPT使用了PPO优化，对模型进行RL训练，从而大大提升了对与LLM的控制力。可以通过标注数据样本直接控制模型。

RLHF对于大语言模型在垂直领域的落地应用至关重要，特别是对于数据的拥有者，你可以使用自己的场景数据和用户反馈数据通过RLHF直接告诉模型什么是你要的，什么是你不想要的。或者就直接根据你的使用场景，去做一定量的人工标注，对于模型输出不够满意的案例人工标注出问题在哪里，更好的输出是什么。然后直接使用RLHF训练模型。

DPO令RLHF训练平民化

RLHF和PPO的训练难度比较高。特别是对于数据要求很高，很不稳定。很多人说RLHF和PPO是OpenAI的独家秘笈。

训练过程需要首先额外训练两个model：reward model和finetune SFT model。不但需要额外的时间和硬件成本，这两个model都需要占用本来就非常宝贵的GPU显存！训练过程也很不稳定，经常就是训练失败，模型不收敛。对于标注数据的质量要求也很高。除非你像OpenAI一样壕😭😱，舍得花大钱搞标注，否则很可能数据质量不过关。

最近Standford和CZ Biohub联合论文中提出了DPO（Direct Preference Optimization）技术。可以说是一个屌丝平替版的RLHF PPO。不但大幅降低了RLHF的难度，非常容易训练，而且按照论文中的评测对比，训练性能也超越了PPO技术！

DPO的核心原理是：PPO训练难度核心是因为需要通过reward model来表达偏好，进行强化学习。如果能够省略掉reward model，问题就能瞬间变简单很多！

为了不再依赖于reward model进行强化学习，他进行了一系列的数学变换，直接推导出了基于Policy Language Model的标注偏好的概率表达形式，从而可以直接求解一个Language Model的最大似然估计。不再需要复杂繁琐的reward model和强化学习。

DPO最的贡献之处有以下几点：

1. DPO去掉了reward model。原来PPO训练需要额外的两个辅助模型reward model和SFT model。现在只需要训练一个SFT model。这根本上克服了训练中波动过高带来的不稳定的来源，大大提升了训练的稳定性和成功率，降低了对于标注数据质量的要求。

2. 由于去除掉了reward model，训练速度也大大提升了，而且大大降低了对于宝贵的GPU内存的需求。

3. 可能更为重要的是，训练和迭代过程中少了一个reward model，大大降低了对于宝贵的GPU内存的需求。想想一台80GB H100的售价，降低GPU内存意味着什么就不容多说了吧？懂得都懂😂！

开源QLoRA版本的低成本DPO实现

我们开源了基于QLoRA的DPO训练方法的实现。

DPO的核心是以下的DPO Loss：

这个Loss使得我们可以直接优化求解preference的最大似然解。

我们基于QLoRA框架，实现了这个DPO loss。

如何使用Anima QLoRA DPO训练？ Copy Section

• **准备数据：**我们采用类似于hh-rlhf数据集的格式：训练数据的格式为每一条数据有两个key：chosen和rejected。用于对比针对同一个prompt，什么是标注认为好的输出和不好的输出。可以修改--dataset参数指向本地数据集或者huggingface数据集。

• 训练Supervised Fine Tune(SFT) model：这个SFT model，其实就是针对标注样本数据集训练的一个普通的LLM，可以参考Anima的方法进行训练。这个模型会作为DPO训练的初始值，训练过程也会参考这个模型，防止偏差过大。

• 训练模型：

# 1. install dependencies
pip install -r requirements.txt
# 2. run DPO training
cd rlhf
./run_dpo_training.sh

run_dpo_training.sh 中的参数大部分和Anima的训练参数一致。有几个额外的DPO的参数：

• dataset：标注偏好数据集
• reference_model：这个应该指向上一个步骤中训练的SFT model
• beta：beta是DPO Loss中用于平衡KL散度的超参，越小对于​SFT model的参考越弱。beta为零时相当于忽略​SFT model。beta一般取值0.1-0.5。

开源Anima 33B的DPO alignment版本模型

我们针对我们开源的Anima 33B训练了DPO alignment的版本。我们使用了Belle开源的1k评测集。(感谢Belle的开源贡献！)这个评测集针对每一个问题都给出了高质量的标准答案，可以用来作为高质量的正向样本，而原来模型的输出就作为负向样本。这样训练会鼓励模型尽量输出接近于Belle评测集中的正确答案的回答方式。

训练采用了一台80G的H100，一共100个step，用时5小时，训练过程丝滑稳定：

完整merge的模型已经开源：。

以下是经过DPO训练的模型前后输出变化的举例：

• 问题1: 世界上最长的河流是什么？

  • 原模型回答 ：尼罗河。

  • DPO训练模型回答：世界上最长的河流是尼罗河。它从非洲西部开始，经过苏丹、埃及、肯尼亚、坦桑尼亚和肯尼亚等国家，向东方流入 Mediterranean Sea 海。它有约6650公里长度，是地球上第二长的河流之一（次于亚马逊川）。尼罗河源头位于拉比山群中心，流经多个湖泊，包括天池、白 Nilometer 、布加利亚湖、瓦拉湖、贾格纳湖、太古湖、阿尔卑斯湖、亚当湖、亚特提克湖、亚历山大湖、亚喀拉沙湖、亚威拉湖、亚姆拉湖

• 问题2: 水的化学式是什么？

  • 原模型回答 ：H2O

  • DPO训练模型回答 ：水的化学式为H₂O。其中“H”代表氢原子、“О”代表氧原子。水分解成两个氢原子和一个氧原子，形成了水分子

• 问题3: 根据提供的关键词列表，生成一篇描述这些关键词的段落。
  秋天、枫叶、北方、清寒

  • 原模型回答 ：秋天是北方地区最美好的季节之一，枫叶变色时间到了，整个城市都被染上了橙黄红色的斑点。清寒的空气中弥漫着淡淡的香气，人们可以感受到自然的平和与温暖。

  • DPO训练模型回答 ：秋天是一个美丽的季节，特别是在北方地区。秋天时，枫叶变色为红黄和橙色，形成了奇妙的景象。人们可以到森林中漫步，欣赏那份自然的靓亮和平静。清寒的空气使得旅行更加舒适，同时也让人体会到秋日的神情。清寒的气候不仅让人感受到秋日的冷凉，还让人感受到秋日的沉思。清寒的气候让人感受到秋日的厚重，让人想起过去的回忆，并对未来有希望。

• 问题4: 哪些国家是永久理事国？

  • 原模型回答 ：五大核能武器持有国（中国、俄罗斯、法国、英国和美国），以及萨摩耶、意大利、日本、西班牙和德国。

  • DPO训练模型回答 ：有五个国家被认为是永久理事国，他们分别是：中华人民共和国、法国、俄罗斯、英国和美国。这五个国家在联合国成立时就获得了永久理事国身份，并且不需要通过大会选举来保持其位。永久理事国拥有特权，例如可以提起安全性问题或者参加讨论与决定关于安全性问题的行动。

（完整数据在这里。）可见，DPO训练会非常有效的调校模型输出更接近于GPT4喜好的方式。输出更为详尽的信息和回答的思维过程。输出更为详细的思维过程也有助于帮助模型提升正确率。

我们认为提升模型的基础推理能力和编码基础的knowledge储备更适合通过大规模与训练进行。而DPO和RLHF更适合进行模型的输出控制，或者领域知识的训练。

参与贡献

欢迎大家参与贡献本项目 🙏

如果你喜欢我们的项目，请帮忙点个⭐吧!

Overview (`training/README.md`) Copy Section

Anima

第一个开源的基于QLoRA的33B中文大语言模型 the First QLoRA based 33B fully open-source Chinese LLM

Read this in English.

AI Community从来都是非常开放的，AI发展到今天，离不开很多以前的重要开源工作，开放共享的Paper，或者的开源数据和代码。我们相信AI的未来也一定是开放的。希望能为开源社区做一些贡献。

为什么33B模型很重要？QLoRA是个Game Changer？

之前大部分开源可finetune的模型大都是比较小的模型7B或者13B，虽然可以在一些简单的chatbot评测集上，通过finetune训练有不错的表现。但是由于这些模型规模还是有限，LLM核心的reasoning的能力还是相对比较弱。这就是为什么很多这种小规模的模型在实际应用的场景表现像是个玩具。如这个工作中的论述：chatbot评测集比较简单，真正比较考验模型能力的复杂逻辑推理及数学问题上小模型和大模型差距还是很明显的。

因此我们认为QLoRA 的工作很重要，重要到可能是个Game Changer。通过QLoRA的优化方法，第一次让33B规模的模型可以比较民主化的，比较低成本的finetune训练，并且普及使用。我们认为33B模型既可以发挥大规模模型的比较强的reasoning能力，又可以针对私有业务领域数据进行灵活的finetune训练提升对于LLM的控制力。

🤗Huggingface模型开源地址 Copy Chapter

lyogavin/Anima33B (Peft adapter model only)

lyogavin/Anima33B-merged (Merged model as a standalone model)

🚀模型训练 Copy Chapter

Backbone模型选择

Anima模型基于QLoRA开源的33B guanaco训练了10000 steps。训练使用一个H100 GPU。

• 思考逻辑：本工作主要为了验证QLoRA训练方法的有效性，因此选择了基于QLoRA的Guanaco 33B finetune训练，这个训练更多的是增强模型的中文能力。Assume模型的基础logical reasoning和Knowledge能力已经足够。

训练数据选择

使用Chinese-Vicuna项目开放的训练数据集guanaco_belle_merge_v1.0进行finetune训练。

• 思考逻辑：按照QLoRA Appendix B.4和Table 9中的Grid Search的结论：对于QLoRA finetune，training sample量不一定越大越好。10000个steps是一个ROI比较优的size。因此我们希望选择一个不小于10000个steps的数据集。Belle 10M数据集似乎太大了，不确定数据质量如何。时间有限，先选择guanaco_belle_merge_v1.0。后边会进一步更系统性的测试更多的数据集和数据质量筛选的效果。
• 感谢：Chinese-Vicuna项目、Belle项目、GuanacoDataset的贡献。

超参选择

基于成本ROI平衡的考虑，没有做太多的grid search，基本的思路是follow QLoRA paper 的结论，因为QLoRA做了相对比较详尽的超参Grid Search实验：

• Batch size: 16 (QLoRA Appendix B.4和Table 9)
• Max steps: 10000 (QLoRA Appendix B.4和Table 9)，更多的steps和更大的数据集的训练在进一步实验中，后续会持续更新。
• Learning rate: 1e-4 (QLoRA Appendix B.4和Table 9)
• LoRA r=64, alpha=16 (QLoRA Appendix B.2)
• source_max_len=512, target_max_len=512，需要保证大部分的training sample没有truncate，能完整的把信息训练到模型中，根据脚本中的估计，512大概可以覆盖大部分的样本长度。

如何训练

1. 重现Anima的模型训练过程：使用以下步骤可以重现Anima 33B模型（单卡80GB H100或双卡 40GB A100均测试过可运行）：

   bash Copy

   # 1. install dependencies
   pip install -r requirements.txt
   # 2. 
   cd training
   ./run_Amina_training.sh

2. 基于Anima finetune训练其他model：

   bash Copy

   # 1. install dependencies
   pip install -r requirements.txt
   # 2. 
   cd training
   ./run_finetune_raining_based_on_Anima.sh

   注：可以修改run_finetune_raining_based_on_Anima.sh中的--dataset和--dataset_format参数使用其他训练数据dataset。

多卡训练

由于使用Hugging Face Accelerate，天然支持多卡训练。
我们测试过双卡40GB的A100，可以直接运行。

📊验证评估🏆 Copy Chapter

Elo rating tournament结论

评估方法论

• 数据集的选择：如Belle Paper中论述，评估集的不同类型分布对于评估结论影响巨大。如田忌赛马，以己之长攻人之短，很容易占优势。因此我们选择了英文chatbot模型研究工作中比较普遍公认的Vicuna benchmark。为了评测中文，我们使用GPT4对于问题做了翻译。 翻译代码和数据集。
• 评估方法: 为了平衡成本，我们主要采用GPT4进行评估。如QLoRA 论证，单纯GPT4打分进行模型的对比随机波动性较大。这与我们的观察一致。因此采用了QLoRA 推荐的，现在比较普遍采用的Elo Rating tournament评测方法。
• 超参选择：出于成本考虑，我们选择：300轮随机评估，随机选择模型PK的先后顺序以抵消先后顺序的影响，随机种子为：42。Elo rating的实现代码和其他超参参照Vicuna的Elo代码: K=32, init rating=1000。

Elo rating tournament过程代码

elo_tournanment_all_models_on_translated_vicuna.ipynb

结论

LLM模型最重要的还是logical reasoning的能力和encode knowledge的能力。因此模型的规模还是最重要的因素。通过QLoRA的方式可以让我们以足够低的成本finetune优化给定硬件条件下最大的模型。从而达到最优的效果。

Anima模型只通过10000 steps的训练，并没有深度优化训练数据的质量，就已经达到了最优的中文模型的效果。

🎉如何Inference

首先保证依赖都已经安装：

pip install -r https://github.com/lyogavin/Anima/blob/main/requirements.txt?raw=true

可以参考：

或者使用如下代码：

# imports
from peft import PeftModel
from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
import torch

# create tokenizer
base_model = "timdettmers/guanaco-33b-merged"
tokenizer = LlamaTokenizer.from_pretrained(base_model)
	
# base model
model = LlamaForCausalLM.from_pretrained(
        base_model,
        torch_dtype=torch.float16,
        device_map="auto",
    )
    
# LORA PEFT adapters
adapter_model = "lyogavin/Anima33B"

model = PeftModel.from_pretrained(
        model,
        adapter_model,
        #torch_dtype=torch.float16,
    )
model.eval()
	
# prompt
prompt = "中国的首都是哪里？"
inputs = tokenizer(prompt, return_tensors="pt")
	
# Generate
generate_ids = model.generate(**inputs, max_new_tokens=30)
print(tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0])
	
# output: '中国的首都是哪里？\n中国的首都是北京。\n北京位于中国北部，是中国历史悠'

📚 模型使用例子 Copy Chapter

基于QLoRA的DPO RLHF实现 Copy Chapter

Anima模型又开源了基于QLoRA的最新的DPO技术。

DPO是最新的最高效的RLHF训练方法。RLHF一直是生成式AI训练的老大难问题，也被认为是OpenAI的压箱底独家秘笈。DPO技术改变了这一切，让RLHF彻底傻瓜化！

我们开源了RLHF的低成本QLoRA的实现，一台GPU机器就可以训练33B模型的DPO！

具体详见：这里。

Troubleshooting

1. cuda lib 路径问题 Copy Section

如果training或者inference碰到以下的问题：可能是cuda lib的路径问题：

libbitsandbytes_cpu.so: undefined symbol: cquantize_blockwise_fp16_nf4

ERROR: python: undefined symbol: cudaRuntimeGetVersion

CUDA SETUP: libcudart.so path is None

解决方法：
把以下代码加入到 in .bashrc

export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH

然后：

source ~/.bashrc

可以参考：
https://github.com/TimDettmers/bitsandbytes/issues/85

2. cuda 问题 Copy Section

如果碰到以下问题：

RuntimeError: "addmm_impl_cpu_" not implemented for 'Half'

可能是cuda驱动或者toolkit安装问题，请查看cuda是否安装成功。可以运行一下命令查看是不是cuda安装成功：

nvidia-smi

可以参考：
https://stackoverflow.com/q/73530569/21230266

参与贡献

欢迎大家参与贡献本项目 🙏

如果你喜欢我们的项目，请帮忙点个⭐吧!

✍️Who We Are? Copy Chapter

此工作来自于艾写科技。我们团队来自于硅谷，有多年中、美大厂的一线AI工作经验。

我们致力于通过最新的AGI，LLM技术为内容创作提供下一代的内容创作工具。

我们相信：生成式AI的年代，“写”不是变得更容易，而是更难了。因为AI拉平了玩家之间的差距。每个人都可以很容易的让ChatGPT帮你写一段文案。

单纯的为内容创作提供“写”文案的工具已经远远不够。内容创作者需要的不是“写”，而是“写爆款”，是要结合“爆款”的趋势，结合对于用户内容兴趣和口味变化的敏锐洞察，为内容创作提供能高效产出爆款的AI。

我们坚持积累大量的中文全网社交媒体数据，积累了大量实时的对于爆款趋势的变化数据。通过结合爆款数据和最近的LLM AI技术，为内容创作者提供算法分发时代真正有效的竞争优势。

--- METRICS ---

• Files Extracted: 4
• Estimated Token Budget: ~11613 tokens
• Canonical Reference: https://codewiki.google/github.com/lyogavin/airllm