diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index f72715d410c..638490c41dc 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/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: diff --git a/tests/unit/torch/opt/plugins/test_transformers_fsdp.py b/tests/unit/torch/opt/plugins/test_transformers_fsdp.py new file mode 100644 index 00000000000..08dbcb7116b --- /dev/null +++ b/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))