多机多卡训练pytorch大模型的配置

配置 PyTorch 多机多卡分布式训练大模型的最佳实践
使用 DistributedDataParallel (DDP)
为了高效地利用多个 GPU 和多台机器来加速大型模型的训练，推荐采用 DDP 方案。相比于 DataParallel，DDP 提供更优性能并支持真正的分布式环境。
在构建分布式应用之前，需先安装必要的依赖库，并确认集群中的每台计算节点已正确配置网络连接以便相互通讯[^1]。
对于单机场景下的多卡训练：

初始化进程组
创建 Python 文件作为入口脚本，在其中调用 dist.init_process_group() 函数完成通信后端的选择以及 rank 和 world size 的设定[^2]。

import torch.distributed as dist

def setup(rank, world_size):
# 初始化方法可以是 'tcp'、'env' 或者 'file'
init_method = f'tcp://localhost:{free_port}'

dist.init_process_group(
backend='nccl', # NCCL 后端适合于 NVIDIA GPU
init_method=init_method,
rank=rank,
world_size=world_size
)
登录后复制

数据集划分与加载
应当引入 DistributedSampler 来替代默认的数据迭代方式，从而确保各 worker 获取到不重复且均衡分布的小批量样本。

from torch.utils.data import DataLoader, DistributedSampler

sampler = DistributedSampler(dataset, shuffle=True)
dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)
登录后复制

模型定义与封装
完成上述准备工作之后，即可实例化目标神经网络结构并将其实例传入至 nn.DistributedDataParallel 构造函数内进行包装处理。

model = MyModel()
device = torch.device('cuda:%d'%local_rank)
model = model.to(device)
ddp_model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
登录后复制
针对跨服务器部署的情况，则还需额外考虑容器化技术的应用以简化资源调度流程；比如借助 Docker 实现 GPU 设备映射及 IP 地址共享等功能特性[^3]。
最后一步就是编写启动命令或脚本来触发各个子进程中执行相同的逻辑代码片段，通常会涉及到设置环境变量如 LOCAL_RANK 等指示当前运行位置的信息。