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:
@@ -101,6 +101,10 @@ fn main() {
|
||||
// Optimization knobs.
|
||||
let steps: usize = flag(&args, "--steps", 2000);
|
||||
let batch_size: usize = flag(&args, "--batch", 8);
|
||||
// Micro-batch gradient accumulation (Phase T16): effective batch =
|
||||
// accum_steps × batch, at one micro-batch's activation-memory cost. Default 1
|
||||
// = no accumulation (bit-identical to the pre-T16 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);
|
||||
@@ -201,6 +205,7 @@ fn main() {
|
||||
let tcfg = TrainConfig {
|
||||
seq_len,
|
||||
batch_size,
|
||||
accum_steps,
|
||||
steps,
|
||||
schedule: LrSchedule {
|
||||
max_lr,
|
||||
@@ -219,10 +224,13 @@ fn main() {
|
||||
};
|
||||
|
||||
println!(
|
||||
"training: {} steps, seq {}, batch {}, lr {:.1e}→{:.1e}, eval every {}",
|
||||
"training: {} steps, seq {}, batch {} × accum {} = effective batch {}, \
|
||||
lr {:.1e}→{:.1e}, eval every {}",
|
||||
tcfg.steps,
|
||||
tcfg.seq_len,
|
||||
tcfg.batch_size,
|
||||
tcfg.accum_steps,
|
||||
tcfg.batch_size * tcfg.accum_steps,
|
||||
tcfg.schedule.max_lr,
|
||||
tcfg.schedule.min_lr,
|
||||
tcfg.eval_every
|
||||
|
||||
Reference in New Issue
Block a user