【训练营】小模型训练支持 - #225
Open
ShenYouSOTA wants to merge 5 commits into
Open
【训练营】小模型训练支持#225ShenYouSOTA wants to merge 5 commits into
ShenYouSOTA wants to merge 5 commits into
Conversation
…MNIST demo, DDP) - Conv2d v1: FP32/NCHW-OIHW, scalar kernel/stride/padding API, Kh/Kw-independent internals; Conv2dMeta centralizes shape arithmetic; selective-save autograd - Backends: CPU direct reference + CUDA im2col + framework GEMM (gather col2im, cub bias reduction) - DataLoader shape-preserving Stack ([B,*sample_dims]); MNIST sample [1,28,28]; stride bug fix; MLP Flatten(1) adaptation - MNIST: MnistCnn + --model (default mlp unchanged); shared_ptr migration for enable_shared_from_this Modules - DDP: infini_run multiprocess, Broadcast-before-wrap, sharded eval + epoch-level AllReduce(SUM), no per-step logging collectives - Tests: kernel exact-value, gradcheck, torch golden parity, im2col/CUDA grids, dataset/dataloader, network parity dumper
Reducer::PrepareForBackward (inside DDP Forward, gradient_as_bucket_view=true) binds param.grad to the bucket view; ZeroGrad(set_to_none=true) after Forward reset that binding, so backward accumulated into a standalone grad the reducer never all-reduced (silent no-sync, cross-rank weight fork). Same order as gpt2/mixtral/llama3 demos.
…NoGradGuard Call modules via operator() (the hook entry, same as gpt2/llama3 and PRs 219/220) instead of Forward() directly. Wrap evaluation in NoGradGuard so forward-only graphs never prime grad accumulators (PR 220 acfbe87 hazard); resolves TODO(dcj) no_grad().
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
MNIST CNN 两层卷积训练支持：Conv2d / ReLU（CPU direct 基准 + CUDA im2col+GEMM，双后端互为 cross-check）、
--model cnn|mlp（默认 mlp 不变）、infini_run多进程 DDP（Broadcast-before-wrap、分片 eval + epoch 级 AllReduce SUM）。Scope
Conv2dMeta集中 shape 算术；selective-save autograd（三个独立 backward op）[B,*sample_dims]，MNIST 样本[1,28,28]）；MNIST stride bug 修复（FLOAT32 步长）ZeroGrad → Forward（DDP bucket-view 要求，与 gpt2/mixtral/llama3 对齐）、Module::operator()调用、evalNoGradGuardVerification
Step后 cross-rank 权重一致（post-epoch w0 两 rank 均为-0.001326）；单卡基线 bit 无回归USE_NCCL=ON一次编过；MNIST 官方数据细节见分支内
docs/small_model_patch.md§7。