AirLLM：单GPU运行大语言模型

AirLLM 项目概述

AirLLM 通过逐层流式加载的方式，大幅降低大语言模型推理所需的显存，使得 70B 量级模型可在单卡 4GB 显存上运行，且无需量化、蒸馏或剪枝。更大规模的模型同样支持：DeepSeek-V3（671B）可在约 12GB 显存上运行，Qwen3.8-Flash-Next（125B）可在 5.95GB 上运行，甚至 Kimi K3（2.8T）也能在 4GB 以内运行。

核心思路

传统加载方式需将全部参数常驻 GPU，显存占用与模型总参数量成正比。AirLLM 只在 GPU 上保留单层权重，其余层在推理过程中按需从磁盘流式读取，所需显存仅取决于单层大小，与模型总规模无强关联。

快速开始

安装

pip install airllm

推理示例

from airllm import AutoModel

# 一行代码加载任意模型（接受 HuggingFace repo id 或本地路径）
model = AutoModel.from_pretrained("Qwen/Qwen3-32B")

input_text = ['What is the capital of United States?']
input_tokens = model.tokenizer(
    input_text,
    return_tensors="pt",
    return_attention_mask=False,
    truncation=True,
    max_length=128,
    padding=False)

output = model.generate(
    input_tokens['input_ids'].cuda(),
    max_new_tokens=20,
    use_cache=True,
    return_dict_in_generate=True)

print(model.tokenizer.decode(output.sequences[0]))

首次推理时，模型会按层切分保存至 HuggingFace 缓存目录，请确保磁盘空间充足。

模型压缩加速

基于分块量化（block-wise quantization）的压缩方案可带来最高约 3 倍推理加速，精度损失极小。配置方法：

pip install -U bitsandbytes
pip install -U airllm

# 在初始化时传入 compression 参数
model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    compression='4bit'  # 也可设为 '8bit'
)

配置参数

• compression：可选 4bit / 8bit / None
• profiling_mode：是否输出各阶段耗时
• layer_shards_saving_path：自定义分层模型保存路径
• hf_token：用于下载 gated 模型
• prefetching：预取下一层以重叠加载与计算（默认开启）
• delete_original：磁盘空间不足时设为 True，删除原始模型以节省一半空间

MacOS 支持

仅支持 Apple Silicon 芯片。安装 mlx 与 torch 后，代码与 Linux 一致。

已验证显存占用

模型	参数量	显存
Qwen3 / Mistral / Phi	~8B	~1–2 GB
Qwen3-235B	235B	~3 GB
Llama 3.x 70B	70B	~4 GB
Llama 3.1 405B	405B	~8 GB
DeepSeek-V3	671B	~12 GB

训练支持

除推理外，AirLLM 现已支持在低显存设备上微调超大模型。冻结的基座权重按层从磁盘流式读取，仅适配器常驻显存。Qwen3.8-Flash-Next（125B）可在 6GB 显存下训练，Qwen3.8-27B 可在约 2GB 下训练。

数据格式支持 JSONL（每行一个 JSON 对象）或 TXT：

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

启动训练：

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

Python API：

from airllm import AirLLMLoRAQwen4Exp

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

# ...分词、train_step()、trainer.save_adapter()

常见问题

• MetadataIncompleteBuffer：通常是磁盘空间不足，分层过程消耗大量磁盘。
• ValueError: max() arg is an empty sequence：模型类与权重不匹配，统一使用 AutoModel.from_pretrained 即可。
• 401 Client Error (gated repo)：在初始化时传入 hf_token。
• tokenizer 无 pad_token：将 padding 设为 False，或为 tokenizer 显式设置 pad_token。