transformers

b3187c7f - [`distributed`] Divide the expert gradients once when `efsdp_size == 1` (#50292)

[`distributed`] Divide the expert gradients once when `efsdp_size == 1` (#50292) * Divide the expert gradients once when efsdp_size == 1 With token dispatch, the experts are FSDP-sharded over efsdp_size ranks with a gradient divide factor of fsdp_size and a forced sum reduction. When efsdp_size == 1 (for example fsdp_size == ep_size without TP), FSDP2 has no reduce-scatter to run: it divides the gradient while copying it, then the forced sum divides it again. Expert gradients came out fsdp_size times too small. Only force the sum when there is a reduce-scatter. * Shorten the efsdp_size == 1 comment * Move the efsdp_size == 1 comment under the if * Simplify the bf16 expert grads test docstrings