[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.