train+ddp: micro-batch gradient accumulation (--accum-steps)

Accumulate grads over N micro-batches, then one AdamW step + zero_grad,
for an effective batch of N×micro at one micro-batch's activation cost.
Each micro-loss is scaled by 1/N before backward (the tape SUM-accumulates
the scaled grads) so the boundary grad equals a single step over an N×
batch. accum==1 skips the scale → bit-identical to the pre-T16 path.

DDP: the cross-rank all-reduce fires ONLY at the accumulation boundary
(intermediate micro-steps are local-only, no NCCL); the /world average is
orthogonal to the per-micro 1/N, so the boundary grad is the effective
global-batch mean. New --accum-steps flag in both train binaries; effective
batch is printed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-06-17 23:45:33 +08:00
parent d01fec6639
commit 7a03b0054a
4 changed files with 101 additions and 43 deletions

View File

@@ -74,6 +74,10 @@ fn main() {
// Optimization knobs (mirror bin/train).
let steps: usize = flag(&args, "--steps", 100);
let batch: usize = flag(&args, "--batch", 16);
// Micro-batch gradient accumulation (Phase T16): effective global batch =
// accum_steps × batch, all-reducing only at the accumulation boundary. Default
// 1 = no accumulation (bit-identical to the pre-T16 DDP path).
let accum_steps: usize = flag(&args, "--accum-steps", 1).max(1);
let seq_len: usize = flag(&args, "--seq", 64);
let max_lr: f32 = flag(&args, "--max-lr", 3e-3);
let min_lr: f32 = flag(&args, "--min-lr", max_lr * 0.1);
@@ -147,6 +151,7 @@ fn main() {
let dcfg = DdpConfig {
seq_len,
batch_size: batch,
accum_steps,
steps,
schedule: LrSchedule {
max_lr,
@@ -164,8 +169,9 @@ fn main() {
};
println!(
"training: {steps} steps, seq {seq_len}, global batch {batch}, lr {max_lr:.1e}{min_lr:.1e}, \
eval every {eval_every}"
"training: {steps} steps, seq {seq_len}, global batch {batch} × accum {accum_steps} = \
effective global batch {}, lr {max_lr:.1e}{min_lr:.1e}, eval every {eval_every}",
batch * accum_steps
);
if bf16 {