——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之十九
深度学习训练长期面临一个”不可能三角”：算得快、算得准、占得少，三者似乎很难同时满足。FP32 精度高、动态范围大，但显存占用大和计算密度低；FP16 显存减半、Tensor Core 吞吐高，却容易因为动态范围不足而把梯度”吞”成零。自动混合精度训练（AMP，Automatic Mixed Precision）就是在两者之间搭一座桥：让大部分前向/反向计算走 FP16，把优化器更新和精度敏感的中间状态留在 FP32。
从 2018 年 NVIDIA 的 Apex，到 PyTorch 1.6 正式引入 torch.cuda.amp，AMP 已经成为现代训练框架的标配能力。但 AMP 并不只是把权重类型改一下那么简单——它涉及主权重双轨存储、损失缩放、NaN/Inf 检测、梯度裁剪、FP16↔FP32 批量转换，以及与底层算子融合（如 Conv+BN+ReLU）的深度耦合。
Tech-Renaissance 选择了一条与 PyTorch 动态 AMP 不同的路：图级、静态、显式的 AMP。所有精度转换、缩放检测、类型转换在编译期就写入计算图，运行期按固定顺序启动 CUDA Graph。这篇文章就来拆解这条流水线是怎么工作的，以及为什么我们坚持使用 FP16 和固定损失缩放。
一、为什么需要混合精度训练
NVIDIA 从 Volta 架构开始在 GPU 中集成 Tensor Core，专门加速半精度矩阵乘加。以 A100 为例，FP16 Tensor Core 的稠密峰值算力约为 312 TFLOPS，而 FP32 只有约 19.5 TFLOPS，相差整整一个数量级。对卷积、全连接这类以 GEMM/卷积为核心的算子，使用 FP16 计算意味着理论上可以把峰值算力提升数倍。
除了峰值算力，FP16 还带来两个更实际的好处：
- 显存占用减半：同样尺寸的权重、激活、梯度，FP16 只需要 FP32 一半的空间；
- 显存带宽减半：数据在 HBM 与计算单元之间搬运时，FP16 传输量也只有 FP32 的一半。
对 ResNet-50、VGG16BN 这类经典 CNN，训练过程在 A100 上往往是显存带宽受限而非纯算力受限。减半带宽能直接提升有效吞吐，这也是 AMP 在 CNN 训练中收益明显的原因之一。
但 FP16 的代价同样明显。它只有 1 位符号位、5 位指数位、10 位尾数位，最小正正规数约 6.1×10⁻⁵，最大可表示值约 6.55×10⁴。训练中的很多梯度值在 10⁻⁵ 甚至 10⁻⁷ 量级，已经接近或低于 FP16 的正规数范围，可能进入 subnormal 区间，精度大幅下降，甚至在启用 flush-to-zero 的路径下被直接归零，导致有效梯度信息丢失；偶尔出现的大梯度又可能上溢成 Inf 或 NaN，导致训练崩溃。
因此混合精度训练的核心思想是两条规则：
- FP32 主权重（master weights）：优化器在 FP32 权重上更新，避免微小更新被 FP16 精度”吞掉”；
- 损失缩放（loss scaling）：前向时把损失值乘以一个较大的缩放因子
S，让反向传播的梯度也放大S倍，离开下溢区；优化器更新前再除以S还原。
主流框架 PyTorch 的 torch.cuda.amp 采用动态损失缩放：初始 scale 通常设为 65536，每轮检测梯度是否溢出，无溢出则连续成功若干步后翻倍，溢出则跳过更新并减半。这种策略对用户透明，但维护 per-device 的 scale 状态、FoundInf 标志和增长计数器，对 eager 模式很自然，对追求确定性、静态化的框架却不是最优解。
二、Tech-Renaissance 的 AMP 选择：FP16 + 固定缩放
Tech-Renaissance 选择固定损失缩放，初始值硬编码在 include/renaissance/core/global_config.h：
#define TR_AMP_INITIAL_SCALING 8192.0f
这个值是 2 的幂（2¹³）。取 2 的幂有两个好处：
- 放大和还原本质上是指数位加减，舍入误差更小；
- 对典型的 CNN 训练，8192 足够把大部分小梯度抬升到 FP16 可表示范围，同时又不至于把损失放大到接近 FP16 上限。
与 PyTorch 的动态策略不同，我们的缩放因子只减不增。当 RANGE_CHECK_NAN 检测到梯度中存在 NaN 时，RANGE_GRAD_SCALING 会把 scaling 乘以 0.5，下限截断到 1.0：
// grad_scaling_op.cu
__global__ void grad_scaling_kernel(const int32_t* __restrict__ has_nan,
float* __restrict__ scaling)
{
if (*has_nan != 0) {
float new_scaling = (*scaling) * 0.5f;
*scaling = (new_scaling < 1.0f) ? 1.0f : new_scaling;
}
}
has_nan：RANGE_CHECK_NAN写入的 GPU 标量标志，0 表示无异常，非 0 表示发现异常；scaling：当前损失缩放因子，保存在S_SCALAR_FP32区域，可被 CUDA Graph 动态更新。
这种”只减不增”的策略看起来会损失一些缩放空间，但它有两个明确优势：
- CUDA Graph 友好：缩放因子的调整不需要 CPU 重新生成 kernel 参数或控制流，完全在 GPU 侧一个 kernel 内完成；
- 确定性：在固定随机种子、固定算子实现的前提下，是否出现 NaN 是确定的，缩放路径也是确定的。
它的代价是鲁棒性不如动态策略：如果训练早期因为异常把 scale 降到很低，后续不会自动恢复。但在我们实测的 CNN 训练场景中，溢出通常只发生在权重初始化不稳定的最初几个 batch，scale 从 8192 降到 4096 或 2048 后就稳定下来，从未观察到一路降到 1 的情况。
为什么不是 BF16？
BF16 用 8 位指数位、7 位尾数位，动态范围与 FP32 相同，通常不需要 loss scaling，训练更稳定。但 Tech-Renaissance 当前版本只支持 FP16，原因非常具体：cuDNN 的 Conv+GenStats 算子目前只支持 FP16。
这个融合算子是下一篇要讲的 CBR（Conv+BN+ReLU）融合的基础。它可以在卷积计算的同时，直接生成 BN 需要的通道和与通道平方和（sum 和 sq_sum），避免 BN 执行时重新读取一次卷积输出这个大张量。对带宽敏感的 CNN 训练来说，省掉的这一次大张量读写非常可观。BF16 目前做不到这一点，所以框架做出了明确的工程取舍：先拿下 FP16 的融合红利，而不是为了 BF16 的稳定性放弃 CBR 融合。
三、AMP 显存布局：W、A、G、E 四条线
在 Tech-Renaissance 里，AMP 不是运行时临时决定精度，而是在编译期就写进了显存分区。打开 include/renaissance/core/types.h，可以看到 68 个命名语义 Region 中专门有四条线服务于 AMP：
- W 系列（005–012）：FP32 主权重，包括 BN 参数、FC 参数、卷积参数；
- E 系列（013–018）：FP32 EMA 权重；019–021 是对应的 FP16 EMA 副本；
- A 系列（022–024）：FP16 AMP 计算权重；
- G 系列（025–030）：FP32 梯度；032–034 是对应的 FP16 Conv/CBR 梯度。
// include/renaissance/core/types.h 中 AMP 相关 Region 节选 A_FC_WEIGHT, // 022 A_FIRST_CONV, // 023 A_DEEP_CONV, // 024 G_BN_BIAS, // 025 G_BN_WEIGHT, // 026 G_FC_BIAS, // 027 G_FC_WEIGHT, // 028 G_FIRST_CONV, // 029 G_DEEP_CONV, // 030 G_FC_WEIGHT_FP16, // 032 G_FIRST_CONV_FP16, // 033 G_DEEP_CONV_FP16, // 034
这种分区设计的关键在于：同类型张量被连续摆放。所有深层卷积的 FP32 权重连续放在 W_DEEP_CONV，所有 FP16 计算权重连续放在 A_DEEP_CONV，所有 FP16 Conv 梯度连续放在 G_DEEP_CONV_FP16。于是框架可以用一次 RANGE_CAST_FP32_TO_FP16 就把整个 W_DEEP_CONV 区域 cast 到 A_DEEP_CONV，而不用逐层启动 kernel。
MemoryPlan 在分配时也显式检查了 AMP 与数据类型的一致性。以特征图分配为例：
// src/graph/memory_plan.cpp
DTensor MemoryPlan::alloc_feature(const Shape& shape, DType dtype) {
bool amp = GlobalRegistry::instance().using_amp();
if (dtype == DType::FP16) {
TR_CHECK(amp, ValueError, "FP16 feature requires amp_enabled");
return alloc_impl(shape, DType::FP16, Region::F_FEATURE_FP16);
}
TR_CHECK(dtype == DType::FP32, ValueError,
"Feature dtype must be FP32 or FP16, got ...");
TR_CHECK(!amp, ValueError, "FP32 feature requires !amp_enabled");
return alloc_impl(shape, DType::FP32, Region::F_FEATURE_FP32);
}
amp：全局 AMP 开关，由GlobalRegistry::using_amp()提供；dtype：请求分配的特征图数据类型；- 返回值
DTensor：纯描述符，只记录 shape/offset/stride/Region，不持有实际显存。
类似的断言也存在于 alloc_grad_slot 和 alloc_scalar 中。FP16 张量要求 AMP 开启，FP32 特征图要求 AMP 关闭。这种”编译期决定一切”的偏执，与静态图、CUDA Graph、MemoryPlan 等设计理念一脉相承。
在 MemoryPlan::is_condition_enabled() 中，A 系列、G 系列的 FP16 梯度区、E 系列的 FP16 副本、F 系列的 FP16 特征图，以及 FP16 标量区，都被显式地门控在 GlobalRegistry::using_amp() 上；而 FP32 特征图区只在 AMP 关闭时才启用。这意味着 AMP 开关一旦确定，整套显存布局就是静态且可验证的。
为什么一次 RangeOp 能转换全模型？
这依赖于 DTensor::compute_slot_bytes 的设计。对权重/梯度区（W、A、G、E 系列），C 通道对齐因子为 1，因此同一层同一 region 的 FP32 槽位大小恰好是 FP16 槽位大小的两倍：
// 简化自 include/renaissance/tensor/distributed_tensor.h
if (dtype == DType::FP16) {
return align_up_256(elems * 2 + 16); // 每个元素 2 字节
} else if (dtype == DType::FP32) {
return 2 * align_up_256(elems * 2 + 16); // FP32 槽位 = 2 × FP16 槽位
}
elems：张量元素总数；align_up_256：向上对齐到 256 字节；+16：末尾保留 16 字节，满足 XNNPACK 等后端的对齐要求。
由于 FP32 与 FP16 槽位大小呈固定 2:1 关系，且同一层在 W/A 或 G/G_FP16 中的顺序完全一致，编译器可以把 W_FC_WEIGHT..W_DEEP_CONV 到 A_FC_WEIGHT..A_DEEP_CONV 的转换表达为一个 RANGE_CAST_FP32_TO_FP16 节点，一个 kernel 调用完成全模型权重转换。
四、训练一步的 AMP 流水线
在 src/backend/graph_executor.cpp 的 run_train_step() 中，一次常规训练迭代的执行顺序大致如下（省略部分 A/B 双缓冲细节）：
// src/backend/graph_executor.cpp
launch(GraphId::TRANSFER_A); // H2D 异步传输
launch(GraphId::ZERO_GRAD); // 梯度区清零
launch(GraphId::FIRST_LAYER_FWD_A); // 首层前向
launch_dual(next_xfer, GraphId::DEEP_FWD_BWD); // 下一 batch 传输 + 深层正反向
sync_all();
launch(GraphId::FIRST_LAYER_BWD_A); // 首层反向
sync_all();
launch(GraphId::CAST_DEEP_GRAD_FP16_TO_FP32); // 深层卷积梯度 FP16 → FP32
launch(GraphId::DEEP_COMM); // 深层梯度 AllReduce
sync_all();
launch(GraphId::CAST_FIRST_GRAD_FP16_TO_FP32); // 首层卷积梯度 FP16 → FP32
launch(GraphId::FIRST_COMM); // 首层梯度 AllReduce
launch(GraphId::NAN_CHECK_AND_GRAD_SCALING); // NaN 检查 + 缩放调整
launch(GraphId::STATS_COMM); // BN 统计量同步
launch(GraphId::UPDATE_STATS); // BN next → prev 复制
bool has_nan = check_nan_flag();
if (!has_nan) {
launch(GraphId::OPTIMIZER); // FP32 主权重更新
launch(GraphId::EMA_UPDATE); // EMA 更新（含 FP32→FP16 cast）
} else {
on_nan_detected();
}
launch_dual(GraphId::CAST_MAIN_FP32_TO_FP16,
GraphId::CAST_EMA_FP32_TO_FP16); // 主权重 / EMA FP32 → FP16
sync_all();
launch：在指定 CUDA Stream 上启动一个 CUDA Graph；launch_dual：同时在两个 stream 上启动两个图，实现并发；sync_all：等待所有流完成当前阶段。
这里有两个实现细节值得补充。第一，所有 FP32↔FP16 的 cast、优化器更新、EMA 更新、NAN_CHECK_AND_GRAD_SCALING 都运行在 StreamKind::UPDATE 上，与主计算流 COMP_1 互不阻塞，便于和计算阶段重叠。第二，EMA_UPDATE 图中已经包含了 EMA FP32→FP16 的 RANGE_CAST_FP32_TO_FP16 节点；GraphExecutor 末尾同时启动的 CAST_EMA_FP32_TO_FP16 图当前为空操作，保留它只是为了保持执行器结构的统一性。
这个流程里有几个值得单独展开的细节。
1. FC 与 Conv 的梯度路径不同
FC 层的反向走 cuBLAS，可以直接要求输出 FP32。在 src/backend/ops/dtensor/fc_op.cpp 中，cublasGemmEx 的输出数据类型指定为 CUDA_R_32F：
// 简化自 fc_op.cpp
float* dw = ...;
cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_T, ...,
x, CUDA_R_16F, x_ns, // 输入 X：FP16
dy, CUDA_R_16F, dy_ns, // 输入 dY：FP16
dw, CUDA_R_32F, dw_ns, // 输出 dW：FP32
CUBLAS_COMPUTE_32F,
CUBLAS_GEMM_DEFAULT_TENSOR_OP);
因此 FC 梯度天然就是 FP32，直接写入 G_FC_WEIGHT / G_FC_BIAS，不需要后续转换。
Conv 和 CBR 的情况更复杂。cuDNN 的卷积反向 filter（wgrad）在 FP16 模式下通常输出 FP16。CBR 反向图在 cuDNN Frontend Graph 中构建时，dY、W、dX 都是 FP16，wgrad 输出自然也是 FP16，先被写到 G_DEEP_CONV_FP16 或 G_FIRST_CONV_FP16，然后由 RANGE_CAST_FP16_TO_FP32 一次性批量转成 FP32：
// src/graph/compiler.cpp
if (amp_on && memory_plan.is_region_populated(Region::G_DEEP_CONV_FP16)) {
MemRange in_deep = memory_plan.region_range(Region::G_DEEP_CONV_FP16);
MemRange out_deep = memory_plan.region_range(Region::G_DEEP_CONV);
train_cg.append_range(GraphId::CAST_DEEP_GRAD_FP16_TO_FP32,
RangeOp::RANGE_CAST_FP16_TO_FP32, {in_deep}, {out_deep});
}
为什么不直接在 cuDNN 里要求 wgrad 输出 FP32？因为 cuDNN Frontend Graph 的某些融合路径对输出数据类型有约束，强制指定 FP32 输出可能导致找不到合法执行计划，或者性能反而下降。接受 FP16 中间结果，再用一次高效的 RangeOp 批量转换，是更务实的选择。
2. 损失缩放到底发生在哪里？
损失缩放不是独立的外挂，而是分散在 SoftmaxCE 反向和优化器更新中：
- SoftmaxCE BWD 输出的梯度已经带有
scaling / batch的缩放； - Optimizer 内部再把梯度乘以
1 / scaling，完成还原。
net effect 是梯度按 batch 平均，与无缩放时一致。把 1/scaling 放到优化器里做，可以省下一次对梯度区的单独显存遍历：
// optimizer_op.cu 中 SGD 内核节选
float _inv_scaling = (scaling && *scaling != 0.0f) ? (1.0f / *scaling) : 1.0f;
for (...) {
float g_i = g[i] * _inv_scaling;
w[i] = w_i * (1.0f - _lr * _wd) - _lr * g_i;
}
g[i]：第i个 FP32 梯度元素；scaling：当前损失缩放因子；_inv_scaling：1 / scaling；w[i]：第i个 FP32 主权重。
如果 has_nan 标志非零，优化器内核第一行就直接 return，不会修改任何主权重。这保证了 NaN 梯度被隔离在单次迭代内。
3. NaN 检测与梯度裁剪
RANGE_CHECK_NAN 扫描整个 FP32 梯度区 G_BN_BIAS..G_DEEP_CONV。由于 MemoryPlan 的语义分区，所有参与训练的梯度在内存中连续排布，一个 RangeOp 就能覆盖全部。
// src/graph/compiler.cpp
node.range_op = RangeOp::RANGE_CHECK_NAN;
node.input_ranges.push_back(
memory_plan.region_range(Region::G_BN_BIAS, Region::G_DEEP_CONV));
node.output_ids.push_back(nan_flag_id);
check_op.cu 实现了两个版本：
// 仅检测
__global__ void check_nan_kernel(
volatile int32_t* __restrict__ has_nan_flag,
const float* __restrict__ data, size_t n)
{
// 任一元素为 NaN 或 Inf，设置 has_nan_flag
}
// 检测 + 裁剪
__global__ void check_nan_and_clip_kernel(
volatile int32_t* __restrict__ has_nan_flag,
float* __restrict__ data, size_t n, float clip_val)
{
// NaN：设置标志，不修改值；
// 其他值：clamp 到 [-clip_val, +clip_val]，Inf 也会被 clamp。
}
- 在不开启梯度裁剪时，
isnan和isinf都会设置has_nan； - 在开启梯度裁剪时，只有
NaN设置标志，Inf和超出边界的值会被clamp到边界，因为 Inf 通常意味着梯度过大而非计算错误，裁剪后仍可更新。
需要说明的是，上述把 Inf 当作可裁剪超大值的处理是 Tech-Renaissance 当前实现的选择；在更保守的训练配置中，Inf 常与 NaN 一样被视为数值溢出，统一触发损失缩放下降并跳过该步更新，以避免可能的方向失真。
如果用户配置了 GradClipParams.max_abs > 0，编译器会把阈值嵌入节点参数，同一个 kernel 同时完成检测与裁剪，避免第二次读取梯度区。
五、CBR 融合：为什么只支持 AMP？
第 19 篇会专门讲 CBR（Conv + BatchNorm + ReLU）融合算子，但这里必须先提一句：Tech-Renaissance 的 CBR 融合目前只实现了 AMP 版本，没有 FP32 版本。
原因出在 cuDNN Frontend Graph 的 GenStats 节点上。CBR 前向要把卷积输出、BN 统计量生成、ReLU 激活尽量融合在一起。其中 BN 需要的通道和与通道平方和，可以通过 genstats 节点在卷积输出上直接生成：
// cbr_op.cpp 中 CBR 前向图构建节选
auto genstats_opts = Genstats_attributes()
.set_name("genstats")
.set_compute_data_type(fe::DataType_t::FLOAT); // 内部计算精度仍为 FP32
auto genstats_outputs = graph->genstats(conv_out, genstats_opts);
auto sum = genstats_outputs[0];
auto sq_sum = genstats_outputs[1];
conv_out：卷积输出，数据类型为HALF（FP16）；compute_data_type：内部累加精度为FLOAT（FP32），保证数值稳定性；sum/sq_sum：BN 所需的通道统计量。
genstats 当前只支持 FP16 输入。这意味着如果想用 CBR 融合，卷积输入、输出、权重都必须是 FP16。FP32 模式下即使实现类似融合，也享受不到 Conv+GenStats 这条优化路径，性能反而未必比分立算子更好。因此框架做了一个务实的选择：CBR 融合算子只服务于 AMP 训练。
CBR 融合带来的收益非常直接：省去了一次完整的 conv_output 读写。以 ResNet-50 典型层为例，FP16 的 conv_output 大约几 MB，训练中有数十个这样的层，每个 batch 都重复——节省的显存带宽累加起来相当可观。
六、与 PyTorch AMP 的本质差异
到这里，可以清晰看出 Tech-Renaissance 的 AMP 和 PyTorch AMP 走的是两条不同的路。
PyTorch AMP 是算子级、动态、自动的：在 autocast 上下文里写模型代码，PyTorch 运行时自动决定每个算子的精度，并用 GradScaler 动态调整 loss scale。好处是对用户透明、几乎零侵入；代价是运行时要做大量判断、状态维护和 CPU-GPU 同步。
Tech-Renaissance AMP 是图级、静态、显式的：精度转换、NaN 检查、梯度缩放、AllReduce 都是编译期生成的独立子图，运行时按固定顺序 launch。用户通过 GlobalRegistry::amp(true) 打开 AMP 后，后续一切由框架自动处理，不需要手动写 cast 或 scale；但所有选择都是预先决定的，无法在中途动态插入自定义精度策略。
这两条路线没有绝对的优劣，只是适合不同的设计目标。PyTorch AMP 适合研究阶段的灵活性；Tech-Renaissance 的 AMP 适合把训练循环 captured 成 CUDA Graph、把显存布局静态化、把每一点 CPU 开销都榨干的场景。如果说 PyTorch 的 AMP 是”优雅的自动挡”，那 Tech-Renaissance 的 AMP 更像是”为赛道调校的手动挡”——你需要在出发前挂好挡，但跑起来后每个挡位都咬合得严丝合缝。
结语
AMP 自动混合精度训练，本质上是把 FP16 的速度、FP32 的精度和一个聪明的梯度缩放策略缝合在一起。Tech-Renaissance 的实现思路可以概括为：
- 用显存分区把 FP32 主权重、FP16 计算权重、FP16 梯度、EMA 权重物理隔离；
- 用 RangeOp 把精度转换、NaN 检查、梯度缩放、AllReduce 表达为批量图阶段；
- 用固定损失缩放换取确定性、低开销和 CUDA Graph 友好性；
- 用 CBR 融合把 AMP 的性能红利发挥到极致。
它不像 PyTorch AMP 那样对用户完全透明，但它在静态图、静态显存、静态多流的框架语境里，找到了一条自洽且高效的实现路径。下一篇，我们就来聊聊 AMP 模式下最具代表性的算子优化——CBR 融合：如何把卷积、批归一化、ReLU 三个算子揉成一次 cuDNN Graph 执行。
