Transformers Fixes Expert-Parallel Gradient Scaling to Prevent Double Division in Specific Configurations

Hugging Face Transformers merged PR #50292 into its main branch on November 5, fixing a gradient-scaling error in mixture-of-experts (MoE) models under specific expert-parallel configurations. The issue involves the interaction between token dispatch and FSDP2. According to the maintainers, when efsdp_size == 1, expert gradients are divided twice, leaving them smaller than expected by a factor of fsdp_size. Patch proposal

This condition is tied to the parallel topology. The official documentation defines efsdp_size = fsdp_size * tp_size / ep_size, representing the number of replicas of the full expert set. For example, with tensor parallelism disabled, fsdp_size=8, and ep_size=8, experts are distributed across eight processes and efsdp_size is one. Token dispatch lets each process hold its own batch, then uses two all-to-all operations to send tokens to the corresponding experts and retrieve the results. Expert parallelism documentation

The error occurs at a gradient-reduction boundary case. The usual path requires FSDP2 to sum gradients first and then divide by the data-parallel size. But when the expert side has only one FSDP shard, there is no reduce-scatter to perform. FSDP2 has already divided the gradients while replicating them, and the setting that forces summation causes them to be divided again. The patch forces summation only when reduce-scatter is needed, so expert gradients are scaled just once. Maintainers’ explanation of the root cause

For training teams, this could affect the relative updates to expert parameters and other modules, so it is worth investigating when comparing different parallel configurations. This is an engineering inference based on the scaling mechanism. The public proposal provides no full convergence curves or downstream quality measurements, so an eightfold reduction in gradient magnitude should not be interpreted as an eightfold drop in model quality, nor does it imply that all MoE training is affected.

Engineers can start by checking the expert dispatch method in use and the three parallelism parameters above, then compare expert gradient norms and single-step updates before and after the patch on fixed data. The official documentation also notes that the main branch documentation corresponds to the source code version, and that Trainer support for token-dispatch batch and token counting still needs follow-up work. This merge addresses a specific gradient issue; before adopting it, verify whether the installed version includes the patch and validate the full training workflow. Official usage limitations