[Levius-Fubuki]

Fixes #8913.

With BF16 AutoEP/AutoTP folding and FP32 accumulation, _update_hp_grad currently rounds the FP32 folding result through BF16 before adding it to the FP32 accumulator. Apply the existing folding correction to temporary FP32 staging and accumulate it directly. Mirror the corrected value into retained BF16 gradients so the existing correction marker and boundary reducer keep their current contract. No new persistent buffer or marker lifecycle is introduced; BF16-only accumulation is unchanged.

On upstream bf222a2, two exact BF16 inputs (1 and 1.0078125) should average to 1.00390625 in FP32; the new two-rank regression fails with absolute error 0.00390625. Device validation: 2 x NVIDIA GeForce RTX 3090 24 GB, driver 580.105.08, torch 2.12.1+cu130, CUDA 13.0, NCCL 2.29.7, Python 3.12.3.

Validation:

DEEPSPEED_RUN_H100_TESTS=1 python -m pytest --forked tests/unit/v1/moe/test_autoep_autotp_bf16_folding_parity.py --torch_ver=2.12.1+cu130 --cuda_ver=13.0 -q
13 passed

The same new engine test fails on the unmodified optimizer with explicit FP32 communication. Changed-file pre-commit, syntax and diff checks pass. Independent spec/code reviews found no remaining issue. A two-rank NCCL boundary microbenchmark with 1,048,576 elements, 10 warmups and 40 timed calls records 6 MiB → 4 MiB temporary peak allocation per rank and preserves the exact FP32 midpoint; measured medians are about 0.85–0.87 ms baseline and 0.82–0.83 ms fixed. This is an operator-boundary measurement, not end-to-end model throughput. New tests use real NCCL reductions and real BF16 ZeRO 1 TP2/DP1 training, analytic gradients and clipped-SGD master updates, both accumulation modes, GAS 1/2, and two optimizer steps. An auxiliary sharded parameter keeps norm contributions nonempty, so this change has no dependency on #8912.

The non-immediate GAS>1 correction-marker limitation already present upstream is outside this change. Later microbatches in that test use equal rank inputs to isolate this precision boundary; arbitrary unequal later gradients remain unvalidated. Full AutoEP MoE model conversion, TP2/EP2/DP2 four-rank execution, other wrappers/accelerators, production convergence and performance remain unvalidated. Existing ZeRO 1/2/3 and engine buffer paths were inspected; this PR changes only the dedicated BF16 accumulation boundary.

Duplicate check: no open PR found for this FP32 staging fix. Merged #8654 manages gradient release/correction marker cleanup; open #8823 changes expert clipping; neither preserves this FP32 folding boundary.

AI assistance (OpenAI Codex) was used for investigation, implementation and device tests. The human submitter confirmed manual diff review and test-result verification and authorized submission. Device tests above were executed on the requested server with Codex; no additional human test execution is asserted.

Local recheck on 2026-10-28, unchanged head 05e16530: Apple M3/macOS ARM64, Python 3.12.13, torch 2.14.1, CPU/Gloo. The complete folding-parity file above reports 13 passed with DEEPSPEED_RUN_H100_TESTS=1; all applicable changed-file pre-commit hooks pass. This is CPU distributed validation, not an H100 or new CUDA/NCCL run. The non-immediate GAS>1 correction-marker limitation described above remains outside this change. Keeping this PR in draft while the author's ready-PR limit is reached.

[0z5a]

Read at 05e16530.

This looks right to me. The correction marker lives on the parameter, not on the gradient tensor, so correcting a temporary fp32 copy and mirroring it back into lp.grad keeps the marker contract. The TP reduction was already done in fp32 on master: reduce_autoep_folding_gradient upcasts bf16 gradients before the all-reduce, then rounds back. So the communication volume is unchanged, and the only difference is that the fp32 accumulator now gets the unrounded value.

On CPU (gloo), test_autoep_autotp_bf16_folding_parity.py passes at this head (13 tests). With the same test file on master, all 8 new cases fail, off by 0.0039 or 0.0078 (the bf16 rounding step), and the 5 existing cases still pass.

One nit, inline.

[0z5a]

The copy-back is unconditional. It's wasted in two common cases. Under immediate_grad_update, accumulate_hp_grads_and_remove_lp sets lp.grad = None right after this call. For parameters whose folding strategy is SKIP, grad holds the unchanged value. apply_folding_correction_to_grad_buffer returns the strategy, so the copy could be skipped there, and also when the caller is about to drop lp.grad.