AirLLM：4GB 显存运行 70B 模型的分层推理

AirLLM 是一个面向 Hugging Face 大语言模型的推理库，项目的核心目标是降低推理阶段的显存占用。官方 README 给出的典型配置是：70B 模型使用单张 4GB GPU，Llama 3.1 405B 使用 8GB，DeepSeek-V3 671B 使用约 12GB。它依靠分层拆分、按需加载和权重释放完成这件事，模型参数总量与显存容量之间的关系因此被重新组织。

显存压力来自哪里

Transformer 模型推理时，显存里通常需要容纳模型权重、运行中的中间激活、注意力缓存以及输入输出张量。模型越大，权重占用越高。一次性把完整模型加载到 GPU，显存很快成为硬限制。

AirLLM 采用了另一种执行顺序。模型结构仍由 Transformers 负责，AirLLM 把模型构建在 meta 设备上，让参数占位而不分配完整权重内存；随后将 embedding、每个 decoder layer、最终归一化层和 lm_head 视为可流式加载的模块。某个模块真正执行前，程序从磁盘读入它的权重，将其移动到 GPU，完成计算后释放对应权重，再处理下一个模块。

这套机制把显存的主要压力限定在当前模块加上运行时状态。因而，70B 模型能否运行，关键取决于单层权重、激活和缓存能否放入设备内存，完整权重无需一次性装入显存。

分层拆分如何工作

第一次加载模型时，AirLLM 会先定位 Hugging Face 模型目录，建立本地分层缓存。Safetensors 持久化器把每个模块的状态字典写入独立文件，并创建 .done 标记。后续启动会检查这些文件，已经完成拆分的模型可以直接进入分层读取流程。

执行阶段的模块序列通常是 embedding、多个 decoder layer、final norm 和 lm_head。模型本身保持完整的 forward 和 generation 逻辑，AirLLM 通过 hook 在模块执行前后插入权重搬运。这个设计减少了对每一种模型架构分别维护注意力、旋转位置编码和缓存逻辑的需要。只要 Transformers 能识别模型架构，AirLLM 就可以复用这部分执行能力。

代码中的关键流程可以概括为：

下载完整 checkpoint
        |
        v
按模块拆分为本地 safetensors
        |
        v
在 meta 设备创建模型结构
        |
        v
加载当前模块 -> CPU -> GPU
        |
        v
执行当前模块并释放权重
        |
        v
继续下一个模块

对于稀疏 MoE 模型，AirLLM 还可以进一步按 expert 流式读取。一个 MoE 层包含大量 expert，但每个 token 只会路由到其中一部分。逐个 expert 读取可以避免把整个 MoE 层同时物化到 GPU。README 中给出的 Kimi K3 说明正是基于这一点：单个 expert 的实际载入规模远低于完整层。

预取与压缩

磁盘读取会成为分层推理的主要成本之一。AirLLM 的默认预取机制使用一个 worker 线程，在当前模块计算时提前加载下一个模块。下一层权重准备完成后，计算线程可以直接使用，减少计算和磁盘 I/O 的串行等待。

代码对预取有明确限制：开启压缩时，预取会被自动关闭。原因是压缩权重的读取和解压流程需要额外的资源管理，当前实现没有把它与下一层预取同时启用。配置上应在两种路径中选择一种：优先降低磁盘读取等待，或者优先降低分层文件体积与加载成本。

AirLLM 支持 4bit 和 8bit 的 block-wise 权重压缩。该选项只处理磁盘上的权重，依赖 bitsandbytes。项目 README 将压缩模式描述为最高约 3 倍的推理速度提升，并给出“几乎可忽略的精度损失”这一项目方表述。实际速度仍会受到磁盘、CPU、GPU、模型架构和上下文长度影响，不能把这个数字当成所有模型和设备上的固定结果。

模型规模与设备内存

官方 README 给出了一组用于理解分层推理的参考表：

这张表应当按“推理内存参考”理解，不能等同于完整模型文件大小。模型仍然需要下载并在磁盘上保存，第一次运行还要额外生成分层缓存。README 明确提醒，模型拆分阶段会消耗较多磁盘空间；磁盘不足时，常见错误包括 Safetensors 头部不完整。

MoE 模型还需要区分总参数和每次激活的 expert。总参数可以非常大，但单个 token 的路由只访问部分 expert。AirLLM 在这类模型上按 expert 流式加载，显存需求由当前活跃的权重片段和运行时张量共同决定。

最小运行示例

安装包和基础依赖由项目的 3.1.0 setup 配置声明，核心依赖包括 PyTorch 2.4 以上、Transformers 4.49 以上且低于 5.13、Accelerate、Safetensors、Hugging Face Hub、SciPy 和 SentencePiece。安装命令如下：

pip install airllm

使用 AutoModel 加载 Hugging Face 模型时，代码结构与 Transformers 推理相近：

from airllm import AutoModel

MAX_LENGTH = 128
model = AutoModel.from_pretrained("Qwen/Qwen3-32B")

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

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

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

切换到更大模型时，通常只需要替换 from_pretrained 的模型 ID。例如项目 README 给出了 Qwen3-235B-A22B 和 DeepSeek-V3 的调用示例。对于需要授权的 Hugging Face 模型，可以通过 hf_token 传入访问令牌。

常用配置包括：

• layer_shards_saving_path：把分层缓存放到指定目录。
• profiling_mode：记录模型加载和压缩耗时。
• compression：选择 4bit 或 8bit 权重压缩。
• prefetching：在当前层计算时读取下一层，默认开启。
• delete_original：拆分完成后删除原始下载文件，以节省磁盘空间。
• max_seq_len：设置最大序列长度，默认值为 512。

Apple Silicon 与 CUDA 环境

项目 README 单独列出了 macOS 路径，要求 Apple Silicon，并提示安装 MLX 和 PyTorch。Apple Silicon 版本使用 MLX 实现对应的模型执行逻辑，权重仍按层从本地缓存读取。这个路径适合验证分层加载和低内存推理，但可用模型架构、MLX 支持范围和统一内存余量都需要单独检查。

CUDA 路径的默认设备是 cuda:0。对于现代深层模型，代码优先使用模型配置中的原生数据类型，通常是 bfloat16；项目代码特别避免强制所有模型使用 float16，因为深层模型可能出现溢出、无穷值或 NaN。使用压缩时还需要安装 bitsandbytes，而部分采用压缩张量格式的模型可能需要额外安装 compressed-tensors 和相匹配的 CUDA、Transformers 版本。

这些条件决定了 AirLLM 的适用边界。它能降低显存门槛，不能消除模型下载、分层转换、磁盘吞吐、CPU 搬运和生成时延。低显存设备运行 70B 模型，通常要接受更高的首轮准备成本和更慢的逐 token 速度；它更适合本地实验、模型结构验证和低并发推理，不能直接推导出生产服务吞吐。

与常规量化的关系

AirLLM 的无压缩路径通过“分层加载”降低峰值显存，保持完整权重格式。量化路径则进一步压缩磁盘上的权重，通过 block-wise 方式减少读取量。项目代码将这两项能力分开配置，compression=None 表示不启用压缩，compression='4bit' 或 compression='8bit' 表示启用对应模式。

传统量化往往需要同时处理权重和激活，容易受到异常值影响。AirLLM README 对自身方案的描述是只压缩权重部分，因为项目观察到瓶颈主要位于磁盘加载。这个设计有利于控制压缩带来的精度变化，但速度收益与模型、存储和设备组合有关。压缩还会关闭当前实现中的预取，因此需要以实际 profile 结果决定是否开启。

适合什么场景

AirLLM 的价值集中在“模型规模超过显存，但仍希望本地运行”的场景。开发者可以用它验证新的 Hugging Face 模型，研究模型输出，做低频离线任务，或者在有限 GPU 上进行原型测试。通过 profiling_mode 记录磁盘读取、GPU 加载和压缩时间，可以进一步判断瓶颈位于存储、CPU 还是设备计算。

如果任务要求稳定的高并发、较低首 token 延迟或持续吞吐，应该把 AirLLM 的结果和常规量化、CPU offload、分布式推理方案放在同一组实测条件下比较。分层推理解决的是峰值内存问题，服务级性能还取决于上下文长度、批处理、缓存策略、存储设备和模型生成配置。

AirLLM 的工程思路很有代表性：当完整对象无法放进一个资源池时，可以把执行过程改成“按需取用、及时释放、必要时提前准备”。它把大模型推理拆成了可观察的 I/O、搬运和计算阶段，开发者也因此可以用 profile 数据调整缓存路径、预取开关和压缩策略。

来源：

• AirLLM GitHub 仓库
• AirLLM README
• AirLLM 分层推理实现