[fsdp] fix: wrap embed_tokens/lm_head by name for peft models (#5516)

Author: cavities12

tests/utils/test_fsdp2_peft_wrapping.py

@@ -0,0 +1,131 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# 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.
+"""Test that apply_fsdp2's module selection handles peft-wrapped models.
+
+peft wraps embed_tokens in a ModulesToSaveWrapper, so isinstance(module, nn.Embedding)
+fails. Without name-based matching, embed_tokens + lm_head land in the root FSDP unit,
+causing OOM from oversized allgather. These tests verify the module selection logic
+works for: (1) vanilla models, (2) peft-wrapped models, (3) tied embeddings.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch.nn as nn
+
+from verl.utils.fsdp_utils import _select_fsdp2_wrap_targets
+
+
+class MockDecoderLayer(nn.Module):
+    """Simulates a transformer decoder layer (e.g. Qwen3DecoderLayer)."""
+
+    def __init__(self, hidden_size=64):
+        super().__init__()
+        self.self_attn = nn.Linear(hidden_size, hidden_size)
+        self.mlp = nn.Linear(hidden_size, hidden_size)
+
+
+class MockModulesToSaveWrapper(nn.Module):
+    """Simulates peft's ModulesToSaveWrapper around nn.Embedding.
+
+    peft wraps modules listed in modules_to_save (like embed_tokens) in this wrapper,
+    which breaks isinstance(module, nn.Embedding) checks.
+    """
+
+    def __init__(self, original_module):
+        super().__init__()
+        self.original_module = original_module
+        self.weight = original_module.weight  # peft exposes weight
+
+
+class MockCausalLM(nn.Module):
+    """Simulates a causal LM with embed_tokens, decoder layers, and lm_head."""
+
+    _no_split_modules = ["MockDecoderLayer"]
+
+    def __init__(self, vocab_size=1000, hidden_size=64, num_layers=2, tie_word_embeddings=False):
+        super().__init__()
+        self.config = SimpleNamespace(tie_word_embeddings=tie_word_embeddings)
+        self.model = nn.Module()
+        self.model.embed_tokens = nn.Embedding(vocab_size, hidden_size)
+        self.model.layers = nn.ModuleList([MockDecoderLayer(hidden_size) for _ in range(num_layers)])
+        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
+
+        if tie_word_embeddings:
+            self.lm_head.weight = self.model.embed_tokens.weight
+
+
+class TestFSDP2PeftWrapping(unittest.TestCase):
+    """Test module selection in apply_fsdp2 for vanilla and peft-wrapped models."""
+
+    def _get_wrapped_names(self, model, cls_names):
+        """Return names of modules selected for wrapping."""
+        selected = _select_fsdp2_wrap_targets(model, cls_names)
+        # _select_fsdp2_wrap_targets returns module objects; map back to names
+        module_to_name = {id(m): n for n, m in model.named_modules()}
+        return [module_to_name[id(m)] for m in selected]
+
+    def test_vanilla_model_wraps_layers_and_embedding(self):
+        """Vanilla model (no peft): embed_tokens matched by isinstance, layers by class name."""
+        model = MockCausalLM(tie_word_embeddings=False)
+        names = self._get_wrapped_names(model, ["MockDecoderLayer"])
+
+        self.assertIn("model.embed_tokens", names)
+        self.assertIn("lm_head", names)
+        self.assertTrue(any("layers.0" in n for n in names))
+        self.assertTrue(any("layers.1" in n for n in names))
+
+    def test_peft_wrapped_model_wraps_embed_tokens_by_name(self):
+        """peft-wrapped model: embed_tokens fails isinstance but is matched by name."""
+        model = MockCausalLM(tie_word_embeddings=False)
+        original_embed = model.model.embed_tokens
+        model.model.embed_tokens = MockModulesToSaveWrapper(original_embed)
+
+        names = self._get_wrapped_names(model, ["MockDecoderLayer"])
+
+        self.assertIn("model.embed_tokens", names)
+        self.assertIn("lm_head", names)
+        self.assertTrue(any("layers.0" in n for n in names))
+
+    def test_tied_embeddings_skips_name_based_wrapping(self):
+        """With tie_word_embeddings=True, embed_tokens/lm_head are NOT wrapped separately."""
+        model = MockCausalLM(tie_word_embeddings=True)
+        names = self._get_wrapped_names(model, ["MockDecoderLayer"])
+
+        self.assertNotIn("model.embed_tokens", names)
+        self.assertNotIn("lm_head", names)
+        self.assertTrue(any("layers.0" in n for n in names))
+
+    def test_peft_wrapped_tied_embeddings_skips_wrapping(self):
+        """peft + tied embeddings: name-based matching is disabled, no wrapping."""
+        model = MockCausalLM(tie_word_embeddings=True)
+        original_embed = model.model.embed_tokens
+        model.model.embed_tokens = MockModulesToSaveWrapper(original_embed)
+
+        names = self._get_wrapped_names(model, ["MockDecoderLayer"])
+
+        self.assertNotIn("model.embed_tokens", names)
+        self.assertNotIn("lm_head", names)
+
+    def test_no_duplicate_wrapping_for_vanilla_embedding(self):
+        """Vanilla nn.Embedding should not be wrapped twice (by isinstance AND by name)."""
+        model = MockCausalLM(tie_word_embeddings=False)
+        names = self._get_wrapped_names(model, ["MockDecoderLayer"])
+
+        embed_count = sum(1 for n in names if n == "model.embed_tokens")
+        self.assertEqual(embed_count, 1, f"embed_tokens wrapped {embed_count} times, expected 1")
+
+
+if __name__ == "__main__":
+    unittest.main()

verl/utils/fsdp_utils.py

@@ -507,6 +507,30 @@ class FSDPModuleABC(ABC, orig_fsdp_module):
         fully_shard_module.FSDPModule = orig_fsdp_module
 
 
+def _select_fsdp2_wrap_targets(model, fsdp_transformer_layer_cls_to_wrap):
+    """Select modules to wrap individually with fully_shard in FSDP2.
+
+    Matches transformer layers by class name, and embed_tokens/lm_head by name
+    (with isinstance fallback). Name-based matching is needed because peft wraps
+    embed_tokens in ModulesToSaveWrapper, breaking isinstance(module, nn.Embedding).
+    When tie_word_embeddings is True, embed_tokens and lm_head share weights and
+    must not be wrapped separately.
+    """
+    _tie = getattr(model.config, "tie_word_embeddings", False)
+    _wrap_by_name = set() if _tie else {"embed_tokens", "lm_head"}
+
+    modules = []
+    for name, module in model.named_modules():
+        leaf_name = name.rsplit(".", 1)[-1] if "." in name else name
+        if (
+            module.__class__.__name__ in fsdp_transformer_layer_cls_to_wrap
+            or (isinstance(module, nn.Embedding) and not _tie)
+            or (leaf_name in _wrap_by_name and hasattr(module, "weight"))
+        ):
+            modules.append(module)
+    return modules
+
+
 def apply_fsdp2(model, fsdp_kwargs, config):
     """model: AutoModelForCausalLM"""
     assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)"
@@ -521,12 +545,7 @@ def apply_fsdp2(model, fsdp_kwargs, config):
 
     assert len(fsdp_transformer_layer_cls_to_wrap) > 0 and fsdp_transformer_layer_cls_to_wrap[0] is not None
 
-    modules = []
-    for name, module in model.named_modules():
-        if module.__class__.__name__ in fsdp_transformer_layer_cls_to_wrap or (
-            isinstance(module, nn.Embedding) and not model.config.tie_word_embeddings
-        ):
-            modules.append(module)
+    modules = _select_fsdp2_wrap_targets(model, fsdp_transformer_layer_cls_to_wrap)
 
     for idx, module in enumerate(modules):
         # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0: