AirLLM 非分片模型实战：低显存 GPU 跑通小模型的完整路径

AirLLM 非分片模型实战：低显存 GPU 跑通小模型的完整路径

【免费下载链接】airllmAirLLM 70B inference with single 4GB GPU项目地址: https://gitcode.com/GitHub_Trending/ai/airllm

手里只有一张 6GB 的卡，想跑个 7B 模型，却在torch.cuda分配那一步直接 OOM？这正是 AirLLM 非分片模型要解决的问题。它不把整个模型塞进显存，而是逐层从磁盘流式换入，所以 7B 模型的峰值显存只有"一层 + KV cache"的量，而不是 fp16 完整模型的 14GB。对低显存 GPU 推理来说，这条路径绕开了"显存装不下"的死结 ⚡

分片 vs 非分片：小模型为什么该走另一条路

先分清两种加载方式。分片路径是给 70B 这类超大模型准备的：checkpoint 本身是多份 safetensors，AirLLM 还要在磁盘上按层切出 layer shards，再一层一层流式换入显存，你得额外管layer_shards_saving_path，切分过程还很吃磁盘。非分片是 v2.10.1 加进来的：checkpoint 是单文件或小目录时，框架直接对原始文件做逐层流式，省掉切分这步预处理。

小模型（7B 以下）走非分片更划算，原因很具体。Llama 系 7B 大约 32 层，单层放进显存的峰值本就低，切分能省下的显存几乎为零，却平添了写盘开销和更多可能出错的环节。层数少、单卡显存又够用，直接流式加载原始 checkpoint，路径更短、加载更稳。换句话说，分片是为"装不下"服务的，而小模型本来就装得下，没必要为它切一刀。

从零跑通：三步完成 AutoModel 非分片加载

入口统一是AutoModel.from_pretrained，实现在 airllm/auto_model.py 里：它先读config判断架构，命中特殊布局（如 ChatGLM、Qwen、Kimi K3）走对应子类，其余架构统统落到通用的AirLLMBaseModel，后者能流式处理任何标准*ForCausalLM，所以新架构基本不用改代码。

Step 1｜先跑通非分片加载。小模型直接给本地路径或 repo id，不传任何切分参数即可：

from airllm import AutoModel # 小模型直接给本地路径或 repo id，走非分片流式加载 model = AutoModel.from_pretrained("Qwen/Qwen-7B")

Step 2｜上小模型 4bit 量化压缩磁盘分片。AirLLM 的瓶颈在磁盘 IO 而非显存，所以量化压的是 on-disk shards 的体积，用来加快加载。注意 compression 依赖bitsandbytes：

model = AutoModel.from_pretrained( "Qwen/Qwen-7B", compression="4bit", # 或 "8bit"，压缩磁盘上的 shards )

Step 3｜挂 profiling 并回收磁盘。调参阶段打开计时，确认无误后让框架删掉原始 checkpoint 省空间：

model = AutoModel.from_pretrained( "Qwen/Qwen-7B", profiling_mode=True, # 打印每层加载/压缩耗时 delete_original=True, # 用完删原始 checkpoint 省盘 )

跨平台落地：Linux、macOS、CPU 各自怎么跑

Linux 是默认场景，device取"cuda:0"，标准 GPU 流式推理；要用量化就先把bitsandbytes装上，否则 compression 参数不生效。

macOS 走的是另一条路：auto_model.py在platform == "darwin"时直接改调AirLLMLlamaMlx，底层交给 MLX 框架，贴合 Apple Silicon 的统一内存，省去了 CUDA 依赖。

CPU 推理是 v2.10.1 同时带来的能力，适合完全没独显的机器——把device指到 CPU 即可跑，代价是吞吐明显下降，适合验证逻辑或低并发批处理，不适合追求低延迟的在线服务。

开启profiling_mode调参、或量化后核对评估侧指标时，可以像上图这样跟踪曲线变化，确认压缩没有把精度打崩。

踩坑清单：这些配置项最容易翻车

磁盘不足时开delete_original=True会翻车。切分本身很吃盘，原始 checkpoint 和 shards 并存那段时间磁盘压力最大，盘不够就中途失败。解法是先给足余量，或手动清一下缓存目录再重跑，别在半途把空间吃光。

量化级别和显存/加载速度的匹配是第二坑。很多人以为 4bit 就是"更快"，但在 AirLLM 里 compression 会直接禁用 prefetching（源码里明确两者互斥），而 prefetching 正是用下一层的磁盘读取去重叠当前层计算的关键。所以 4bit 省了盘、却可能让加载变慢；显存和盘都够时优先 8bit，保精度、少踩这个互斥，追求极致省盘再上 4bit，并且别指望 prefetching 帮忙。

最后是profiling_mode的隐形开销。它会在每次load_layer里插桩计时（add_profiling_time），逐层累加。调参、排障时开着很有用，但生产环境长期开着只会平添计时成本，跑稳了就关掉。

🔧 配置不复杂，关键是分清"瓶颈在哪"。想确认自己的模型命中了哪个子类、compression和profiling_mode具体怎么接线，直接翻 airllm/auto_model.py 里from_pretrained的分支，或者拿你手头那个 7B 以下的小模型，按上面三步试跑一次，显存曲线会替你说话。

【免费下载链接】airllmAirLLM 70B inference with single 4GB GPU项目地址: https://gitcode.com/GitHub_Trending/ai/airllm