AirLLM 使单个 4GB GPU 能够进行 70B LLM 推理

概览

AirLLM 通过每次仅在设备上保留一个模型层,使得在消费级 GPU 上运行超大型语言模型成为可能。

AirLLM 的工作原理

该库将模型拆分为单个层,并仅在计算需要时才将每一层从磁盘流式传输到 GPU。这意味着 VRAM 的需求量由最大层的尺寸决定,而不是参数的总数。

支持的模型和 VRAM 需求

AirLLM 支持广泛的开源 LLM。下表显示了在使用默认(无压缩)模式时,各种模型所需的近似 VRAM:

  • Qwen3, Mistral, Phi (≈8 B) – ~1–2 GB
  • Qwen3‑30B, Mixtral (MoE) – ~1–3 GB
  • Qwen3‑235B (MoE) – ~3 GB
  • Llama 3.x 70B (full precision) – ~4 GB
  • Llama 3.1 405B – ~8 GB
  • DeepSeek‑V3 (671 B) – ~12 GB
  • Kimi K3 (2.8 T) – ~3.72 GB VRAM measured on an RTX 6000 Ada (requires additional dependencies)

性能特征

在默认模式下运行模型会产生高延迟,因为每个 token 会触发从磁盘加载一层。该库提供可选的分块压缩功能,可以在精度损失微乎其微的情况下将推理速度提高多达 3 倍。用户报告的关于 Kimi K3 在 RTX 6000 Ada 上的测量结果显示,每个 token 大约需要 292 秒。

设置与使用

  1. 安装包:pip install airllm
  2. 使用 AutoModel.from_pretrained(<repo_id>) 初始化模型。
  3. (可选)启用压缩:添加 compression='4bit'compression='8bit' 并确保已安装 bitsandbytes
  4. (可选)通过 hf_token 参数为受限模型提供 Hugging Face token。
  5. (可选)设置 layer_shards_saving_path 以控制拆分后的层存储在哪里。
  6. (可选)开启 profiling_mode=True 以查看时间分解。
  7. 按照快速入门示例所示运行生成。

注意:第一次运行会分解模型并保存层分片,因此需要在 Hugging Face 缓存目录中有足够的空闲磁盘空间。

平台支持

  • Linux 和 Windows:标准 PyTorch 安装。
  • macOS:仅支持 Apple Silicon;您必须按照 macOS 部分的说明安装 mlxtorch

局限性与权衡

  • 延迟:逐层流式传输使得交互式聊天变得缓慢;该技术更适合批处理或隔夜任务。
  • 磁盘使用:存储所有层分片可能会消耗大量空间;您可以启用 delete_original=True 以仅保留转换后的版本。
  • 默认不进行量化:基础方法避免了量化以保持精度;压缩是可选的附加功能。
  • 某些模型(例如 Kimi K3)有额外要求,例如 compressed-tensors, flash-attn, 一个 CUDA 12 构建的 PyTorch,以及 transformers 4.56.x。

社区观点

Hacker News 的评论者强调了兴奋感和实际担忧:

  • "Kimi K3 on RTX 6000 Ada (48GB) takes 292 s/token" – @imenani
  • "Seeing a lot of these ‘run 1TB models with 1GB RAM’ projects recently. Most seem vibe coded and probably won’t be maintained." – @roger_
  • "I love it how the rampocalypse is pushing people to squash all the performance they can. I hope that this also leads to rethinking model architecture." – @seu
  • "I’m still slightly confused on what this adds… Does this basically load layers in and out on demand? So I still have to download the full model to disk..." – @cpfohl
  • "You can run any frontier model on your PC if you just wait long enough..." – @xg15
  • "The layer‑by‑layer streaming is clever. Curious how the throughput compares to running a quantized model on the same GPU — seems like a speed might still win on quantization?" – @junsu22
  • "At this latency, I can still imagine batch or overnight jobs being interesting; for chat, time to first useful answer matters much more..." – @Alisaqqt

这些评论强调,虽然 AirLLM 实现了令人印象深刻的显存效率,但由此产生的延迟限制了交互式使用,并邀请了与 llama.cpp 或 unsloth 等量化推理引擎的比较。

Sources

相关