PyTorch 实验环境本地跑通：固定依赖、随机种子与设备探测

PyTorch 环境第一次跑通，目标不是把所有 GPU 都压满，而是确认依赖能导入、设备能识别、固定小输入能完成前向和反向。

1. 先验证单进程最小闭环

本地环境验证先看依赖导入、设备探测和单进程数值正确性。多进程只在前一步通过后开启，并把 NCCL 初始化与训练逻辑分开记录。

保存 Python、PyTorch、CUDA 与驱动版本摘要，使用合成张量运行一个小模型。只有单进程结果正确，才进入torchrun和 NCCL 路径。

2. 再检查多进程启动

环境摘要至少包含 Python、PyTorch、CUDA、驱动和设备类型。随机种子与小批量形状写进自检命令，失败时先比较各 rank 的退出原因。

本地测试只验证进程组初始化、各 rank 输出形状和梯度同步，不用真实训练数据。示例日志中的 rank 与数值是构造输入，执行时应输出当前环境的实际结果。

3. 环境探测与分布式自检

[ERROR] torch.distributed.elastic.multiprocessing.api: failed (exitcode: -9) [Rank 2] RuntimeError: NCCL error in: /pytorch/torch/csrc/distributed/c10d/NCCLUtils.hpp:275, Unhandled system error / internal error or browser network connection lost. NCCL WARN: Call to connect returned Connection refused [Rank 0] RuntimeError: Timed out waiting for ranks to synchronize.

import os import sys import socket import torch import torch.nn as nn import torch.distributed as dist import torch.multiprocessing as mp from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader, Dataset, DistributedSampler def find_free_port() -> int: """寻找本地空闲的网络端口，避免 Port 冲突""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind(('', 0)) return s.getsockname()[1] class SyntheticDataset(Dataset): """用于本地验证的伪造数据集""" def __init__(self, size: int = 100): self.data = torch.randn(size, 16) self.labels = torch.randint(0, 2, (size,)) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx], self.labels[idx] def run_distributed_worker(rank: int, world_size: int, master_port: int): """ 单节点分布式 Worker 进程入口。 自动检测 GPU 可用性，无 GPU 时降级为 GLOO + CPU 仿真。 """ os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(master_port) use_cuda = torch.cuda.is_available() and torch.cuda.device_count() >= world_size backend = "nccl" if use_cuda else "gloo" device = torch.device(f"cuda:{rank}" if use_cuda else "cpu") if use_cuda: torch.cuda.set_device(device) try: # 1. 初始化分布式进程组 dist.init_process_group( backend=backend, rank=rank, world_size=world_size ) print(f"[RANK {rank}/{world_size}] 初始化成功! 使用后端: {backend}, 设备: {device}") # 2. 构建模型与 DDP 包装 model = nn.Sequential( nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 2) ).to(device) # CPU 模式下无需 specify device_ids ddp_model = DDP(model, device_ids=[rank] if use_cuda else None) # 3. 构造数据分片加载器 dataset = SyntheticDataset() sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True) loader = DataLoader(dataset, batch_size=4, sampler=sampler) # 4. 模拟 2 个 Epoch 的训练流程 optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01) criterion = nn.CrossEntropyLoss() for epoch in range(2): sampler.set_epoch(epoch) # 确保随机种子跨 Epoch 递进 for step, (inputs, targets) in enumerate(loader): inputs, targets = inputs.to(device), targets.to(device) optimizer.zero_grad() outputs = ddp_model(inputs) loss = criterion(outputs, targets) loss.backward() # 梯度裁剪与优化更新 torch.nn.utils.clip_grad_norm_(ddp_model.parameters(), max_norm=1.0) optimizer.step() # 5. 梯度同步验证：全 Rank 计算 Loss 均值 loss_tensor = torch.tensor([loss.item()], device=device) dist.all_reduce(loss_tensor, op=dist.ReduceOp.SUM) avg_loss = loss_tensor.item() / world_size if rank == 0: print(f"[RANK 0 VERIFY] 分布式跑通成功! 最后一个 Epoch 均化 Loss: {avg_loss:.4f}") except Exception as err: print(f"[RANK {rank} FATAL] 运行过程抛出异常: {str(err)}") raise finally: # 优雅清理分布式进程组 if dist.is_initialized(): dist.destroy_process_group() def launch_local_test(world_size: int = 4): """一键拉起本地多进程分布式仿真""" port = find_free_port() print(f"[LAUNCH] 启动本地分布式仿真测试, World Size: {world_size}, 自动分配端口: {port}") mp.spawn( run_distributed_worker, args=(world_size, port), nprocs=world_size, join=True ) if __name__ == "__main__": launch_local_test(world_size=4)

[LAUNCH] 启动本地分布式仿真测试, World Size: 4, 自动分配端口: 49152 [RANK 0/4] 初始化成功! 使用后端: gloo, 设备: cpu [RANK 1/4] 初始化成功! 使用后端: gloo, 设备: cpu [RANK 2/4] 初始化成功! 使用后端: gloo, 设备: cpu [RANK 3/4] 初始化成功! 使用后端: gloo, 设备: cpu [RANK 0 VERIFY] 分布式跑通成功! 最后一个 Epoch 均化 Loss: 0.6842

4. 复核清单

• 单进程前向与反向是否先于多进程验证。
• Python、PyTorch、CUDA 和驱动摘要是否保存。
• 各 rank 的形状、梯度与退出原因是否可对齐。
• 自检是否只使用合成张量和非敏感日志。

总结

“本地环境怎样一次跑通”应以清晰的条件和脚本复核。先记录边界，再解释结果。