【训练营】小模型训练支持 by accelerator-llc · Pull Request #220 · InfiniTensor/InfiniTrain

accelerator-llc · 2026-09-09T08:22:34Z

Summary

This PR adds small-model (CNN) training support to InfiniTrain: Conv2d / ReLU
operators for both CPU and CUDA backends (kernel + autograd + nn::Module, three
coherent layers), an MNIST CNN training demo, and distributed data-parallel
training wired through the framework's existing DDP components. The operator
interfaces follow PyTorch conventions; the implementation scope covers what the
CNN demo needs (Conv2d with stride=1 / padding=0, leaving room for future
extensions).

What is included

Conv2d: im2col + GEMM (Eigen on CPU, cuBLAS on CUDA; forward and
backward-input issued as single strided-batched GEMMs).
ReLU: CPU/CUDA elementwise kernels, bit-faithful to PyTorch semantics
(NaN propagation on forward and backward).
Linear bias-gradient regression tests: a CUDA LinearBackwardBias defect
was found independently while developing this branch (LinearBackwardBias read
the row-major (bs, out_features) gradient transposed and reduced over the
wrong dimension; for constant-valued inputs the row sums happen to equal the
column sums, and the existing test asserted only the output count, so the bug
was never exposed). Master has since fixed the same defect (a684567), so this
PR carries no kernel change for it; what it contributes is three numeric
regression tests (non-square repro shape, random pattern, BF16 branch) pinning
the sample-dimension reduction.
MNIST CNN demo: MnistCnn
(Conv2d(1,16,3)→ReLU→Conv2d(16,32,3)→ReLU→Flatten→Linear); --model cnn|mlp
keeps the existing MLP path. DDP is enabled automatically by the environment
variables that infini_run injects (one GPU per process via LOCAL_RANK, rank-0
parameter broadcast at construction, per-step loss AllReduce for global
logging). Test metrics are reported after every epoch (evaluation runs under
no_grad), and the progress line prints an explicit step counter.
Performance: conv GEMM scheduling fix (per-image loop → strided-batched)
and removal of redundant im2col scratch initialization; single-GPU MNIST CNN
throughput 1.98× (16582 → 32802 samples/s on a 4090D).
Tests: Conv2d 15 cases (device-parameterized single files, covering
empty batch, kernel == spatial size, non-square regression shapes, torch
golden), ReLU cases, Linear backward numeric assertions; full suite
CPU 276 / CUDA 10, all passing on the current tree.
CMake integration: no CMakeLists changes are needed — the CNN network is
a header-only addition to the existing mnist target, and tests/ collects
its sources via glob, so re-running cmake picks up the new test files.
README: documents --model usage and the multi-process infini_run
launch example.

Verification

Network-level numerical alignment against PyTorch (same weights, inputs and
hyper-parameters): forward logits 5.2e-8, loss 1.3e-9, gradients 4.3e-8,
one optimizer step 1.5e-8, 10-step trajectory 2.4e-7 (thresholds 6e-5 / 2e-6).
DDP equivalence: single-process bs=128 gradients vs two-process bs=64 with
AllReduce averaging — overall max_abs 1.863e-8; two-process end-to-end 3-epoch
run shows bitwise-identical loss curves and identical accuracy (97.57%).
Distributed runs were taken on a 2×4090D machine on this branch's pre-rebase
tree. Rebasing onto current master picked up master's reworked sharded loader
(per-rank batch count derived from the global batch size, trailing partial
batch dropped per epoch), so the demo's old divisibility guard was removed —
per-rank step equality is now a construction-time invariant of the loader
itself. The distributed path has not been re-run on 2 GPUs since the rebase
and the evaluation-timing change.
End-to-end: CNN CPU 3ep 97.79% / CUDA 98.06% / 10ep 98.23% (per-epoch
accuracy sweep in the project report, submitted separately).

@kilinchange please review, thank you!

Implemented inline (im2col + GEMM) rather than exposing a public operator since the MNIST CNN is the only consumer. Scoped to square kernel, stride 1, padding 0, optional bias, and FP32 to meet the project need without covering the full Conv2d parameter space. Uses PyTorch cross-correlation semantics (no kernel flip) to keep numerical alignment with torch.nn.functional.conv2d.

Add autograd::ReLU function (forward with clamp_min semantics, backward with threshold_backward semantics) and nn::ReLU module, following the existing activation pattern. CPU and CUDA elementwise kernels preserve the exact NaN and negative-zero behavior of PyTorch, verified bit-exact against PyTorch (fixed seed). Includes forward, backward, end-to-end training chain and Flatten gradient tests.

The MNIST dataset loader computed the per-sample byte stride from the on-disk UINT8 element size, but the image tensor is normalized to FLOAT32 in the constructor. Sample views after the first therefore read from a wrong, overlapping byte range, so most training images are misaligned mixes of real pixels and the network cannot learn. Use the FLOAT32 element size when computing the stride. Validated: MNIST MLP and CNN training reach >=95% test accuracy.

Add MnistCnn (Conv2d(1,16,3)->ReLU->Conv2d(16,32,3)->ReLU->Flatten->Linear(18432,10)) alongside the existing MLP. A --model=mlp|cnn flag (default cnn) selects the network; the CNN restores the (N,1,28,28) spatial layout from the flattened [N,784] input with Tensor::View. Hold the loss in a shared_ptr to match the other examples. Tune the default --num_epoch/--lr (3 / 0.1) so the default CNN demo reaches ~97.8% test accuracy. Training converges after the MNIST dataset loader's per-sample stride is fixed to use the normalized FLOAT32 element size (older stride misaligned the image views).

Launch with infini_run to shard MNIST batches across processes: each process selects its GPU from LOCAL_RANK, wraps the network with DistributedDataParallel, broadcasts rank-0 parameters before training, and averages the per-step loss for global logging. Single-process behavior is unchanged; document --model and the multi-process launch in the README.

Cover non-square grad_output shapes whose per-element values separate the sample-dimension sum from a row sum, and the bf16 branch whose grad_bias is promoted to fp32.

Route the training and test loops through Module::operator() instead of calling Forward directly, matching the gpt2 example and the hook design doc: operator() is the entry that runs module hooks. No hook consumers exist yet, so outputs are unchanged.

Fold the *_cuda_* conv test files into their CPU counterparts so each test definition instantiates on every available device, matching the test infrastructure design and leaving no ONLY_ macros in the merged files. The CUDA-only empty-batch cases are device-independent and carry over, and the negative-value check in the extreme-values forward case now also runs on CUDA.

Forward and backward-input dispatched one cuBLAS call per image with batch_count=1, while the Gemm interface already exposes a strided-batched path that Matmul uses. Both now issue a single strided-batched GEMM with the weight shared via stride 0, and backward-input runs one col2im launch for the whole batch. The CPU kernels zero-initialized their im2col scratch right before overwriting every element; allocate it uninitialized instead. Single-GPU MNIST CNN throughput: 16582 -> 32802 samples/s per epoch.

The train pair duplicated the same end-to-end SGD step per device, leaving a *_cuda_* file behind after the forward/backward merge. Fold them into one parameterized body that inspects loss, gradients and parameters through host copies, with parameter snapshots taken into freshly allocated host buffers so they survive the in-place update.

Drop the unused Gemm declaration header from the CUDA conv kernel, add the missing <cstdint> to the CPU ReLU kernel, and remove the copy-and- keep-in-sync note above the CUDA ReLU block-size helper.

Run the full test set after each epoch instead of only once at the end of training, so a single run shows the test loss falling and the test accuracy rising. The evaluation runs under no_grad: a forward-only graph leaves each parameter's grad accumulator with a dependency count that the next backward pass cannot satisfy, silently stopping its gradient accumulation. The training progress line now carries an explicit 1-based [step k/total] field next to the samples consumed so far.

accelerator-llc added 14 commits September 11, 2026 22:01

fix: use shared_ptr for MNIST network to avoid bad_weak_ptr

c1c9028

feat(conv): add CUDA Conv2d kernels with autograd and unit tests

51a3490

fix(conv): guard Conv2d CPU kernels against non-FP32 inputs

57e95fe

test(linear): add bias-gradient regression tests

10923eb

Cover non-square grad_output shapes whose per-element values separate the sample-dimension sum from a row sum, and the bf16 branch whose grad_bias is promoted to fp32.

chore(kernels): fix include hygiene in conv and relu

badbb9f

Drop the unused Gemm declaration header from the CUDA conv kernel, add the missing <cstdint> to the CPU ReLU kernel, and remove the copy-and- keep-in-sync note above the CUDA ReLU block-size helper.

accelerator-llc force-pushed the feat/small-model-cnn branch from 502f2bd to badbb9f Compare September 11, 2026 15:24

Navigation Menu

Uh oh!

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

【训练营】小模型训练支持#220
accelerator-llc wants to merge 15 commits into
InfiniTensor:masterfrom
accelerator-llc:feat/small-model-cnn

accelerator-llc commented Sep 9, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

accelerator-llc commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What is included

Verification

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

accelerator-llc commented Sep 9, 2026 •

edited

Loading