ChatGLM-6B模型并行推理：多GPU加速方案

如果你手头有多张GPU，但每张卡的显存都不足以完整放下ChatGLM-6B模型，或者你想让推理速度再快一点，那么模型并行推理就是你需要的技术。简单来说，就是把一个模型“切开”，分别放到不同的GPU上运行，让它们协同工作。

这听起来有点复杂，但实际操作起来，借助一些成熟的工具，并没有想象中那么难。今天，我就带你一步步实现ChatGLM-6B在多GPU环境下的并行推理，让你手里的多张显卡真正“物尽其用”。

1. 为什么需要多GPU并行推理？

在开始动手之前，我们先聊聊为什么需要这么做。ChatGLM-6B模型如果以FP16精度加载，大概需要13GB的显存。如果你的单张显卡只有8GB或11GB，那就跑不起来。

这时候你有几个选择：

• 模型量化：把模型精度降低（比如INT4），减少显存占用，但可能会损失一些生成质量。
• CPU推理：用内存代替显存，但速度会慢很多。
• 多GPU并行：把模型拆开，让多张卡一起扛，既能跑起来，速度还可能更快。

多GPU并行的好处很明显：它能突破单卡显存的限制，并且通过并行计算提升吞吐量。对于需要部署服务、处理并发请求的场景，这个优势尤其突出。

2. 环境准备与核心工具

要实现多GPU并行，我们主要依赖两个工具：Hugging Face Transformers库和Accelerate库。Accelerate是Hugging Face出品的一个神器，它能帮我们轻松地把模型和计算分布到多个设备上，而不用写复杂的分布式代码。

首先，确保你的环境已经准备好了：

# 安装必要的Python包
pip install transformers==4.27.1  # 推荐使用这个版本，兼容性好
pip install accelerate
pip install torch  # 请根据你的CUDA版本安装对应的PyTorch

如果你的网络环境访问Hugging Face比较慢，可以考虑先提前下载好模型。但为了教程的连贯性，我们这里会直接从网络加载，你可以根据实际情况调整。

3. 基础的单GPU加载回顾

在进入多GPU之前，我们先看看正常的单卡加载是怎么写的，这样对比起来更清楚：

from transformers import AutoTokenizer, AutoModel

# 单GPU标准加载方式
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
model = model.eval()

这段代码会把整个模型加载到第一张GPU（cuda:0）上。如果你的显存不够，这里就会报错。

4. 实现多GPU并行加载

现在到了关键部分。ChatGLM-6B的官方代码库里其实已经为我们准备好了多GPU加载的工具函数。我们不需要自己从头写模型切分的逻辑，直接调用就行。

4.1 使用官方工具函数

在ChatGLM-6B的GitHub仓库中，有一个utils.py文件，里面提供了load_model_on_gpus函数。我们来看看怎么用：

from transformers import AutoTokenizer
# 注意这里导入方式的变化
from utils import load_model_on_gpus

# 初始化tokenizer
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)

# 将模型加载到2张GPU上
model = load_model_on_gpus("THUDM/chatglm-6b", num_gpus=2)
model = model.eval()

就这么简单！num_gpus=2表示使用2张GPU。函数会自动把模型均匀地切分到两张卡上。如果你的机器有4张卡，改成num_gpus=4就行。

4.2 如果找不到utils模块怎么办？

有时候，你可能不想克隆整个ChatGLM-6B仓库，或者环境配置有问题。别担心，我们也可以直接用Accelerate库来实现类似的功能，虽然稍微多写几行代码，但更灵活：

from transformers import AutoTokenizer, AutoModel
from accelerate import init_empty_weights, load_checkpoint_and_dispatch
import torch

# 1. 初始化tokenizer
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)

# 2. 创建空模型（不立即加载权重）
with init_empty_weights():
    model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)

# 3. 计算每层应该放在哪个设备上
# 这里我们手动指定一个简单的切分策略：前一半层在GPU 0，后一半在GPU 1
device_map = {}
num_layers = len(model.transformer.encoder.layers)
half_point = num_layers // 2

for i in range(num_layers):
    if i < half_point:
        device_map[f"transformer.encoder.layers.{i}"] = 0  # GPU 0
    else:
        device_map[f"transformer.encoder.layers.{i}"] = 1  # GPU 1

# 其他部分放在GPU 0上
device_map["transformer.embedding"] = 0
device_map["transformer.rotary_pos_emb"] = 0
device_map["lm_head"] = 0

# 4. 加载并分发模型
model = load_checkpoint_and_dispatch(
    model,
    "THUDM/chatglm-6b",
    device_map=device_map,
    no_split_module_classes=["GLMBlock"],  # 告诉accelerate不要拆分GLMBlock内部
    offload_folder="offload",  # 如果显存还不够，可以设置offload到CPU
)
model = model.eval()

这种方法的好处是你可以精细控制每一层放在哪张卡上，对于负载均衡调优特别有用。

5. 进行对话测试

模型加载好了，我们来测试一下它是否正常工作。多GPU并行的模型在使用上和单GPU模型完全一样，这要归功于Accelerate的封装：

def chat_with_model():
    history = []
    print("开始与ChatGLM-6B对话（输入'quit'退出）")
    
    while True:
        user_input = input("\n你: ")
        if user_input.lower() == 'quit':
            break
            
        response, history = model.chat(tokenizer, user_input, history=history)
        print(f"ChatGLM: {response}")

if __name__ == "__main__":
    chat_with_model()

运行这段代码，你应该能看到模型正常回复。你可以问它一些问题，比如“你好”、“介绍一下你自己”、“晚上睡不着怎么办”等等。

6. 负载均衡策略与性能调优

默认的均匀切分不一定是最优的。不同的模型层计算量不同，有些层更“重”，有些更“轻”。如果简单均匀切分，可能会导致一张卡算完了在等另一张卡，这就是负载不均衡。

6.1 如何查看当前的设备分布？

在调试时，我们可以先看看模型是怎么分布在各个设备上的：

# 打印模型各层所在的设备
for name, param in model.named_parameters():
    print(f"{name}: {param.device}")

这会输出每一层参数所在的设备，帮助你了解当前的分布情况。

6.2 手动优化设备映射

基于上面的信息，你可以调整device_map，把计算密集的层分散到不同的卡上，或者根据你GPU的性能差异来分配。比如，如果你有一张3090和一张2080Ti，可以把更多的层放在3090上。

# 示例：自定义设备映射
custom_device_map = {
    # 嵌入层和开头几层放在GPU 0
    "transformer.embedding": 0,
    "transformer.encoder.layers.0": 0,
    "transformer.encoder.layers.1": 0,
    "transformer.encoder.layers.2": 0,
    
    # 中间层交替放置
    "transformer.encoder.layers.3": 1,
    "transformer.encoder.layers.4": 0,
    "transformer.encoder.layers.5": 1,
    
    # 最后几层和输出层放在GPU 1
    "transformer.encoder.layers.28": 1,
    "transformer.encoder.layers.29": 1,
    "transformer.rotary_pos_emb": 1,
    "lm_head": 1,
}

6.3 使用自动平衡策略

如果你不想手动调，也可以让Accelerate自动尝试平衡：

from accelerate import infer_auto_device_map

# 自动推断设备映射，尽量平衡显存使用
device_map = infer_auto_device_map(
    model,
    max_memory={0: "10GB", 1: "10GB"},  # 每张卡最多用10GB
    no_split_module_classes=["GLMBlock"]
)

7. 实际效果与注意事项

在实际使用中，多GPU并行会带来一些开销，主要是GPU之间的通信。但对于ChatGLM-6B这样规模的模型，这个开销通常是值得的。

几点实用建议：

1. GPU选择：尽量使用相同型号的GPU，避免因性能差异导致等待。
2. 通信带宽：如果有多条PCIe通道，确保每张卡都有足够的带宽。NVLink连接的GPU之间通信效率更高。
3. 监控工具：使用nvidia-smi或gpustat监控每张卡的显存使用和利用率，帮助调优。
4. 批量处理：在处理多个请求时，多GPU的优势更明显，可以考虑批量处理输入。

8. 完整示例代码

这里给一个完整的、可以直接运行的示例：

#!/usr/bin/env python3
"""
ChatGLM-6B多GPU并行推理示例
"""

import torch
from transformers import AutoTokenizer
from utils import load_model_on_gpus  # 来自ChatGLM-6B仓库

def main():
    # 设置使用的GPU数量
    num_gpus = 2  # 根据你的实际情况修改
    
    print(f"正在加载ChatGLM-6B模型到{num_gpus}张GPU上...")
    
    # 加载tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        "THUDM/chatglm-6b", 
        trust_remote_code=True
    )
    
    # 多GPU加载模型
    model = load_model_on_gpus(
        "THUDM/chatglm-6b", 
        num_gpus=num_gpus
    )
    model = model.eval()
    
    print("模型加载完成！开始对话测试...")
    
    # 测试对话
    test_prompts = [
        "你好，请介绍一下你自己",
        "深度学习和机器学习有什么区别？",
        "用Python写一个快速排序算法"
    ]
    
    history = []
    for prompt in test_prompts:
        print(f"\n用户: {prompt}")
        response, history = model.chat(tokenizer, prompt, history=history)
        print(f"AI: {response}")
    
    print("\n测试完成！")

if __name__ == "__main__":
    # 设置PyTorch可以看到所有GPU
    torch.cuda.device_count()
    main()

9. 总结

多GPU并行推理是释放ChatGLM-6B潜力的有效方法，特别适合那些拥有多张显卡但单卡显存不足的场景。通过Accelerate库和ChatGLM官方提供的工具，实现起来比想象中简单。

实际用下来，这种方案确实能解决显存不够的问题，而且当处理多个并发请求时，吞吐量的提升比较明显。不过也要注意，GPU之间的通信会带来一些额外开销，所以并不是GPU越多就一定越快，需要根据实际情况找到平衡点。

如果你刚开始尝试，建议先从2张GPU开始，熟悉了整个流程和性能表现后，再考虑扩展到更多卡或者尝试更复杂的负载均衡策略。毕竟，能让模型跑起来才是第一步，优化可以慢慢来。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。