Fix FSDP2 handling for tied embeddings

modelopt/torch/opt/plugins/transformers.py

@@ -64,6 +64,44 @@ def is_liger_available():
     return True
 
 
+def _fully_shard_tied_embeddings(model, accelerator):
+    """Put tied input/output embeddings in the same FSDP2 parameter group.
+
+    Accelerate otherwise visits the two owner modules independently. PyTorch rejects
+    the second ``fully_shard`` call because their shared weight is already managed by
+    the first group.
+    """
+    if not getattr(accelerator, "is_fsdp2", False):
+        return
+
+    input_embedding = getattr(model, "get_input_embeddings", lambda: None)()
+    output_embedding = getattr(model, "get_output_embeddings", lambda: None)()
+    if (
+        input_embedding is None
+        or output_embedding is None
+        or input_embedding is output_embedding
+        or getattr(input_embedding, "weight", None) is not getattr(output_embedding, "weight", None)
+    ):
+        return
+
+    from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard
+
+    if isinstance(input_embedding, FSDPModule) or isinstance(output_embedding, FSDPModule):
+        return
+
+    fsdp_plugin = accelerator.state.fsdp_plugin
+    mesh = getattr(accelerator, "torch_device_mesh", None)
+    fully_shard(
+        [input_embedding, output_embedding],
+        reshard_after_forward=fsdp_plugin.reshard_after_forward,
+        offload_policy=fsdp_plugin.cpu_offload,
+        mp_policy=fsdp_plugin.mixed_precision_policy or MixedPrecisionPolicy(),
+        mesh=(
+            mesh[tuple(accelerator.parallelism_config.fsdp_dim_names)] if mesh is not None else None
+        ),
+    )
+
+
 @contextmanager
 def _undo_torch_init_override_by_transformers():
     if not hasattr(tf_modeling_utils, "TORCH_INIT_FUNCTIONS"):
@@ -523,12 +561,18 @@ def _prepare_model(self, model):
         trainable_param_groups; in that case the caller is responsible for gathering
         ``zero.Init``-partitioned params around forward passes.
         """
+        _fully_shard_tied_embeddings(model, self.accelerator)
         if self.is_deepspeed_enabled and not any(p.requires_grad for p in model.parameters()):
             return self.accelerator.prepare_model(model, evaluation_mode=True)
         dummy_optimizer = torch.optim.SGD([next(model.parameters())], lr=0.0)
         model, _ = self.accelerator.prepare(model, dummy_optimizer)
         return model
 
+    def train(self, *args, **kwargs):
+        """Prepare tied embeddings before Trainer applies FSDP2 auto wrapping."""
+        _fully_shard_tied_embeddings(self.model, self.accelerator)
+        return super().train(*args, **kwargs)
+
     def training_step(self, *args, **kwargs):
         """Run gc.collect() before the training step if manual_gc is enabled."""
         if self.trainer_args.manual_gc:

tests/unit/torch/opt/plugins/test_transformers_fsdp.py

@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for Hugging Face Trainer FSDP integration."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+pytest.importorskip("transformers")
+
+from modelopt.torch.opt.plugins.transformers import _fully_shard_tied_embeddings
+
+
+class _TiedModel(nn.Module):
+    def __init__(self, tied=True):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(16, 8)
+        self.lm_head = nn.Linear(8, 16, bias=False)
+        if tied:
+            self.lm_head.weight = self.embed_tokens.weight
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+
+def _accelerator(is_fsdp2=True):
+    plugin = SimpleNamespace(
+        reshard_after_forward=True,
+        cpu_offload=None,
+        mixed_precision_policy=None,
+    )
+    return SimpleNamespace(
+        is_fsdp2=is_fsdp2,
+        state=SimpleNamespace(fsdp_plugin=plugin),
+        torch_device_mesh=None,
+    )
+
+
+def test_fully_shard_tied_embeddings_as_one_group(monkeypatch):
+    model = _TiedModel()
+    calls = []
+
+    def _record_fully_shard(modules, **kwargs):
+        calls.append((modules, kwargs))
+
+    monkeypatch.setattr(torch.distributed.fsdp, "fully_shard", _record_fully_shard)
+
+    _fully_shard_tied_embeddings(model, _accelerator())
+
+    assert len(calls) == 1
+    modules, kwargs = calls[0]
+    assert modules == [model.embed_tokens, model.lm_head]
+    assert kwargs["reshard_after_forward"] is True
+    assert kwargs["offload_policy"] is None
+    assert isinstance(kwargs["mp_policy"], torch.distributed.fsdp.MixedPrecisionPolicy)
+    assert kwargs["mesh"] is None
+
+
+@pytest.mark.parametrize(("tied", "is_fsdp2"), [(False, True), (True, False)])
+def test_fully_shard_tied_embeddings_skips_unsupported_cases(monkeypatch, tied, is_fsdp2):
+    model = _TiedModel(tied=tied)
+
+    def _unexpected_fully_shard(*args, **kwargs):
+        pytest.fail("fully_shard should not be called")
+
+    monkeypatch.setattr(torch.distributed.fsdp, "fully_shard", _unexpected_fully_shard)
+
+    _fully_shard_tied_embeddings(model, _accelerator(is_fsdp2=is_fsdp2))