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 3f1e4ab.

* 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>

Author: 3 people

src/transformers/distributed/mixin.py

@@ -21,6 +21,7 @@
 
 import torch
 
+from ..core_model_loading import build_glob_alternation
 from ..dynamic_module_utils import custom_object_save
 from ..utils import is_torch_greater_or_equal, logging
 from ..utils.hub import create_and_tag_model_card
@@ -237,7 +238,9 @@ def _resolve_parallel_plans(
         Each returned plan maps one concrete module or parameter name to a style, e.g. `model.layers.3.mlp.up_proj`
         instead of `model.layers.*.mlp.up_proj`. A module that has an EP rule keeps no TP rule.
         """
-        names = {replace_layer_number_by_wildcard(n) for n, _ in chain(self.named_modules(), self.named_parameters())}
+        module_names = {replace_layer_number_by_wildcard(n) for n, _ in self.named_modules()}
+        param_names = {replace_layer_number_by_wildcard(n) for n, _ in self.named_parameters()}
+        names = module_names | param_names
         for plan_name in ("tp_plan", "ep_plan", "fsdp_plan"):
             override = getattr(distributed_config, plan_name)
             if isinstance(override, dict):
@@ -276,9 +279,17 @@ def _resolve_parallel_plans(
         # Tied embed/lm_head weights are wrapped once: fold the head's FSDP rule into the embedding's.
         fsdp_plan = _resolve_tied_embed_lm_head_plan(fsdp_plan, self)
 
+        # FSDP2 needs one group per dtype: give the modules kept in float32 their own group.
+        if fsdp_plan and (dtype_plan := self._get_dtype_plan(self.dtype)):
+            kept_in_fp32, _, _ = build_glob_alternation(list(dtype_plan))
+            fp32_modules = {name.rpartition(".")[0] for name in param_names if kept_in_fp32.search(name)}
+            fsdp_plan = dict.fromkeys(fp32_modules, "free_full_weight") | fsdp_plan
+
         tp_plan = self._expand_parallel_plan(tp_plan)
         ep_plan = self._expand_parallel_plan(ep_plan)
         fsdp_plan = self._expand_parallel_plan(fsdp_plan)
+        # FSDP2 needs a submodule sharded before the module containing it, e.g. `mlp.gate` before its layer.
+        fsdp_plan = dict(sorted(fsdp_plan.items(), key=lambda item: item[0].count("."), reverse=True))
 
         return tp_plan, ep_plan, fsdp_plan
 

src/transformers/models/ernie4_5_moe/modeling_ernie4_5_moe.py

@@ -467,6 +467,8 @@ def _init_weights(self, module):
         elif isinstance(module, Ernie4_5_MoeExperts):
             init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range)
             init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, Ernie4_5_MoeTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
 
 
 @auto_docstring

src/transformers/models/ernie4_5_moe/modular_ernie4_5_moe.py

@@ -199,6 +199,8 @@ def _init_weights(self, module):
         elif isinstance(module, Ernie4_5_MoeExperts):
             init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range)
             init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, Ernie4_5_MoeTopKRouter):
+            init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
 
 
 @auto_docstring

tests/test_fsdp_mixin.py

@@ -79,6 +79,8 @@
 # Set to None to run distributed FSDP tests for every model with a plan.
 FSDP_DISTRIBUTED_TEST_MODEL_TYPES = {
     "cohere2_moe",
+    # Keeps its router weight in float32 (`_keep_in_fp32_modules_strict`)
+    "ernie4_5_moe",
     "llama",
     "mistral",
     "mixtral",
@@ -359,6 +361,24 @@ def _test_fsdp2_save_load_impl(rank, config_class, config_dict):
             logger.debug(f"FSDP2 save/load test passed: all {len(state_dict_before)} parameters match exactly.")
 
 
+def _test_fsdp2_bf16_training_impl(rank, config_class, config_dict):
+    """Train in bfloat16 with the model sharded at load time, float32-kept parameters included."""
+    init_test_logger()
+
+    config = config_class.from_dict(config_dict)
+    distributed_config = DistributedConfig(fsdp_size=dist.get_world_size())
+
+    with _deterministic_init_model_dir(rank, config, torch.bfloat16) as init_dir:
+        _set_determinism(SEED)
+        model = AutoModelForCausalLM.from_pretrained(
+            init_dir, torch_dtype=torch.bfloat16, distributed_config=distributed_config
+        )
+        optimizer = torch.optim.AdamW(model.parameters(), lr=LR)
+        batches = _build_repeated_training_batches(config, _get_rank_device(rank), num_steps=2)
+        losses, _ = _run_training_steps(model, optimizer, batches, track_grad_norms=False)
+        assert all(torch.isfinite(torch.tensor(losses))), f"Non-finite losses: {losses}"
+
+
 def _test_fsdp2_save_load_dcp_impl(rank, config_class, config_dict):
     """Save FSDP2 model via save_pretrained(distributed_checkpoint=True), reload, compare state dicts."""
     init_test_logger()
@@ -597,7 +617,7 @@ def _test_fsdp2_expert_parallel_bf16_grads_impl(rank, config_class, config_dict)
         reference = AutoModelForCausalLM.from_pretrained(init_model_dir, torch_dtype=dtype).to(device)
         reference_loss = reference(input_ids=input_ids, labels=input_ids, use_cache=False).loss
         reference_loss.backward()
-        expected = {name: param.grad.float() for name, param in reference.named_parameters()}
+        expected = {name: param.grad.float() for name, param in reference.named_parameters() if param.grad is not None}
         del reference
 
         model = AutoModelForCausalLM.from_pretrained(
@@ -610,7 +630,11 @@ def _test_fsdp2_expert_parallel_bf16_grads_impl(rank, config_class, config_dict)
         rows = slice(dp_rank * BATCH_SIZE, (dp_rank + 1) * BATCH_SIZE)
         loss = model(input_ids=input_ids[rows], labels=input_ids[rows], use_cache=False).loss
         loss.backward()
-        actual = {name: param.grad.full_tensor().float() for name, param in model.named_parameters()}
+        actual = {
+            name: param.grad.full_tensor().float()
+            for name, param in model.named_parameters()
+            if param.grad is not None
+        }
         loss = loss.detach()
         dist.all_reduce(loss, group=dp_group)
         loss = loss / world_size
@@ -777,6 +801,11 @@ def test_fsdp2_sharding_structure(self, label):
     def test_fsdp2_save_load(self):
         self._run_fsdp2_distributed_test("test_fsdp2_save_load", _test_fsdp2_save_load_impl)
 
+    @require_torch_greater_or_equal("2.7")
+    @is_fsdp_test
+    def test_fsdp2_bf16_training(self):
+        self._run_fsdp2_distributed_test("test_fsdp2_bf16_training", _test_fsdp2_bf16_training_impl)
+
     @require_torch_greater_or_equal("2.7")
     @is_fsdp_test
     def test_fsdp2_save_load_dcp(self):