AirLLM逐层推理：4GB显存运行70B大模型的技术原理与实践

1. 项目概述：当大模型遇见“显存焦虑”

最近在折腾本地大模型的朋友，估计没少为显存发愁。想跑个70B参数的模型，动辄要求几十GB显存，手里的消费级显卡（比如RTX 4060的8GB，或者RTX 3090的24GB）看着就力不从心。更别提那些动辄数百亿、上万亿参数的“巨无霸”了，那简直是云端计算资源的专属游戏。但就在这个节骨眼上，一个叫AirLLM的开源项目火了，它的口号相当震撼：4GB显存跑70B模型，8GB显存跑405B模型，甚至能用3.7GB显存去推理一个有2.8万亿参数的模型（比如Kimi的K3）。

这听起来有点“违背物理定律”，对吧？显存就像模型运行时的“工作台”，模型参数和计算中间结果都得放在上面。一个70B的FP16模型，光参数就要占掉140GB，4GB显存连零头都装不下。AirLLM的魔力就在于，它根本就没打算把整个模型都搬到显存里。它的核心思路叫做“逐层推理”。你可以想象成处理一个超长的流水线作业：整个模型（原材料）存放在你的硬盘或者内存里，GPU（工作站）每次只从流水线上取下一层模型（一个零件）到显存（工作台）上进行计算，算完立刻把结果传给下一层，同时把这一层的参数从显存里清出去，再加载下一层。

这种方法彻底颠覆了传统大模型推理需要将整个模型加载到显存的范式。对于绝大多数个人开发者、研究者，或者预算有限的小团队来说，这无疑打开了一扇新的大门。你不再需要苦苦等待云厂商昂贵的API，或者斥巨资购买多张A100/H100，用现有的、甚至是一些老旧的显卡，就能在本地探索和部署曾经遥不可及的超大规模模型。接下来，我们就深入拆解AirLLM是如何实现这一“魔法”的，以及在实际操作中你会遇到什么，又该如何避开那些坑。

2. 核心原理拆解：逐层推理与显存优化的艺术

AirLLM的技术核心并不复杂，但实现得非常巧妙。要理解它，我们需要先看看标准的大模型推理是怎么“吃”显存的。

2.1 传统全量加载的显存瓶颈

在标准的PyTorch或Hugging Face Transformers库加载模型时，通常采用全量加载模式。以Llama 2 70B模型为例，如果使用半精度（FP16）存储，其参数所占空间约为70B * 2 bytes = 140 GB。这140GB需要一次性加载到GPU显存中，才能开始前向传播计算。这还没算上计算过程中产生的激活值（Activations）、优化器状态（如果微调）以及KV Cache（用于加速自回归生成）所占用的显存。所以，实际需求往往远超140GB。这就是为什么像70B这样的模型，通常需要至少2张80GB显存的A100才能流畅运行。

这种模式对硬件的要求形成了极高的门槛，将很多有趣的实验和创新挡在了门外。

2.2 AirLLM的“化整为零”策略

AirLLM的解决方案是Layer-wise Inference（逐层推理）。它将整个神经网络模型视为一个由许多层（如Transformer Block）堆叠起来的结构。推理时，它不再一次性加载所有层，而是：

1. 单层加载：将模型权重以层为单位，存储在速度较慢但容量巨大的介质上，如系统内存（RAM）甚至固态硬盘（SSD）。当需要计算某一层时，仅将该层所需的权重从内存/硬盘加载到GPU显存。
2. 即时计算与卸载：在GPU上完成该层的计算后，立即将这一层的权重从显存中移除，释放空间。计算得到的激活值（该层的输出）会保留，作为下一层的输入。
3. 流水线推进：重复步骤1和2，像流水线一样，一层接一层地完成整个模型的前向传播。

这个过程听起来简单，但实现起来有几个关键的技术挑战和优化点：

• 权重存储格式：为了加快单层权重的加载速度，AirLLM通常会将模型权重预先转换成一种更适合快速读取的格式，并对权重进行量化（如INT4、INT8），进一步减少单次需要传输的数据量。例如，一个70B的模型，如果量化到INT4，那么单层参数可能只有几十到几百MB，从内存加载到显存的速度极快。
• 激活值管理：虽然权重被逐层卸载，但层与层之间传递的激活值必须保留在显存中。对于生成任务，随着生成序列变长，激活值也会增长。AirLLM需要精细管理这部分显存，有时可能需要对长序列进行切块处理。
• IO与计算重叠：为了不让硬盘或内存的IO（读取权重）成为性能瓶颈，AirLLM会使用预取（Prefetching）技术。即在计算当前层时，异步地将下一层所需的权重提前加载到CPU内存的缓冲区中，当需要时再快速送入GPU，从而实现IO和GPU计算的重叠，最大化GPU利用率。

一个生活化的比喻：想象你要阅读一本1000页的巨著（大模型）。传统方法要求你必须有能同时摊开1000页的巨大桌子（大显存）。而AirLLM给你的是一张只能放一页纸的小书桌（小显存），但你有一个高效的书架（内存/硬盘）。你的阅读流程是：从书架上取下一页，放在书桌上阅读，读完把这一页放回书架，再取下下一页。只要你的“取书-放书”动作够快，你就能用一张小书桌读完整个巨著。AirLLM的核心就是优化“取书-放书”这个流程，让它快到几乎不影响“阅读”体验。

3. 环境部署与模型准备实战

理解了原理，我们来看看如何亲手把AirLLM跑起来。这里我以在Ubuntu 22.04系统，搭配一张RTX 4060（8GB显存）的显卡上，运行一个量化后的70B参数模型为例。

3.1 基础环境搭建

首先确保你的Python环境（建议3.9以上）和CUDA工具包（>=11.7）已经就绪。然后安装AirLLM：

pip install airllm

AirLLM的依赖相对干净，主要基于PyTorch。如果你的PyTorch还没有安装GPU版本，需要先安装对应CUDA版本的PyTorch。

注意：这里有一个新手常踩的坑。如果你之前用pip install torch安装了CPU版本，需要先卸载pip uninstall torch torchvision torchaudio，然后去 PyTorch官网 根据你的CUDA版本，复制对应的安装命令。例如，对于CUDA 11.8，命令可能是pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118。

3.2 模型获取与转换

AirLLM本身不提供模型，它支持加载Hugging Face格式的模型。但直接加载原始模型是行不通的，因为我们需要的是量化后并按层存储的格式。社区通常已经为我们准备好了这些模型。

以NousResearch/Llama-2-70b-hf这个模型为例，我们无法直接使用。我们需要寻找已经用airllm工具量化并分片好的版本。这些模型通常会在Hugging Face Hub上以*-airllm或类似后缀命名。

实操步骤：使用内置工具转换模型

AirLLM提供了一个命令行工具，可以将标准的Hugging Face模型转换为它需要的格式。但请注意，转换一个70B模型需要大量的CPU内存和硬盘空间（可能需要200GB+的临时空间），过程也非常漫长。

# 假设我们已经下载了Llama-2-70b-hf到本地目录 ./llama-2-70b-hf # 使用airllm的转换命令，进行INT4量化并分片存储 airllm convert ./llama-2-70b-hf ./llama-2-70b-airllm-int4 --quantization int4

这个过程会持续数小时甚至更久。对于大多数用户，我强烈建议直接下载社区已经转换好的模型，而不是自己转换。你可以在Hugging Face上搜索关键词如llama-2-70b-airllm-int4。

例如，找到一个模型后，你可以使用huggingface-cli下载：

pip install huggingface-hub huggingface-cli download TheBloke/Llama-2-70B-AirLLM-INT4 --local-dir ./llama-2-70b-airllm-int4

下载完成后，检查模型目录，你会看到很多以.bin结尾的分片文件，每个文件对应模型的一层或几层，这就是AirLLM能逐层加载的关键。

3.3 首次运行与验证

模型准备好后，写一个最简单的Python脚本来验证一切是否正常：

from airllm import AutoModel # 指定你下载的AirLLM格式模型路径 model_path = "./llama-2-70b-airllm-int4" # 加载模型。注意，这里不会把整个模型读入显存，只是初始化了模型结构。 model = AutoModel.from_pretrained(model_path) # 准备输入 input_text = "请用中文介绍一下AirLLM这个项目。" input_tokens = model.tokenizer(input_text, return_tensors="pt").input_ids.to('cuda') # 生成文本 # max_length控制生成总长度，airllm内部会自动管理显存 generated_ids = model.generate(input_tokens, max_length=100) output_text = model.tokenizer.decode(generated_ids[0], skip_special_tokens=True) print("模型回复：", output_text)

第一次运行AutoModel.from_pretrained时，它会加载模型的配置文件（如config.json）和分词器，速度很快。当你调用generate时，才会开始真正的逐层推理。观察你的GPU显存占用（可以用nvidia-smi命令），你会发现显存占用会在一个较低的水平波动，而不是一下子被撑满。

4. 关键参数调优与性能深度解析

让AirLLM跑起来只是第一步，让它跑得又快又好，就需要理解并调整一些关键参数。这部分是区分“能用”和“好用”的关键。

4.1 影响性能的核心参数

在初始化模型和生成时，有几个参数至关重要：

model = AutoModel.from_pretrained( model_path, device_map='cuda', # 指定使用GPU max_cache_size=4.0, # **关键参数**：单位GB，控制用于缓存模型层的“显存池”大小 offload_folder='./offload' # 如果单层模型超过max_cache_size，临时卸载到硬盘的路径 ) generated_ids = model.generate( input_tokens, max_length=200, min_length=50, do_sample=True, temperature=0.8, top_p=0.95, repetition_penalty=1.1, # AirLLM特有参数 use_beam_search=False, # 束搜索会大幅增加激活值显存，慎用 )

• max_cache_size：这是AirLLM的灵魂参数。它定义了一个显存缓存池。AirLLM会尝试将最常用或即将用到的模型层保留在这个缓存里，避免反复从内存加载。例如，设为4.0（GB），那么AirLLM会努力让缓存的层权重总量不超过4GB。这个值设置得太小，会导致缓存命中率低，频繁IO，速度慢；设置得太大，可能会挤占激活值所需的空间，导致OOM（内存溢出）。经验法则：对于8GB显存的卡，跑70B INT4模型，可以尝试设置为3.0到4.0。你需要根据任务和模型大小进行微调。
• offload_folder：当某一层模型的大小超过了max_cache_size（对于非常大的层），或者系统内存不足时，AirLLM会启用第二级卸载，将数据临时存到硬盘。务必指向一个高速SSD路径，否则性能会急剧下降。
• 生成参数中的use_beam_search：束搜索（Beam Search）在生成时会在内存中保留多个候选序列，这对于逐层推理来说是显存杀手，因为它会使每一层的激活值成倍增加。在显存紧张的情况下，建议保持为False，使用采样（Sampling）策略。

4.2 性能监控与瓶颈分析

如何判断你的AirLLM运行是否高效？主要看两个指标：

1. GPU利用率（GPU-Util）：通过nvidia-smi -l 1持续观察。理想状态下，在生成文本时，GPU利用率应该持续在高位（如70%以上）。如果利用率频繁地、大幅度地波动（如从90%骤降到10%，再升回去），这通常意味着IO瓶颈——GPU在等待从内存或硬盘加载下一层权重。
2. 每Token生成时间：计算生成一段文本的总时间除以生成的token数量。这是最直观的体验指标。

如果发现性能不佳，可以按以下步骤排查：

• 瓶颈在IO：表现为GPU利用率周期性波动。解决方案：
  • 增加max_cache_size，让更多层留在显存缓存。
  • 确保模型文件存放在NVMe SSD上，而不是机械硬盘。
  • 检查系统内存是否充足，如果内存不足导致频繁使用offload_folder到硬盘，速度会慢很多。
• 瓶颈在计算：表现为GPU利用率持续很高，但每Token生成时间依然很慢。这可能是模型本身的计算量太大（例如使用了未量化的模型）。解决方案：
  • 使用更低比特的量化模型（如从INT8换到INT4）。
  • 如果支持，考虑开启PyTorch的torch.compile（需要PyTorch 2.0+）对模型进行图优化，可能获得一定的加速。

4.3 与Ollama、vLLM等方案的对比

你可能也听说过Ollama、vLLM等优秀的本地大模型部署工具。这里简单对比一下：

总结一下：如果你的唯一约束是显存太小，想挑战在4GB或8GB卡上运行百亿、千亿模型，AirLLM是目前几乎唯一的选择。如果你追求开箱即用的流畅体验，Ollama很棒。如果你要搭建一个高并发的模型服务，vLLM是专业之选。

5. 高级应用场景与避坑指南

掌握了基础运行和调优后，我们可以探索一些更实际、也更容易踩坑的应用场景。

5.1 长文本处理与上下文长度扩展

大模型的一个魅力是长上下文。但逐层推理下，长序列会带来巨大的激活值显存占用。假设序列长度为L，模型隐藏层维度为H，那么一层的激活值（fp16）就占L * H * 2字节。对于Llama 2 (H=8192)，一个4096长度的序列，单层激活值就要占4096 * 8192 * 2 ≈ 67 MB。70B模型有80层，光保存所有层的激活值就需要超过5GB显存，这还没算KV Cache。

AirLLM的应对策略与实操： AirLLM在处理长文本时，可能需要结合序列切分技术。它不是一次性处理整个长序列，而是将其分成重叠的块（chunks），逐块进行前向传播，再合并结果。这需要模型本身支持这种“滑动窗口”注意力机制，或者使用外挂的上下文扩展方法。

避坑提示：

• 不要盲目追求极限上下文长度。先从1024、2048开始测试，监控显存占用。
• 如果遇到生成长文本时显存溢出，首先尝试减少生成时的max_new_tokens，或者降低max_cache_size，为激活值腾出空间。
• 查阅模型文档，确认其是否原生支持长上下文，或是否有配套的扩展方案（如LongLoRA）。

5.2 与LangChain等应用框架集成

AirLLM本身是一个推理引擎，要构建复杂的AI应用，需要与LangChain、LlamaIndex等框架集成。好消息是，由于AirLLM提供了标准的generate接口，集成起来并不复杂。

示例：创建AirLLM的LangChain LLM封装器

from langchain.llms.base import LLM from typing import Optional, List, Any, Mapping from airllm import AutoModel class AirLLMWrapper(LLM): model_path: str model: Any = None tokenizer: Any = None max_cache_size: float = 4.0 def __init__(self, model_path: str, max_cache_size: float = 4.0, **kwargs): super().__init__(**kwargs) self.model_path = model_path self.max_cache_size = max_cache_size # 延迟加载，避免在初始化时就加载模型 self._model = None self._tokenizer = None @property def _llm_type(self) -> str: return "airllm" def _call(self, prompt: str, stop: Optional[List[str]] = None, **kwargs) -> str: if self._model is None: # 首次调用时加载模型 self._model = AutoModel.from_pretrained( self.model_path, max_cache_size=self.max_cache_size, offload_folder='./offload' ) self._tokenizer = self._model.tokenizer inputs = self._tokenizer(prompt, return_tensors="pt").input_ids.to('cuda') outputs = self._model.generate(inputs, max_length=kwargs.get('max_length', 100)) response = self._tokenizer.decode(outputs[0], skip_special_tokens=True) # 去除输入提示部分，只返回新生成的文本 return response[len(prompt):] # 使用示例 llm = AirLLMWrapper(model_path="./llama-2-70b-airllm-int4", max_cache_size=3.5) print(llm("什么是机器学习？"))

这样，你就可以把这个llm对象代入到LangChain的任何Chain中使用了，比如做检索增强生成（RAG）。

集成时的坑：

• 线程安全：AirLLM模型对象可能不是线程安全的。如果在多线程Web服务（如FastAPI）中直接使用，可能会出错。建议为每个请求创建独立的模型实例，或使用线程锁进行保护。更好的方式是将AirLLM作为一个独立的后端服务，通过HTTP接口调用。
• 内存泄漏：长时间运行后，注意监控系统内存。由于Python的垃圾回收和CUDA内存管理机制，在反复加载卸载模型层时，可能会有内存碎片或未及时释放的内存。定期重启进程是一个简单粗暴但有效的办法。

5.3 模型微调的可能性探讨

一个很自然的问题是：能用AirLLM的方式微调大模型吗？答案是：理论上可行，但非常复杂且不推荐。

训练/微调需要保存优化器状态、参数梯度，并且需要反向传播，这要求所有层的参数在计算梯度时都必须可访问。逐层加载会使得反向传播的链路断裂。虽然有一些研究致力于“逐层训练”或“内存高效的训练”，如Zero-Offload、FSDP等，但它们的设计远比AirLLM的推理方案复杂。

AirLLM的核心定位是推理。对于微调，如果你的目标是在有限资源下微调大模型，应该去关注LoRA、QLoRA这类低秩适配技术。QLoRA甚至可以在单张24GB的3090上微调65B的模型。你可以用QLoRA进行微调，然后将合并后的模型再用AirLLM的方式转换用于推理，这才是更高效的组合拳。

6. 常见问题排查与实战心得

最后，分享一些我在折腾AirLLM过程中遇到的实际问题和解决心得，希望能帮你少走弯路。

6.1 问题速查表

6.2 实战心得与技巧

1. 模型来源是成功的第一步：自己转换模型耗时耗力，失败率高。优先在Hugging Face Hub上搜索[模型名]-airllm-int4或咨询社区。像TheBloke这样的用户经常会上传各种模型的AirLLM量化版。
2. 从“小”模型开始试水：不要一上来就用70B模型。先找个7B或13B的AirLLM版本模型，确保你的环境、流程全部跑通，对显存占用和速度有个感性认识，再挑战大家伙。
3. 监控是你的眼睛：打开两个终端窗口，一个跑程序，一个用watch -n 0.5 nvidia-smi实时监控显存和GPU利用率的变化。用htop监控内存和CPU。数据比感觉更可靠。
4. offload_folder一定要用SSD：这个参数看起来是备胎，但一旦被用到，如果指向机械硬盘，速度会立刻降到令人无法忍受的程度。确保它在一个高速的NVMe SSD分区上。
5. 理解“速度”的代价：AirLLM让你能用小显存跑大模型，代价就是生成速度（延迟）。它的吞吐量可能不高，尤其是首次加载（冷启动）时。这适合做研究、实验、或者对实时性要求不高的批处理任务，不适合需要毫秒级响应的对话场景。
6. 社区是宝藏：AirLLM是一个快速发展的项目，遇到奇怪的问题，去GitHub的Issues区搜一搜，很可能已经有人遇到并解决了。主动提问时，附上你的环境信息、错误日志和已经尝试过的步骤，能更快获得帮助。

AirLLM就像给大模型推理领域扔下的一颗“技术震撼弹”，它用一种聪明而直接的方式，打破了显存资源的硬约束。虽然它在性能上做了妥协，但极大地提升了大型语言模型的可及性。随着模型量化技术、IO优化技术的不断进步，相信这类“小显存跑大模型”的工具会越来越成熟、易用。对于每一个渴望在本地探索AI前沿的个人开发者来说，这无疑是一个值得投入时间学习和尝试的强大工具。