【训练营】小模型训练支持 by gavin-richie · Pull Request #232 · InfiniTensor/InfiniTrain

gavin-richie · 2026-09-20T13:52:27Z

【训练营】小模型训练支持 — PR 描述

标题:【训练营】小模型训练支持
分支:feat/small-model-support → master(14 commits)
预览链接:https://github.com/gavin-richie/InfiniTrain/pull/new/feat/small-model-support

概述

扩展 InfiniTrain 的小模型训练能力:补齐 CNN 所需的 Module / 算子(从零实现 Conv2d 前向/反向,
CPU 与 CUDA 双后端),新增 ReLU 算子与 Flatten 模块,并在 MNIST 上完成端到端 CNN 训练 Demo,
支持 单卡 CPU / CUDA 训练与 DDP 多卡分布式训练,提供完整单元测试、SwanLab 训练可视化,
以及与 PyTorch 的逐点数值对齐验证。

最终训练流程:MNIST Dataset → DataLoader → CNN → CrossEntropyLoss → Backward → Optimizer → Accuracy

主要变更

1. Conv2d 算子(核心)

autograd::Conv2d Function:支持 stride / padding 与 input / weight / bias 三路梯度,
依据 needs_input_grad 裁剪保存张量;接口对齐 torch.nn.Conv2d 的最小集合
(dilation / groups 按项目要求不在范围内,非 fp32 显式报错)。
CPU kernel:im2col 展开 + 行主序 Eigen GEMM;grad_input 走 col2im 散射累加。
CUDA kernel:im2col kernel + 复用框架 cuBLAS Gemm 封装(strided-batch,stride_b=0
广播共享权重)+ col2im(atomicAdd)+ 按 channel 的 bias 规约。
nn::Conv2d 模块:默认初始化与 PyTorch 一致(KaimingUniform(a=√5) + bias U(±1/√fan_in))。

2. ReLU / Flatten

autograd::Relu(CPU / CUDA fp32/bf16)+ nn::ReLU + nn::functional::Relu。
nn::Flatten(start_dim=1):基于 Tensor::Flatten(View/NoOp,零拷贝,梯度透传)。

3. MNIST CNN Demo(`example/mnist`)

新增 --model mlp|cnn(默认 mlp,保持既有行为):
Conv2d(1,16,3) → ReLU → Conv2d(16,32,3) → ReLU → Flatten → Linear(18432,10)。
--init_weights:加载 InfiniTrain Checkpoint 格式的初始权重(PyTorch 导出,用于对齐)。
--metrics_file:JSON Lines 指标输出(逐 step train loss + 逐 epoch train/test loss 与
accuracy),配合 scripts/swanlab_upload.py 上传 SwanLab 做训练可视化。
每个 epoch 输出 test loss / test accuracy(评估在 NoGradGuard 下);日志仅 rank 0。

4. DDP 分布式训练

Demo 接入框架既有 DDP 栈:InitAllEnv(读取 infini_run 注入的 RANK/WORLD_SIZE)→
每 rank 绑定各自 GPU → DP 进程组 → DistributedDataParallel(zero_stage=0)(桶化梯度
AllReduce,Backward() 零改动)→ DistributedDataLoader 数据分片 → 启动时 rank0 参数广播
→ 训练 loss 跨 rank AllReduce 平均;仅 rank0 输出日志与指标。
按 rank 下限封顶每 epoch 步数:batch 数不整除 world size 时,DistributedDataLoader 的交错
分片会让一个 rank 多跑一次 collective 导致 NCCL 死锁,故丢弃尾部不完整的全局 batch。
启动:./build/infini_run --nnodes=1 --nproc_per_node=2 ./build/mnist --model cnn --device cuda ...

5. 修复的两个既有缺陷(对齐排查中发现)

example/mnist/dataset.cc:图像转 float32 后样本视图步长仍按 uint8 计算(784B vs 3136B),
图像与标签错位,训练精度停滞在随机水平(约 10%)→ 转换后按 float32 重算步长。
infini_train/src/kernels/cuda/linear.cu LinearBackwardBias:ReduceColumnsKernel 按转置
索引读取梯度,bs=1 时恰好正确、bs>1 时 bias 梯度错误 → 修正归约轴,并在
test_autograd_linear_backward.cc 增加 bias 梯度数值断言防止回归。

6. CMake / 文档

CUDA_ARCHITECTURES 增加 sm_86(RTX 3090)与 sm_89(RTX 4060)。
README:MNIST 的 --model / --init_weights / --metrics_file 用法与 DDP 启动命令。

单元测试(设备参数化,CPU / CUDA 各执行一遍)

测试文件	覆盖
`test_autograd_conv2d_forward.cc` / `..._backward.cc`	stride/padding/no-bias 组合的前向与三路梯度,期望值由 PyTorch 生成后硬编码(容差 1e-4,实测 < 1e-5)
`test_autograd_relu.cc`	ReLU 前向/反向
`test_module_conv2d.cc`	参数命名/形状/初始化、Flatten、Conv→ReLU→Flatten→Linear→CE→SGD 全链路一步训练
`test_autograd_linear_backward.cc`(增强)	Linear bias 梯度数值断言(修复回归防护)

回归:CPU 全量 251 项 100% 通过;CUDA 套件除 master 既有缺陷
AutogradElementwiseBackwardTest.ExpBackward(Release 构建段错误,已在 master 工作树复现确认,
与本次改动无关)外全部通过。

结果

端到端训练(MNIST,CNN:Conv2d(1,16,3)→ReLU→Conv2d(16,32,3)→ReLU→Flatten→Linear)

配置	epoch	test loss	test accuracy
CUDA(4060),bs64,lr0.05	3	0.064	97.79%
CPU,bs64,lr0.05	1	0.120	96.30%
DDP 2×3090,bs64/rank,lr0.1	3	0.086	97.08%(236k samples/s)

Loss 随训练明显下降、Accuracy 逐步提升(SwanLab 曲线截图见项目报告)。

与 PyTorch 的数值对齐

方法:PyTorch 固定种子初始化后,将权重导出为 InfiniTrain Checkpoint 二进制格式;InfiniTrain 以
--init_weights 加载;双方相同顺序 batch、相同超参。对齐阈值约定 |Δloss| ≤ 1e-3:

对比项	实测
Forward logits(batch 0,20 值)	逐元素一致(< 1e-6)
Loss(batch 0)	2.300651 vs 2.300651(CNN);2.308022 vs 2.308022(MLP)
梯度范数(step 0,CNN 6 参数)	全部一致(最大差 2.7e-6,cuBLAS 归约顺序)
一次 Optimizer Step 后参数范数	全部一致
多步 loss 曲线	CNN 200 步 max \|Δ\| = 2.0e-6;MLP 938 步 max \|Δ\| = 1.0e-6
DDP 梯度等价性	DDP(2×bs32) 与单卡 bs64 的 938 步轨迹前 40 步逐点一致,全程 max \|Δ\| = 2.6e-5
端到端精度(CNN 1 epoch,lr=0.05)	96.30% vs 96.30%

复现

git submodule update --init
cmake -B build -DUSE_CUDA=ON -DBUILD_TEST=ON -DUSE_NCCL=OFF -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_POLICY_VERSION_MINIMUM=3.5
cmake --build build -j$(nproc) --target mnist infini_run

./scripts/assets/prepare-infinitrain-assets.sh mnist        # 数据准备
./build/mnist --model cnn --device cuda --dataset data/mnist --num_epoch 3 --lr 0.05

# DDP 双卡
./build/infini_run --nnodes=1 --nproc_per_node=2 ./build/mnist \
  --model cnn --device cuda --dataset data/mnist --num_epoch 3 --lr 0.1

# 单元测试
ctest -L cpu -j4

详细复现日志、对齐数据与排障记录见随报告提交的《小模型训练支持-详细结果日志》。

Add the Relu autograd function with forward/backward kernels for CPU and CUDA (fp32/bf16 via the elementwise dispatch helpers), a functional wrapper, and an nn::ReLU module, mirroring the existing Sigmoid operator.

Add the Conv2d autograd function supporting symmetric stride and padding with input, weight, and bias gradients, plus an nn::Conv2d module using the same default initialization as torch.nn.Conv2d (kaiming_uniform with a=sqrt(5)). CPU: im2col unfolding with row-major Eigen GEMMs and a col2im scatter for grad_input. CUDA: an im2col kernel, strided-batched cuBLAS GEMMs through the existing Gemm wrapper (weight broadcast via stride_b=0), an atomic col2im kernel, and per-channel bias reductions. Kernels cover the fp32 path required by the MNIST CNN training use case and reject other dtypes explicitly.

image_size_in_bytes_ was computed in the member initializer list against the original uint8 IDX payload, but the constructor then replaces the image tensor with a float32 copy. Every sample view therefore advanced by 784 bytes instead of 3136, so images were misaligned with their labels and training silently stagnated around chance accuracy. Recompute the per-sample stride once the conversion to float32 is done.

Add a --model flag (mlp|cnn, defaulting to the existing MLP) with a CNN classifier following Conv2d(1,16,3) -> ReLU -> Conv2d(16,32,3) -> ReLU -> Flatten -> Linear(18432, 10). The evaluation loop now runs under no_grad, per-step losses are reported, and the final test loss and accuracy are logged explicitly. A new --init_weights flag loads an InfiniTrain checkpoint as initial weights, enabling PyTorch-alignment runs.

Device-parameterized gtest coverage with PyTorch-derived reference vectors: ReLU elementwise forward/backward, Conv2d forward with and without bias across stride/padding combinations, Conv2d input, weight, and bias gradients, and module-level parameter, shape, and conv-relu-flatten-linear-cross-entropy training integration checks.

ReduceColumnsKernel summed along the wrong axis: it indexed the (bs, out_features) grad_output as if it were (out_features, bs), which happens to be correct only for bs == 1. Bias gradients on CUDA were therefore transposed sums and diverged from PyTorch for larger batches. The kernel now reduces rows within each column block for a row-major (num_rows, num_cols) input, and LinearBackwardBias passes (bs, out_features). Extend the linear backward test to check the bias gradient values, which previously only asserted the result size.

Add a --metrics_file flag that appends train and test metrics as JSON lines for visualization: a train_step line every 10 steps and an epoch_end line with train loss, test loss, and test accuracy after each epoch. Evaluation is extracted into an Evaluate() helper and now runs after every epoch under no_grad.

Upload the JSON lines written by the mnist demo's --metrics_file flag to SwanLab via the official swanlab package. The API key is read from the SWANLAB_API_KEY environment variable and is never stored in the repository.

Wire the mnist demo into the existing DDP stack: initialize the parallel environment from the RANK/LOCAL_RANK/WORLD_SIZE variables set by tools/infini_run, place each rank on its own CUDA device, create the data- parallel process group, wrap the model with DistributedDataParallel (zero_stage 0, bucketed gradient all-reduce), and shard the training data with DistributedDataLoader. Per-rank training losses are averaged with an AllReduce for logging, parameters are broadcast from the global-rank-0 process at startup, and rank-0-only logging/metrics keep single-process behavior unchanged when no parallel environment is set. The per-epoch step count is capped at the per-rank floor so every rank runs the same number of collectives; the strided DistributedDataLoader would otherwise hand one rank an extra batch and deadlock NCCL on the last all-reduce. Single-process 2-GPU verification (2x RTX 3090, CUDA 12.9/NCCL cu12): DDP with 2 ranks x batch 32 reproduces the single-process batch-64 loss trajectory exactly (max |diff| 2.6e-5 over 938 steps), and a 3-epoch DDP run reaches 97.08% test accuracy at ~236k samples/s.

gavin-richie added 14 commits September 7, 2026 06:37

feat: add ReLU activation operator and module

76e6c98

Add the Relu autograd function with forward/backward kernels for CPU and CUDA (fp32/bf16 via the elementwise dispatch helpers), a functional wrapper, and an nn::ReLU module, mirroring the existing Sigmoid operator.

feat: add nn::Flatten module

5f4742f

build: add sm_89 to the CUDA architecture list

5805794

docs: document the MNIST CNN demo options

16c92e7

chore: add swanlab metrics upload script

76c9600

Upload the JSON lines written by the mnist demo's --metrics_file flag to SwanLab via the official swanlab package. The API key is read from the SWANLAB_API_KEY environment variable and is never stored in the repository.

build: add sm_86 to the CUDA architecture list

15a6265

docs: document mnist DDP launch

a8674d7

Navigation Menu

【训练营】小模型训练支持 - #232

【训练营】小模型训练支持#232
gavin-richie wants to merge 14 commits into
InfiniTensor:masterfrom
gavin-richie:feat/small-model-support

gavin-richie commented Sep 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

gavin-richie commented Sep 20, 2026

【训练营】小模型训练支持 — PR 描述

概述

主要变更

1. Conv2d 算子(核心)

2. ReLU / Flatten

3. MNIST CNN Demo(example/mnist)

4. DDP 分布式训练

5. 修复的两个既有缺陷(对齐排查中发现)

6. CMake / 文档

单元测试(设备参数化,CPU / CUDA 各执行一遍)

结果

端到端训练(MNIST,CNN:Conv2d(1,16,3)→ReLU→Conv2d(16,32,3)→ReLU→Flatten→Linear)

与 PyTorch 的数值对齐

复现

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

3. MNIST CNN Demo(`example/mnist`)