Add gradient accumulation for training (single-GPU/DDP/FSDP2)

[eldarkurtic]

Summary

Adds a --gradient-accumulation-steps / TrainerConfig.gradient_accumulation_steps knob (default 1) that accumulates gradients over N microbatches before each optimizer step, growing the effective batch size (per-step batch × N) without extra per-microbatch memory. Matches the convention used by peer speculative-decoding frameworks (DeepSpec, TorchSpec, SpecForge): loss scaled by 1/N, optimizer + LR scheduler stepped once per accumulation window, counting in optimizer (macro) steps.

Behavior

• Macro-step counting: global_step, --log-freq, --checkpoint-freq (when < 1), and the LR scheduler total all advance once per accumulation window (an accum=N run does 1/N the scheduler steps per epoch).
• Drop remainder: trailing microbatches that don't fill a full window are dropped each epoch. _optimizer_steps_per_epoch is the single source of truth for this policy (used by the scheduler math, the __init__ guard, and the loop boundary).
• _maybe_no_sync: on DDP, skips the gradient all-reduce on non-boundary microbatches and syncs the accumulated gradient on the boundary. Single-GPU (raw module) and FSDP2 (fully_shard, DTensor params) fall back to a no-op context and stay correct because gradients accumulate additively into .grad.
• Guardrail: Trainer raises early if N exceeds the batches-per-epoch (which would otherwise run zero optimizer steps).
• accum=1 is behavior-identical to the previous one-step-per-batch loop (regression-safe).

Verification

Numerically confirmed equivalent to the mean of per-microbatch gradients (accumulated grad vs. reference) on real hardware:

Path	Launch	cadence (accum=1 / 4)	max abs grad diff
Single-GPU	python	13 / 3	3.7e-9
DDP	torchrun ×8	13 / 3	9.3e-9
FSDP2 (fully_shard)	torchrun ×4	13 / 3	2.3e-9

All three exercise the real Trainer.train_epoch → _accumulate_and_step → _maybe_no_sync path. DDP/FSDP2 use different data per rank, so cross-rank averaging composes with accumulation.

Unit tests added: scheduler-step math (accum, remainder, accum=1 regression), numerical grad-equivalence on a toy model, _maybe_no_sync behavior (single-process DDP), the accum-too-large guard, and CLI parsing/rejection.

Test plan

• python -m pytest tests/unit/train/test_gradient_accumulation.py tests/unit/train/test_trainer_scheduler.py tests/unit/train/test_cli_args.py — passing.
• make quality (ruff + format + mypy) clean on changed files.

🤖 Generated with Claude Code

[eldarkurtic]

Superseded by the upstream PR vllm-project#859.