AMP 混合精度训练中的动态缩放机制: grad_scaler.py函数解析（ torch._amp_update_scale

AMP 混合精度训练中的动态缩放机制: grad_scaler.py函数解析（ torch._amp_update_scale_）

AMP 混合精度训练中的动态缩放机制

在深度学习中，混合精度训练（AMP, Automatic Mixed Precision）是一种常用的技术，它利用半精度浮点（FP16）计算来加速训练，同时使用单精度浮点（FP32）来保持数值稳定性。为了在混合精度训练中避免数值溢出，PyTorch 提供了一种动态缩放机制来调整 “loss scale”（损失缩放值）。本文将详细解析动态缩放机制的实现原理，并通过代码展示其内部逻辑。

动态缩放机制简介

动态缩放机制的核心思想是通过一个可动态调整的缩放因子（scale factor）放大 FP16 的梯度，从而降低舍入误差对训练的影响。当检测到数值不稳定（例如 NaN 或无穷大）时，缩放因子会被降低；当连续多步未检测到数值问题时，缩放因子会被提高。其调整策略基于以下两个参数：

growth_factor: 连续成功步骤后用于增加缩放因子的乘数（通常大于 1，如 2.0）。
backoff_factor: 检测到数值溢出时用于减少缩放因子的乘数（通常小于 1，如 0.5）。

此外，动态缩放还使用 growth_interval 参数控制连续成功步骤的计数阈值。当达到这个阈值时，缩放因子才会增加。

AMP 缩放更新核心代码解析

PyTorch 实现了一个用于更新缩放因子的 CUDA 核函数以及相关的 Python 包装函数。以下是核心代码解析：

CUDA 核函数实现

// amp_update_scale_cuda_kernel 核函数实现
__global__ void amp_update_scale_cuda_kernel(float* current_scale,
                                             int* growth_tracker,
                                             const float* found_inf,
                                             double growth_factor,
                                             double backoff_factor,
                                             int growth_interval) {
  if (*found_inf) {
    // 如果发现梯度中存在 NaN 或 Inf，缩放因子乘以 backoff_factor，并重置 growth_tracker。
    *current_scale = (*current_scale) * backoff_factor;
    *growth_tracker = 0;
  } else {
    // 未发现数值问题，增加 growth_tracker 的计数。
    auto successful = (*growth_tracker) + 1;
    if (successful == growth_interval) {
      // 当 growth_tracker 达到 growth_interval，尝试增长缩放因子。
      auto new_scale = static_cast<float>((*current_scale) * growth_factor);
      if (isfinite_ensure_cuda_math(new_scale)) {
        *current_scale = new_scale;
      }
      *growth_tracker = 0;
    } else {
      *growth_tracker = successful;
    }
  }
}

核函数逻辑

发现数值溢出（found_inf > 0）：
- 缩放因子 current_scale 乘以 backoff_factor。
- 重置成功计数器 growth_tracker 为 0。
未发现数值溢出：
- 增加成功计数器 growth_tracker。
- 如果 growth_tracker 达到 growth_interval，则将缩放因子乘以 growth_factor。
- 保证缩放因子不会超过 FP32 的数值上限。

C++ 包装函数实现

在 PyTorch 中，这一 CUDA 核函数通过 C++ 包装函数 _amp_update_scale_cuda_ 被调用。以下是实现代码：

Tensor& _amp_update_scale_cuda_(Tensor& current_scale,
                                Tensor& growth_tracker,
                                const Tensor& found_inf,
                                double growth_factor,
                                double backoff_factor,
                                int64_t growth_interval) {
  TORCH_CHECK(growth_tracker.is_cuda(), "growth_tracker must be a CUDA tensor.");
  TORCH_CHECK(current_scale.is_cuda(), "current_scale must be a CUDA tensor.");
  TORCH_CHECK(found_inf.is_cuda(), "found_inf must be a CUDA tensor.");
  
  // 核函数调用
  amp_update_scale_cuda_kernel<<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(
    current_scale.mutable_data_ptr<float>(),
    growth_tracker.mutable_data_ptr<int>(),
    found_inf.const_data_ptr<float>(),
    growth_factor,
    backoff_factor,
    growth_interval);
  C10_CUDA_KERNEL_LAUNCH_CHECK();

  return current_scale;
}

Python 调用入口

AMP 的 GradScaler 类通过 _amp_update_scale_ 函数更新缩放因子，以下是相关代码：
代码来源：anaconda3/envs/xxxx/lib/python3.10/site-packages/torch/amp/grad_scaler.py

具体调用过程可以参考笔者的另一篇博文：PyTorch到C++再到 CUDA 的调用链（C++ ATen 层）：以torch._amp_update_scale_调用为例

def update(self, new_scale: Optional[Union[float, torch.Tensor]] = None) -> None:
    """更新缩放因子"""
    if not self._enabled:
        return

    _scale, _growth_tracker = self._check_scale_growth_tracker("update")

    if new_scale is not None:
        # 设置用户定义的新缩放因子。
        self._scale.fill_(new_scale)
    else:
        # 收集所有优化器中的 found_inf 数据。
        found_infs = [
            found_inf.to(device=_scale.device, non_blocking=True)
            for state in self._per_optimizer_states.values()
            for found_inf in state["found_inf_per_device"].values()
        ]

        found_inf_combined = found_infs[0]
        if len(found_infs) > 1:
            for i in range(1, len(found_infs)):
                found_inf_combined += found_infs[i]

        # 更新缩放因子。
        torch._amp_update_scale_(
            _scale,
            _growth_tracker,
            found_inf_combined,
            self._growth_factor,
            self._backoff_factor,
            self._growth_interval,
        )