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%）。