6. 标准训练流程

标准训练流程

• 标准训练流程 = 训练/验证模式切换 → 取 batch → 前向 → loss → 反向 → 更新 → 清理状态 → 记录

• 最外层结构（epoch 级）

  # epoch 是数据遍历次数
  # Dataset 不知道 eposh
  # Sampler（DDP）时需要知道 epoch
  for epoch in range(num_epochs):
      if sampler is not None:
          sampler.set_epoch(epoch)
  
      train_one_epoch(...)
      val_loss = validate(...)
  
      if rank == 0:
          print(f"epoch {epoch}, val_loss={val_loss:.4f}")

• train_one_epoch

  def move_to_device(batch, device):
      if isinstance(batch, torch.Tensor):
          return batch.to(device, non_blocking=True)
      if isinstance(batch, dict):
          return {k: move_to_device(v, device) for k, v in batch.items()}
      if isinstance(batch, list):
          return [move_to_device(v, device) for v in batch]
      return batch
  
  def train_one_epoch(
      model,
      dataloader,
      optimizer,
      device,
      scheduler=None
  ):
      model.train() # 切换到训练模式
  
      # Sampler 已决定 index、DataLoader 已构造 batch、collate_fn 已对齐结构
      for batch in dataloader:
          batch = move_to_device(batch, device)
  
          outputs = model(**batch["inputs"])
          # loss 必须是标量，loss 必须在当前 device，loss 必须有 grad
          loss = outputs["loss"]
  		# autograd 从 loss 出发，沿计算图反向传播，每个参数的 grad 被累积
          loss.backward()
          # 使用当前 grad，更新模型的参数
          optimizer.step()
          # 在更新之后，清空梯度历史，防止下一个 batch 梯度叠加
          optimizer.zero_grad()
  
          if scheduler is not None:
              # 学习率的时间轴，可能会是 epoch 级别的更新，即
              # for epoch:
              #     train_one_epoch()
              #     scheduler.step()
              scheduler.step()

• validate

  def validate(model, dataloader, device):
      model.eval() # 切换至验证模式
      total_loss = 0
  
      with torch.no_grad(): # 无需计算梯度
          for batch in dataloader:
              batch = move_to_device(batch, device)
              outputs = model(**batch["inputs"])
              total_loss += outputs["loss"].item()
  
      return total_loss / len(dataloader)

• 在一个完整的训练流程中，不应该出现 Dataset 逻辑、DataLoader 配置、Model 定义、分布式初始化

训练流程增强

• 增强后的训练流程

  def train_one_epoch(
      model,
      dataloader,
      optimizer,
      device,
      scheduler=None,
      scaler=None,
      max_grad_norm=None
  ):
      model.train()
  
      for batch in dataloader:
          batch = move_to_device(batch, device)
  
          optimizer.zero_grad()
  
          with torch.cuda.amp.autocast(enabled=(scaler is not None)):
              outputs = model(**batch["inputs"])
              loss = outputs["loss"]
  
          if scaler is not None:
              scaler.scale(loss).backward()
              # AMP 的梯度裁剪
              if max_grad_norm is not None:
                  scaler.unscale_(optimizer)
                  torch.nn.utils.clip_grad_norm_(
                      model.parameters(),
                      max_grad_norm
                  )
              scaler.step(optimizer)
              scaler.update()
          else:
              loss.backward()
              # fp32 的梯度裁剪
              if max_grad_norm is not None:
                  torch.nn.utils.clip_grad_norm_(
                      model.parameters(),
                      max_grad_norm
                  )
              optimizer.step()
  
          if scheduler is not None:
              scheduler.step()

• 增强后的验证流程

  def validate(model, dataloader, device):
      model.eval()
      total_loss = 0
  
      with torch.no_grad():
          for batch in dataloader:
              batch = move_to_device(batch, device)
              outputs = model(**batch["inputs"])
              total_loss += outputs["loss"].item()
  
      return total_loss / len(dataloader)

AMP

• AMP（Automatic Mixed Precision） 要实现的是：用更低精度算得更快，但又不把数值稳定性搞崩

• 深度学习里的计算任务可以分为两类

  • 计算密集型：matmul、conv、attention

  • 数值敏感型：loss、梯度累积、optimizer 更新

• AMP 的核心思想是不同类型的算子，使用不同精度

• fp16 的致命缺陷是表示范围小、精度低，在反向传播中梯度往往非常小（$10^{-8}$、$10^{-10}$），而 fp16 表示不了 → underflow → 梯度直接变 0，从而导致梯度消失，模型无法学习

• AMP 由两个组件构成

  • autocast → 控制 forward 的计算精度，只影响 forward，不改变模型参数的存储精度（仍是 fp32），而是在 forward 时临时切换算子精度

    # forward 才是计算密集区，backward 不在 autocast 控制范围
    # 使用 BF16 训练更稳定，autocast(dtype=torch.bfloat16)
    # bf16 表示范围 ≈ fp32，精度稍低但不容易 underflow，通常不需要 GradScaler
    with torch.cuda.amp.autocast():
        outputs = model(x)
        loss = criterion(outputs, y)
    # matmul / conv      → fp16 / bf16
    # softmax / norm     → fp32
    # loss               → fp32

  • GradScaler → 控制 backward 的数值稳定性，解决梯度 underflow，其先把 loss 放大，再反向传播，反向传播后再把梯度缩小

    # 初始化 scaler，scaler 内部维护当前 scale 值、是否出现 inf / nan、动态调整策略
    scaler = torch.cuda.amp.GradScaler()
    
    # backward 前先放大 loss
    # 等价于 (loss * scale).backward()
    scaler.scale(loss).backward()
    
    # backward 之后缩小梯度
    # 把梯度除以 scale，写回 parameter.grad
    scaler.unscale_(optimizer)
    
    # 通过 scaler 进行 optimizer.step
    # 如果梯度有 inf / nan → 跳过 step
    # 否则 → 正常 optimizer.step()
    scaler.step(optimizer)
    
    # 更新 scaler
    # 若多次稳定 → 增大 scale
    # 若出现溢出 → 减小 scale
    scaler.update()

• AMP 的完整顺序

  with torch.cuda.amp.autocast():
      outputs = model(...)
      loss = outputs["loss"]
  
  scaler.scale(loss).backward()
  scaler.unscale_(optimizer)
  # clip 需要真实梯度，必须放在 unscale 之后
  torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
  scaler.step(optimizer)
  scaler.update()

• AMP 并不会减少显存里的参数大小，其参数仍是 fp32，而显存节省主要来自激活/中间结果

• AMP 不保证一定更快，当 batch 很小或者 CPU/IO 瓶颈时收益有限

• 在 loss 本身极不稳定或自定义 CUDA op 不支持 fp16 时，不能使用 AMP

• AMP + DDP 为什么不会冲突

  • DDP 的 AllReduce 用的是 parameter.grad，而 AMP 在 AllReduce 之前已经把它 unscale 成 fp32

  • 一轮 iteration 的完整顺序

    forward (autocast)
    ↓
    loss (fp32)
    ↓
    scaler.scale(loss).backward()
    ↓
    grad still scaled (临时状态)
    ↓
    scaler.unscale_(optimizer)
    ↓
    parameter.grad = fp32 的真实梯度
    ↓
    DDP AllReduce (fp32 梯度)
    ↓
    optimizer.step()

  • DDP 同步发生在 backward 结束之后，且看到的是已经 unscale 的梯度

  • 在 PyTorch 中，DDP 会在 backward 时给每个参数注册一个 autograd hook，当梯度计算完成时，hook 触发，对该参数的梯度做 AllReduce

  • 但这个 hook 看到的是 grad 张量本身，而 AMP 的流程是 backward 结束、scaler.unscale_() 改写 grad 的数值，此时 DDP hook 同步的是 unscale 后的结果

• GradScaler 不会影响多卡一致性，因为所有 rank 用的是同一个 loss scale 策略，如果某个 rank 出现 inf / nan，scaler.step() 会在所有 rank 跳过更新

• 推理阶段的混合精度与训练阶段的 AMP

  • 训练的 AMP 是“数值安全工程”，而推理的混合精度是“纯性能工程”

  • 推理阶段的流程

    with torch.autocast(device_type="cuda", dtype=torch.float16):
        y = model(x)
    # 更激进地
    model.half()

  • 推理阶段无需计算梯度，数值问题主要来自 softmax / normalization，而这些算子会由 autocast 保留为 fp32，因此不需要 GradScaler

  • 对比

    # 训练阶段
    参数        → fp32
    forward     → mixed (fp16 / fp32)
    backward    → fp32
    optimizer   → fp32
    需要 scaler → 是
    # 推理阶段
    参数        → fp16 或 bf16
    forward     → 低精度为主
    backward    → 不存在
    需要 scaler → 否

学习率调度（Scheduler）

• Step / Cosine

  scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
      optimizer,
      T_max=total_steps
  )
  # 放在 optimizer.step() 之后，每个 step 调用一次
  # scheduler 依赖当前 step 的 lr 状态，不能在 backward 前

• Warmup 是 scheduler 的一部分，不应散落在训练循环各处

  # 前 N 个 step 线性上升学习率
  # 后续交给 Cosine / Step
  if global_step < warmup_steps:
      lr = base_lr * global_step / warmup_steps
      set_lr(optimizer, lr)
  else:
      scheduler.step()

Checkpoint

• 保存 checkpoint （只在 rank 0）：保存参数状态、优化器动量、scheduler 位置、AMP 缩放状态、当前 epoch

  def save_checkpoint(path, model, optimizer, scheduler, scaler, epoch):
      torch.save({
          "model": model.state_dict(),
          "optimizer": optimizer.state_dict(),
          "scheduler": scheduler.state_dict() if scheduler else None,
          "scaler": scaler.state_dict() if scaler else None,
          "epoch": epoch
      }, path)

• 恢复 checkpoint

  # load 在 model.to(device) 之后
  # optimizer / scaler 必须与当初结构一致
  ckpt = torch.load(path, map_location="cpu")
  
  model.load_state_dict(ckpt["model"])
  optimizer.load_state_dict(ckpt["optimizer"])
  
  if scheduler and ckpt["scheduler"]:
      scheduler.load_state_dict(ckpt["scheduler"])
  
  if scaler and ckpt["scaler"]:
      scaler.load_state_dict(ckpt["scaler"])
  
  start_epoch = ckpt["epoch"] + 1

梯度裁剪

• 位置放在 backward 之后，optimizer.step 之前

  # fp32
  loss.backward()
  torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
  optimizer.step()
  # AMP
  scaler.scale(loss).backward()
  scaler.unscale_(optimizer)
  torch.nn.utils.clip_grad_norm_(...)
  scaler.step(optimizer)

• clip 是训练稳定性工具，不改变 loss，只限制更新幅度

早停（Early Stopping）

• 早停不在 train_one_epoch 里，而在 epoch 外

  best_loss = float("inf")
  patience = 5
  bad_epochs = 0
  
  for epoch in range(num_epochs):
      train_one_epoch(...)
      val_loss = validate(...)
  
      if val_loss < best_loss:
          best_loss = val_loss
          bad_epochs = 0
          save_checkpoint(...)
      else:
          bad_epochs += 1
  
      if bad_epochs >= patience:
          print("early stopping")
          break

完整示例

• 一轮训练 iteration 中，真实发生的顺序是：

  • Dataset 定义「样本是什么」

  • DataLoader 定义「batch 如何产生」

  • DistributedSampler 决定「每张卡拿哪些样本」

  • model = DDP(model)

  • forward 在 autocast 里

  • loss 计算（fp32）

  • scaler.scale(loss).backward()

  • scaler.unscale_(optimizer)

  • 可选：gradient clipping

  • optimizer.step()

• 第一步，实现 Dataset

  import torch
  from torch.utils.data import Dataset
  
  class ToyDataset(Dataset):
      def __init__(self, data, labels):
          self.data = data          # Tensor or numpy
          self.labels = labels
  
      def __len__(self):
          return len(self.data)
  
      def __getitem__(self, idx):
          x = self.data[idx]
          y = self.labels[idx]
          return x, y

• 第二步，初始化分布式

  import torch
  import torch.distributed as dist
  
  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

• 第三步，实现 DataLoader

  from torch.utils.data import DataLoader
  from torch.utils.data.distributed import DistributedSampler
  
  dataset = ToyDataset(data, labels)
  
  sampler = DistributedSampler(
      dataset,
      shuffle=True,     # 注意：shuffle 在 sampler，不在 DataLoader
      drop_last=True
  )
  
  dataloader = DataLoader(
      dataset,
      batch_size=32,
      sampler=sampler,
      num_workers=4,
      pin_memory=True,
      persistent_workers=True
  )

• 第四步，实现模型

  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  class SimpleMLP(nn.Module):
      def __init__(self, input_dim, hidden_dim, num_classes):
          super().__init__()
  
          self.fc1 = nn.Linear(input_dim, hidden_dim)
          self.fc2 = nn.Linear(hidden_dim, hidden_dim)
          self.fc3 = nn.Linear(hidden_dim, num_classes)
  
          self.norm = nn.LayerNorm(hidden_dim)
  
      def forward(self, x):
          # x: [B, D]
          x = self.fc1(x)
          x = F.relu(x)
  
          x = self.fc2(x)
          x = self.norm(x)
          x = F.relu(x)
  
          x = self.fc3(x)
          return x

• 第五步，初始化相关组件

  import torch.nn as nn
  from torch.nn.parallel import DistributedDataParallel as DDP
  
  model = MyModel().cuda()
  model = DDP(model, device_ids=[local_rank])
  
  optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
  
  scaler = torch.cuda.amp.GradScaler()

• 第六步，实现训练流程

  from torch.cuda.amp import autocast
  
  def train_one_epoch(
      model,
      dataloader,
      optimizer,
      scaler,
      loss_fn,
      device,
      epoch,
      grad_accum_steps
  ):
      model.train()
      optimizer.zero_grad(set_to_none=True)
  
      total_loss = 0.0
  
      for step, (x, y) in enumerate(dataloader):
          x = x.to(device, non_blocking=True)
          y = y.to(device, non_blocking=True)
  
          with autocast(device_type="cuda", dtype=torch.float16):
              logits = model(x)
              loss = loss_fn(logits, y)
              loss = loss / grad_accum_steps   # 关键：缩放 loss
  
          scaler.scale(loss).backward()
  
          is_accum_step = (step + 1) % grad_accum_steps == 0
  
          if is_accum_step:
              # 反向传播已经结束，开始处理真实梯度
              scaler.unscale_(optimizer)
  
              torch.nn.utils.clip_grad_norm_(
                  model.parameters(),
                  max_norm=1.0
              )
  
              scaler.step(optimizer)
              scaler.update()
  
              optimizer.zero_grad(set_to_none=True)
  
          total_loss += loss.item() * grad_accum_steps
  
      return total_loss / len(dataloader)

• 第七步，实现验证流程

  @torch.no_grad()
  def validate(
      model,
      dataloader,
      loss_fn,
      device
  ):
      model.eval()
  
      total_loss = 0.0
      total_correct = 0
      total_samples = 0
  
      with autocast(device_type="cuda", dtype=torch.float16):
          for x, y in dataloader:
              x = x.to(device, non_blocking=True)
              y = y.to(device, non_blocking=True)
  
              logits = model(x)
              loss = loss_fn(logits, y)
  
              total_loss += loss.item() * x.size(0)
  
              preds = logits.argmax(dim=-1)
              total_correct += (preds == y).sum().item()
              total_samples += x.size(0)
  
      avg_loss = total_loss / total_samples
      accuracy = total_correct / total_samples
  
      return avg_loss, accuracy

• 第八步，完整代码流程

  def main():
      local_rank = setup_ddp()
      device = torch.device("cuda", local_rank)
  
      # --------------------
      # 超参数
      # --------------------
      input_dim = 128
      hidden_dim = 256
      num_classes = 10
      batch_size = 32
      grad_accum_steps = 4
      num_epochs = 10
      lr = 3e-4
  
      # --------------------
      # 构造假数据（示例）
      # --------------------
      num_samples = 10000
      data = torch.randn(num_samples, input_dim)
      labels = torch.randint(0, num_classes, (num_samples,))
  
      # --------------------
      # Dataset / Sampler / DataLoader
      # --------------------
      train_dataset = ToyDataset(data, labels)
      val_dataset = ToyDataset(data[:2000], labels[:2000])
  
      train_sampler = DistributedSampler(
          train_dataset,
          shuffle=True,
          drop_last=True
      )
  
      val_sampler = DistributedSampler(
          val_dataset,
          shuffle=False
      )
  
      train_loader = DataLoader(
          train_dataset,
          batch_size=batch_size,
          sampler=train_sampler,
          num_workers=4,
          pin_memory=True,
          persistent_workers=True
      )
  
      val_loader = DataLoader(
          val_dataset,
          batch_size=batch_size,
          sampler=val_sampler,
          num_workers=4,
          pin_memory=True
      )
  
      # --------------------
      # 模型 / DDP
      # --------------------
      model = SimpleMLP(input_dim, hidden_dim, num_classes)
      model.to(device)
      model = DDP(model, device_ids=[local_rank])
  
      # --------------------
      # Optimizer / Loss / AMP
      # --------------------
      optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
      loss_fn = nn.CrossEntropyLoss()
      scaler = GradScaler()
  
      # --------------------
      # 训练循环
      # --------------------
      for epoch in range(num_epochs):
          train_sampler.set_epoch(epoch)
  
          train_loss = train_one_epoch(
              model,
              train_loader,
              optimizer,
              scaler,
              loss_fn,
              device,
              grad_accum_steps
          )
  
          if dist.get_rank() == 0:
              val_loss, val_acc = validate(
                  model,
                  val_loader,
                  loss_fn,
                  device
              )
  
              print(
                  f"Epoch {epoch} | "
                  f"Train Loss {train_loss:.4f} | "
                  f"Val Loss {val_loss:.4f} | "
                  f"Val Acc {val_acc:.4f}"
              )