[Precision] BF16 AutoEP folding rounds FP32 TP reduction before FP32 gradient accumulation #8913

Summary And Scope

At master fc7bc301f8301e72fed37e2caa3c1d40dae3bbb4, AutoEP+AutoTP folding in BF16_Optimizer._update_hp_grad reduces the low-precision gradient through a helper that promotes to FP32, performs the TP correction, and copies the result back to BF16. The wrapper then casts this rounded BF16 result to FP32 for accumulation.

Under an FP32 gradient-accumulation policy this introduces an avoidable BF16 rounding boundary between FP32 reduction and FP32 accumulation.

This is a precision-policy report, not a claim that every BF16-only reduction policy is incorrect or that cross-framework bitwise equality is an unconditional DeepSpeed contract.

Sources:

• https://github.com/deepspeedai/DeepSpeed/blob/fc7bc301f8301e72fed37e2caa3c1d40dae3bbb4/deepspeed/runtime/bf16_optimizer.py
• https://github.com/deepspeedai/DeepSpeed/blob/fc7bc301f8301e72fed37e2caa3c1d40dae3bbb4/deepspeed/module_inject/auto_ep_folding.py

Observed Boundary

With BF16 TP2/EP2/ETP1 folding and FP32 accumulation, input and expert gradients matched the numerical reference, while the first router accumulated gradient, FP32 master weight and Adam moments differed.

The relevant path is:

BF16 local gradient
  -> FP32 TP folding reduction/correction
  -> BF16 writeback
  -> FP32 accumulation

The intended matched FP32 policy instead preserves the FP32 reduction result directly into the accumulation buffer.

Reproduction Guidance

In an AutoEP+AutoTP folding training workload selecting BF16_Optimizer and FP32 accumulation, capture:

1. The BF16 local router gradient.
2. The helper's FP32 result before low-precision writeback.
3. The high-precision gradient after _update_hp_grad.

Use the same topology, correction rule, local gradients and DP normalization in the control. Compare the third value with FP32 accumulation of the second value, without inserting a BF16 conversion. A difference at this boundary isolates the precision loss without relying on different model forward implementations.

The full external training harness used in the investigation is not an upstream minimal reproduction attached to this issue; the source-level boundary and validation scope are stated here rather than claiming otherwise.

Proposed Fix

In _update_hp_grad, create FP32 staging from lp.grad, apply the existing folding correction to that staging tensor, and add it directly to the high-precision gradient buffer. Preserve readiness flags, correction-once bookkeeping, parameter-name/topology mapping and existing gradient-clear semantics.

Keep BF16-only policies explicit rather than silently changing their contract. Other wrappers and bucket paths need separate auditing to avoid re-rounding FP32 staging into low-precision storage.

Validation

A local unsubmitted fix using FP32 staging passed three BF16 updates on four ranks for TP2/EP2/ETP1, comparing 66 tensors per rank, including router/expert gradients, FP32 masters and Adam moments.

The accepted numerical oracle is native Megatron model/optimizer code with a harness's explicit FP32 TP-then-DP/per-parameter reducer. It is not the production Megatron distributed-optimizer entry or a claim that default collective layouts should naturally match. Production shapes, other wrappers, GAS, checkpoint compatibility and performance need additional validation.