大模型多卡并行的一些尝试

前言

需要评估一个大模型 ，不涉及任何训练微调。我的a100有40G显存，8B模型以fp16加载只需要16G显存。在单卡的小型模型很简单，如图直接加载模型，移动上设备，然后传入评估函数即可。
但是，我需要测试一个更大的模型，27B以fp16加载需要54G，根本放不下。于是我申请了一个三卡的集群（A100(40G)*3）。后续就是考虑使用啥方法进行模型配置。

尝试1torch.nn.DataParallel

OOM（out of memery）失败。

• torch.nn.DataParallel 仅适合多卡“复制模型+数据并行”，不适合超大模型拆分。
• 对于超大模型，torch.nn.parallel.DistributedDataParallel + 张量/模型并行更合适，但这需要代码显著改动。
• PyTorch 并没有开箱即用的“模型分片（tensor/model parallel）”功能，需要额外库（如 Megatron-LM）配合。

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "/path/to/your/27B-model"
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,  # 优化CPU内存占用
)
model = model.to(device)

# 多卡包装，适合模型能完整放入单卡的情况
if torch.cuda.device_count() > 1:
    model = torch.nn.DataParallel(model)

def evaluate(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to(device)
    model.eval()
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=100)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))


input_text = "Hello world"
evaluate(input_text)

尝试2accelerate

accelerate 是Huggingface提供的轻量多卡分布式与混合精度训练/推理工具，支持自动拆分模型权重到多卡，适合大模型推理。
成功但是速度些许慢。

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import torch
from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map

model_path = "/home/.../model/gemma_27B"

tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)

# 1. 加载配置
config = AutoConfig.from_pretrained(model_path)

# 2. 用init_empty_weights创建空模型
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

# 3. 绑权重，避免警告
model.tie_weights()

# 4. 指定模型层不拆分的类名，避免层内跨设备切分
# 请替换成你模型中层的实际类名，或者打印确认
no_split_module_classes = ["Gemma2DecoderLayer"]

# 5. 自动推断设备映射，带上 no_split_module_classes 参数
device_map = infer_auto_device_map(
    model,
    max_memory={0: "30GB", 1: "40GB", 2: "40GB"},
    dtype=torch.float16,
    no_split_module_classes=no_split_module_classes,
)

print("Device map:", device_map)
print(type(model.model.layers[24]))

# 6. 加载权重并分配设备，带上 no_split_module_classes
model = load_checkpoint_and_dispatch(
    model,
    model_path,
    device_map=device_map,
    no_split_module_classes=no_split_module_classes,
    dtype=torch.float16,
)

model.eval()

def evaluate(input_text):
    inputs = tokenizer(input_text, return_tensors="pt").to("cuda:0")
    with torch.no_grad():
        outputs = model.generate(**inputs, max_length=100)
    print(tokenizer.decode(outputs[0], skip_special_tokens=True))

input_text = "Hello world"
evaluate(input_text)

尝试3vllm

因为我的模型不需要微调，只需要评估。我就使用了一个很出名的推理框架vllm.python代码是自己写的，脚本代码是我问同门说这个脚本启动更常见。
结果可以顺利运行，但是vllm封装得太好了，我的需求是要保留model供我调用。因此pass。

import os
from vllm import LLM, SamplingParams

# 设置可见的 CUDA 设备
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"

def infer(model_path, prompts):
    # 初始化 LLM，使用 bfloat16 或 float32
    llm = LLM(
        model=model_path,
        tensor_parallel_size=2,
        dtype="bfloat16",  # 或者使用 "float32"
        seed=42,
        disable_log_stats=True,
        trust_remote_code=True,
        gpu_memory_utilization=0.95,
        enforce_eager=True,
    )

    # 设置采样参数
    sampling_params = SamplingParams(
        top_p=1.0,
        top_k=-1,
        n=3,
        temperature=0.5,
        max_tokens=2048
    )

    # 生成输出
    outputs = llm.generate(prompts=prompts, sampling_params=sampling_params)

    # 打印输出
    for output in outputs:
        print(output)

    return None

# 初始化输入
model_path = "/home/.../model/gemma_27B"  # 替换为您的模型路径
prompts = ["这是一个示例输入。"]  # 输入提示，可以是多个

# 调用 infer 函数
_ = infer(model_path, prompts)

#!/bin/bash

# 参数设置
CUDA_VISIBLE_DEVICES=0,1 # GPU编号
TOKENIZERS_PARALLELISM=false # 是否使用多线程分词器
MODEL_NAME="Qwen2.5-VL-7B-Instruct"  # 模型名称
PORT=8000  # 端口号
GPU_MEMORY_UTILIZATION=0.9  # GPU 内存利用率

MODELS_FOLDER="/home/Disk/llm_models/"  # 模型存放路径
TENSOR_PARALLEL_SIZE=2  # 模型并行度
MAX_MODEL_LEN=30000  # 最大模型长度

MODEL_PATH="${MODELS_FOLDER}/${MODEL_NAME}"
SERVED_MODEL_NAME=$(echo $MODEL_NAME | awk -F'/' '{print $NF}')
LOG_FILE="./logs/vllm_${SERVED_MODEL_NAME}.log"  # 日志文件

echo "Starting vLLM server..."
echo "  CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES}"
echo "  MODEL_PATH: ${MODEL_PATH}"
echo "  SERVED_MODEL_NAME: ${SERVED_MODEL_NAME}"
echo "  PORT: ${PORT}"
echo "  GPU_MEMORY_UTILIZATION: ${GPU_MEMORY_UTILIZATION}"
echo "  LOG_FILE: ${LOG_FILE}"

nohup bash -c "CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} vllm serve ${MODEL_PATH} --served-model-name ${SERVED_MODEL_NAME} --port ${PORT} --trust-remote-code --tensor-parallel-size ${TENSOR_PARALLEL_SIZE} --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION} --max_model_len ${MAX_MODEL_LEN}" > ${LOG_FILE} 2>&1 &

尝试4 deepspeed

类似3
DeepSpeed的Inference功能支持大模型零冗余优化(ZeRO)、张量并行、流水线并行等，推理效率和显存利用率都不错。
使用门槛比accelerate稍高，需配置相关参数。

import os
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import deepspeed

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"
model_path = "/home/.../model/gemma_27B"

tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)

print("Loading model with DeepSpeed inference...")

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="cpu",   # 先加载到 CPU
    low_cpu_mem_usage=True  # 尽量节省 CPU 内存
)

model = deepspeed.init_inference(
    model,
    mp_size=3,
    dtype=torch.float16,
    replace_method="auto",
    replace_with_kernel_inject=True,
)


model.eval()


def eval(model, tokenizer, input_text, dataset_name):
    print(f"Evaluating on dataset: {dataset_name}")
    inputs = tokenizer(input_text, return_tensors="pt", padding=True)
    device = next(model.module.parameters()).device
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=100,
            do_sample=True,
            top_p=0.9,
            temperature=0.7,
            num_return_sequences=1
        )
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print("Input text:", input_text)
    print("Generated text:", generated_text)


input_text = "Hello, how are you?"
dataset = "ptb_text_only"

eval(model, tokenizer, input_text, dataset)

结论

大模型多卡推理，最简单可用、灵活的方案是 accelerate，适合评估和开发；追求极致性能或服务部署可考虑 DeepSpeed 或 vllm。