基于 AirLLM 实现 4GB 显存环境下的 Qwen2.5-72B 本地推理

背景与挑战

随着开源大语言模型的发展，通义千问 2.5（Qwen2.5）系列在多项评测指标上表现优异。其中，72B 参数版本在 MMLU、MATH、MBPP 等基准测试中展现了强大的能力，部分指标甚至超越了 Llama3 405B 和 GPT-4o。

然而，如此庞大的参数量对硬件资源提出了极高要求。以 FP16 精度为例，72B 模型仅权重就需要约 144GB 显存。即便使用 INT4 量化，也需要约 40GB 以上的显存才能流畅加载。对于拥有 4GB 显存的消费级显卡用户而言，直接运行该模型几乎是不可能的任务。

解决方案：AirLLM 分层推理

为了解决低显存环境下的超大模型部署问题，AirLLM 项目提供了一种创新的解决方案。其核心思想是分层推理（Layer-wise Inference）。

传统的推理方式通常需要将整个模型一次性加载到 GPU 显存中。而 AirLLM 允许将模型的不同层动态地加载到显存中。每次计算时，仅将当前需要的层载入 GPU，计算完成后卸载或替换为下一层。这种机制使得模型可以像流水一样通过有限的显存进行推理，从而突破了显存容量的限制。

原理简述

CPU-GPU 内存协同：利用系统内存（RAM）存储完整的模型权重，仅在需要计算时将特定层传输至 GPU。
动态调度：根据推理流程，自动管理层的加载与卸载，减少 PCIe 带宽的瓶颈影响。
兼容性：支持多种主流大模型架构，包括 Qwen2.5、Llama3、Mixtral 等。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

操作系统：Linux 或 Windows (WSL)
Python 版本：3.8 及以上
PyTorch：安装对应 CUDA 版本的 PyTorch
CUDA 驱动：确保显卡驱动已正确安装
硬件：至少 4GB 显存的 NVIDIA 显卡，建议配备较大容量的系统内存（RAM）

安装 AirLLM

可以通过 pip 直接安装 AirLLM 库：

pip install airllm

同时，确保安装了必要的依赖包：

pip install transformers accelerate torch

代码实现

以下是使用 AirLLM 加载并运行 Qwen2.5-72B-Instruct 模型的完整示例。该代码展示了如何初始化模型、处理输入以及生成输出。

from airllm import AutoModel
import torch

# 设置最大生成长度
MAX_LENGTH = 128

# 初始化模型
# 注意：首次运行会下载模型权重，请耐心等待
model = AutoModel.from_pretrained("Qwen/Qwen2.5-72B-Instruct")

# 准备输入文本
input_text = [
    "What is the capital of United States?",
]

# Tokenize 输入
input_tokens = model.tokenizer(
    input_text,
    return_tensors=,
    return_attention_mask=,
    truncation=,
    max_length=MAX_LENGTH,
    padding=
)



generation_output = model.generate(
    input_tokens[].cuda(), 
    max_new_tokens=,
    use_cache=,
    return_dict_in_generate=
)


output = model.tokenizer.decode(generation_output.sequences[])

(output)

基于 AirLLM 实现 4GB 显存环境下的 Qwen2.5-72B 本地推理