AirLLM：单卡4GB显存运行700亿大模型，革命性轻量化框架

Introduction

运行大型语言模型（LLM）传统上是一项资源密集型任务，通常需要昂贵的高端GPU和大量显存。这种硬件壁垒限制了许多开发者和研究人员接触和试验最先进模型的机会。AirLLM作为一项突破性解决方案应运而生。它通过创新的内存优化技术，使得在显存低至4GB的单张GPU上运行大规模模型（例如700亿参数模型）成为可能。本文提供了AirLLM的完整指南，涵盖安装、核心概念和高级优化技巧。

Quick Start

Environment Setup and Installation

• Python 3.8+
• PyTorch 1.13+

pip install airllm
pip install transformers peft accelerate bitsandbytes einops sentencepiece

开始使用AirLLM前，请确保系统满足以下先决条件：

• Python 3.8+
• PyTorch 1.13+
• CUDA 11+ (如果使用NVIDIA GPU)
• 充足的磁盘空间（模型拆分过程需要大量存储空间）
  使用pip安装AirLLM核心包及其依赖项：

pip install airllm
pip install transformers peft accelerate bitsandbytes einops sentencepiece

Basic Inference Example

from airllm import AutoModel

# Initialize the model (supports Hugging Face model IDs or local paths)
model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct")

# Prepare input text
input_text = ['What is the capital of the United States?',]

# Tokenization
input_tokens = model.tokenizer(input_text,
    return_tensors="pt",
    return_attention_mask=False,
    truncation=True,
    max_length=128,
    padding=False)

# Generate output
generation_output = model.generate(
    input_tokens['input_ids'].cuda(),
    max_new_tokens=20,
    use_cache=True,
    return_dict_in_generate=True)

# Decode and print the result
output = model.tokenizer.decode(generation_output.sequences[0])
print(output)

以下代码片段展示了使用AirLLM进行推理的最简单方法。其API设计让熟悉Hugging Face transformers库的用户感到亲切。

from airllm import AutoModel
>
# 初始化模型（支持Hugging Face模型ID或本地路径）
model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct")
>
# 准备输入文本
input_text = ['What is the capital of the United States?',]
>
# 分词处理
input_tokens = model.tokenizer(input_text,
    return_tensors="pt",
    return_attention_mask=False,
    truncation=True,
    max_length=128,
    padding=False)
>
# 生成输出
generation_output = model.generate(
    input_tokens['input_ids'].cuda(),
    max_new_tokens=20,
    use_cache=True,
    return_dict_in_generate=True)
>
# 解码并输出结果
output = model.tokenizer.decode(generation_output.sequences[0])
print(output)

Core Technical Principles

AirLLM之所以能在有限的硬件上运行大模型，并非魔法，而是多项专注于内存管理的关键技术创新的成果。其核心思想围绕着智能的模型分区和流式处理。AirLLM并非一次性将整个数十GB的模型加载到GPU显存中，而是将模型分割成层或更小的分片。在推理过程中，它动态地将仅需要的组件（例如，正在计算的当前层）加载到GPU显存中，而将其余部分保留在CPU内存甚至磁盘上。这种“交换”机制，结合预取和缓存策略，极大地降低了GPU显存的峰值占用。

Technical Feature Comparison

下表突显了AirLLM与传统推理方法相比所带来的变革性优势。

Advanced Configuration

1. Model Compression and Acceleration

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    compression='4bit'  # or '8bit'
)

您可以启用4位或8位量化来显著减小模型大小并加速推理，通常可实现高达3倍的速度提升。

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    compression='4bit'  # 或 '8bit'
)

2. Performance Profiling Mode

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    profiling_mode=True  # Outputs detailed timing information
)

启用性能分析模式可以深入了解推理不同阶段的时间消耗，这对于优化至关重要。

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    profiling_mode=True  # 输出详细的时间消耗信息
)

3. Custom Storage Path

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    layer_shards_saving_path="/path/to/your/custom/save/directory"
)

指定一个自定义目录来保存模型层/分片文件，这对于管理磁盘空间或使用网络存储非常有用。

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    layer_shards_saving_path="/path/to/your/custom/save/directory"
)

4. Prefetching Optimization

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    prefetching=True  # Enabled by default
)

预取功能默认启用，它有助于将后续模型层的加载与当前层的计算重叠进行，从而隐藏I/O延迟并提高吞吐量。

model = AutoModel.from_pretrained(
    "garage-bAInd/Platypus2-70B-instruct",
    prefetching=True  # 默认启用
)

Multi-Model Support

ChatGLM Model

from airllm import AutoModel
model = AutoModel.from_pretrained("THUDM/chatglm3-6b-base")
# ... subsequent processing is identical to the Llama2 example

from airllm import AutoModel
model = AutoModel.from_pretrained("THUDM/chatglm3-6b-base")
# ... 后续处理与Llama2示例相同

QWen Model

from airllm import AutoModel
model = AutoModel.from_pretrained("Qwen/Qwen-7B")
# ... subsequent processing

from airllm import AutoModel
model = AutoModel.from_pretrained("Qwen/Qwen-7B")
# ... 后续处理

Other Supported Models

# Baichuan
model = AutoModel.from_pretrained("baichuan-inc/Baichuan2-7B-Base")

# InternLM
model = AutoModel.from_pretrained("internlm/internlm-20b")

# Mistral
model = AutoModel.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")

AirLLM的支持扩展到许多其他架构，包括但不限于：

# Baichuan
model = AutoModel.from_pretrained("baichuan-inc/Baichuan2-7B-Base")
>
# InternLM
model = AutoModel.from_pretrained("internlm/internlm-20b")
>
# Mistral
model = AutoModel.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")