remove compile dict

Author: wenxin0319

fastgen/configs/config.py

@@ -160,8 +160,8 @@ class BaseModelConfig:
     # - however, it is required if the model has a discriminator or the net initializes unused modules (e.g., for logvar predictions)
     ddp_find_unused_parameters: bool = True
 
-    # torch.compile mode for the networks in `compile_dict` ("default", "reduce-overhead", "max-autotune")
-    # None disables torch.compile.
+    # torch.compile mode for the training networks ("default", "reduce-overhead", "max-autotune")
+    # applied in apply_torch_compile. None disables torch.compile.
     torch_compile_mode: Optional[str] = None
 
     # precision variables (choose from "float64", "float32", "bfloat16", or "float16")

fastgen/methods/model.py

@@ -25,7 +25,7 @@
 
 class FastGenModel(torch.nn.Module):
     # Preprocessor sub-objects of ``net`` (handled specially for device/dtype placement
-    # in on_train_begin and for torch.compile in compile_dict). Single source of truth.
+    # in on_train_begin and for torch.compile in apply_torch_compile). Single source of truth.
     _PREPROCESSOR_ATTRS = ("vae", "text_encoder", "image_encoder")
 
     def __init__(self, config: BaseModelConfig):
@@ -268,22 +268,28 @@ def build_model(self):
         if hasattr(self.net, "init_preprocessors") and self.config.enable_preprocessors:
             self.net.init_preprocessors()
 
-    @property
-    def compile_dict(self) -> dict:
-        """Return dict of modules to compile with torch.compile, keyed by name.
-
-        Built from model_dict (e.g. net, plus fake_score/discriminator for DMD2)
-        minus the EMA networks, plus the teacher (if any, cf. fsdp_dict) and the
-        net's preprocessors. Compilation is applied in place (via
-        ``nn.Module.compile``) by ``apply_torch_compile``, which the trainer calls
-        *after* DDP/FSDP wrapping so torch.compile composes with the wrappers.
+    def apply_torch_compile(self):
+        """Compile the training networks in place with torch.compile.
+
+        Called by the trainer after DDP/FSDP wrapping (and after the networks
+        have been moved to their device in ``on_train_begin``) so torch.compile
+        composes with the distributed wrappers. No-op when
+        ``config.torch_compile_mode`` is None.
+
+        The modules compiled are those in model_dict (e.g. net, plus
+        fake_score/discriminator for DMD2) minus the EMA networks, plus the
+        teacher (if any, cf. fsdp_dict) and the net's preprocessors.
         """
+        mode = self.config.torch_compile_mode
+        if mode is None:
+            return
+
         # model_dict contains the trainable networks (incl. EMA); EMA networks are
         # weight-averaged copies that aren't run during training, so drop them.
-        compile_dict = {name: net for name, net in self.model_dict.items() if name not in self.ema_dict}
+        modules = {name: net for name, net in self.model_dict.items() if name not in self.ema_dict}
         # The teacher is not part of model_dict; add it when present (cf. fsdp_dict).
         if getattr(self, "teacher", None) is not None:
-            compile_dict["teacher"] = self.teacher
+            modules["teacher"] = self.teacher
         # Preprocessors (VAE, text/image encoders) are often lightweight wrappers
         # (e.g. WanVideoEncoder, SDVAE) that are not nn.Modules themselves but hold
         # the actual nn.Module under an attribute; compile those submodules too.
@@ -292,24 +298,13 @@ def compile_dict(self) -> dict:
             if obj is None:
                 continue
             if isinstance(obj, torch.nn.Module):
-                compile_dict[name] = obj
+                modules[name] = obj
             else:
                 for attr, submodule in getattr(obj, "__dict__", {}).items():
                     if isinstance(submodule, torch.nn.Module):
-                        compile_dict[f"{name}.{attr}"] = submodule
-        return compile_dict
-
-    def apply_torch_compile(self):
-        """Compile the modules in ``compile_dict`` in place with torch.compile.
+                        modules[f"{name}.{attr}"] = submodule
 
-        Called by the trainer after DDP/FSDP wrapping (and after the networks
-        have been moved to their device in ``on_train_begin``). No-op when
-        ``config.torch_compile_mode`` is None.
-        """
-        mode = self.config.torch_compile_mode
-        if mode is None:
-            return
-        for name, module in self.compile_dict.items():
+        for name, module in modules.items():
             logger.info(f"Applying torch.compile (mode={mode}) to {name}")
             module.compile(mode=mode)
 

tests/test_torch_compile.py

@@ -110,9 +110,11 @@ def test_sft_compile_disabled(sft_model_not_compiled):
 
 
 def test_dmd2_compile_enabled(dmd2_model_compiled):
+    # apply_torch_compile draws from model_dict (net, fake_score, discriminator) plus the teacher.
     assert _is_compiled(dmd2_model_compiled.net)
     assert _is_compiled(dmd2_model_compiled.teacher)
     assert _is_compiled(dmd2_model_compiled.fake_score)
+    assert _is_compiled(dmd2_model_compiled.discriminator)
 
 
 def test_dmd2_compile_disabled(dmd2_model_not_compiled):
@@ -121,26 +123,24 @@ def test_dmd2_compile_disabled(dmd2_model_not_compiled):
     assert not _is_compiled(dmd2_model_not_compiled.fake_score)
 
 
-def test_compile_dict_contains_expected_modules(sft_model_not_compiled, dmd2_model_not_compiled):
-    assert set(sft_model_not_compiled.compile_dict) == {"net"}
-    # DMD2's model_dict contributes fake_score (+ discriminator); the base adds the teacher.
-    assert {"net", "teacher", "fake_score"} <= set(dmd2_model_not_compiled.compile_dict)
-
-
-def test_compile_dict_excludes_ema(sft_model_not_compiled):
+def test_compile_excludes_ema(sft_model_not_compiled):
     # EMA networks live in model_dict but are weight-averaged copies that are not run
-    # during training, so compile_dict must drop them.
+    # during training, so apply_torch_compile must not compile them.
     model = sft_model_not_compiled
     model.use_ema = ["ema"]
     model.ema = torch.nn.Linear(4, 4)  # any nn.Module suffices for ema_dict/model_dict
-    assert "ema" in model.ema_dict
-    assert "ema" in model.model_dict
-    assert "ema" not in model.compile_dict
+    assert "ema" in model.ema_dict and "ema" in model.model_dict
+
+    model.config.torch_compile_mode = "default"
+    model.apply_torch_compile()
+    assert _is_compiled(model.net)
+    assert not _is_compiled(model.ema)
 
 
-def test_compile_dict_discovers_preprocessor_submodules(sft_model_not_compiled):
+def test_compile_discovers_preprocessor_submodules(sft_model_not_compiled):
     # Preprocessor wrappers (VAE, text/image encoders) are not nn.Modules themselves but
-    # hold the actual nn.Module under an attribute; compile_dict must find those submodules.
+    # hold the actual nn.Module under an attribute; apply_torch_compile must find and
+    # compile those submodules.
     model = sft_model_not_compiled
 
     class _DummyVAEWrapper:  # mimics WanVideoEncoder/SDVAE (not an nn.Module)
@@ -149,13 +149,13 @@ def __init__(self):
             self.scaling_factor = 0.18  # non-module attributes are ignored
 
     model.net.vae = _DummyVAEWrapper()
-    keys = set(model.compile_dict)
-    assert "vae.vae" in keys
-    assert isinstance(model.compile_dict["vae.vae"], torch.nn.Module)
-
     # An attribute that is itself an nn.Module is compiled directly under its own name.
     model.net.text_encoder = torch.nn.Linear(4, 4)
-    assert "text_encoder" in model.compile_dict
+
+    model.config.torch_compile_mode = "default"
+    model.apply_torch_compile()
+    assert _is_compiled(model.net.vae.vae)
+    assert _is_compiled(model.net.text_encoder)
 
 
 def test_sft_compiled_train_step(sft_model_compiled):