使用 device_map 优化 Transformers 模型的多 GPU 显存分配

在部署大型语言模型（如 Hugging Face Transformers 模型）时，显存管理是优化性能和避免内存溢出的关键。device_map 参数提供了一种灵活的方式，帮助用户在多 GPU 环境下分配模型的不同部分。本文将详细介绍 device_map 的用法、适用场景以及常见参数的作用。

---

为什么需要 device_map？

当模型大小超过单个 GPU 的显存容量时，模型并行化（Model Parallelism）成为必要手段。通过将模型的不同层分配到不同的 GPU 上，可以显著降低单个设备的显存压力。device_map 参数允许用户控制模型在多个 GPU 上的分布策略，从而实现显存的高效利用。

---

device_map 的四种模式

以下是 device_map 支持的四种模式及其适用场景：

1. device_map="auto" 或 device_map="balanced"

• 作用：在所有可用的 GPU 上均匀分配模型。 • 适用场景：当 GPU 的显存容量相近时，这是最推荐的方式。例如，两个 24GB 显存的 GPU。 • 示例代码：

python
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-VL-7B",
    device_map="auto",  # 或 "balanced"
    torch_dtype=torch.float16
)

2. device_map="balanced_low_0"

• 作用：优先在其他 GPU 上均匀分配模型，仅在显存不足时使用 GPU 0。 • 适用场景：当 GPU 0 需要预留显存用于其他任务（如数据预处理或梯度计算）时。例如，在多任务系统中，GPU 0 可能负责推理以外的计算。 • 示例代码：

python
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-VL-7B",
    device_map="balanced_low_0",
    torch_dtype=torch.float16
)

3. device_map="sequential"

• 作用：按顺序填充 GPU（从 GPU 0 开始，依次使用 GPU 1、GPU 2...），仅当当前 GPU 显存不足时才使用下一个 GPU。 • 适用场景：当 GPU 的显存容量差异较大时。例如，GPU 0 有 24GB，GPU 1 有 48GB。 • 示例代码：

python
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-VL-7B",
    device_map="sequential",
    torch_dtype=torch.float16
)

4. 手动分配（device_map={...}）

• 作用：自定义模型各部分的设备分配。 • 适用场景：需要精准控制显存分布时（例如，将视觉模块和语言模块分离到不同 GPU）。 • 示例代码：

python
device_map = {
    "model.layers.0": 0,      # 第 0 层分配到 GPU 0
    "model.layers.1": 1,      # 第 1 层分配到 GPU 1
    "lm_head": 1,             # 语言模型头分配到 GPU 1
    "visual.blocks": 0,       # 视觉模块分配到 GPU 0
}
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-VL-7B",
    device_map=device_map,
    torch_dtype=torch.float16
)

举个例子，比如这个模型：https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct/blob/main/model.safetensors.index.json

可以这么写：

python
# 手动分配设备映射
device_map = {
    # 将 model.layers 的前 10 层分配到 GPU0，后 n 层分配到 GPU1
    **{f"model.layers.{i}": 0 for i in range(10)},  # GPU0
    **{f"model.layers.{i}": 1 for i in range(10, 28)},  # GPU1

    # 将 visual.blocks 的前 10 层分配到 GPU0，后 n 层分配到 GPU1
    **{f"visual.blocks.{i}": 0 for i in range(10)},  # GPU0
    **{f"visual.blocks.{i}": 1 for i in range(10, 32)},  # GPU1

    # 其他模块分配到 GPU0 或 GPU1
    "lm_head": 1,  # GPU1
    "model.embed_tokens": 0,  # GPU0
    "model.norm": 1,  # GPU1
    "visual.merger": 0,  # GPU0
    "visual.patch_embed": 0,  # GPU0
}
# 加载模型
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    model_path, torch_dtype="auto", device_map=device_map
)

---

如何选择模式？

模式	优点	缺点
auto/balanced	显存利用均衡	要求 GPU 显存容量相近
balanced_low_0	保留 GPU 0 的显存灵活性	其他 GPU 显存需足够
sequential	适配显存差异大的硬件环境	可能导致显存利用率不均
手动分配	完全控制显存分布	需要熟悉模型结构

推荐策略：

1. 优先尝试 device_map="auto"。
2. 如果 GPU 0 需要预留显存，使用 balanced_low_0。
3. 仅在需要精细控制时手动分配。