PyTorch 分布式训练部署指南

适用环境: K3s + HAMi vGPU 集群
最后更新: 2026-08-22

1. 分布式训练简介

为什么需要分布式训练？

问题	分布式方案
单卡显存不足	模型并行 (Model Parallelism)
训练时间太长	数据并行 (Data Parallelism)
数据集太大	数据并行 + 分布式数据加载

并行策略对比

策略	原理	适用场景	本集群支持
DataParallel (DP)	单进程多线程，GIL 限制	单机 2-4 卡	✅
DistributedDataParallel (DDP)	多进程，AllReduce 通信	单机/多机	✅
FSDP	模型参数分存储	超大模型	✅ (PyTorch 2.0+)
Pipeline Parallel	模型按层切分	超大模型	❌ (需特殊框架)
Tensor Parallel	单层层内切分	超大模型	❌ (需 Megatron)

2. PyTorch 分布式架构

2.1 DDP 工作流程

┌─────────────────────────────────────────────────────┐
│                   Master (Rank 0)                    │
│  - 参数初始化                                         │
│  - 梯度聚合 (AllReduce)                               │
│  - 参数更新                                           │
└─────────────────────────────────────────────────────┘
                          ↕ NCCL/gloo
┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐
│  Worker Rank 1  │  │  Worker Rank 2  │  │  Worker Rank 3  │
│  GPU 0          │  │  GPU 1          │  │  GPU 2          │
│  数据子集 A     │  │  数据子集 B     │  │  数据子集 C     │
└─────────────────┘  └─────────────────┘  └─────────────────┘

2.2 环境变量

变量	说明	示例
`MASTER_ADDR`	Master 节点地址	`localhost` 或 Pod IP
`MASTER_PORT`	Master 端口	`12355`
`WORLD_SIZE`	总进程数	`4` (4 GPU)
`RANK`	当前进程序号	`0`, `1`, `2`, `3`
`LOCAL_RANK`	当前节点内序号	`0`, `1`

3. Kubernetes 部署方案

3.1 方案选择

方案	复杂度	适用场景	本集群推荐
单 Pod 多 GPU	低	单机训练	✅
多 Pod + Service	中	多机训练	✅
PyTorchJob (Kubeflow)	高	生产级训练	❌ (需额外安装)

3.2 单 Pod 多 GPU (推荐)

适用于单节点多卡训练（如 RTX 3090 节点）：

yaml

apiVersion: v1
kind: Pod
metadata:
  name: pytorch-training
spec:
  schedulerName: hami-scheduler
  nodeSelector:
    kubernetes.io/hostname: "10-60-18-8"
  containers:
  - name: trainer
    image: 117.50.188.237:30000/pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime
    imagePullPolicy: IfNotPresent
    command: ["torchrun", "--nproc_per_node=2", "/workspace/train.py"]
    env:
    - name: HF_ENDPOINT
      value: "https://hf-mirror.com"
    - name: MASTER_ADDR
      value: "localhost"
    - name: MASTER_PORT
      value: "12355"
    - name: NCCL_DEBUG
      value: "INFO"
    resources:
      requests:
        cpu: "2"
        memory: "8Gi"
        nvidia.com/gpu: "2"
        nvidia.com/gpumem: "12k"
      limits:
        cpu: "8"
        memory: "16Gi"
        nvidia.com/gpu: "2"
        nvidia.com/gpumem: "12k"
    volumeMounts:
    - name: code
      mountPath: /workspace
    - name: models
      mountPath: /models
      readOnly: true
    - name: data
      mountPath: /data
  volumes:
  - name: code
    hostPath:
      path: /opt/training-code
      type: DirectoryOrCreate
  - name: models
    persistentVolumeClaim:
      claimName: nfs-modelscope
  - name: data
    persistentVolumeClaim:
      claimName: nfs-training-data
  restartPolicy: Never

3.3 部署与运行

bash

# 准备训练代码（在节点上）
sudo mkdir -p /opt/training-code
sudo cp train.py /opt/training-code/

# 创建训练数据 PVC（如需要）
kubectl apply -f training-data-pvc.yaml

# 启动训练
kubectl apply -f pytorch-training.yaml

# 查看训练日志
kubectl logs -f pytorch-training

4. DataParallel (单节点多卡)

4.1 简介

最简单的多卡方案，但受 Python GIL 限制，性能较差：

python

import torch
import torch.nn as nn

# 创建模型
model = MyModel()

# 包装为 DataParallel
if torch.cuda.device_count() > 1:
    print(f"Using {torch.cuda.device_count()} GPUs!")
    model = nn.DataParallel(model)

model = model.cuda()

# 正常训练
output = model(input_data.cuda())
loss = criterion(output, targets.cuda())
loss.backward()
optimizer.step()

4.2 优缺点

优点：

代码改动最小
无需设置环境变量

缺点：

GIL 限制，多线程效率低
仅支持单机
负载不均衡（GPU 0 承担更多工作）

5. DistributedDataParallel (DDP)

5.1 简介

推荐的多卡训练方案，支持单机/多机：

python

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup(rank, world_size):
    """初始化分布式环境"""
    dist.init_process_group(
        backend="nccl",  # GPU 用 nccl, CPU 用 gloo
        rank=rank,
        world_size=world_size
    )
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train(rank, world_size):
    setup(rank, world_size)
    
    # 创建模型
    model = MyModel().to(rank)
    model = DDP(model, device_ids=[rank])
    
    # 创建分布式数据加载器
    train_dataset = MyDataset()
    train_sampler = DistributedSampler(
        train_dataset,
        num_replicas=world_size,
        rank=rank
    )
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=32,
        sampler=train_sampler
    )
    
    # 训练循环
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)  # 必须调用！
        for batch in train_loader:
            output = model(batch["input"].to(rank))
            loss = criterion(output, batch["target"].to(rank))
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    cleanup()

5.2 启动方式

bash

# 方式一：torchrun (推荐)
torchrun --nproc_per_node=2 train.py

# 方式二：torch.distributed.launch
python -m torch.distributed.launch \
  --nproc_per_node=2 \
  --master_addr=localhost \
  --master_port=12355 \
  train.py

# 方式三：环境变量
export MASTER_ADDR=localhost
export MASTER_PORT=12355
export WORLD_SIZE=2
export RANK=0
python train.py

6. 训练代码示例

6.1 完整 DDP 训练脚本

python

#!/usr/bin/env python3
"""
DDP 分布式训练示例
用法: torchrun --nproc_per_node=2 train_ddp.py
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler
from torchvision import datasets, transforms

def setup():
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    torch.cuda.set_device(local_rank)
    return local_rank

def cleanup():
    dist.destroy_process_group()

def get_dataloader(rank, world_size, batch_size=32):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,))
    ])
    
    dataset = datasets.MNIST(
        root="/data/mnist",
        train=True,
        download=True,
        transform=transform
    )
    
    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank
    )
    
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=2
    )
    
    return loader

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 512)
        self.fc2 = nn.Linear(512, 10)
        self.relu = nn.ReLU()
    
    def forward(self, x):
        x = x.view(-1, 784)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def main():
    rank = setup()
    world_size = dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    
    model = SimpleNet().to(local_rank)
    model = DDP(model, device_ids=[local_rank])
    
    loader = get_dataloader(rank, world_size, batch_size=64)
    
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    num_epochs = 5
    for epoch in range(num_epochs):
        loader.sampler.set_epoch(epoch)
        
        for batch_idx, (data, target) in enumerate(loader):
            data, target = data.to(local_rank), target.to(local_rank)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            if rank == 0 and batch_idx % 10 == 0:
                print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")
    
    if rank == 0:
        torch.save(model.module.state_dict(), "/workspace/model.pth")
        print("Model saved!")
    
    cleanup()

if __name__ == "__main__":
    main()

6.2 运行脚本

bash

# 在 Pod 内执行
torchrun --nproc_per_node=2 /workspace/train_ddp.py

7. 性能调优

7.1 通信后端选择

后端	适用设备	性能
`nccl`	GPU	⚡ 最快
`gloo`	CPU/GPU	🐢 较慢
`mpi`	多机	🚀 高速网络

7.2 Batch Size 调整

python

# 分布式训练时，总 batch size = per_device_batch_size × world_size
per_device_batch_size = 32
world_size = 4
total_batch_size = per_device_batch_size * world_size  # 128

# 学习率也需要相应调整
base_lr = 0.001
scaled_lr = base_lr * world_size  # 0.004

7.3 混合精度训练

python

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in loader:
    optimizer.zero_grad()
    
    with autocast():  # 自动混合精度
        output = model(batch["input"])
        loss = criterion(output, batch["target"])
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

7.4 梯度累积

python

accumulation_steps = 4

for i, batch in enumerate(loader):
    output = model(batch["input"])
    loss = criterion(output, batch["target"])
    loss = loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

7.5 NCCL 优化

bash

# 环境变量优化
export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0
export NCCL_IB_DISABLE=1  # 禁用 InfiniBand（如果没有）
export NCCL_P2P_DISABLE=1  # 禁用 P2P（如果 GPU 间不支持）

8. 故障排查

8.1 NCCL 初始化失败

RuntimeError: NCCL error: unhandled system error

解决：

bash

export NCCL_DEBUG=INFO
export NCCL_SOCKET_IFNAME=eth0

8.2 GPU 不可见

RuntimeError: No CUDA GPUs are available

解决：确保使用 schedulerName: hami-scheduler 并请求 nvidia.com/gpu

8.3 端口冲突

RuntimeError: Address already in use

解决：更换 MASTER_PORT

bash

export MASTER_PORT=12356

8.4 梯度不同步

RuntimeError: Expected to have finished reduction from the prior iteration

解决：确保每个 epoch 都调用 sampler.set_epoch(epoch)

8.5 OOM (Out of Memory)

CUDA out of memory

解决：

减小 batch size
使用混合精度训练
启用梯度检查点：model.gradient_checkpointing_enable()

附录

A. 分布式训练检查清单

[ ] 使用 DDP 而非 DataParallel
[ ] 设置 backend="nccl"
[ ] 每个 epoch 调用 sampler.set_epoch(epoch)
[ ] 只在 rank 0 保存模型和日志
[ ] 调整学习率：lr = base_lr * world_size
[ ] 使用 dist.barrier() 同步关键操作

B. 常用命令

bash

# 查看 GPU 状态
watch -n 1 nvidia-smi

# 查看 NCCL 通信
nvidia-smi nvlink -s

# 查看进程资源使用
nvidia-smi pmon

# 查看分布式环境
python -c "import torch; print(torch.distributed.is_available())"

C. 参考链接

文档版本: v1.0
更新时间: 2026-08-22
适用版本: PyTorch 2.1+ CUDA 12 | NCCL_DEBUG | NCCL 日志级别 | INFO |

PyTorch 分布式训练部署指南

目录

1. 分布式训练简介

为什么需要分布式训练？

并行策略对比

2. PyTorch 分布式架构

2.1 DDP 工作流程

2.2 环境变量

3. Kubernetes 部署方案

3.1 方案选择

3.2 单 Pod 多 GPU (推荐)

3.3 部署与运行

4. DataParallel (单节点多卡)

4.1 简介

4.2 优缺点

推荐：仅用于快速原型验证，生产环境使用 DDP。

5. DistributedDataParallel (DDP)

5.1 简介

5.2 启动方式

6. 训练代码示例

6.1 完整 DDP 训练脚本

6.2 运行脚本

7. 性能调优

7.1 通信后端选择

7.2 Batch Size 调整

7.3 混合精度训练

7.4 梯度累积

7.5 NCCL 优化

8. 故障排查

8.1 NCCL 初始化失败

8.2 GPU 不可见

8.3 端口冲突

8.4 梯度不同步

8.5 OOM (Out of Memory)

附录

A. 分布式训练检查清单

B. 常用命令

C. 参考链接

PyTorch 分布式训练部署指南 ​

目录 ​

1. 分布式训练简介 ​

为什么需要分布式训练？ ​

并行策略对比 ​

2. PyTorch 分布式架构 ​

2.1 DDP 工作流程 ​

2.2 环境变量 ​

3. Kubernetes 部署方案 ​

3.1 方案选择 ​

3.2 单 Pod 多 GPU (推荐) ​

3.3 部署与运行 ​

4. DataParallel (单节点多卡) ​

4.1 简介 ​

4.2 优缺点 ​

推荐：仅用于快速原型验证，生产环境使用 DDP。 ​

5. DistributedDataParallel (DDP) ​

5.1 简介 ​

5.2 启动方式 ​

6. 训练代码示例 ​

6.1 完整 DDP 训练脚本 ​

6.2 运行脚本 ​

7. 性能调优 ​

7.1 通信后端选择 ​

7.2 Batch Size 调整 ​

7.3 混合精度训练 ​

7.4 梯度累积 ​

7.5 NCCL 优化 ​

8. 故障排查 ​

8.1 NCCL 初始化失败 ​

8.2 GPU 不可见 ​

8.3 端口冲突 ​

8.4 梯度不同步 ​

8.5 OOM (Out of Memory) ​

附录 ​

A. 分布式训练检查清单 ​

B. 常用命令 ​

C. 参考链接 ​

PyTorch 分布式训练部署指南

目录

1. 分布式训练简介

为什么需要分布式训练？

并行策略对比

2. PyTorch 分布式架构

2.1 DDP 工作流程

2.2 环境变量

3. Kubernetes 部署方案

3.1 方案选择

3.2 单 Pod 多 GPU (推荐)

3.3 部署与运行

4. DataParallel (单节点多卡)

4.1 简介

4.2 优缺点

推荐：仅用于快速原型验证，生产环境使用 DDP。

5. DistributedDataParallel (DDP)

5.1 简介

5.2 启动方式

6. 训练代码示例

6.1 完整 DDP 训练脚本

6.2 运行脚本

7. 性能调优

7.1 通信后端选择

7.2 Batch Size 调整

7.3 混合精度训练

7.4 梯度累积

7.5 NCCL 优化

8. 故障排查

8.1 NCCL 初始化失败

8.2 GPU 不可见

8.3 端口冲突

8.4 梯度不同步

8.5 OOM (Out of Memory)

附录

A. 分布式训练检查清单

B. 常用命令

C. 参考链接