【训练营】小模型训练支持 - #220
Open
accelerator-llc wants to merge 14 commits into
Open
Conversation
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.
LinearBackwardBias read the row-major (bs, out_features) grad_output as (out_features, bs) and summed rows, which only matches the column-sum semantics for constant-valued inputs. Reduce the sample dimension per output feature instead (64-bit indexing, one block per feature, both fp32 and bf16 paths) and require the last dim to match out_features. Add numeric gradient assertions to the linear backward tests: exact values on a non-square reproduction shape, a pseudo-random pattern checked against a double-precision host reference, and the bf16 branch.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
backward-input issued as single strided-batched GEMMs).
(NaN propagation on forward and backward).
LinearBackwardBiasread 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). Fixed to reduce
over the sample dimension, with three new numeric regression tests
(non-square repro shape, random pattern, BF16 branch).
MnistCnn(Conv2d(1,16,3)→ReLU→Conv2d(16,32,3)→ReLU→Flatten→Linear);
--model cnn|mlpkeeps the existing MLP path unchanged. DDP is enabled automatically by the
environment variables that
infini_runinjects (one GPU per process viaLOCAL_RANK, rank-0 parameter broadcast at construction, per-step loss
AllReduce for global logging); the single-process path is untouched.
and removal of redundant im2col scratch initialization; single-GPU MNIST CNN
throughput 1.98× (16582 → 32802 samples/s on a 4090D).
empty batch, kernel == spatial size, non-square regression shapes, torch
golden), ReLU cases, Linear backward numeric assertions; full suite
CPU 272 / CUDA 11 with no regressions introduced by this PR (one
environment-sensitive pre-existing upstream test failure is described in the
project report).
example/andtests/use glob-based collection, so the new sources are picked up byre-running cmake.
--modelusage and the multi-processinfini_runlaunch example.
Verification
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).
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%).
accuracy sweep in the project report, submitted separately).
@kilinchange please review, thank you!