单节点多GPU下基于Hugging Face Trainer的DDP端到端实现指南

核心问题解答 #

1. 是否需要将HF模型用DDP包裹？

• 若使用HF原生Trainer类：不需要手动包裹，Trainer会自动完成DDP封装与分布式逻辑处理。
• 若使用原生PyTorch训练循环：必须用torch.nn.parallel.DistributedDataParallel包裹HF模型，否则无法实现分布式训练。

2. 是否需要修改Trainer参数或对优化器进行分布式封装？

• 用Trainer时：只需在TrainingArguments中设置好per_device_train_batch_size等参数，无需手动封装优化器——Trainer会自动创建兼容分布式的优化器和学习率调度器。
• 用原生循环时：优化器需在DDP包裹模型后初始化（或直接传入model.parameters()，DDP会自动处理分布式参数），无需额外封装。

3. 是否需要执行DDP常规的进程组初始化？

• 用Trainer时：不需要手动调用dist.init_process_group，Trainer会通过环境变量自动读取local_rank等信息并完成初始化。
• 用原生循环时：必须手动初始化进程组，示例代码：

import os
import torch.distributed as dist

def setup_ddp():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

4. local_rank的作用是什么？
   local_rank是当前进程在单台机器GPU集群中的索引（比如单台4卡机器，范围是0-3），核心作用：

• 指定当前进程独占的GPU设备，避免多进程抢占同一GPU；
• 配合全局rank参数，帮助DDP识别进程在分布式集群中的位置，实现跨进程/跨机器通信。

5. 终端启动脚本的正确写法？
   无论是Trainer还是原生循环，都推荐用torchrun启动（自动注入LOCAL_RANK等环境变量）：

• 单机器多卡：

torchrun --nproc_per_node=4 train.py

• 多机器多卡（以2台4卡机器为例）：

# 主机器（node_rank=0）
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="你的主机器IP" --master_port=29500 train.py

# 从机器（node_rank=1）
torchrun --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="你的主机器IP" --master_port=29500 train.py

6. 如何利用world_size在训练循环中分片数据？

• 用Trainer时：无需手动分片，Trainer会自动使用DistributedSampler将数据集按world_size（总进程数）均分，每个进程只处理自己的分片数据。
• 用原生循环时：必须手动配置DistributedSampler，示例：

from torch.utils.data import DistributedSampler, DataLoader

train_dataset = ... # 你的HF数据集
train_sampler = DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, batch_size=8, sampler=train_sampler)

# 训练循环中，每个epoch需更新sampler的epoch值，保证跨epoch数据打乱的随机性
for epoch in range(num_epochs):
    train_sampler.set_epoch(epoch)
    for batch in train_loader:
        # 训练逻辑

完整端到端示例 #

示例1：使用HF Trainer实现DDP

from transformers import AutoModelForSequenceClassification, AutoTokenizer, TrainingArguments, Trainer
from datasets import load_dataset

# 加载并预处理数据
dataset = load_dataset("imdb")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize_func(examples):
    return tokenizer(examples["text"], padding="max_length", truncation=True)

tokenized_ds = dataset.map(tokenize_func, batched=True)
small_train_ds = tokenized_ds["train"].shuffle(seed=42).select(range(1000))
small_eval_ds = tokenized_ds["test"].shuffle(seed=42).select(range(1000))

# 加载模型
model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2)

# 设置训练参数
training_args = TrainingArguments(
    output_dir="./ddp_results",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    logging_dir="./ddp_logs",
    logging_steps=10,
    report_to="none",  # 禁用外部报告工具，避免分布式冲突
    fp16=True,  # 可选：开启混合精度加速
)

# 初始化Trainer，自动处理DDP
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=small_train_ds,
    eval_dataset=small_eval_ds,
)

# 启动训练
trainer.train()

示例2：原生PyTorch循环实现DDP

import os
import torch
import torch.distributed as dist
from torch.utils.data import DistributedSampler, DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AdamW
from datasets import load_dataset

def setup_ddp():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup_ddp():
    dist.destroy_process_group()

def main():
    local_rank = setup_ddp()
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 加载并预处理数据
    dataset = load_dataset("imdb")
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

    def tokenize_func(examples):
        return tokenizer(examples["text"], padding="max_length", truncation=True)

    tokenized_ds = dataset.map(tokenize_func, batched=True)
    train_ds = tokenized_ds["train"].shuffle(seed=42).select(range(1000))
    train_ds.set_format("torch", columns=["input_ids", "attention_mask", "label"])

    # 分布式采样器
    train_sampler = DistributedSampler(train_ds)
    train_loader = DataLoader(train_ds, batch_size=8, sampler=train_sampler)

    # 加载模型并包裹DDP
    model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=2).to(local_rank)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])

    # 优化器
    optimizer = AdamW(model.parameters(), lr=5e-5)

    # 训练循环
    num_epochs = 3
    model.train()
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)
        for step, batch in enumerate(train_loader):
            batch = {k: v.to(local_rank) for k, v in batch.items()}
            outputs = model(**batch)
            loss = outputs.loss
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

            # 仅在主进程（rank=0）打印日志
            if rank == 0 and step % 10 == 0:
                print(f"Epoch {epoch+1}/{num_epochs}, Step {step}/{len(train_loader)}, Loss: {loss.item():.4f}")

    cleanup_ddp()

if __name__ == "__main__":
    main()

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