feat: pack 2 fc layers in mlp with swiglu

[Chamberlain0w0]

背景

当前 SwiGLU MLP 使用两个独立的 ColumnParallelLinear 分别计算 gate projection 和 up projection，随后执行 SiLU 和逐元素乘法：

gate = c_fc2(x)
up   = c_fc(x)
out  = SiLU(gate) * up

该实现会产生两次独立 GEMM 调用，并分别保存中间结果。对于 MoE，大 batch 下各 expert 的中间激活会带来较明显的性能和显存开销。Megatron 中常将两个 FC1 投影打包为单个 GEMM，本 PR 也做了类似的修改，并新增融合的 SwiGLU kernel。

主要改动

Packed FC1

将 SwiGLU 的 gate/up projection 合并为一个输出维度为 2H 的 ColumnParallelLinear：

packed = c_fc(x)        # [..., 2H]，布局为 [gate, up]
out    = SwiGLU(packed) # SiLU(gate) * up，输出 [..., H]

packed tensor 采用 Megatron-LM 的 [gate, up] 布局，便于后续对齐 Megatron checkpoint 和相关 fused kernel。

Dense MLP 和 MoE expert 共用该实现，因此均切换为 packed SwiGLU。GELU MLP 路径保持不变。

SwiGLU kernel

新增独立的 SwiGLU kernel 文件：

• CPU：infini_train/src/kernels/cpu/swiglu.cc
• CUDA：infini_train/src/kernels/cuda/swiglu.cu

kernel 融合执行 SiLU 和逐元素乘法，并实现对应反向传播：

output = SiLU(gate) * up

支持范围：

• CPU：FP32
• CUDA：FP32、BF16
• CUDA backward 支持 BF16/FP32 混合梯度类型

本 PR 合并了两个 FC1 GEMM，并融合了 SiLU + multiply。

Autograd 与 Module

• 新增 SwiGLU autograd Function。
• forward 接收末维布局为 [gate, up] 的 packed tensor。
• backward 分别计算 gate 和 up 梯度，并按 [gate, up] 顺序写回。
• SwiGLU module 改为调用 dispatcher 中的 CPU/CUDA kernel。
• 补充输入末维必须为偶数等合法性检查。
• 补充 MLP forward 中各阶段的 shape 注释。

Checkpoint loader

同步调整 LLMC checkpoint loader：

• Llama3 loader 将 gate_proj 写入 packed FC1 前半部分，将 up_proj 写入后半部分。
• Mixtral loader 将每个 expert 的 w1/gate_proj 写入前半部分，将 w3/up_proj 写入后半部分。
• TP 场景下按每个 rank 的本地 [2H_local] 布局分别加载 gate/up shard。
• 抽象通用的 projection-to-packed-weight 加载逻辑。

参数布局变化

SwiGLU MLP 的 state dict 从：

c_fc.weight  [H, C] # up
c_fc2.weight [H, C] # gate

调整为：

c_fc.weight [2H, C] # [gate, up]

c_fc2.weight 不再存在。

并且将 Packed SwiGLU 设置为 SwiGLU MLP 的默认实现，不需要额外配置开关。

MoE 性能提升

统计时排除 step 1 warm-up，使用 step 2–10 的平均耗时。baseline 和 packed 使用相同启动参数及构建配置。

用例	平均 step time（baseline → packed）	耗时下降	Peak used（baseline → packed）	Peak reserved（baseline → packed）
默认 batch，FP32	273.20 → 220.14 ms	19.4%	13308 → 13308 MB	14848 → 14656 MB
默认 batch，BF16	301.80 → 268.08 ms	11.2%	13246 → 13246 MB	15744 → 15360 MB
GBS=MBS=80，FP32	1095.47 → 969.08 ms	11.5%	35472 → 30992 MB	35712 → 31264 MB
GBS=MBS=80，BF16	602.52 → 511.31 ms	15.1%	30367 → 27807 MB	30880 → 28288 MB

结果表明：

• 4 组 MoE 用例的平均 step time 均有改善，降低约 11%–19%。
• 小 batch 下主要体现为计算性能提升。
• GBS=MBS=80、FP32 下 peak used 减少 4480 MB（12.6%）。
• GBS=MBS=80、BF16 下 peak used 减少 2560 MB（8.4%）。

[kilinchange]

这个抽到 infini_train/include/common/cuda/common_cuda.h 里吧，elementwise.cu 和 swiglu.cu 统一调同一个公共函数。

[kilinchange]

留个 TODO，后续考虑新增向量化访存优化。

[kilinchange]

这部分 lora 相关改动麻烦 @chen2021673 也看一下

[kilinchange]

这部分代码和上面 QKV 的处理逻辑上相似度比较高，能不能抽出通用逻辑，合并成一个循环：

for (const auto &[module_name, module] : named_modules) {
    if (dynamic_cast<CausalSelfAttention *>(module.get())) {
        MarkPackedQKVLoRASharding(module_name, module, shardings);
    } else if (dynamic_cast<MLP *>(module.get())) {
        MarkPackedSwiGLULoRASharding(module_name, module, shardings);
    }
}

给指定 projection 的 LoRA-B 设置 sharding 的通用 helper：

void MarkLoRAColumnParallelBSharding(
    const std::string &module_name,
    const std::shared_ptr<Module> &module,
    const std::string &projection_name,
    LoRATensorSharding sharding,
    std::unordered_map<std::string, LoRATensorSharding> &shardings) {

    auto projection = module->mutable_module(projection_name);

    if (!dynamic_cast<LoRAColumnParallelLinear *>(projection.get())) {
        return;
    }

    const auto projection_module_name =
        QualifiedParamName(module_name, projection_name);

    shardings[QualifiedParamName(
        projection_module_name,
        LoRAColumnParallelLinear::kParamLoraBName)] = sharding;
}

QKV helper：

void MarkPackedQKVLoRASharding(
    const std::string &module_name,
    const std::shared_ptr<Module> &module,
    std::unordered_map<std::string, LoRATensorSharding> &shardings) {

    MarkLoRAColumnParallelBSharding(
        module_name,
        module,
        CausalSelfAttention::kCAttnLayerName,
        LoRATensorSharding::kPackedQKVColumnParallelDim0,
        shardings);
}

SwiGLU helper：

void MarkPackedSwiGLULoRASharding(
    const std::string &module_name,
    const std::shared_ptr<Module> &module,
    std::unordered_map<std::string, LoRATensorSharding> &shardings) {

    bool is_swiglu = false;

    for (const auto &child : module->modules()) {
        if (dynamic_cast<SwiGLU *>(child.get())) {
            is_swiglu = true;
            break;
        }
    }

    if (!is_swiglu) {
        return;
    }

    MarkLoRAColumnParallelBSharding(
        module_name,
        module,
        MLP::kFcLayerName,
        LoRATensorSharding::kPackedSwiGLUColumnParallelDim0,
        shardings);
}

[kilinchange]

这两个 slice/restore 函数跟 QKV 的处理逻辑也比较类似，可以看下是否方便抽出通用逻辑，避免后续 fused projection 越来越多，都各自写一套重复的逻辑。

[kilinchange]

这里可能需要适配下 #206 的改动，不能再在测试里使用 ONLY_CUDA 了，应该使用 SKIP_CPU 跳过 cpu 执行，麻烦 @chen2021673 确认下。