(22) 损失函数、学习率调度与参数初始化：训练算法配置层

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

前面的文章花了很大篇幅讲 Tech-Renaissance 的“骨架”：张量、计算图、编译器、显存规划、多流执行、CUDA Graph 捕获、融合算子。骨架再强壮，如果没有训练算法的正确配置，模型也跑不起来——就好比一辆赛车，发动机、变速箱、悬挂都调好了，但如果没有合适的燃油标号、点火时机和轮胎气压，它依然到不了终点。

训练算法的配置层，通常被称为深度学习框架里“最不起眼、却最影响结果”的部分。它包括损失函数、学习率调度、参数初始化，再配合优化器（上一篇已经专门讲过）和随机数生成（下一篇会细讲），共同决定了一次训练能否收敛、收敛到多好的位置、以及收敛得多快。

这篇文章就来聊聊 Tech-Renaissance 里这一层是怎么设计的。

一、训练算法配置层到底在配什么

1.1 损失函数：把预测变成可优化的目标

深度学习模型最后输出的一般是 logits 或概率分布，而损失函数负责衡量这个分布与真实标签之间的差距。对于分类任务，最常用的是交叉熵损失（Cross-Entropy Loss）。它的核心思想很朴素：模型对正确类别的预测概率越高，损失越低；如果模型对错误答案很自信，损失会迅速放大。

具体地，对于单个样本，若真实标签为 one-hot 向量 y，模型输出概率为 p，则交叉熵可以写成下面的伪代码：

# y[i] : 第 i 类的 one-hot 标签（真实类别为 1，其余为 0）
# p[i] : 模型预测的第 i 类概率，由 softmax 输出
loss = -sum(y[i] * log(p[i]) for i in range(num_classes))
# 由于 one-hot 中只有正确类别 y[target] == 1，上式等价于
loss = -log(p[target])

这个式子鼓励模型把正确类别的概率推向 1；同时由于 log 在接近 0 时下降很快，也让“自信但错误”的样本付出极高代价。

交叉熵通常和 softmax 一起使用。为了让数值更稳定，现代框架不会先算 softmax 再取 log，而是直接做fused log-softmax + NLL，利用 log-sum-exp 技巧避免指数爆炸或下溢。PyTorch 的 nn.CrossEntropyLoss 和 TensorFlow 的 SparseCategoricalCrossentropy 都遵循这一惯例，要求用户传入原始 logits 而不是已经归一化的概率。

在分类任务中，一个常用的技巧是标签平滑（Label Smoothing）。传统 one-hot 标签会把全部概率放在正确类别上，模型为了拟合这种“硬”目标，容易对训练样本过度自信，泛化能力下降。标签平滑把目标分布改成：

# eps : 平滑系数，常见取值为 0.1
# num_classes : 类别数
# one_hot : 真实类别的 one-hot 分布
smoothed = (1 - eps) * one_hot + eps / num_classes
# 等价于：把正确类别的概率从 1 降到 1 - eps + eps/num_classes，
#         其余类别各分 eps/num_classes

其中 ε 通常取 0.1，相当于告诉模型“正确答案的置信度不要拉满”。在 ImageNet 等大规模视觉任务里，0.1 的标签平滑几乎是标准配方，能够让 top-1/top-5 轻微提升，同时让模型输出的概率分布更加校准。

交叉熵还有一个重要优势：当它和 softmax 一起求导时，梯度形式非常简洁，等于 softmax(logits) - target。这意味着即使模型在早期输出极差的 logits，梯度也不会像均方误差那样因为链式求导而迅速衰减，深层网络因此更容易获得有效的更新信号。这也是分类任务几乎不用 MSE 作为损失函数的原因之一。

1.2 学习率调度：控制每一步的步长

学习率大概是深度学习训练里最让人又爱又恨的超参数。太大，损失震荡甚至发散；太小，收敛慢还容易掉进鞍点。更复杂的是，训练不同阶段需要不同的学习率：刚开始时梯度方向不稳定，需要Warmup——从小学习率线性爬升到目标学习率；中后期则需要逐步衰减，让模型在低谷附近精细搜索。

Warmup 在大 batch 训练中尤为重要。当 batch size 很大时，梯度估计更精确但方差也更大，一开始就用高学习率容易导致参数更新过猛。先用几个 epoch 把学习率从小慢慢 ramp 上来，相当于给优化器一个“暖机”过程。之后再接各种衰减策略：

• StepLR：每隔固定 epoch 把学习率乘一个因子 gamma，简单直接但衰减不连续；
• MultiStepLR：在指定里程碑处打折，适合有明确阶段划分的训练；
• ExponentialLR：按 lr = base_lr * gamma^epoch 连续衰减；
• PolynomialLR：按 (1 - progress)^power 衰减，power 控制曲线形状；
• CosineAnnealingLR：按余弦曲线从 base_lr 降到 eta_min，平滑自然；
• CosineAnnealingWarmRestarts：余弦退火周期重启，常用于 Transformer 训练；
• WSD（Warmup-Stable-Decay）：前段保持 base_lr 不变，后段线性衰减到 end_lr，近年来在大模型训练里很流行。

PyTorch 里常用 LambdaLR 把这些策略像积木一样拼接，再用 SequentialLR 按阶段切换。但不管公式怎么变，本质上都是给定“当前步数”，返回一个“学习率”。学习率调度器本身不修改模型参数，它只是告诉优化器：这一步该用多大的步长。

以线性 Warmup 为例，其学习率计算公式可以写成：

# t          : 当前步数（从 0 开始计数）
# W          : Warmup 总步数
# start_lr   : Warmup 起始学习率
# base_lr    : Warmup 结束后的目标学习率
# 注意：当 t == W 时，lr 恰好等于 base_lr
if t <= W:
    lr = start_lr + (base_lr - start_lr) * t / W
else:
    lr = decay_schedule(t - W)  # 进入后续衰减阶段

这个简单公式背后是大批量训练的经验：在 Warmup 结束时，参数已经移动到一个相对稳定的区域，此时再用 base_lr 就不容易造成震荡。

1.3 参数初始化：给训练一个合理的起点

神经网络不是从任意状态都能训练成功的。如果权重全设为 0，所有神经元会学到同样的东西，对称性无法打破；如果权重太大或太小，信号会在深层网络里爆炸或消失。参数初始化的目标，就是让前向传播和反向传播中的信号方差保持相对稳定。

最经典的两种思路：

• Xavier / Glorot 初始化：假设激活函数近似线性，通过 fan_in 和 fan_out 的平均值来控制权重方差，适用于 tanh、sigmoid 等对称激活。
• Kaiming / He 初始化：针对 ReLU 类激活函数，把方差乘以 2 来补偿 ReLU 的“截断一半”效应。

用伪代码表示它们的采样参数：

# fan_in  : 权重张量输入侧的连接数
# fan_out : 权重张量输出侧的连接数
# gain    : 由激活函数决定的额外缩放因子；ReLU 默认 gain = sqrt(2)

# Xavier / Glorot 正态分布
xavier_std = gain * sqrt(2.0 / (fan_in + fan_out))

# Kaiming / He 正态分布
kaiming_std = gain / sqrt(fan_in)              # ReLU 默认 gain = sqrt(2)

# Kaiming / He 均匀分布：区间 [-bound, bound]
kaiming_bound = gain * sqrt(3.0 / fan_in)

选择哪种初始化，关键看激活函数。ReLU 及其变体几乎总是用 Kaiming；tanh、sigmoid、swish 等对称激活更适合 Xavier。如果选错，深层网络前几轮就会出现激活值迅速塌陷或爆炸。

除此之外，视觉领域还常用截断正态分布（Truncated Normal），把采样值限制在 ±2σ 以内，避免极端值破坏早期训练；批归一化的 gamma 通常初始化为 1、bias 为 0；ResNet 里还会把部分 BN 的 gamma 初始化为 0，形成所谓的 Zero Gamma 初始化，让深层网络在训练初期近似恒等映射，缓解梯度传播困难。

PyTorch 的 nn.init 模块和 TensorFlow 的 tf.keras.initializers 都实现了上述方案。它们看起来只是几行随机数生成，但背后直接影响的是训练能否启动、收敛速度和最终精度。

批归一化层的初始化也有讲究。γ（weight）初始化为 1、β（bias）初始化为 0，是为了让 BN 在训练初期等价于恒等变换，不破坏前一层已经初始化好的信号分布。ResNet 进一步提出 Zero Gamma：对特定残差分支末尾的 BN，把 γ 初始化为 0，这样整个分支在训练刚开始时输出几乎为 0，网络先学习浅层恒等映射，再逐步激活深层路径，从而缓解极深网络的优化困难。

二、Tech-Renaissance 的设计哲学：无状态、纯函数、配置进图

在 Tech-Renaissance 里，损失函数、学习率调度器、参数初始化器都被设计成无状态的纯配置对象。它们不维护 current_step、current_lr 之类的可变状态，也不参与运行时的逐步更新。这种设计不是偶然，而是服务于整个框架的一个核心约束：多卡分布式训练必须保持一致，且不能依赖运行时同步。

在 PyTorch 里，调度器通常和优化器绑定，scheduler.step() 每轮修改优化器内部的学习率。如果多卡训练时某个 rank 的 step 数因为数据分片不均衡而不一致，学习率就会分叉。Tech-Renaissance 的做法是反过来：给定 epoch 和 batch，直接算出学习率。所有 rank 只要看到相同的 (epoch, batch)，就会得到相同的结果，无需任何同步。

这种无状态设计也简化了 CUDA Graph 的捕获。如果学习率是一个对象内部的 mutable 状态，每次变化都可能触发图的重构；而把它变成 (epoch, batch) → lr 的纯函数，学习率就可以被当作一个普通标量输入，运行时只需要更新这个标量即可。

这三个配置对象最终都会进入编译管线：损失参数写入 MemoryPlan 的标量张量，学习率在每个 batch 开头通过一次轻量 H2D 拷贝注入显存，初始化策略则以 InitConfig 的形式绑定到每一个 DTensor。接下来我们逐一展开。

三、损失函数：一期聚焦交叉熵

Tech-Renaissance 当前版本在损失函数上非常聚焦，只实现了 CrossEntropyLoss。这不是能力上限，而是训练场景的选择：本框架目前主要面向图像分类和类似监督学习任务，交叉熵加标签平滑已经覆盖了从 MNIST 到 ImageNet 的主流需求。

代码层面，CrossEntropyLoss 是一个纯配置类：

// include/renaissance/algo/loss.h
class CrossEntropyLoss {
public:
    CrossEntropyLoss() = default;

    CrossEntropyLoss& label_smoothing(float value) {
        TR_CHECK(value >= 0.0f && value <= 0.20001f, ValueError,
                 "label_smoothing must be in [0, 0.2], got " << value);
        label_smoothing_ = value;
        GlobalRegistry::instance().set_label_smoothing(value);
        return *this;
    }

    [[nodiscard]] float label_smoothing() const noexcept {
        return label_smoothing_;
    }

private:
    float label_smoothing_ = 0.0f;
};

注意这里的一个关键细节：CrossEntropyLoss::label_smoothing() 一被调用，就把值写进了 GlobalRegistry，因此它是全局唯一的权威值；而 DeepLearningTask::loss() 本身只做一个阶段检查，确认用户还在 PLANNING 阶段。对框架来说，真正被编译器读取的是注册表里的值，而不是配置对象内部的副本。

编译阶段，Compiler::compile() 从 GlobalRegistry 读取 label_smoothing，把它写入 MemoryPlan 中 Region::S_SCALAR_FP32 区域的一个标量 DTensor：

// src/graph/compiler.cpp
float ls_val = GlobalRegistry::instance().label_smoothing();
memory_plans[s]->set_init_config(
    memory_plans[s]->baseline().label_smoothing, kInitConstant(ls_val));

这个标量随后被 SOFTMAX_CE_FP32_FWD、SOFTMAX_CE_AMP_FWD 以及对应的反向算子作为输入读取。实际的前向 softmax、交叉熵、反向求导都被融合成一个后端算子执行，中间不会把完整 logits 或概率矩阵写回全局显存。这个设计与 PyTorch 的 fused cross entropy 在数学上等价，但更符合 Tech-Renaissance 静态图编译、减少中间张量的整体风格。

标签平滑的范围在代码里被限制在约 [0, 0.2]（上界写成 0.20001f，保留了一点浮点容差），注释里也提到 MLPerf Closed Division 的常见取值是 0 或 0.1。这个限制看起来很小，但它反映了框架的一个态度：训练配置不是任意值都可以，而是要和主流基准、可复现性要求对齐。

有人可能会问，为什么不多实现几种损失函数？答案是 Tech-Renaissance 目前专注于分类任务的端到端吞吐优化，损失函数本身并不是瓶颈。真正影响性能的是损失与前向、反向算子之间的融合方式。先把交叉熵做透、和 softmax 一起融进后端，比仓促支持十种损失函数更有价值。未来如果需要回归、检测或语言模型任务，再在此基础上扩展也水到渠成。

四、学习率调度器：八种策略，一套纯函数框架

Tech-Renaissance 实现了八种学习率调度器，都继承自同一个抽象基类 LRScheduler：

调度器	核心行为	关键参数
ConstantLR	全程保持 base_lr	—
StepLR	每 step_size 个 epoch 把学习率乘以 gamma	step_size, gamma
MultiStepLR	在预设的 milestones epoch 处乘以 gamma	milestones, gamma
ExponentialLR	每个 epoch 把学习率乘以 gamma	gamma
PolynomialLR	按 (1 - progress)^power 从 base_lr 衰减到 end_lr	power, end_lr
CosineAnnealingLR	按余弦曲线从 base_lr 衰减到 eta_min	eta_min
CosineAnnealingWithWarmRestartsLR	周期性余弦退火，周期长度按 T_mult 递增	T_0, T_mult, eta_min
WSDLR	前 decay_start 比例保持稳定，之后线性衰减到 end_lr	decay_start, end_lr

基类的设计非常明确：

// include/renaissance/algo/scheduler.h
class LRScheduler {
public:
    virtual ~LRScheduler() = default;

    virtual LRScheduler& base_lr(float lr);
    virtual LRScheduler& warmup(int epochs);
    virtual LRScheduler& warmup_start_lr(float start_lr);
    virtual LRScheduler& warmup_start_factor(float factor);
    virtual LRScheduler& step_by_batch(bool v = true);
    virtual LRScheduler& step_by_epoch();

    void prepare(int total_epochs, int steps_per_epoch);
    float get_lr_by_batch(int batch_id) const;
    float get_lr_by_epoch(int epoch_id) const;
    bool is_step_by_batch() const noexcept { return step_by_batch_; }

protected:
    virtual float compute_decay_lr(int decay_step, int total_decay) const = 0;
    float compute_lr_at_step(int effective_step) const;
    // ...
};

基类负责统一裁剪边界、执行线性 Warmup，然后把衰减阶段委托给派生类实现：

// src/algo/scheduler.cpp
float LRScheduler::compute_lr_at_step(int effective_step) const {
    // effective_step : 从训练开始累计的步数（已被裁剪到 [0, total_steps_]）
    // total_steps_   : 训练总步数 = total_epochs_ * steps_per_epoch_
    // warmup_steps_  : Warmup 总步数 = warmup_epochs_ * steps_per_epoch_
    if (effective_step < 0) effective_step = 0;
    if (effective_step >= total_steps_) effective_step = total_steps_;

    // 线性 Warmup：峰值出现在 effective_step == warmup_steps_
    if (warmup_steps_ > 0 && effective_step <= warmup_steps_) {
        float progress = static_cast<float>(effective_step)
                       / static_cast<float>(warmup_steps_);
        float start_lr = resolve_warmup_start_lr();
        return start_lr + (base_lr_ - start_lr) * progress;
    }

    // 衰减阶段：派生类只需提供 compute_decay_lr
    int decay_step  = effective_step - warmup_steps_;
    int total_decay = total_steps_ - warmup_steps_;
    float lr = compute_decay_lr(decay_step, total_decay);
    if (lr < 0.0f || std::isnan(lr)) lr = 0.0f;
    return lr;
}

派生类只需要实现一个纯函数。例如 CosineAnnealingLR：

// src/algo/scheduler.cpp
float CosineAnnealingLR::compute_decay_lr(int decay_step, int total_decay) const {
    if (total_decay <= 0) return base_lr_;

    int effective_total = step_by_batch_ ? total_decay : (total_decay - steps_per_epoch_);
    if (effective_total <= 0) return base_lr_;

    float progress = static_cast<float>(decay_step)
                   / static_cast<float>(effective_total);
    if (progress > 1.0f) progress = 1.0f;

    return eta_min_ + (base_lr_ - eta_min_)
           * (1.0f + std::cos(static_cast<float>(M_PI) * progress)) * 0.5f;
}

PolynomialLR 则对齐了 TensorFlow 的 polynomial_decay 行为，分母使用 total_decay + 1；StepLR 和 MultiStepLR 按 epoch 数打折；WSDLR 先稳定后线性衰减。每一种策略都是纯函数，没有任何内部状态推进。

4.1 逐 epoch 还是逐 batch

LRScheduler 支持两种步进模式。step_by_epoch() 表示每个 epoch 内学习率不变；step_by_batch() 则把 epoch * steps_per_epoch + batch_id 作为全局步数，让学习率在每个 batch 都变化。DeepLearningTask::fetch_lr_for_batch() 里的逻辑很直接：

// src/task/deep_learning_task.cpp
float DeepLearningTask::fetch_lr_for_batch(int batch_id) const {
    return std::visit([this, batch_id](auto&& sch) -> float {
        using T = std::decay_t<decltype(sch)>;
        if constexpr (std::is_same_v<T, std::monostate>) {
            return 0.0f;
        } else {
            if (sch.is_step_by_batch()) {
                int global_step = current_epoch_ * sch.steps_per_epoch() + batch_id;
                return sch.get_lr_by_batch(global_step);
            } else {
                return sch.get_lr_by_epoch(current_epoch_);
            }
        }
    }, sched_cfg_);
}

这里 std::variant 保存了具体调度器类型。编译期确定类型，运行期无虚函数开销。

另外，step_by_batch() 和 step_by_epoch() 一旦调用就会锁定步进模式，后续冲突调用会直接抛错。这是一种防御性设计，防止用户链式配置时先写 .step_by_batch() 又写 .step_by_epoch()，导致学习率更新节奏出现歧义。

4.2 如何把学习率送进 CUDA Graph

Tech-Renaissance 的训练循环整体被捕获成 CUDA Graph。如果学习率变化需要重构整个图，那开销就太大了。框架的解决办法是：学习率本身是一个 S_SCALAR_FP32 标量 DTensor，运行时通过锁页内存做一次 4 字节的 cudaMemcpyAsync。

// src/task/deep_learning_task.cpp
float* lr_dev_ptr = static_cast<float*>(ctx.ptr_at(lr_dtensor_id_));
// ...
bool need_lr = is_step_by_batch_mode() || batch == 0;
if (need_lr) {
    lr = fetch_lr_for_batch(batch);
    *lr_pinned_[rank] = lr;
    cudaMemcpyAsync(lr_dev_ptr, lr_pinned_[rank], sizeof(float),
                    cudaMemcpyHostToDevice, s_trans);
}

lr_pinned_ 在 compile() 阶段通过 cudaMallocHost 分配。这个小拷贝发生在传输流上，和计算流并行，4 字节的数据量几乎可以忽略。优化器算子从同一个标量张量读取当前学习率，因此图结构无需改变。这是静态图框架下实现动态学习率的一种务实方案。

值得补充的是，框架里还存在一个 StagingParamPool 小参数区（默认 256 字节，64 个 FP32 的锁页 staging 区），它配合 RANGE_H2D_COPY_DTENSOR 等范围算子，可以把极小参数的 H2D 拷贝也表达为图中的一个节点。不过在当前 DeepLearningTask 的 GPU 训练循环里，学习率走的是直接的 cudaMemcpyAsync 路径；StagingParamPool 在 SimpleTask 等其它任务路径中发挥作用。

五、参数初始化：从策略到每个 Region 的 8 字节配置

参数初始化在 Tech-Renaissance 里由三个文件共同承担：init_config.h 定义策略枚举和紧凑配置，initializer.h/.cpp 负责策略推导和数学实现，TaskBase::init_all() 负责实际填充。

5.1 InitConfig：8 字节的初始化策略描述符

// include/renaissance/core/init_config.h
enum class InitKind : uint8_t {
    NONE, ZEROS, CONSTANTS,
    KAIMING_NORMAL, KAIMING_UNIFORM,
    XAVIER_NORMAL, XAVIER_UNIFORM,
    TRUNC_NORMAL, FIXED_NORMAL,
    STANDARD, ZERO_GAMMA
};

enum class FanMode : uint8_t { FAN_IN, FAN_OUT, FAN_AVG };

struct InitConfig {
    float    scale = 1.0f;           // CONSTANTS=填充值，Kaiming/Xavier=gain，FIXED_NORMAL=标准差
    InitKind kind  = InitKind::NONE; // 初始化方法种类
    FanMode  fan   = FanMode::FAN_IN; // fan 计算模式
    // 由于对齐，编译器会在尾部补 2 字节，整体恰好 8 bytes
};

static_assert(sizeof(InitConfig) == 8, "InitConfig must be exactly 8 bytes");

InitConfig 被设计成 8 字节的值类型，可以在 CPU 侧作为普通值自由拷贝。scale 字段根据 kind 有不同含义：对 CONSTANTS 它是填充值，对 KAIMING/XAVIER 它是 gain，对 FIXED_NORMAL 它是标准差。

5.2 Initializer：链式配置与 Region 推导

用户侧 API 是典型的链式风格：

Initializer init;
init.conv(InitKind::TRUNC_NORMAL)
    .fc(InitKind::KAIMING_UNIFORM)
    .bn(InitKind::STANDARD)
    .fan(FanMode::FAN_IN)
    .scale(1.0f);

Initializer 本身只记录层类型级别的策略。真正决定某个 DTensor 如何初始化，是在编译期通过 derive(Region) 完成的。Region 是 MemoryPlan 里每个张量的语义分区，Initializer::derive() 采用三段式结构：

// src/core/initializer.cpp
InitConfig Initializer::derive(Region region) const {
    // 第一段：偏置区 → ZEROS
    if (is_bias_region(region)) {
        return InitConfig{0.0f, InitKind::ZEROS, FanMode::FAN_IN};
    }

    // 动量/速度权重区 → ZEROS
    if (region == Region::M_BN_WEIGHT  || region == Region::M_FC_WEIGHT  ||
        region == Region::M_FIRST_CONV || region == Region::M_DEEP_CONV  ||
        region == Region::V_BN_WEIGHT  || region == Region::V_FC_WEIGHT  ||
        region == Region::V_FIRST_CONV || region == Region::V_DEEP_CONV) {
        return InitConfig{0.0f, InitKind::ZEROS, FanMode::FAN_IN};
    }

    // BN running mean / variance 的初始值
    if (region == Region::B_PREV_MEAN || region == Region::B_NEXT_MEAN) {
        return InitConfig{0.0f, InitKind::CONSTANTS, FanMode::FAN_IN};
    }
    if (region == Region::B_PREV_VAR || region == Region::B_NEXT_VAR) {
        return InitConfig{1.0f, InitKind::CONSTANTS, FanMode::FAN_IN};
    }

    // 非参数区 → NONE（显存池已全局置零）
    if (!is_param_region(region)) {
        return InitConfig{1.0f, InitKind::NONE, FanMode::FAN_IN};
    }

    // 第二段：权重区按层类型分发
    if (region == Region::W_EQ_SCALE) {
        return InitConfig{1.0f, InitKind::CONSTANTS, FanMode::FAN_IN};
    }
    if (is_bn_weight(region)) {
        return InitConfig{1.0f, InitKind::CONSTANTS, FanMode::FAN_IN};
    }
    if (is_conv_weight(region)) {
        float gain = global_scale_ * (conv_kind_ == InitKind::KAIMING_NORMAL ||
                                      conv_kind_ == InitKind::KAIMING_UNIFORM
                                      ? std::sqrt(2.0f / (1.0f + kaiming_a_ * kaiming_a_))
                                      : 1.0f);
        return InitConfig{gain, conv_kind_, fan_mode_};
    }
    if (is_fc_weight(region)) {
        if (fc_kind_ == InitKind::FIXED_NORMAL) {
            return InitConfig{fc_param_, InitKind::FIXED_NORMAL, fan_mode_};
        }
        float gain = global_scale_ * (fc_kind_ == InitKind::KAIMING_NORMAL ||
                                      fc_kind_ == InitKind::KAIMING_UNIFORM
                                      ? std::sqrt(2.0f / (1.0f + kaiming_a_ * kaiming_a_))
                                      : 1.0f);
        return InitConfig{gain, fc_kind_, fan_mode_};
    }
    // ...
}

这段代码体现了 Tech-Renaissance 初始化系统的核心思路：不是按“层”去初始化，而是按“Region”去初始化。因为 MemoryPlan 已经把同语义的张量集中排放，偏置、BN weight、Conv weight、FC weight 各自落在不同的 Region 上。derive() 负责覆盖所有需要显式初始化的参数区；像梯度、动量、EMA 等区要么在全局 memset 中已经为零，要么由后续的范围算子（RangeOp）批量处理。

顺带说明：Region 枚举目前共有 69 个槽位（含边界哨兵），其中实际命名的语义 Region 为 68 个，但 derive() 只需处理真正进入初始化管线的参数区即可。

对于 ResNet 中的 Zero Gamma 策略，框架已经预留了 mark_bn3() 接口：编译器在展开残差块时会记录主分支最后一个 BN 的 weight DTensor id。理论上 init_all() 阶段应把这些特定 id 的 BN weight 从默认的 1.0 覆盖为 0.0，从而让整个残差分支在训练初期近似恒等映射。不过当前版本中，Zero Gamma 的最终覆盖逻辑因为验证尚未完全通过而被显式禁用；mark_bn3 的标记链路已经就绪，后续验证完成后即可打开。

5.3 数学实现：从 fan 到随机数

Initializer::apply_to_tensor() 是实际生成随机数的入口，全部在 CPU 端完成：

// src/core/initializer.cpp
void Initializer::apply_to_tensor(Tensor& t, const Shape& shape, InitConfig cfg) {
    switch (cfg.kind) {
        case InitKind::TRUNC_NORMAL: {
            int64_t fan = compute_fan(shape, cfg.fan);
            float std = std::sqrt(cfg.scale / static_cast<float>(fan));
            t.truncated_normal(0.0f, std, -2.0f * std, 2.0f * std);
            return;
        }
        case InitKind::KAIMING_NORMAL: {
            int64_t fan = compute_fan(shape, cfg.fan);
            float std = cfg.scale / std::sqrt(static_cast<float>(fan));
            t.normal(0.0f, std);
            return;
        }
        case InitKind::KAIMING_UNIFORM: {
            int64_t fan = compute_fan(shape, cfg.fan);
            float bound = cfg.scale * std::sqrt(3.0f / static_cast<float>(fan));
            t.uniform(-bound, bound);
            return;
        }
        case InitKind::XAVIER_NORMAL: {
            int64_t fi = compute_fan(shape, FanMode::FAN_IN);
            int64_t fo = compute_fan(shape, FanMode::FAN_OUT);
            float std = cfg.scale * std::sqrt(2.0f / static_cast<float>(fi + fo));
            t.normal(0.0f, std);
            return;
        }
        case InitKind::XAVIER_UNIFORM: {
            int64_t fi = compute_fan(shape, FanMode::FAN_IN);
            int64_t fo = compute_fan(shape, FanMode::FAN_OUT);
            float bound = cfg.scale * std::sqrt(6.0f / static_cast<float>(fi + fo));
            t.uniform(-bound, bound);
            return;
        }
        // ...
    }
}

卷积权重的 fan 计算基于 Tech-Renaissance 统一的 KRSC 布局：[K=outC, R=kH, S=kW, C=inC]，因此 fan_in = C × R × S，fan_out = K × R × S。FC 层则按普通二维权重处理。这些公式和 PyTorch nn.init.kaiming_normal_、nn.init.xavier_uniform_ 的定义一致，只是实现被整合进了静态初始化管线。

5.4 初始化流程：rank 0 生成，H2D，广播

编译完成后，TaskBase::init_all() 会遍历所有 DTensor，根据 init_config 逐个初始化：

// src/task/task_base.cpp
void TaskBase::init_all() {
    for (const auto& dtensor : active_memory_plan_->dtensors()) {
        init(dtensor);
    }
}

TaskBase::init() 的实际流程是：

1. 在 CPU 端申请一个同样 shape 的 Tensor；
2. 调用 Initializer::apply_to_tensor() 生成随机数；
3. 把这个 CPU 张量传输到 rank 0 的 GPU；
4. 如果有多卡，通过 NCCL broadcast_from_rank0 把 rank 0 的权重广播到其他 rank。

这样做的好处是：所有 GPU 上的初始权重完全一致，避免分布式训练因为初始值不同导致梯度 AllReduce 后仍然不一致。随机数生成由 Philox 计数器 RNG 负责，下一篇文章会专门讲它如何保证多线程、多卡下的可复现性。

需要强调的是，初始化虽然发生在训练循环之前，但它对整个分布式系统的正确性至关重要。如果每个 rank 各自独立生成随机初始值，即使后续 AllReduce 同步梯度，各 rank 的参数更新起点也不同，训练结果在数学上就不再等价于单卡大 batch。Tech-Renaissance 选择 rank 0 生成、全卡广播的方案，正是为了确保这个起点的一致性。

六、从用户配置到训练图：一条完整的链路

把这些东西串起来看，Tech-Renaissance 的训练算法配置层大致是这样工作的：

// tests/example/mlp_mnist.cpp
DeepLearningTask task;
task.model(mlp)
    .loss(CrossEntropyLoss().label_smoothing(0.1f))
    .optimizer(AdamW().weight_decay(1e-4f))
    .scheduler(CosineAnnealingLR().base_lr(0.001f).warmup(5))
    .initializer(Initializer()
        .fc(InitKind::KAIMING_UNIFORM)
        .fan(FanMode::FAN_IN))
    .total_epochs(kTotalEpochs);

task.compile();
auto result = task.run();

用户面对的 API 非常薄。DeepLearningTask 在 on_prepare() 里完成几件事：

1. 如果没有显式指定初始化器，就使用默认 Initializer()；
2. 根据优化器类型设置 PlanConfig；
3. 从 BluePrint 生成 ArchPlan，再调用 Compiler::compile(..., initializer_, variant_specs)；
4. 编译完成后，查找 Region::S_SCALAR_FP32 中的学习率标量 DTensor；
5. 为优化器标量（momentum、weight_decay、beta2、eps 等）设置 InitConfig；
6. 把训练图和推理图加入 TaskBase。

Compiler 在 create_memory_plans() 阶段会对每个分配出的 DTensor 调用 initializer.derive(desc.region)，把初始化策略写入 MemoryPlan。随后 TaskBase::init_all() 在训练开始前一次性执行初始化。整个过程没有动态分配，没有运行时状态竞争，所有配置都在编译期落进图和内存布局。

七、小结

损失函数、学习率调度、参数初始化，这些常常被初学者当作“调参三板斧”的东西，在深度学习框架内部其实是一整套需要精心设计的基础设施。它们必须和计算图、显存规划、CUDA Graph 捕获、分布式一致性紧密配合，才能在保证正确的前提下不成为性能瓶颈。

Tech-Renaissance 的做法可以概括为三点：

1. 配置对象无状态：损失、调度器、初始化器都不维护运行时状态，给定输入直接返回结果，天然适合多卡并行。
2. 配置进图、进 Region：损失参数和学习率变成 MemoryPlan 标量张量，初始化策略变成每个 DTensor 的 8 字节 InitConfig，避免运行时反复查询对象属性。
3. 保持与主流做法的数学一致性：交叉熵加标签平滑、八种常见学习率调度、Kaiming/Xavier/Truncated Normal 初始化，公式都与 PyTorch、TensorFlow 对齐，确保迁移模型时不需要重新调参。

下一篇，我们将进入训练流程中另一个容易被忽视却至关重要的主题——随机数生成，以及 Tech-Renaissance 如何用 Philox 计数器实现端到端的确定性训练。