Add torch_compile flag for training/inference (#28)

* Add torch_compile_mode for training networks

Add an optional torch_compile_mode config flag ("default",
"reduce-overhead", "max-autotune"; None disables). Networks to compile
are collected in a compile_dict (derived from model_dict minus EMA, plus
the teacher and the net's preprocessors) and compiled in place via
nn.Module.compile by the trainer, after DDP/FSDP wrapping.

Author: wenxin0319

fastgen/configs/config.py

@@ -160,6 +160,10 @@ 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 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")
     # (precision of the time steps is handled in the noise scheduler, defaulting to float64 for numerical stability)
 

fastgen/methods/model.py

@@ -24,6 +24,10 @@
 
 
 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 apply_torch_compile). Single source of truth.
+    _PREPROCESSOR_ATTRS = ("vae", "text_encoder", "image_encoder")
+
     def __init__(self, config: BaseModelConfig):
         """FastGenModel class for implementing training interface for all fastgen networks.
 
@@ -264,6 +268,50 @@ def build_model(self):
         if hasattr(self.net, "init_preprocessors") and self.config.enable_preprocessors:
             self.net.init_preprocessors()
 
+    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.
+        # None entries arise when cleanup_unused_modules has already been called.
+        modules = {name: net for name, net in self.model_dict.items() if name not in self.ema_dict and net is not None}
+        # The teacher is not part of model_dict; add it when present (cf. fsdp_dict).
+        if getattr(self, "teacher", None) is not None:
+            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.
+        for name in self._PREPROCESSOR_ATTRS:
+            obj = getattr(self.net, name, None)
+            if obj is None:
+                continue
+            if isinstance(obj, torch.nn.Module):
+                modules[name] = obj
+            else:
+                for attr, submodule in getattr(obj, "__dict__", {}).items():
+                    if isinstance(submodule, torch.nn.Module):
+                        modules[f"{name}.{attr}"] = submodule
+
+        for name, module in modules.items():
+            if getattr(module, "_compiled_call_impl", None) is not None:
+                logger.info(f"Skipping torch.compile for {name} (already compiled)")
+                continue
+            logger.info(f"Applying torch.compile (mode={mode}) to {name}")
+            module.compile(mode=mode)
+
     def on_train_begin(self, is_fsdp=False):
         self._is_fsdp = is_fsdp  # Store for later use (e.g., to skip EMA during inference)
         ctx = dict(dtype=self.precision, device=self.device)
@@ -306,15 +354,11 @@ def on_train_begin(self, is_fsdp=False):
         # For networks that don't need gradients, we always manually handle casting and device management
         if hasattr(self.net, "init_preprocessors") and self.config.enable_preprocessors:
             logger.debug(f"Starting moving preprocessors to context: {ctx}.")
-            if hasattr(self.net, "vae"):
-                self.net.vae.to(**ctx)
-                synchronize()
-            if hasattr(self.net, "text_encoder"):
-                self.net.text_encoder.to(**ctx)
-                synchronize()
-            if hasattr(self.net, "image_encoder"):
-                self.net.image_encoder.to(**ctx)
-                synchronize()
+            for name in self._PREPROCESSOR_ATTRS:
+                preprocessor = getattr(self.net, name, None)
+                if preprocessor is not None:
+                    preprocessor.to(**ctx)
+                    synchronize()
             logger.debug(f"Completed moving preprocessors to context: {ctx}.")
 
         synchronize()

fastgen/trainer.py

@@ -118,6 +118,11 @@ def run(
             logger.info("FSDP wrapping completed")
         else:
             model_ddp = model
+
+        # Compile networks after DDP/FSDP wrapping so torch.compile composes
+        # with the distributed wrappers (no-op if torch_compile_mode is None).
+        model.apply_torch_compile()
+
         self.callbacks.on_model_init_end(model_ddp)
         synchronize()
 

scripts/inference/image_model_inference.py

@@ -127,7 +127,7 @@ def main(args, config: BaseConfig):
     # Remove unused modules to free memory
     cleanup_unused_modules(model, args.do_teacher_sampling)
 
-    # Set up inference modules
+    # Set up inference modules (also calls apply_torch_compile internally)
     teacher, student, vae = setup_inference_modules(
         model, config, args.do_teacher_sampling, args.do_student_sampling, model.precision
     )

scripts/inference/inference_utils.py

@@ -124,18 +124,22 @@ def load_checkpoint(
 
 
 def cleanup_unused_modules(model: Any, do_teacher_sampling: bool) -> None:
-    """Remove unused modules to free memory.
+    """Free GPU memory held by modules that are not needed for inference.
+
+    Sets attributes to None rather than deleting them so that model_dict
+    property accesses (which read e.g. self.fake_score) return None instead
+    of raising AttributeError.
 
     Args:
         model: Model to clean up
         do_teacher_sampling: Whether teacher sampling will be performed
     """
     if hasattr(model, "fake_score"):
-        del model.fake_score
+        model.fake_score = None
     if hasattr(model, "discriminator"):
-        del model.discriminator
+        model.discriminator = None
     if (not do_teacher_sampling) and hasattr(model, "teacher"):
-        del model.teacher
+        model.teacher = None
 
 
 def setup_inference_modules(
@@ -177,6 +181,8 @@ def setup_inference_modules(
         vae = model.net.vae
         vae.to(device=model.device, dtype=precision)
 
+    model.apply_torch_compile()  # no-op if torch_compile_mode is None; must run after init_preprocessors
+
     return teacher, student, vae
 
 

scripts/inference/video_model_inference.py

@@ -567,7 +567,7 @@ def main(args, config: BaseConfig):
     # Remove unused modules
     cleanup_unused_modules(model, args.do_teacher_sampling)
 
-    # Get precision and set up inference modules
+    # Get precision and set up inference modules (also calls apply_torch_compile internally)
     teacher, student, vae = setup_inference_modules(
         model, config, args.do_teacher_sampling, args.do_student_sampling, model.precision
     )

tests/test_torch_compile.py

@@ -0,0 +1,207 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import gc
+import torch
+import pytest
+from fastgen.methods import DMD2Model
+from fastgen.configs.methods.config_sft import ModelConfig as SFTModelConfig
+from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig
+from fastgen.configs.config_utils import override_config_with_opts
+from fastgen.methods.fine_tuning.sft import SFTModel
+
+
+def _is_compiled(module):
+    # nn.Module.compile() compiles the module in place: it stores the compiled
+    # callable on `_compiled_call_impl` rather than replacing the module.
+    return getattr(module, "_compiled_call_impl", None) is not None
+
+
+@pytest.fixture
+def sft_model_compiled():
+    gc.collect()
+    instance = SFTModelConfig()
+    opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=False"]
+    instance.net = override_config_with_opts(instance.net, opts)
+    instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16"
+    instance.pretrained_model_path = ""
+    instance.input_shape = [3, 8, 8]
+    instance.torch_compile_mode = "default"
+    instance.cond_dropout_prob = 0.1
+    instance.cond_keys_no_dropout = []
+    instance.guidance_scale = None
+    model = SFTModel(instance)
+    # Mirror the trainer order: on_train_begin() moves parameters to device/dtype and
+    # initialises preprocessors; apply_torch_compile() must come after so that
+    # preprocessors exist and are on the right device when compiled.
+    model.on_train_begin()
+    model.apply_torch_compile()
+    return model
+
+
+@pytest.fixture
+def sft_model_not_compiled():
+    gc.collect()
+    instance = SFTModelConfig()
+    opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=False"]
+    instance.net = override_config_with_opts(instance.net, opts)
+    instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16"
+    instance.pretrained_model_path = ""
+    instance.input_shape = [3, 8, 8]
+    instance.torch_compile_mode = None
+    instance.cond_dropout_prob = 0.1
+    instance.cond_keys_no_dropout = []
+    instance.guidance_scale = None
+    model = SFTModel(instance)
+    model.on_train_begin()
+    model.apply_torch_compile()
+    return model
+
+
+@pytest.fixture
+def dmd2_model_compiled():
+    gc.collect()
+    instance = DMD2ModelConfig()
+    opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1"]
+    instance.net = override_config_with_opts(instance.net, opts)
+    opts_discriminator = ["-", "feature_indices=[0]", "all_res=[8]", "in_channels=128"]
+    instance.discriminator = override_config_with_opts(instance.discriminator, opts_discriminator)
+    instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16"
+    instance.pretrained_model_path = ""
+    instance.student_update_freq = 2
+    instance.input_shape = [3, 8, 8]
+    instance.torch_compile_mode = "default"
+    model = DMD2Model(instance)
+    model.on_train_begin()
+    model.apply_torch_compile()
+    return model
+
+
+@pytest.fixture
+def dmd2_model_not_compiled():
+    gc.collect()
+    instance = DMD2ModelConfig()
+    opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1"]
+    instance.net = override_config_with_opts(instance.net, opts)
+    opts_discriminator = ["-", "feature_indices=[0]", "all_res=[8]", "in_channels=128"]
+    instance.discriminator = override_config_with_opts(instance.discriminator, opts_discriminator)
+    instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16"
+    instance.pretrained_model_path = ""
+    instance.student_update_freq = 2
+    instance.input_shape = [3, 8, 8]
+    instance.torch_compile_mode = None
+    model = DMD2Model(instance)
+    model.on_train_begin()
+    model.apply_torch_compile()
+    return model
+
+
+def test_default_torch_compile_mode_is_none():
+    from fastgen.configs.config import BaseModelConfig
+
+    config = BaseModelConfig()
+    assert config.torch_compile_mode is None
+
+
+def test_sft_compile_enabled(sft_model_compiled):
+    assert _is_compiled(sft_model_compiled.net)
+
+
+def test_sft_compile_disabled(sft_model_not_compiled):
+    assert not _is_compiled(sft_model_not_compiled.net)
+
+
+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):
+    assert not _is_compiled(dmd2_model_not_compiled.net)
+    assert not _is_compiled(dmd2_model_not_compiled.teacher)
+    assert not _is_compiled(dmd2_model_not_compiled.fake_score)
+    assert not _is_compiled(dmd2_model_not_compiled.discriminator)
+
+
+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 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 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_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; apply_torch_compile must find and
+    # compile those submodules.
+    model = sft_model_not_compiled
+
+    class _DummyVAEWrapper:  # mimics WanVideoEncoder/SDVAE (not an nn.Module)
+        def __init__(self):
+            self.vae = torch.nn.Linear(4, 4)
+            self.scaling_factor = 0.18  # non-module attributes are ignored
+
+    model.net.vae = _DummyVAEWrapper()
+    # An attribute that is itself an nn.Module is compiled directly under its own name.
+    model.net.text_encoder = torch.nn.Linear(4, 4)
+
+    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):
+    model = sft_model_compiled
+    model.init_optimizers()
+
+    batch_size = 1
+    labels = torch.nn.functional.one_hot(torch.randint(0, 10, (batch_size,)), num_classes=10).float()
+    data = {
+        "real": torch.randn(batch_size, 3, 8, 8).to(model.device, model.precision),
+        "condition": labels.to(model.device, model.precision),
+        "neg_condition": torch.zeros(batch_size, 10).to(model.device, model.precision),
+    }
+
+    loss_map, _ = model.single_train_step(data, 0)
+    assert "total_loss" in loss_map
+    assert not torch.isnan(loss_map["total_loss"])
+    loss_map["total_loss"].backward()
+
+
+def test_dmd2_compiled_train_step(dmd2_model_compiled):
+    model = dmd2_model_compiled
+    model.init_optimizers()
+
+    batch_size = 1
+    labels = torch.nn.functional.one_hot(torch.randint(0, 10, (batch_size,)), num_classes=10)
+    data = {
+        "real": torch.randn(batch_size, 3, 8, 8).to(model.device, model.precision),
+        "condition": labels.to(model.device, model.precision),
+        "neg_condition": torch.zeros(batch_size, 10).to(model.device, model.precision),
+    }
+
+    # Student update step
+    loss_map, _ = model.single_train_step(data, 0)
+    assert "total_loss" in loss_map
+    assert not torch.isnan(loss_map["total_loss"])
+
+    # Fake score update step
+    model.optimizers_zero_grad(1)
+    loss_map, _ = model.single_train_step(data, 1)
+    assert "total_loss" in loss_map
+    assert not torch.isnan(loss_map["total_loss"])