数据并行

# 伪代码示例
for batch in data_loader:
    # 1. 分割数据到各个GPU
    mini_batch = split_batch(batch, num_gpus)
    
    # 2. 各GPU独立计算
    loss = model(mini_batch)
    gradients = backward(loss)
    
    # 3. 梯度同步
    all_reduce(gradients)  # 求平均
    
    # 4. 更新参数
    optimizer.step()

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# 初始化分布式环境
dist.init_process_group("nccl")

# 创建模型并移到GPU
model = MyModel().cuda()
model = DDP(model)

# 训练循环
for data, target in train_loader:
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()  # DDP自动处理梯度同步
    optimizer.step()

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

# FSDP配置
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    cpu_offload=CPUOffload(offload_params=True),
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
    )
)