fix: update on complete gradient accumulation boundaries

[Liuchenbing-2026]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

With gradient accumulation set to 4, Transformers backward increments the micro-step counter before checking the boundary, but the shared predicate tests (cur_step - 1) % 4. The first optimizer update therefore occurs at micro-step 5 rather than 4, and eight micro-batches produce only one update rather than two.

Check the completed positive step count directly. Add a numerical regression calling the real Transformers backward, optimizer step and zero-grad core methods, with a small local wrapper fixture and real PyTorch autograd/AdamW. It checks accumulation factors 1, 2 and 4 against complete-batch updates, continuation of the counter and unchanged single-step behavior. RPC dispatch and distributed training are outside this test.

AI assistance was used for implementation and testing.

Experiment results

Compared the same tests on main 9953ea1ebbbdbe59c3831702157d3c147ca0ee26 and the fixed implementation:

Source	Result
Main	3 failed, 2 passed
Fixed	5 passed, no skips

The baseline failures are incorrect update boundaries for accumulation 2/4 and continuation to the next complete boundary. Fixed parameters match the full-batch reference at rtol=atol=1e-12 in float64.

PYTHONPATH=src PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 OMP_NUM_THREADS=2 python -m pytest \
  tests/transformers/test_accumulation_boundaries.py -q --confcutdir=tests/transformers
pre-commit run --all-files

All pre-commit hooks passed; the test file also passed manual isort/yapf/flake8 checks because repository hooks exclude tests. Environment: Linux/aarch64, Python 3.12.13, Torch 2.10.0, CPU tensors only. This rerun does not claim NPU throughput, full-model quality or multi-rank acceptance.

[Liuchenbing-2026]

@tastelikefeet could you please review this focused gradient-accumulation fix? The same five regressions fail three cases on current main and pass all five on this head. They call the actual Transformers backward/step/zero_grad core methods with real CPU autograd and AdamW, checking accumulation 1/2/4 against complete batches at 1e-12 tolerance. The final published files were verified byte-for-byte against the rerun snapshot. Commands and validation limits are in the description. Thanks!