(21) 融合优化器：SGD、AdamW与LARS的整区批量更新

技术觉醒

——“一个人用AI如何写出比PyTorch更快的自研深度学习框架”系列文章之二十一

训练神经网络，本质上就是无数次地“算梯度、改权重”。前向、反向、通信这些阶段固然热闹，但真正决定参数怎么走的，是优化器这一步。很多人对优化器的印象还停留在 optimizer.step() 这一行代码上，觉得它不过是在梯度后面乘个学习率、再减到权重上而已。但如果你真的去 profile 一次训练，就会发现优化器 step 常常是容易被低估的性能热点：原因不在计算量，而在调度开销。

现代网络动辄几十、上百个参数张量。在主流框架里，每个张量都要单独启动一个 CUDA kernel，从 GPU 上读权重、读梯度、写权重、写动量，几十次 kernel launch 排队等 CPU 派发。PyTorch 的 torch.optim.AdamW 虽然提供了 fused=True，能把单个参数张量内部的读写合并到一个 kernel 里，且借助 multi-tensor apply 把许多参数张量打包进少量 kernel，但这些张量在显存中并不连续，每次 launch 都要传入逐张量的指针元数据，也无法像整区连续内存那样一次线性扫过；参数组间的循环与 Python 侧的调度开销依然存在。对于一个拥有大量小参数张量的网络（比如 ResNet-50 里大量的 BN gamma/beta 和卷积权重），这些开销加起来相当可观。

Tech-Renaissance 没有走这条路。得益于 MemoryPlan 的显存分区设计，优化器更新被表达成对连续显存区域的批量操作：不管模型有多少层，权重更新往往只需一两个 kernel，就能把全部可训练参数遍历一遍。这篇文章就来聊聊 Tech-Renaissance 的融合优化器设计——SGD、AdamW、LARS 三种算法，如何被映射到同一张“整区更新”的图里。

一、优化器到底在算什么？

在深入实现之前，先快速回顾一下我们要支持的三种优化器。下面的伪代码都用 w 表示权重、g 表示梯度、lr 表示学习率、wd 表示 weight decay。

SGD（随机梯度下降） 是最基础的参数更新策略，每一步沿梯度反方向走一个步长：

// SGD: w = w - lr * g
for (size_t i = 0; i < n; ++i) {
    w[i] = w[i] - lr * g[i];
}

引入动量之后，SGD 不再只看当前梯度，而是维护一个一阶动量缓冲区 m，把历史梯度方向以系数 beta 累积进来：

// Momentum:
//   m = beta * m + g
//   w = w - lr * m
for (size_t i = 0; i < n; ++i) {
    m[i] = beta * m[i] + g[i];
    w[i] = w[i] - lr * m[i];
}

动量让更新更平滑，有助于穿越平缓区域、抑制震荡。Nesterov 动量则是把“先看未来一步的动量方向”再求梯度。在我们的 kernel 里，它等价于下面这种更利于融合的形式：

// Nesterov:
//   m_new = beta * m + g
//   w     = w - lr * (m_new * beta + g)
//   m     = m_new
for (size_t i = 0; i < n; ++i) {
    float g_i  = g[i];
    float m_new = beta * m[i] + g_i;
    w[i] = w[i] - lr * (m_new * beta + g_i);
    m[i] = m_new;
}

这段逻辑在代数上完全等价于标准的 Nesterov 动量，只是我们把它展开并重新结合。这样做的妙处在于：它允许在一个 kernel 的单次循环内，利用寄存器（m_new）同时完成前后步动量的推导与权重的写入，极大降低了显存读写的依赖。

Adam 在动量基础上又加了二阶矩估计 v，它记录梯度平方的指数移动平均，从而对每个参数单独缩放学习率：

// Adam:
//   m = beta1 * m + (1 - beta1) * g
//   v = beta2 * v + (1 - beta2) * g * g
//   w = w - lr * m_hat / (sqrt(v_hat) + eps)
// 其中 m_hat = m / (1 - beta1^t), v_hat = v / (1 - beta2^t)
for (size_t i = 0; i < n; ++i) {
    float g_i = g[i];
    m[i] = beta1 * m[i] + (1.0f - beta1) * g_i;
    v[i] = beta2 * v[i] + (1.0f - beta2) * g_i * g_i;
    float m_hat = m[i] * bias_corr1;   // bias_corr1 = 1 / (1 - beta1^t)
    float v_hat = v[i] * bias_corr2;   // bias_corr2 = 1 / (1 - beta2^t)
    w[i] = w[i] - lr * m_hat / (sqrtf(v_hat) + eps);
}

由于 m 和 v 在训练初期都偏向 0，Adam 还需要做偏差修正（bias correction），把 m 和 v 分别除以 1 - beta1^t 和 1 - beta2^t，其中 t 是训练步数。

AdamW 与 Adam 的数学区别很小，却影响深远：它把权重衰减从梯度里解耦出来，直接作用在参数上，而不是加在梯度里再被二阶矩放大。AdamW 的更新可以写成：

// AdamW:
//   w = w * (1 - lr * wd)
//   m = beta1 * m + (1 - beta1) * g
//   v = beta2 * v + (1 - beta2) * g * g
//   w = w - lr * m_hat / (sqrt(v_hat) + eps)
for (size_t i = 0; i < n; ++i) {
    float g_i = g[i];
    w[i] = w[i] * (1.0f - lr * wd);
    m[i] = beta1 * m[i] + (1.0f - beta1) * g_i;
    v[i] = beta2 * v[i] + (1.0f - beta2) * g_i * g_i;
    float m_hat = m[i] * bias_corr1;
    float v_hat = v[i] * bias_corr2;
    w[i] = w[i] - lr * m_hat / (sqrtf(v_hat) + eps);
}

这也是为什么 AdamW 在大模型训练里成了事实标准。

LARS（Layer-wise Adaptive Rate Scaling） 的思路则完全不同。它不是给每个参数单独的学习率，而是给每一层一个 trust ratio。对第 l 层的权重 w_l 和梯度 g_l，先计算：

// LARS trust ratio:
//   eta_l = trust_coefficient * ||w_l|| / (||g_l|| + wd * ||w_l|| + eps)
// 然后用 eta_l * lr 作为该层的有效学习率

LARS 的初衷是解决超大 batch 训练时的优化稳定性问题：不同层的权重和梯度量级差异巨大，统一的全局学习率往往顾此失彼。LARS 的逐层缩放让大 batch 训练变得可行，也是 ResNet 在 batch size 达到数千乃至上万时仍能收敛的关键技术之一。

PyTorch 的 foreach/fused optimizer 已经能把多个张量的更新在一定程度上合并，显著减少逐张量 Python 循环和 kernel launch 开销。但它的组织单位仍然是参数列表和参数组，无法利用 Tech-Renaissance 这种”同语义参数在显存中连续排布”的 Region 布局。因此，在本框架里，优化器可以进一步被表达成整区 RangeOp，用极少数 kernel 覆盖整组参数。

二、Tech-Renaissance 的优化器：Region 即批量

在 Tech-Renaissance 中，优化器不是对着一堆分散的张量做循环，而是直接对显存 Region 做批量操作。这是 MemoryPlan 分区设计的直接产物。

回到 include/renaissance/core/types.h 里的 Region 枚举，与优化器直接相关的有以下几条线：

  • W 系列：W_BN_BIAS、W_BN_WEIGHT、W_FC_BIAS、W_FC_WEIGHT、W_FIRST_CONV、W_DEEP_CONV，存放 FP32 主权重；
  • G 系列：对应同名的梯度区；
  • M 系列：一阶动量缓冲区；
  • V 系列：二阶动量缓冲区（Adam/AdamW 专用）；
  • N 系列：LARS 逐层 trust ratio 的存储区；
  • S 系列：S_SCALAR_FP32、S_SCALAR_INT32，存放 lr、wd、beta、scaling、has_nan 等标量。

由于同一语义的张量被连续摆放，优化器可以用一次 region_range(start, end) 就覆盖全部同类参数。例如，所有 FC 与卷积权重连续放在 W_FC_WEIGHT 到 W_DEEP_CONV 之间，它们的梯度也连续放在 G_FC_WEIGHT 到 G_DEEP_CONV 之间，动量缓冲区连续放在 M_FC_WEIGHT 到 M_DEEP_CONV 之间，二阶矩缓冲区连续放在 V_FC_WEIGHT 到 V_DEEP_CONV 之间。这种“同形同区”的布局不是偶然的，而是编译器和 MemoryPlan 共同保证的：当某一层被分配到 W_FC_WEIGHT 时，它的梯度、动量、二阶矩都会按相同顺序落在对应区域。

用户在 include/renaissance/algo/optimizer.h 中通过纯配置类选择优化器：

// SGD with Momentum
task.optimizer(SGD().momentum(0.9f).weight_decay(5e-5f).nesterov(false));

// AdamW
task.optimizer(AdamW().beta1(0.9f).beta2(0.999f).eps(1e-8f).weight_decay(0.01f));

// LARS
task.optimizer(LARS().momentum(0.9f).weight_decay(5e-5f)
                       .trust_coefficient(0.001f));

这些配置在 task.compile() 时被编译器读取。plan_config_from_optimizer() 函数将 OptimizerKind 翻译为 PlanConfig 的标志位，决定 MemoryPlan 是否分配 M、V、N 系列区域：

inline PlanConfig plan_config_from_optimizer(OptimizerKind kind, bool has_ema = false) {
    PlanConfig cfg;
    cfg.has_ema = has_ema;
    switch (kind) {
        case OptimizerKind::SGD:
            cfg.use_momentum = false;
            cfg.use_adam     = false;
            cfg.use_lars     = false;
            break;
        case OptimizerKind::SGD_MOMENTUM:
        case OptimizerKind::SGD_NESTEROV:
            cfg.use_momentum = true;
            cfg.use_adam     = false;
            cfg.use_lars     = false;
            break;
        case OptimizerKind::LARS:
        case OptimizerKind::LARS_NESTEROV:
            cfg.use_momentum = true;
            cfg.use_adam     = false;
            cfg.use_lars     = true;
            break;
        case OptimizerKind::ADAM:
        case OptimizerKind::ADAMW:
            cfg.use_momentum = true;
            cfg.use_adam     = true;
            cfg.use_lars     = false;
            break;
        default:
            break;
    }
    return cfg;
}

选 SGD 就不分配 M 系列，选 AdamW 就多分配 V 系列，没有浪费。这种设计下，优化器的选择直接影响显存布局。

于是 SGD/Momentum/Adam/AdamW 的 weight 更新可以被表达为一个 RANGE_UPDATE_WEIGHT_* 节点。以 AdamW 为例，src/graph/compiler.cpp 中的构建逻辑大致如下：

// 1. step += 1
GraphNode inc_node;
inc_node.kind = GraphNode::Kind::COMPUTE;
inc_node.compute_op = ComputeOp::SCALAR_INCREMENT;
inc_node.input_ids  = {scalar_ids.step};
inc_node.output_ids = {scalar_ids.step};
train_cg.append(GraphId::OPTIMIZER, inc_node);

// 2. bc1 = 1/(1-beta1^step), bc2 = 1/(1-beta2^step)
GraphNode bc_node;
bc_node.kind = GraphNode::Kind::COMPUTE;
bc_node.compute_op = ComputeOp::ADAM_BIAS_CORRECTION;
bc_node.input_ids  = {scalar_ids.step, scalar_ids.beta, scalar_ids.beta2};
bc_node.output_ids = {scalar_ids.bias_corr1, scalar_ids.bias_corr2};
train_cg.append(GraphId::OPTIMIZER, bc_node);

// 3. 整区 weight 更新
MemRange w_range = memory_plan.region_range(
    Region::W_FC_WEIGHT, Region::W_DEEP_CONV);
MemRange g_range = memory_plan.region_range(
    Region::G_FC_WEIGHT, Region::G_DEEP_CONV);
MemRange m_range = memory_plan.region_range(
    Region::M_FC_WEIGHT, Region::M_DEEP_CONV);
MemRange v_range = memory_plan.region_range(
    Region::V_FC_WEIGHT, Region::V_DEEP_CONV);

GraphNode node;
node.kind = GraphNode::Kind::RANGE;
node.range_op = RangeOp::RANGE_UPDATE_WEIGHT_ADAMW;
node.input_ranges.push_back(w_range);  // W
node.input_ranges.push_back(g_range);  // G
node.input_ranges.push_back(m_range);  // M
node.input_ranges.push_back(v_range);  // V
node.output_ranges.push_back(w_range);
node.output_ranges.push_back(m_range);
node.output_ranges.push_back(v_range);
train_cg.append(GraphId::OPTIMIZER, node);

这个节点最后被放进 GraphId::OPTIMIZER 图里，运行在 StreamKind::UPDATE 流上。运行时只需要 launch 一次图，就能完成整个模型所有权重的更新。更重要的是，这些 Region 的起始地址和大小在编译期就已经确定，w_range 和 g_range 不再是张量描述符列表，而是 (offset, size) 的内存区间。kernel 启动时不需要遍历参数列表、不需要动态计算偏移，只需要两个指针和元素个数。这为 CUDA Graph 全捕获创造了条件：图的拓扑、内存地址、kernel grid 全部静态可知，运行时只是 cudaGraphLaunch 的一次重放。

三、SGD、Momentum、AdamW 的 CUDA 内核

具体的内核实现位于 src/backend/ops/range/optimizer_op.cu。以 AdamW 为例，其核心 kernel 如下：

__global__ void update_adamw_kernel(
    float* __restrict__ w, const float* __restrict__ g,
    float* __restrict__ m, float* __restrict__ v, size_t n,
    const float* __restrict__ lr, const float* __restrict__ wd,
    const float* __restrict__ b1, const float* __restrict__ b2,
    const float* __restrict__ eps,
    const int32_t* __restrict__ has_nan,
    const float* __restrict__ scaling,
    const float* __restrict__ bias_corr1,
    const float* __restrict__ bias_corr2)
{
    if (*has_nan != 0) return;
    float _lr = *lr;
    float _wd = wd ? *wd : 0.0f;
    float _b1 = *b1;
    float _b2 = *b2;
    float _eps = *eps;
    float _inv_scaling = (scaling && *scaling != 0.0f) ? (1.0f / *scaling) : 1.0f;
    float _bc1 = bias_corr1 ? *bias_corr1 : 1.0f;
    float _bc2 = bias_corr2 ? *bias_corr2 : 1.0f;

    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n; i += gridDim.x * blockDim.x) {
        float g_i = g[i] * _inv_scaling;
        w[i] = w[i] * (1.0f - _lr * _wd);
        m[i] = m[i] * _b1 + (1.0f - _b1) * g_i;
        v[i] = v[i] * _b2 + (1.0f - _b2) * g_i * g_i;
        float m_hat = m[i] * _bc1;
        float v_hat = v[i] * _bc2;
        w[i] = w[i] - _lr * m_hat / (sqrtf(v_hat) + _eps);
    }
}

这个 kernel 里融合了四件事：

  1. AMP 反缩放：g_i = g[i] * _inv_scaling，把 AMP 阶段放大过的梯度还原；
  2. 权重衰减：w[i] = w[i] * (1.0f - _lr * _wd)，这是 AdamW 的解耦 weight decay；
  3. 一阶/二阶矩更新；
  4. 偏差修正后的参数更新。

注意所有超参数都不是通过 kernel 参数传入，而是通过设备指针传入。这一点非常关键：如果 lr、wd 等标量以 kernel 参数形式传入，那么每次学习率变化都需要重新编译 CUDA Graph；而以设备指针传入时，kernel 读到的是显存中实时变化的标量，图结构本身无需改变。

const float* lr  = scalar_ptr<0>(mp, node.input_ids.data(), ctx.rank_for_context());
const float* wd  = scalar_ptr<1>(mp, node.input_ids.data(), ctx.rank_for_context());
// ...

这些标量全部存放在 S_SCALAR_FP32 或 S_SCALAR_INT32 区域。运行时只要通过一次 H2D 把小标量写进显存，kernel 读取的就是最新值，CUDA Graph 无需重新捕获。这对于学习率调度至关重要：每个 batch 的学习率可以变化，但训练图保持不变。

has_nan 也以设备指针传入。一旦 RANGE_CHECK_NAN 检测到梯度异常，所有优化器 kernel 会在最开始的 if (*has_nan != 0) return; 处直接退出，等价于跳过本轮更新。这个判断在 GPU 内部完成，不需要 CPU 介入每个参数张量。

SGD 与 Momentum 的 kernel 更简洁，但遵循同样的模式：

// SGD: w = w * (1 - lr * wd) - lr * g
__global__ void update_sgd_kernel(...)
{
    if (*has_nan != 0) return;
    float _lr = *lr;
    float _wd = wd ? *wd : 0.0f;
    float _inv_scaling = (scaling && *scaling != 0.0f) ? (1.0f / *scaling) : 1.0f;
    for (size_t i = ...; i < n; i += ...) {
        float w_i = w[i];
        float g_i = g[i] * _inv_scaling;
        w[i] = w_i * (1.0f - _lr * _wd) - _lr * g_i;
    }
}

// Momentum:
//   m = beta * m + g
//   w = w * (1 - lr * wd) - lr * m
__global__ void update_momentum_kernel(...)
{
    if (*has_nan != 0) return;
    float _beta = *beta;
    for (size_t i = ...; i < n; i += ...) {
        float g_i = g[i] * _inv_scaling;
        m[i] = m[i] * _beta + g_i;
        w[i] = w[i] * (1.0f - _lr * _wd) - _lr * m[i];
    }
}

所有 kernel 都使用 __launch_bounds__(128, 2) 编译提示。它告诉编译器：每块 128 线程、期望每个 SM 至少驻留 2 个 block（minBlocksPerMultiprocessor），从而引导编译器控制寄存器用量以保证占用率，在大多数 GPU 架构上是一个稳健的平衡点。

CPU fallback 版本位于 optimizer_op.cpp 下半部分，用于无 GPU 环境或正确性验证。CUDA 与 CPU 的数学公式保持一致，这是 Tech-Renaissance 保证数值可复现的基础。

四、Bias-like 参数的分组哲学

在 compiler.cpp 中构建优化器图时，Tech-Renaissance 把所有可训练参数分成了两组：

  • Weight 组：W_FC_WEIGHT、W_FIRST_CONV、W_DEEP_CONV，需要 weight decay；
  • Bias-like 组：W_BN_BIAS、W_BN_WEIGHT、W_FC_BIAS，不需要 weight decay。

这个分组不是随意定的，而是基于深度学习训练中的一个常见约定：BatchNorm 的 scale/bias 和全连接层的 bias 通常不参与 L2 正则化。如果强行把 BN 的 gamma 也做 weight decay，在很多标准训练配置里会被认为破坏了归一化层的尺度不变性。

代码里通过 static_assert 强制要求这两组在 Region 枚举中连续排列：

static_assert(
    static_cast<int>(Region::W_BN_BIAS) + 1 == static_cast<int>(Region::W_BN_WEIGHT) &&
    static_cast<int>(Region::W_BN_WEIGHT) + 1 == static_cast<int>(Region::W_FC_BIAS),
    "W_BN_BIAS, W_BN_WEIGHT, W_FC_BIAS must be consecutive in Region enum"
);
static_assert(
    static_cast<int>(Region::G_BN_BIAS) + 1 == static_cast<int>(Region::G_BN_WEIGHT) &&
    static_cast<int>(Region::G_BN_WEIGHT) + 1 == static_cast<int>(Region::G_FC_BIAS),
    "G_BN_BIAS, G_BN_WEIGHT, G_FC_BIAS must be consecutive in Region enum"
);
static_assert(
    static_cast<int>(Region::M_BN_BIAS) + 1 == static_cast<int>(Region::M_BN_WEIGHT) &&
    static_cast<int>(Region::M_BN_WEIGHT) + 1 == static_cast<int>(Region::M_FC_BIAS),
    "M_BN_BIAS, M_BN_WEIGHT, M_FC_BIAS must be consecutive in Region enum"
);
static_assert(
    static_cast<int>(Region::V_BN_BIAS) + 1 == static_cast<int>(Region::V_BN_WEIGHT) &&
    static_cast<int>(Region::V_BN_WEIGHT) + 1 == static_cast<int>(Region::V_FC_BIAS),
    "V_BN_BIAS, V_BN_WEIGHT, V_FC_BIAS must be consecutive in Region enum"
);

于是 Bias-like 更新可以用一个 region_range(W_BN_BIAS, W_FC_BIAS) 覆盖全部，对应 RANGE_UPDATE_BIAS_* 系列 RangeOp。Bias 路径的 kernel 不传 wd 指针，内部 _wd = 0.0f，因此无论 SGD、Momentum、Nesterov 还是 Adam/AdamW，Bias-like 参数都自然得到正确的行为。

Adam 与 AdamW 的 Bias 路径共用同一个 RANGE_UPDATE_BIAS_ADAM，因为当 wd = nullptr 时，update_adam_kernel 和 update_adamw_kernel 数学等价：

// Bias 路径 wd 固定传 nullptr，_wd = 0.0f
float _wd = wd ? *wd : 0.0f;
// Adam 的 g_i = g * inv_scaling + _wd * w = g * inv_scaling
// AdamW 的 g_i = g * inv_scaling
// 当 _wd == 0 时两者完全相同

这种“用分组代替逐层判断”的思路，大幅简化了编译期图构建，也减少了运行时的 kernel 数量。

五、Adam 的偏差修正：两个标量 kernel

Adam/AdamW 还需要在每步更新前计算偏差修正系数 bc1 和 bc2。Tech-Renaissance 没有把这件事放到主更新 kernel 里，而是单独放在 GraphId::OPTIMIZER 图的最前面，用两个极小的 kernel 完成：

// SCALAR_INCREMENT: step += 1
__global__ void scalar_increment_kernel(int32_t* value) {
    if (threadIdx.x == 0 && blockIdx.x == 0) {
        *value += 1;
    }
}

需要说明的是，这个 step += 1 的 kernel 与后续所有参数更新 kernel 同处 GraphId::OPTIMIZER 图；当 has_nan 标志为真时，整段 OPTIMIZER 图不会被 launch（见 src/backend/graph_executor.cpp），因此 step 不会在 NaN/Inf 步错误推进。

// ADAM_BIAS_CORRECTION:
//   bc1 = 1 / (1 - beta1^step)
//   bc2 = 1 / (1 - beta2^step)
__global__ void adam_bias_correction_kernel(
    const int32_t* __restrict__ step,
    const float*   __restrict__ beta1,
    const float*   __restrict__ beta2,
    float*         __restrict__ bias_corr1,
    float*         __restrict__ bias_corr2)
{
    if (threadIdx.x != 0 || blockIdx.x != 0) return;
    int32_t t = *step;
    float b1 = *beta1;
    float b2 = *beta2;
    *bias_corr1 = 1.0f / (1.0f - powf(b1, static_cast<float>(t)));
    *bias_corr2 = 1.0f / (1.0f - powf(b2, static_cast<float>(t)));
}

这两个 kernel 都运行在 UPDATE 流上，与后续的大规模权重更新 kernel 处于同一条流，天然按顺序执行。它们的计算量极小，但把它们独立出来的好处是：偏差修正只需要做一次，随后被所有 weight 和 bias 更新 kernel 共享。没有重复计算，也没有把标量逻辑耦合进大数据 kernel。

注：在训练步数极大的生产场景中，为了避免单精度 powf 的舍入/下溢误差，也可将这一步改在核内用 double 计算，或改为 beta^t 的迭代累乘。

你可能会问：为什么不在主更新 kernel 里直接算 powf(beta, t)？技术上当然可以，但那样每个线程都要重复一次指数运算，而指数运算在 GPU 上并不便宜。更重要的是，把它独立成标量 kernel 后，bc1 和 bc2 可以被多个 RangeOp 节点共享——weight 更新用一次，bias 更新也用一次，确保所有参数的偏差修正在同一步保持一致。

此外，把 step 的递增也做成显式节点，让训练步数这个全局状态完全在 GPU 侧维护。CPU 只负责在每个 epoch 或每个 batch 开始时把学习率等小标量 H2D 写入显存，不参与步数计数。这种“CPU 只做决策、GPU 做执行”的分工，是静态图框架实现极致低延迟的关键。

六、LARS：为什么不能走 RangeOp？

LARS 是三种优化器里最特殊的一个。SGD、AdamW 的更新只依赖元素级别的读写，可以批量覆盖整个 Region；但 LARS 需要为每一层计算一个 trust ratio，而这个 trust ratio 是 ||w_l|| 和 ||g_l|| 两个 L2 范数的比值，本质上是对单个张量的归约操作。

RangeOp 的设计目标是“无视张量边界，一次性处理同类数据”。一旦需要 per-tensor 的归约结果，RangeOp 就不再适用，因为整个 Region 被揉在一起后，你无法区分哪一段属于哪一层。

因此 Tech-Renaissance 为 LARS 单独设计了一套 ComputeOp，编译器按层构建节点。在 src/graph/compiler.cpp 中，LARS 的图被分配到三个独立的 GraphId：LARS_FC_OPT、LARS_FIRST_CONV_OPT、LARS_DEEP_CONV_OPT，分别对应 FC 权重、首层卷积、深层卷积：

auto build_lars_pair = [&](Region w_region, GraphId gid,
                            ComputeOp trust_op, ComputeOp update_op,
                            int32_t temp_partial_id) {
    Region g_region = paired_grad_region(w_region);
    Region m_region = paired_momentum_region(w_region);
    Region n_region = paired_norm_region(w_region);

    const auto& w_ids = memory_plan.get_ids_by_region(w_region);
    const auto& g_ids = memory_plan.get_ids_by_region(g_region);
    const auto& m_ids = memory_plan.get_ids_by_region(m_region);
    const auto& n_ids = memory_plan.get_ids_by_region(n_region);

    for (size_t i = 0; i < w_ids.size(); ++i) {
        // Step 1: 计算 trust ratio
        GraphNode trust_node;
        trust_node.kind = GraphNode::Kind::COMPUTE;
        trust_node.compute_op = trust_op;
        trust_node.input_ids  = {w_ids[i], g_ids[i], scalar_ids.tc,
                                  scalar_ids.wd, scalar_ids.eps,
                                  scalar_ids.scaling, scalar_ids.has_nan,
                                  temp_partial_id};
        trust_node.output_ids = {n_ids[i]};
        train_cg.append(gid, trust_node);

        // Step 2: 用 trust ratio 更新权重
        GraphNode update_node;
        update_node.kind = GraphNode::Kind::COMPUTE;
        update_node.compute_op = update_op;
        update_node.input_ids  = {w_ids[i], g_ids[i], m_ids[i], n_ids[i],
                                   scalar_ids.lr, scalar_ids.beta,
                                   scalar_ids.wd, scalar_ids.scaling,
                                   scalar_ids.has_nan};
        update_node.output_ids = {w_ids[i], m_ids[i]};
        train_cg.append(gid, update_node);
    }
};

这里的关键是三张 LARS 图、三条计算流：

  • LARS_FC_OPT → COMP_1
  • LARS_FIRST_CONV_OPT → COMP_2
  • LARS_DEEP_CONV_OPT → COMP_3

FC、首层卷积、深层卷积的 trust ratio 归约和权重更新分别落在三条独立的计算流上，天然并行。src/backend/op_stream_policy.cpp 中的流分配策略确认了这一点：

case ComputeOp::LARS_COMPUTE_TRUST_RATIO_FC:
case ComputeOp::LARS_UPDATE_FC:
case ComputeOp::LARS_NESTEROV_UPDATE_FC:
    return StreamKind::COMP_1;

case ComputeOp::LARS_COMPUTE_TRUST_RATIO_FIRST:
case ComputeOp::LARS_UPDATE_FIRST:
case ComputeOp::LARS_NESTEROV_UPDATE_FIRST:
    return StreamKind::COMP_2;

case ComputeOp::LARS_COMPUTE_TRUST_RATIO_DEEP:
case ComputeOp::LARS_UPDATE_DEEP:
case ComputeOp::LARS_NESTEROV_UPDATE_DEEP:
    return StreamKind::COMP_3;

lars_op.cu 中的 trust ratio kernel 采用两阶段 reduce：phase1 用多个 block 并行累加局部 sum_w2/sum_g2，phase2 用单个线程完成最终归约并计算 eta：

// Phase 1: 每个 block 归约一段数据，输出局部 sum_w2 / sum_g2
__global__ void lars_trust_ratio_phase1_kernel(...)
{
    float local_w2 = 0.0f, local_g2 = 0.0f;
    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x;
         i < n; i += gridDim.x * blockDim.x) {
        float wv = w[i];
        float gv = g[i];
        local_w2 += wv * wv;
        local_g2 += gv * gv;
    }
    // warp shuffle + shared memory 完成 block 内归约
    // 最终写入 out_w2[blockIdx.x], out_g2[blockIdx.x]
}

// Phase 2: 单线程汇总所有 block 结果，计算 eta
__global__ void lars_trust_ratio_phase2_kernel(...)
{
    float sum_w2 = 0.0f, sum_g2 = 0.0f;
    for (int i = 0; i < num_blocks; ++i) {
        sum_w2 += in_w2[i];
        sum_g2 += in_g2[i];
    }
    float w_norm = sqrtf(sum_w2);
    float g_norm = sqrtf(sum_g2) * inv_scaling;
    float eta = 1.0f;
    if (w_norm >= 1e-12f && g_norm >= 1e-12f) {
        eta = tc * w_norm / (g_norm + wd * w_norm + eps);
        if (eta > 100.0f) eta = 100.0f;  // 数值稳定性钳制
    }
    *out_eta = eta;
}

临时缓冲区 T_TEMP_FP32 的分配顺序是 FC → FirstConv → DeepConv，索引从 0 开始，每个 LARS 层对占用 2 × kLarsMaxPartial 个 float。kLarsMaxPartial 定义为 65535，因此每个层对约 512 KB，远小于一次中间特征图的内存占用。

两阶段归约的核心优势在于并行度。单 block 归约最多只能利用 1024 个线程，但 Phase 1 可以启动多达 65535 个 block，每个 block 256 个线程——对于大型全连接层（数百万参数）的范数计算，这能充分利用 GPU 的计算单元。

LARS 的权重更新 kernel 也有其独特之处：它同时接受 eta（per-layer trust ratio）和 lr（全局学习率），将两者相乘作为实际学习率，并在一次遍历中完成动量更新和权重衰减：

// LARS_UPDATE:
//   gp      = g + wd * w
//   m_new   = beta * m + lr * eta * gp
//   w_new   = w - m_new
//   m       = m_new
for (size_t i = ...; i < n; i += ...) {
    float wv = w[i];
    float gv = g[i] * inv_scaling;
    float gp = gv + _wd * wv;
    float m_new = _beta * m[i] + _lr * _eta * gp;
    w[i] = wv - m_new;
    m[i] = m_new;
}

注意 LARS 的 weight decay 是嵌入在梯度更新中的（gp = gv + _wd * wv），这与 SGD/Adam 的 weight decay 嵌入方式不同，是 LARS 原始论文中定义的公式。

LARS 的 Bias-like 参数则不参与 trust ratio，直接退化为带 Momentum 或 Nesterov 的标准 RangeOp 更新，与 SGD_MOMENTUM / SGD_NESTEROV 共用同一条路径。这再次体现了“按行为分组”的设计哲学。

七、NaN 保护、梯度缩放与 CUDA Graph 兼容性

优化器更新不是孤立执行的——它处在整个训练循环的最后阶段，前面是梯度计算、梯度缩放（AMP 的反向过程）、NaN 检测。这些环节的信息通过一个 has_nan 标志传递给优化器。

在 optimizer_op.cpp 的 launcher 中，每个优化器在启动 kernel 前都会解析 has_nan 标志的指针：

const int32_t* has_nan = static_cast<const int32_t*>(
    ArenaKeeper::instance().ptr_at(ctx.rank_for_context(),
                                   mp.get_dtensor(node.input_ids.back()).offset()));

kernel 内部的第一条指令就是检查这个标志：

if (*has_nan != 0) return;

这个设计的精妙之处在于：即使 kernel 什么都不做，CUDA Graph 的重放依然有效。因为 has_nan 是一个可更新的 GPU 标量，Graph 重放前可以通过 host 端更新其值，kernel 读取后根据值决定是否执行实际计算。这保证了确定性训练场景下的图稳定性——不会因为某步出现 NaN 而导致图结构变化。

梯度缩放（gradient scaling）的处理同样被融合进了优化器 kernel 内部：

float _inv_scaling = (scaling && *scaling != 0.0f) ? (1.0f / *scaling) : 1.0f;
// ...
float g_i = g[i] * _inv_scaling;

这避免了额外的一次全局显存遍历，将 unscale 的开销完全隐藏在了优化器的计算中。一个值得与 PyTorch 对比的细节是：PyTorch 的 AMP 通常需要显式调用 scaler.unscale_(optimizer) 来还原梯度，然后再调用 optimizer.step()。这两个步骤是分离的——unscale 遍历一次所有梯度，step 再遍历一次所有参数。在我们的融合设计中，unscale 和 update 在同一个 kernel 的同一个循环中完成，减少了显存访问次数。对于 memory-bound 的优化器更新来说，这是一项实质性的性能改进。

八、学习率调度：一个小标量 H2D 就够了

在 Tech-Renaissance 中，学习率不是优化器状态的一部分，而是由 Scheduler 计算、通过一次极小的 H2D 拷贝写入 S_SCALAR_FP32 区域的全局标量。由于优化器 kernel 以设备指针读取 lr，学习率的变化不会导致 CUDA Graph 重新捕获。

这个设计在运行时的 GraphExecutor::update_lr_scalar() 中体现：它只把一个 float 值 copy 到显存，耗时通常在微秒级，并且可以和后续计算流重叠。如果是 step-by-batch 的调度策略，每个 batch 都会有一次这样的小 H2D，但因为它只写单个标量，开销几乎可以忽略。相比之下，如果优化器 step 是以 kernel 参数形式接收学习率，那么学习率每变一次就需要重构图，step-by-batch 的调度成本就会高得多。

九、小结

回顾一下本框架融合优化器的核心设计思路：

  1. 从“逐参数循环”到“整区批量更新”：利用 MemoryPlan 的 Region 连续布局，将模型所有参数分为 Weight 组和 Bias 组，每组只需一个 kernel 完成全部更新。无论模型有多少层，优化器更新永远是 1~2 个 kernel。
  2. Weight 与 Bias 的分组设计：基于 weight decay 的语义差异——Weight 组施加 weight decay，Bias 组 weight decay 恒为零。这种分组不是普通的性能技巧，而是深度学习界长期验证的收敛实践（防止正则化过度抹去仿射变换的平移语义）。
  3. 五种优化器、统一框架：SGD、Momentum、Nesterov、Adam、AdamW 通过统一的 RangeOp 路径实现，LARS 通过 ComputeOp 路径实现。两者共享相同的标量参数管理、NaN 保护和图兼容机制。
  4. LARS 的三流并行：将 FC 层、首层卷积、深层卷积的 trust ratio 计算和更新分别映射到三条计算流上并行执行，配合两阶段归约，将 LARS 的额外开销降到最低。
  5. CUDA Graph 全兼容：所有标量参数以指针传入，学习率、动量系数、bias correction 等可以在 Graph 重放前动态更新，无需重新捕获图。

记住：

Tech-Renaissance 的优化器是整区批量的：所有同行为的参数在内存中连续，一次 kernel launch 覆盖全部。SGD/Momentum/AdamW 的 weight 更新通常只需要一个 RANGE_UPDATE_WEIGHT_* 节点，bias 更新只需要一个 RANGE_UPDATE_BIAS_* 节点，再加上 Adam 的两个标量 kernel。无论模型是 10 层还是 100 层，优化器 step 的 kernel 数量是恒定的。

这种设计也带来了更好的显存访问局部性。主权重、梯度、动量缓冲区按相同顺序连续存放，kernel 在 grid-stride loop 中顺序读写，能够更好地利用合并访问和缓存行。相比之下，逐张量更新虽然每个张量内部也是顺序访问，但张量之间的地址不连续，容易在 TLB 和 cache 上产生更多抖动。

当然，这种整区批量的前提是：所有被合并的参数必须共享相同的超参数和行为。Tech-Renaissance 目前把 BN bias/weight、FC bias 分为一组，FC weight、Conv weight 分为另一组，正是因为它们的 weight decay 行为一致。如果你需要给不同层设置完全不同的超参数，这种强批量设计就不太适用——不过对标准 CNN 训练来说，这种分组已经覆盖了绝大多数场景。

优化器融合是显存分区设计带来的直接红利——分区让“批量”成为可能，融合让“批量”发挥出性能价值。下一篇，我们将介绍损失函数、学习率调度与参数初始化，看看框架的训练算法配置层是如何设计的。

Related Posts

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注

ICP备案号：京ICP备2025133467号-1