分布式训练先从一组可复跑的样本和配置开始记录应当能让另一个人复跑因此还要写明数据读取顺序和启动方式。排除了哪些因素同样值得留下。当一次复跑结果不一致时先把变动写成假设再逐项排除。这样留下的是可验证的线索而不是一句模糊的“环境不稳定”。单卡跑通了多卡启动却卡在 rank 0 卡死不动这里记录的是分布式训练从单卡迁到多卡时最容易丢失的运行条件数据如何分片、进程在哪一步同步、checkpoint 怎样保存。文中的现象用于说明排查顺序实际结果仍要在固定版本和配置下复跑。真正让人头疼的是把脚本推上 8 卡 A100 服务器的那一刻。终端打出torchrun --nproc_per_node8后屏幕停在第一个 epoch 初始化的位置显存占满了但 GPU 利用率全挂在 0%。看进程树才发现Rank 0 正在同步执行数据集的大目录扫描而 Rank 1 到 Rank 7 早早就调用了dist.barrier()陷入挂起死锁。单卡脚本包上一层DistributedDataParallel并不等于训练就能稳定跑起来。数据集解析、分布式采样、混合精度和 checkpoint 轮转要各自有清楚的入口出问题时才知道该查哪一段。------------------------------------------------------------------------- | Trainer Engine (调度引擎) | ------------------------------------------------------------------------- | | | v v v ------------------ -------------------- -------------------- | DataLoader (数据) | | Model DDP (模型) | | StateSaver (状态) | | - DistSampler | | - AMP autocast | | - Rank 0 checkpoint| | - Prefetching | | - Gradient Accu | | - Metric log | ------------------ -------------------- --------------------依赖解耦把 DataLoader、Loss 和 Engine 的边界划干净在成熟的分布式训练 Harness 架构中各个组件只对自己的输入输出负责。Dataset只做单纯的 Byte 到 Tensor 转换与增强绝不上揽跨进程的状态共享DataLoader必须交由DistributedSampler来支配索引分片并且必须保证在每个 epoch 开始前由外部 Engine 显式调用sampler.set_epoch(epoch)否则每个 GPU 每一轮读入的数据顺序都会完全一样。Engine则是训练的主控中枢负责初始化进程组、管理梯度累加步数Gradient Accumulation Steps、调度GradScaler防止 FP16 梯度下溢以及拦截任何单卡抛出的 Runtime Exception 并安全退回 NCCL 通信。状态与通信流DDP 初始化与全卡屏障同步在这套结构里所有的 I/O 开销与 Checkpoint 硬盘写入操作必须严格约束在rank 0上执行而梯度的 Reduce 汇总则在后台由 NCCL 隐式完成。生产级单机多卡训练最小 Harness 代码实现下面是一个符合生产环境要求的完整 PyTorch 分布式训练 Harness 模版。包含了 DDP 初始化、自适应混合精度AMP、梯度裁剪以及稳健的屏障异常处理。import os import sys import time import logging import torch import torch.nn as nn import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import Dataset, DataLoader, DistributedSampler from torch.cuda.amp import autocast, GradScaler # 配置日志 def setup_logger(rank: int): logger logging.getLogger(DDP_Trainer) logger.setLevel(logging.INFO if rank 0 else logging.ERROR) handler logging.StreamHandler(sys.stdout) formatter logging.Formatter(f[Rank {rank}] %(asctime)s - %(levelname)s - %(message)s) handler.setFormatter(formatter) if not logger.handlers: logger.addHandler(handler) return logger # 1. 模拟数据集 class SyntheticDataset(Dataset): def __init__(self, size: int 10000, feature_dim: int 128): self.size size # 预先生成随机数据避免训练中做复杂 CPU 运算 self.data torch.randn(size, feature_dim) self.labels torch.randint(0, 2, (size,)) def __len__(self): return self.size def __getitem__(self, idx): return self.data[idx], self.labels[idx] # 2. 简单模型定义 class ToyClassifier(nn.Module): def __init__(self, in_dim: int 128, hidden_dim: int 256): super().__init__() self.net nn.Sequential( nn.Linear(in_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 2) ) def forward(self, x): return self.net(x) # 3. 主训练 Harness def train_ddp(): # 从环境变量读取分布式参数由 torchrun 自动注入 rank int(os.environ.get(RANK, 0)) local_rank int(os.environ.get(LOCAL_RANK, 0)) world_size int(os.environ.get(WORLD_SIZE, 1)) logger setup_logger(rank) # 初始化 NCCL 进程组 dist.init_process_group(backendnccl, init_methodenv://) torch.cuda.set_device(local_rank) device torch.device(fcuda:{local_rank}) logger.info(f成功卡位分布式环境: local_rank{local_rank}, world_size{world_size}) # 构建数据与 Sampler dataset SyntheticDataset() sampler DistributedSampler( dataset, num_replicasworld_size, rankrank, shuffleTrue, seed42 ) dataloader DataLoader( dataset, batch_size64, samplersampler, num_workers2, pin_memoryTrue ) # 模型与 DDP 包装 model ToyClassifier().to(device) model DDP(model, device_ids[local_rank], output_devicelocal_rank) criterion nn.CrossEntropyLoss().to(device) optimizer torch.optim.AdamW(model.parameters(), lr1e-3) scaler GradScaler() # 混合精度 Scaler total_epochs 3 for epoch in range(total_epochs): # 关键步骤设置 sampler 的 epoch保证每轮洗牌随机种子更新 sampler.set_epoch(epoch) model.train() running_loss 0.0 start_time time.perf_counter() for step, (inputs, targets) in enumerate(dataloader): inputs inputs.to(device, non_blockingTrue) targets targets.to(device, non_blockingTrue) optimizer.zero_grad() # 开启前向混合精度上下文 with autocast(): outputs model(inputs) loss criterion(outputs, targets) # 缩放梯度并反向传播 scaler.scale(loss).backward() # 梯度裁剪必须先 unscale scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0) # 步进更新 scaler.step(optimizer) scaler.update() running_loss loss.item() epoch_time time.perf_counter() - start_time avg_loss running_loss / len(dataloader) if rank 0: logger.info(fEpoch [{epoch1}/{total_epochs}] 完成 | Avg Loss: {avg_loss:.4f} | 耗时: {epoch_time:.2f}s) # 仅在 Rank 0 保存模型 checkpoint checkpoint_path fmodel_epoch_{epoch1}.pt torch.save({ epoch: epoch 1, model_state_dict: model.module.state_dict(), optimizer_state_dict: optimizer.state_dict(), }, checkpoint_path) logger.info(f保存检查点至 {checkpoint_path}) # 确保所有进程同步完当前 epoch 再进入下一轮 dist.barrier() # 资源清理 dist.destroy_process_group() logger.info(训练完成已注销分布式进程组。) if __name__ __main__: train_ddp()踩坑记录Sampler 没加 epoch 计数导致分布式训练数据重复吐出在一次线上大模型预训练中团队发现模型在第 5 个 epoch 时 Loss 几乎不再下降且 Validation 结果出现严重的过拟合。调出 Trainer 的代码仔细对查发现循环体内唯独漏掉了sampler.set_epoch(epoch)这一行。由于DistributedSampler默认的种子计算依赖epoch参数默认为 0当省略该调用时虽然每个卡拿到了总数据集的 1/N但卡与卡之间在所有 epoch 中分发的数据分片顺序完全静止不变。整个分布式训练变成了把同一个子集重复刷了 5 遍。结语训练配置、随机种子与数据切分一旦丢失后续结论就很难核验。把首轮失败也留在样本里分布式训练第一次跑不通时我不会先把日志删掉重来。通信超时、某张卡先退出、数据分片为空这些现象都要和当时的启动参数放在一起。后续换了 batch size 或节点数才能知道改变的是训练行为还是只碰巧避开了原来的条件。