diff --git a/fastgen/configs/config.py b/fastgen/configs/config.py index c7cbf3a..81eefdd 100644 --- a/fastgen/configs/config.py +++ b/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) diff --git a/fastgen/methods/model.py b/fastgen/methods/model.py index 72d6f05..9bfba9f 100644 --- a/fastgen/methods/model.py +++ b/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() diff --git a/fastgen/trainer.py b/fastgen/trainer.py index 7f6be53..b30c530 100644 --- a/fastgen/trainer.py +++ b/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() diff --git a/scripts/inference/image_model_inference.py b/scripts/inference/image_model_inference.py index 1286f14..ab26a16 100644 --- a/scripts/inference/image_model_inference.py +++ b/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 ) diff --git a/scripts/inference/inference_utils.py b/scripts/inference/inference_utils.py index 7fd32da..e544a83 100644 --- a/scripts/inference/inference_utils.py +++ b/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 diff --git a/scripts/inference/video_model_inference.py b/scripts/inference/video_model_inference.py index 83b9856..31bce8e 100644 --- a/scripts/inference/video_model_inference.py +++ b/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 ) diff --git a/tests/test_torch_compile.py b/tests/test_torch_compile.py new file mode 100644 index 0000000..650da2b --- /dev/null +++ b/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"])