求助：Mistral-7B在多GPU集群推理仅单GPU占用，如何实现并行？

问题背景 #

在最多4块GPU的集群上运行Mistral-7B-Instruct-v0.2推理，当前仅GPU0被充分利用（占用率98%），其余GPU完全闲置，单GPU推理耗时过长。原代码尝试使用DataParallel但未生效，调用model.module.generate也无法解决多GPU利用问题。

核心原因 #

1. DataParallel的局限性：DataParallel是单进程多线程方案，在大模型生成任务中，generate方法的逻辑难以被有效分发到多GPU，仅能实现简单数据并行，对7B级大模型的适配性极差。
2. 单条prompt处理逻辑：原代码虽设置了batch_size，但实际是循环逐个处理单条prompt，未真正形成批量输入，无法触发多GPU的负载分担机制。
3. 错误调用model.module.generate：model.module会绕过DataParallel的封装，直接调用原始模型，完全丧失多GPU分发能力。

解决方案 #

方案1：使用Hugging Face Accelerate库（推荐） #

Accelerate可自动处理多GPU/分布式配置，无需手动编写复杂的分布式逻辑，代码简洁易维护。

方案2：使用DistributedDataParallel（DDP） #

DDP是多进程分布式方案，更适合大模型的多GPU推理，性能优于DataParallel，但需要手动处理分布式环境初始化和数据分发。

修改后的完整代码（Accelerate版本） #

from accelerate import Accelerator
from transformers import AutoTokenizer, AutoModelForCausalLM
import pandas as pd
from tqdm import tqdm
import torch

model_name = 'mistralai/Mistral-7B-Instruct-v0.2'

# 初始化Accelerator，自动适配多GPU环境
accelerator = Accelerator()

# 加载tokenizer并设置pad token
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 加载模型，Accelerate自动处理多GPU分发
model = AutoModelForCausalLM.from_pretrained(model_name)

# 用Accelerator准备模型和tokenizer
model, tokenizer = accelerator.prepare(model, tokenizer)

# 加载prompt数据集（替换为你的数据源）
prompts_df = pd.read_csv("your_prompts.csv")
prompts = prompts_df["prompt"].tolist()

# 批量生成响应函数
def generate_batch_responses(prompts_batch, max_new_tokens=30):
    # 批量tokenize，统一padding和截断
    inputs = tokenizer(
        prompts_batch,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512  # 根据GPU内存调整
    ).to(accelerator.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            do_sample=False  # 不需要采样时关闭，提升推理速度
        )
    
    # 解码结果并移除原始prompt（可选，根据需求调整）
    responses = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    responses = [resp[len(prompt):].strip() for resp, prompt in zip(responses, prompts_batch)]
    return responses

# 按批次处理所有prompt
batch_size = 8  # 根据单GPU内存调整，16GB GPU建议设为8-12
responses = []

for i in tqdm(range(0, len(prompts), batch_size)):
    batch = prompts[i:i+batch_size]
    batch_responses = generate_batch_responses(batch)
    responses.extend(batch_responses)

# 仅主进程保存结果
if accelerator.is_main_process:
    prompts_df["response"] = responses
    prompts_df.to_csv("inference_results.csv", index=False)

修改后的完整代码（DDP版本） #

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from transformers import AutoTokenizer, AutoModelForCausalLM
import pandas as pd
from tqdm import tqdm
import os

model_name = 'mistralai/Mistral-7B-Instruct-v0.2'

# 初始化DDP分布式环境
def setup_ddp():
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup_ddp():
    dist.destroy_process_group()

local_rank = setup_ddp()
rank = dist.get_rank()
world_size = dist.get_world_size()

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 加载模型并绑定到当前GPU
model = AutoModelForCausalLM.from_pretrained(model_name).to(local_rank)
model = DDP(model, device_ids=[local_rank])

# 主进程加载prompt并拆分到各进程
if rank == 0:
    prompts_df = pd.read_csv("your_prompts.csv")
    prompts = prompts_df["prompt"].tolist()
    # 均分prompt到各个GPU进程
    chunk_size = len(prompts) // world_size
    process_prompts = [prompts[i*chunk_size : (i+1)*chunk_size] for i in range(world_size)]
    # 处理余数
    remaining = len(prompts) % world_size
    for i in range(remaining):
        process_prompts[i].append(prompts[world_size*chunk_size + i])
else:
    process_prompts = None

# 分发prompt到各进程
process_prompts = dist.scatter_object_list(process_prompts, src=0)[0]

# 批量生成响应函数
def generate_batch_responses(prompts_batch, max_new_tokens=30):
    inputs = tokenizer(
        prompts_batch,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512
    ).to(local_rank)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    
    return tokenizer.batch_decode(outputs, skip_special_tokens=True)

# 处理当前进程的prompt
batch_size = 8
process_responses = []
for i in tqdm(range(0, len(process_prompts), batch_size)):
    batch = process_prompts[i:i+batch_size]
    process_responses.extend(generate_batch_responses(batch))

# 收集所有进程的结果到主进程
all_responses = [None for _ in range(world_size)]
dist.gather_object(process_responses, all_responses if rank == 0 else None)

# 主进程合并并保存结果
if rank == 0:
    all_responses = [resp for sublist in all_responses for resp in sublist]
    prompts_df["response"] = all_responses
    prompts_df.to_csv("inference_results.csv", index=False)

cleanup_ddp()

关键注意事项 #

• 启动方式：DDP版本必须用torchrun启动，命令示例：torchrun --nproc_per_node=4 your_script.py（4为GPU数量）。
• batch_size调整：根据单GPU内存设置，16GB GPU运行Mistral-7B时，批量大小建议设为8-12。
• 避免单条处理：必须采用批量tokenize和生成逻辑，单条prompt无法触发多GPU负载分担。

内容的提问来源于stack exchange，提问作者t1nkov