transformers

e82b00dd - Shard the parameters kept in float32 in their own FSDP2 group (#49501)

Shard the parameters kept in float32 in their own FSDP2 group (#49501) * FSDP2: shard the parameters kept in float32 in their own group * Test a bf16 training step under FSDP2, and run the FSDP tests on ERNIE-4.5-MoE * build and expand once the tp/ep_plan * unify * clean fsdp * reorganize to avoid breaking changes * better distinction * revert * add tp plan to deepseek v4 (but attention stays replicated) * Revert "add tp plan to deepseek v4 (but attention stays replicated)" This reverts commit 3f1e4ab990df1fa64cb577695f6602e07610a5ab. * Fix the merge of main: restore moved imports and the empty tp_plan override * cleaner mixin * cleaner * deprecation cycle * cleaner * revert tp plan * Pass the expanded plan in the bare-module colwise_gather_output test * retrigger to verify fetcher * the branch is now pointing to main and not to unify-resolve-plan * Refactor expert parallelism registration and improve documentation - Consolidated the registration of parallel styles into the `ParallelInterface` class for better organization. - Updated the `expert_parallel.py` and `tensor_parallel.py` files to streamline the registration process of expert parallel strategies. - Enhanced documentation for clarity regarding the usage of tensor parallel strategies and their registration. - Removed redundant code related to `MoeExpertsParallel` and `MoeIdentityParallel` classes, which are now properly registered in the new structure. * comment * some models missed the tests EP mixin * fix copyright * move replace_layer_number_by_wildcard and _get_parameter_plan to parallel_utils * Group the float32-kept parameters while resolving the FSDP plan * Find the float32 modules from the names `_resolve_parallel_plans` already collects, and shard submodules first The modules kept in float32 become wildcard rules that `_expand_parallel_plan` expands with the other rules, so the model is not walked again. Sorting the resolved plan by depth puts every submodule before the module containing it, for any nested entry, so the float32 entries no longer need an ordering of their own. A module already in the plan keeps its style. * Initialize the ERNIE-4.5-MoE router, and skip parameters without a gradient in the EP gradient test Enabling the ERNIE-4.5-MoE FSDP tests also runs its two expert-parallel tests, which need 4 ranks: - `test_fsdp2_expert_parallel_2d_vs_ddp`: `_init_weights` left the router at zero, so every expert ties and, after one Adam step, the routing depends on rounding noise. Initialize it like ERNIE-4.5-VL-MoE and Qwen3-MoE do. - `test_fsdp2_expert_parallel_bf16_grads`: `e_score_correction_bias` only feeds the top-k selection, so it gets no gradient. * Reword the comment on the float32 FSDP2 groups * Swap the words in the float32 FSDP2 group comment --------- Co-authored-by: 3outeille <ferdinand.mom@epita.fr> Co-authored-by: Ferdinand Mom <47445085+3outeille@users.noreply.github.com> Co-authored-by: Tarek Ziade <tarek@ziade.org>