混合精度训练与分布式训练（DDP）

import torch
from torch.cuda.amp import GradScaler, autocast

model = model.cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = GradScaler()  # 自动 loss scaling

for x, y in dataloader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()

    # autocast 上下文内的运算用 FP16
    with autocast():
        pred = model(x)
        loss = criterion(pred, y)

    # loss scaling + 反向
    scaler.scale(loss).backward()

    # 更新前 unscale 梯度,可选地做梯度裁剪
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    # 优化器更新;scaler 会自动调整 scale factor
    scaler.step(optimizer)
    scaler.update()

from torch.cuda.amp import autocast

for x, y in dataloader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad()

    # dtype=bfloat16,无需 scaler
    with autocast(dtype=torch.bfloat16):
        pred = model(x)
        loss = criterion(pred, y)

    loss.backward()
    optimizer.step()

import torch
from torch.cuda.amp import autocast, GradScaler

def train_one_epoch(model, loader, optimizer, criterion, device, use_bf16=True):
    model.train()
    scaler = None if use_bf16 else GradScaler()
    total_loss = 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()

        amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
        with autocast(dtype=amp_dtype):
            pred = model(x)
            loss = criterion(pred, y)

        if scaler:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

        total_loss += loss.item()
    return total_loss / len(loader)

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


def setup_ddp():
    """初始化分布式进程组"""
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank


def cleanup_ddp():
    dist.destroy_process_group()


def main():
    rank = setup_ddp()
    is_main = (rank == 0)

    # 1. 模型放到当前 GPU
    model = MyModel().cuda(rank)
    model = DDP(model, device_ids=[rank])

    # 2. 数据集 + DistributedSampler
    dataset = MyDataset()
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=64, sampler=sampler,
                        num_workers=4, pin_memory=True)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    scaler = torch.amp.GradScaler('cuda')

    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # 关键:每个 epoch 重置 shuffle
        for x, y in loader:
            x, y = x.cuda(rank), y.cuda(rank)
            optimizer.zero_grad()

            with torch.amp.autocast('cuda', dtype=torch.bfloat16):
                pred = model(x)
                loss = criterion(pred, y)

            loss.backward()
            optimizer.step()

            if is_main and step % 100 == 0:
                print(f"epoch {epoch} step {step} loss {loss.item():.4f}")

    cleanup_ddp()


if __name__ == "__main__":
    main()

sampler = DistributedSampler(dataset, shuffle=True)
loader = DataLoader(dataset, batch_size=64, sampler=sampler)

# 只有 rank 0 保存
if dist.get_rank() == 0:
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.module.state_dict(),  # 注意 .module
        'optimizer_state_dict': optimizer.state_dict(),
    }, 'checkpoint.pt')

# 所有 rank 加载(各加载到自己的模型)
map_location = {'cuda:%d' % 0: 'cuda:%d' % rank}
ckpt = torch.load('checkpoint.pt', map_location=map_location)
model.module.load_state_dict(ckpt['model_state_dict'])

accumulation_steps = 4
for i, (x, y) in enumerate(loader):
    with autocast(dtype=torch.bfloat16):
        loss = criterion(model(x.cuda()), y.cuda()) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# GPU 显存
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Cached:    {torch.cuda.memory_reserved() / 1e9:.2f} GB")

# 吞吐量
torch.cuda.synchronize()
start = time.time()
# ... train ...
torch.cuda.synchronize()
print(f"Throughput: {num_samples / (time.time() - start):.1f} samples/s")