ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

4GB显存跑70B大模型:AirLLM分层推理原理与实战

4GB显存跑70B大模型:AirLLM分层推理原理与实战 第一次看到 AirLLM 这个项目标题时我第一反应是标题党70B 大模型单卡 4GB 显存就能跑这和我对显存容量的所有直觉是拧着的。70B 参数用 FP16 存权重就得接近 140GB一张消费级小卡连零头都装不下凭什么能跑后来把源码和 README 过了一遍才明白它并不是把 70B 塞进 4GB而是让 GPU 一次只照顾模型的一层跑完就扔逐层接力。这个思路不复杂但确实把能不能跑这件事的门槛拉到了一个以前不敢想的程度。AirLLM 是一个纯开源项目GitHub 上叫 lyogavin/airllm核心卖点就一句话在单张 4GB 显存的 GPU 上用分层推理的方式加载并运行 70B 级别的大模型。它非常适合下面这几类人手头没有 A100、H100 这种大显存卡但又想在本地体验 70B 模型、做离线批量推理的开发者想把私有数据喂给开源模型却不想把模型 API 化、也不想为显存付费的研究者以及单纯想搞明白显存不够时大模型推理到底还能怎么玩的硬核玩家。这篇文章我会把它的原理、实测代码、性能量级和常见坑一次讲清楚。1. 4GB 显存跑 70B到底反常识在哪里1.1 70B 模型的重量级不是 70GB是 140GB先算一笔账。所谓 70B指的是模型参数总量约 700 亿个。如果每个参数用 FP16 存储也就是 2 字节那么光权重文件就有大约 140GB。你要是把加载中间状态、KV Cache、激活值也算进去一张 80GB 显存的 A100 跑 70B 都谈不上轻松更别提 4GB 的消费级卡了。这里的单位换算很多人会搞混70B 这个数字是参数个数不是存储容量。参数个数乘上每个参数的字节数才是最基础的显存下限。很多教程喜欢说70B 量化成 4bit 后约 35GB这确实是压缩后的体量但即使是 35GB4GB 显存依然放不下。所以一开始我默认 AirLLM 走了什么黑魔法量化或者对模型做了蒸馏裁剪但读完实现才发现它走的是完全不同的路线不压缩模型改变的是模型的存放位置和计算时序。1.2 为什么常规推理框架在 4GB 卡上必然 OOM绝大多数深度学习框架的默认行为是先把整个模型权重全部加载到 GPU 显存然后才开始推理。transformers 库也好vLLM 也好在标准配置下都遵循这个逻辑。显存不够的时候会直接报 CUDA out of memory连推理的第一步都迈不出去。你可能会说不是有 device_mapauto 吗这个参数确实能把部分层放到 CPU 上但它的设计目标是充分利用多设备默认更倾向于优先填满 GPU然后溢出到 CPU。对一张 4GB 的小卡来说哪怕只放几层 Transformer也很容易把显存挤爆因为模型初始化时还会有额外的中间开销。而且大部分人在 4GB 卡上跑的是量化版小模型7B 或 13B极少有人敢想 70B。1.3 显存之外还有两道坎内存和带宽这里必须先给4GB 显存跑 70B补一个前提条件4GB 说的是显卡显存但整机 CPU 内存依然要足够大。70B 模型 FP16 权重约 140GB在 AirLLM 的机制里这些权重通常要常驻 CPU 内存或磁盘。如果你的机器只有 16GB 内存就需要用量化权重或者干脆让小模型跑否则依然是跑不起来的。带宽是另一道容易被忽略的坎。GPU 一次只能处理一层但模型有 80 层左右每一层都要从 CPU 内存搬到 GPU 显存算完再搬走。PCIe 总线再快也远不如显存内部带宽更不用说如果机器只有机械硬盘磁盘到内存再到显存这条链路会慢到让人怀疑人生。所以 AirLLM 这类方案的本质是用带宽换容量用时间换空间。它能做到的是跑起来而不是飞快地跑。2. AirLLM 的破局思路逐层加载用完即走2.1 Transformer 推理天然就是按层执行的要理解 AirLLM得先看一眼 Transformer 大模型的基本结构。一个标准的大模型可以粗分成三块最底层的 Embedding、中间一大摞重复的 Transformer Block以及最上层的输出头。Llama-70B 这种规模中间的 Transformer Block 有 80 层。每一层的计算逻辑是线性的输入向量进入第 1 层经过注意力机制和 MLP输出新的向量这个输出再作为第 2 层的输入。第 2 层只关心第 1 层给过来的中间结果并不需要第 1 层的权重还在显存里待命。这个特性是分层推理的根基。AirLLM 正是抓住了这一点既然层与层之间的依赖只是中间结果的传递而不是所有权重同时在线那就完全可以逐层搬运。2.2 核心循环Load-Forward-ReleaseAirLLM 的推理流程可以简化成下面这个循环首次加载时把目标模型的全部权重读入 CPU 内存或者按需从磁盘读入。在 GPU 上只预留一块刚好够放当前层权重和激活值的显存。把第 1 层权重从 CPU 内存拷贝到 GPU执行前向计算。得到第 1 层的输出中间结果后把它传回 CPU 内存保存然后立刻释放 GPU 上第 1 层的显存。加载第 2 层重复同样的过程直到最后一层。最后把输出头的计算在 GPU 或 CPU 上完成得到 logits。因为任意时刻 GPU 上只存在一层权重 当前激活 必要的 KV Cache所以峰值显存需求从整个模型量级骤降到某一层量级。对一个 70B 模型来说单层权重可能只有 1~2GB4GB 显存刚好能拿下一层。这才是4GB 跑 70B的真正含义。2.3 三级存储协作显存、内存和磁盘谁干什么AirLLM 之所以能做到底层是因为它没有把全部东西都压在显存里而是构建了一条显存 - 内存 - 磁盘的存储链路。显存只存放当前正在计算的那一层权重、输入激活和必要的中间缓存。内存保存剩下的所有层权重以及已经算出来的中间激活结果。这里的内存压力很大70B FP16 大约需要 140GB 内存。磁盘当内存也不够时AirLLM 会把一部分暂时用不到的层写回磁盘需要时再搬运到内存。这个行为很像虚拟内存的换页机制会进一步放大时延但保住了能跑的下限。不同版本的 AirLLM 在缓存策略上做过不少优化比如在同一层 GPU 计算的时候预先去加载下一层的数据。这个异步预取的细节很关键它防止 GPU 在大部分时间里空转等数据虽然整体依然慢但不会慢到完全不能用的地步。2.4 和 DeepSpeed ZeRO、llama.cpp 的 offload 有什么本质区别很多人看到 AirLLM 第一反应是这不就是 offload 吗。确实DeepSpeed ZeRO-Inference 也在做 CPU offloadllama.cpp 也在用 mmap 按需读取权重但它们的侧重点完全不同。DeepSpeed ZeRO 的强项是极大规模的训练和推理并行它把模型参数、梯度和优化器状态分片到多张卡配合 CPU offload 会涉及大量通信同步配置成本高通常是为多卡集群设计的。llama.cpp 则默认把计算放在 CPU 上通过修改后的 GGUF 格式做量化用 mmap 把权重映射到内存按需读页它的优势是 CPU 也能跑但如果你有一张 GPU想充分利用 CUDA 加速某一层的计算llama.cpp 的 GPU offload 是大粒度的层分配也不是逐层换入换出的逻辑。AirLLM 的定位非常轻在 PyTorch 生态内以最接近 HuggingFace 使用习惯的 API把GPU 只计算不存储这件事做到极致。它不追求多卡通信不依赖特殊内核只要 CUDA 能用就能用。3. 实操在 4GB 显卡上从零跑通 AirLLM3.1 环境准备先验证 torch.cuda.is_available()AirLLM 的安装本身不复杂最省事的方式是直接走 pippip install airllm但真正决定你能不能跑起来的是 PyTorch 和 CUDA 的版本。建议先把 PyTorch 单独装好确认你的 Python 环境里能正常执行下面这段代码import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果torch.cuda.is_available()返回 False后面 AirLLM 所有流程都白搭。这里有一个容易踩的细节点很多人的机器上其实装了多个 Python 环境pip 安装和 Jupyter Kernel 用的可能是不同的解释器一定要在同一个环境里完成所有操作。另外如果遇到和 transformers 的兼容问题比如某个 API 在新版本里被改名导致报错可以直接从源码安装 AirLLMgit clone https://github.com/lyogavin/airllm.git cd airllm pip install -e .我个人的建议是有条件的话优先在 Linux 环境跑Windows 下不是不能用但环境变量、路径和 CUDA 驱动的问题会多出不少调试成本更高。3.2 模型权重从哪来在线下载或本地路径AirLLM 的底层依赖 HuggingFace 生态from_pretrained会默认去 HuggingFace Hub 拉取权重。要注意的是meta-llama 官方的 Llama-2-70B 属于受限模型需要你先在 Hub 上申请访问权限再用 token 登录才能下载。如果嫌麻烦可以直接用不需要额外申请的微调版本比如garage-bAInd/Platypus2-70B它的基座是 Llama-2-70B但权重文件是非受限状态能省掉不少事。如果你所在环境的网络访问 HuggingFace 不稳定或者公司内网屏蔽了外网也可以先找国内镜像站手动下载模型文件然后落到本地目录。AirLLM 的from_pretrained是支持本地路径的你只要把参数从模型名改成文件夹路径就行比如model_path /data/models/Platypus2-70B这里有个实操技巧即使最后要用 AirLLM 跑 70B也建议先用一个几百 MB 的小模型比如 7B 版本把整条链路跑通确认权重路径、tokenizer 和生成逻辑都没问题再去下载几百 GB 的 70B 权重。3.3 最小可运行代码一段一次能抄走的示例下面这段代码是 AirLLM 跑 70B 模型的最小闭环我在 4GB 显存的机器上用过类似版本只要内存和磁盘够能跑完整个生成过程import os os.environ[AIRLLM_GPU_MEMORY_THRESHOLD] 0.8 from transformers import AutoTokenizer from airllm import AutoModel, InferenceConfig model_path garage-bAInd/Platypus2-70B tokenizer AutoTokenizer.from_pretrained(model_path) config InferenceConfig(max_seq_len512, torch_dtypetorch.float16) model AutoModel.from_pretrained(model_path, inference_configconfig) prompt Explain quantum computing in simple terms. inputs tokenizer(prompt, return_tensorspt).to(cuda) output model.generate(**inputs, max_new_tokens64) print(tokenizer.decode(output[0], skip_special_tokensTrue))代码逻辑很简单但有几个关键点要理解。InferenceConfig(max_seq_len512, torch_dtypetorch.float16)里的max_seq_len不是随便填的它决定了 AirLLM 在显存里预分配的缓存槽位大小如果你的max_new_tokens设置得比max_seq_len小很多生成过程中要预留的空间也会更宽松。torch_dtypetorch.float16能让权重按半精度读取显存压力更小。3.4 控制资源的关键参数max_seq_len、环境变量和 batchAirLLM 的默认配置可以满足能跑但实际使用中你大概率要调几个参数。max_seq_len这是最重要的显存控制旋钮。4GB 显存下512 是一个比较稳的起步值如果 OOM先把它降到 256。AIRLLM_GPU_MEMORY_THRESHOLD环境变量控制 GPU 显存使用阈值。设置为 0.8 之类的值时AirLLM 会尽量把显存留给当前层计算但又不会把所有显存都吃光。AIRLLM_CPU_MEMORY_THRESHOLD控制 CPU 内存使用阈值超过这个阈值后部分层会溢出到磁盘。如果你的内存比较紧张可以把阈值调低让 AirLLM 更早地把层写回磁盘保住进程不崩溃。AIRLLM_TORCH_DEVICE手动指定推断设备。默认会自动选择 cuda但如果机器上有多个设备可以用它固定到某一张卡。AIRLLM_CACHE_DIR修改 AirLLM 的权重缓存目录默认在用户目录下如果系统盘空间不够建议指到一块大容量 SSD。这些环境变量在不同小版本里可能略有差异写进脚本之前花半分钟翻一眼当前版本的 README 是最稳妥的。还有一个很容易被忽略的点AirLLM 这类逐层推理方案最适合的 batch size 就是 1。只要 batch 变大激活值就会成倍增加GPU 峰值显存会瞬间上涨速度也会变得更不可控。设计实验时尽量把请求拆成单条处理不要贪多。3.5 先用小模型验证流程再上 70B我见过不少人在 4GB 显存机器上第一次跑 AirLLM就直接指到 70B 权重然后卡在下载阶段等了半天最后因为磁盘空间不足失败非常打击信心。一个更务实的做法是先用 7B 或 13B 级别的模型验证确认安装没问题确认 model_path 是本地路径时 AirLLM 能正确加载确认 generate 出来的文本质量正常观察模型跑一个完整生成周期时CPU 内存和 GPU 显存的变化曲线。等这些都正常了再切换到 70B你会有把握得多。毕竟 70B 权重少则几十 GB、多则一两百 GB下载和校验本身就是一件不轻松的事情。4. 实测性能与常见坑4.1 速度量级为什么每生成一个 token 都要等半天先给一个直观结论AirLLM 在小显存卡上跑 70B速度大约在每生成一个 token 需要几十秒的量级。这个数字来自社区反馈和我的实际体验不同机器差异很大但不要期望它接近正常的 GPU 推理速度。原因可以算一笔账。假设你的 PCIe 3.0 x16 实际带宽在 10GB/s 左右70B FP16 权重合计约 140GB每生成一个 token都要把所有层完整前向传播一遍意味着要把 140GB 的权重至少从 CPU 内存搬到 GPU 一次。光传输时间就要 14 秒再加上每一层的计算、中间激活的回传、缓存管理、以及可能的磁盘交换30 到 60 秒一个 token 是非常正常的。所以用 AirLLM 跑生成任务时心态要调整过来。它不是让你实时聊天的而是让你在后台慢慢泡咖啡的。我一般用它做离线批处理比如给一组测试 prompt 生成回复样例、在私有数据上跑评估、或者验证某个微调版本在 70B 规模下的回答风格。4.2 CPU 内存不够怎么办量化叠加和磁盘缓存前面反复提到 70B FP16 权重大约需要 140GB 内存。如果你的机器内存只有 64GB连完整 FP16 权重都放不下AirLLM 的磁盘换页机制会开始介入把一部分层丢到磁盘上速度会进一步变慢。想改善这种情况可以引入量化权重。AirLLM 的分层推理和你自己叠加的量化方案并不冲突它的核心循环是在加载一层权重到 GPU 计算层面至于这一层权重本身是 FP16 还是 4bit并不影响流程。社区里比较常见的做法是先把模型权重量化成 4bit 或 8bit再用 AirLLM 加载。量化之后 70B 权重可能压缩到 35GB~70GB普通工作站内存也能撑住速度会比磁盘交换顺畅很多。这里有一个细节要提醒量化本身会带来一定精度损失尤其对复杂推理类任务。如果只是小范围跑几个例子体感可能不明显但批量评估时务必保留一组 FP16 或高精度结果作为对照防止被量化误差带偏结论。4.3 常见 crash 排查OOM、下载中断、版本冲突我把自己和社区里碰到过的坑按频率排了个序GPU OOM。绝大多数情况是max_seq_len设置得太大或者max_new_tokens太长导致 KV Cache 和激活值超过了单层计算预留量。解决办法是把max_seq_len降到 256 或更小同时把 batch 固定为 1。CPU 内存耗尽。70B FP16 需要约 140GB 内存小内存机器会在某个层加载时直接 kill 进程。优先检查系统内存如果不足用量化权重或者直接换小模型。HuggingFace 下载中断。大权重文件下载经常因为网络波动中断表现是进程卡住或者报连接错误。可以先用下载工具把权重完整拉下来再走本地路径加载。transformers 版本冲突。AirLLM 迭代过程中依赖的 transformers API 会变化如果你本地刚好更新了大版本可能出现某个类找不到或者参数不匹配的报错。解决办法就是按照 README 推荐版本安装或者把 AirLLM 升级到最新源码版。磁盘空间不足。70B 模型的权重文件有好几百 GB 的缓存和落盘需求跑之前先确认AIRLLM_CACHE_DIR指向的磁盘有足够空间别把家目录塞爆。4.4 这个方案真正适合的场景AirLLM 不是生产级推理方案这一点必须想清楚。它适合的场景基本都有几个共同点不追求低延迟、数据量大但能异步处理、硬件条件有限但想用大模型。我最常用的场景是两个。一是给一批文本跑自动评估比如对比几个 70B 模型的回答风格差异样本量几百条跑一晚上能出结果完全可接受。二是隐私敏感场景数据不能出内网但内网的 GPU 只是一张老款的 4GB 卡这时候 AirLLM 是少数能让你在本地把 70B 模型跑起来的 PyTorch 方案。反而如果你要做线上聊天机器人、高并发 API或者需要秒级响应AirLLM 帮不了你。5. AirLLM 和其他低配跑大模型方案的取舍5.1 主流方案横向对比为了帮你更快做选型我把 AirLLM、llama.cpp、vLLM 和 DeepSpeed ZeRO-Inference 放在一张表里对比维度AirLLMllama.cppvLLMDeepSpeed ZeRO-Inference最低显存要求约 4GB只需容纳单层可以不依赖大显存GPU offload 可选很高70B 通常需要 80GB 以上面向多卡集群依赖通信环境推理速度很慢受 PCIe 和内存带宽限制中等CPU 优化内核成熟量化支持好极快高吞吐适合生产服务较快但主要面向多卡并行依赖生态PyTorch、HuggingFace APIGGUF 格式C/C 推理PyTorch、PagedAttentionPyTorch、NCCL、多卡管理模型格式HuggingFace 原生权重需要转成 GGUFHuggingFace 权重HuggingFace 权重上手难度低代码量和 transformers 一致中工具链独立中服务端参数较多高配置复杂最佳场景小显存、离线批处理、PyTorch 实验CPU 机器、本地聊天、轻量部署显存充足、高并发生产多卡训练/推理统一管理5.2 什么时候选 AirLLM什么时候选 llama.cpp如果你的设备根本没有独立 GPU或者只有集成显卡那 AirLLM 的优势发挥不出来优先考虑 llama.cpp。llama.cpp 对 CPU 推理做了很多底层优化量化后 70B 模型在内存足够的情况下可以慢慢跑生态也很成熟Ollama 这类工具已经把它包装得很友好。如果你手里有一张 4GB 的小 GPU又想用 PyTorch 和 HuggingFace 生态做实验AirLLM 几乎是当前最顺滑的选择。它不需要转格式代码风格很接近你平时写推理脚本的方式调试和二次开发的门槛都低。尤其当你需要把 AirLLM 的推理过程嵌到自己的数据处理流程里PyTorch 的 API 就是一个巨大的优势。vLLM 是另一个极端性能和吞吐都极好但显存门槛不是随便能绕开的。如果你的机器已经有一张 80GB 的 A100或者预算允许租卡那没必要折腾 AirLLMvLLM 才是正路。最怕的就是拿着 4GB 显卡硬上 vLLM然后因为 OOM 怀疑人生。5.3 多卡扩展和下一步能玩什么AirLLM 并不是只能单卡运行它同样支持把不同层分配到多张 GPU 上通过减少单卡的逐层交换次数来提升速度。如果你手头有若干张 4GB 卡比如一台老机器插了三四张旧显卡多卡模式会比单卡流畅不少。具体开启方式在项目 README 的多卡章节里写得很清楚主要思路就是让 AirLLM 知道每一层应该放在哪张设备上。我个人还比较看好这种分层推理思路的延展空间。当前 AirLLM 解决的是推理阶段显存不够的问题但同样的思想完全可以延伸到微调、评估、甚至多模态模型上。社区里已经有人把它和参数高效微调工具组合在小显存机器上尝试 LoRA 之类的实验虽然速度不快但至少提供了一个不用大规模攒硬件就能做研究的路径。回到我自己的实际感受AirLLM 不是那种装完就让你激动到原地起飞的性能工具它更像是一个门槛消失器。它让你意识到真正限制你本地跑大模型的往往不是显卡多贵而是你对显存、内存、带宽这三件事的理解是否到位。以前我会不假思索地说70B 需要大显存被 AirLLM 折腾过一遍之后我会先问一句你是想飞快跑还是只要能跑想飞快跑老老实实租 A100只要能跑4GB 显卡配一个大内存的旧工作站AirLLM 就能给你开一条路。
返回列表