diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index c31aaef9c..b3e8a05b2 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -43,6 +43,9 @@ class ComposedPipelineBase(ABC): training_args: TrainingArgs | None = None fastvideo_args: FastVideoArgs | TrainingArgs | None = None modules: dict[str, Any] = {} + # do not need to include moe related transformers + trainable_transformer_names: list[str] = ["transformer"] + trainable_transformer_modules: dict[str, torch.nn.Module] = {} post_init_called: bool = False # TODO(will): args should support both inference args and training args @@ -87,13 +90,14 @@ class ComposedPipelineBase(ABC): def set_trainable(self) -> None: # Only train DiT if getattr(self.fastvideo_args, "training_mode", False): - for name, module in self.modules.items(): + for name, module in self.trainable_transformer_modules.items(): + logger.info("Setting %s to requires_grad=True", name) if not isinstance(module, torch.nn.Module): + logger.info( + "Skipping %s because it is not a torch.nn.Module", name) continue - if "transformer" in name: - module.requires_grad_(True) - else: - module.requires_grad_(False) + module.requires_grad_(True) + module.train() def post_init(self) -> None: assert self.fastvideo_args is not None, "fastvideo_args must be set" @@ -122,18 +126,32 @@ class ComposedPipelineBase(ABC): fsdp_module_cls = FSDPModule except Exception: # pragma: no cover - FSDP not always available fsdp_module_cls = None + + compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {} if fsdp_module_cls is not None and isinstance( transformer_module, fsdp_module_cls): logger.info( "Transformer is already FSDP-wrapped; skipping torch.compile in pipeline" ) else: - compile_kwargs = self.fastvideo_args.torch_compile_kwargs or {} logger.info("Enabling torch.compile for DiT with kwargs=%s", compile_kwargs) self.modules["transformer"] = torch.compile( transformer_module, **compile_kwargs) - logger.info("Torch Compile enabled for DiT") + if "transformer_2" in self.modules: + transformer_module_2 = self.modules["transformer_2"] + if fsdp_module_cls is not None and isinstance( + transformer_module_2, fsdp_module_cls): + logger.info( + "Transformer_2 is already FSDP-wrapped; skipping torch.compile in pipeline" + ) + else: + logger.info( + "Enabling torch.compile for Transformer_2 with kwargs=%s", + compile_kwargs) + self.modules["transformer_2"] = torch.compile( + transformer_module_2, **compile_kwargs) + logger.info("Torch Compile enabled for DiT") if not self.fastvideo_args.training_mode: logger.info("Creating pipeline stages...") diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index b3babed32..d9fd4b424 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/fastvideo/pipelines/lora_pipeline.py @@ -32,10 +32,9 @@ class LoRAPipeline(ComposedPipelineBase): ) # state dicts of loaded lora adapters (includes lora_A, lora_B, and lora_alpha) cur_adapter_name: str = "" cur_adapter_path: str = "" - lora_layers: dict[str, BaseLayerWithLoRA] = {} - lora_layers_critic: dict[str, BaseLayerWithLoRA] = {} + lora_layers: dict[str, dict[str, BaseLayerWithLoRA]] = {} fastvideo_args: FastVideoArgs | TrainingArgs - exclude_lora_layers: list[str] = [] + exclude_lora_layers: dict[str, list[str]] = {} device: torch.device = get_local_torch_device() lora_target_modules: list[str] | None = None lora_path: str | None = None @@ -47,8 +46,32 @@ class LoRAPipeline(ComposedPipelineBase): def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self.device = get_local_torch_device() - self.exclude_lora_layers = self.modules[ - "transformer"].config.arch_config.exclude_lora_layers + # build list of trainable transformers + for transformer_name in self.trainable_transformer_names: + if transformer_name in self.modules and self.modules[ + transformer_name] is not None: + self.trainable_transformer_modules[ + transformer_name] = self.modules[transformer_name] + # check for transformer_2 in case of Wan2.2 MoE or fake_score_transformer_2 + if transformer_name.endswith("_2"): + raise ValueError( + f"trainable_transformer_name override in pipelines should not include _2 suffix: {transformer_name}" + ) + + secondary_transformer_name = transformer_name + "_2" + if secondary_transformer_name in self.modules and self.modules[ + secondary_transformer_name] is not None: + self.trainable_transformer_modules[ + secondary_transformer_name] = self.modules[ + secondary_transformer_name] + + logger.info("trainable_transformer_modules: %s", + self.trainable_transformer_modules.keys()) + + for transformer_name, transformer_module in self.trainable_transformer_modules.items( + ): + self.exclude_lora_layers[ + transformer_name] = transformer_module.config.arch_config.exclude_lora_layers self.lora_target_modules = self.fastvideo_args.lora_target_modules self.lora_path = self.fastvideo_args.lora_path self.lora_nickname = self.fastvideo_args.lora_nickname @@ -65,8 +88,16 @@ class LoRAPipeline(ComposedPipelineBase): if self.lora_target_modules is None: self.lora_target_modules = [ "q_proj", "k_proj", "v_proj", "o_proj", "to_q", "to_k", - "to_v", "to_out", "to_qkv" + "to_v", "to_out", "to_qkv", "to_gate_compress" ] + logger.info( + "Using default lora_target_modules for all transformers: %s", + self.lora_target_modules) + else: + logger.warning( + "Using custom lora_target_modules for all transformers, which may not be intended: %s", + self.lora_target_modules) + self.convert_to_lora_layers() # Inference elif not self.training_mode and self.lora_path is not None: @@ -100,13 +131,18 @@ class LoRAPipeline(ComposedPipelineBase): super().set_trainable() return - self.modules["transformer"].requires_grad_(False) - if "fake_score_transformer" in self.modules: - self.modules["fake_score_transformer"].requires_grad_(False) device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1), mesh_dim_names=["fake", "replicate"]) - set_lora_grads(self.lora_layers, device_mesh) - set_lora_grads(self.lora_layers_critic, device_mesh) + for transformer_name, transformer_module in self.trainable_transformer_modules.items( + ): + transformer_module.train() + transformer_module.requires_grad_(False) + if transformer_name in self.lora_layers: + set_lora_grads(self.lora_layers[transformer_name], device_mesh) + else: + raise ValueError( + f"Transformer {transformer_name} should be trainable but not found in lora_layers" + ) def convert_to_lora_layers(self) -> None: """ @@ -115,46 +151,33 @@ class LoRAPipeline(ComposedPipelineBase): if self.lora_initialized: return self.lora_initialized = True - converted_count = 0 - for name, layer in self.modules["transformer"].named_modules(): - if not self.is_target_layer(name): - continue - - excluded = False - for exclude_layer in self.exclude_lora_layers: - if exclude_layer in name: - excluded = True - break - if excluded: - continue - - layer = get_lora_layer(layer, - lora_rank=self.lora_rank, - lora_alpha=self.lora_alpha, - training_mode=self.training_mode) - if layer is not None: - self.lora_layers[name] = layer - replace_submodule(self.modules["transformer"], name, layer) - converted_count += 1 - logger.info("Converted %d layers to LoRA layers", converted_count) - - if "fake_score_transformer" in self.modules: - for name, layer in self.modules[ - "fake_score_transformer"].named_modules(): + for transformer_name, transformer_module in self.trainable_transformer_modules.items( + ): + converted_count = 0 + if transformer_name not in self.lora_layers: + self.lora_layers[transformer_name] = {} + logger.info("Converting %s to LoRA Transformer", transformer_name) + for name, layer in transformer_module.named_modules(): if not self.is_target_layer(name): continue + + excluded = False + for exclude_layer in self.exclude_lora_layers[transformer_name]: + if exclude_layer in name: + excluded = True + break + if excluded: + continue + layer = get_lora_layer(layer, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, training_mode=self.training_mode) if layer is not None: - self.lora_layers_critic[name] = layer - replace_submodule(self.modules["fake_score_transformer"], - name, layer) + self.lora_layers[transformer_name][name] = layer + replace_submodule(transformer_module, name, layer) converted_count += 1 - logger.info( - "Converted %d layers to LoRA layers in the critic model", - converted_count) + logger.info("Converted %d layers to LoRA layers", converted_count) def set_lora_adapter(self, lora_nickname: str, @@ -238,39 +261,45 @@ class LoRAPipeline(ComposedPipelineBase): # Merge the new adapter adapted_count = 0 - for name, layer in self.lora_layers.items(): - lora_A_name = name + ".lora_A" - lora_B_name = name + ".lora_B" - lora_alpha_name = name + ".lora_alpha" - if lora_A_name in self.lora_adapters[lora_nickname]\ - and lora_B_name in self.lora_adapters[lora_nickname]: - # Get alpha value for this layer (defaults to None if not present) - lora_A = self.lora_adapters[lora_nickname][lora_A_name] - lora_B = self.lora_adapters[lora_nickname][lora_B_name] - # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B - alpha = self.lora_adapters[lora_nickname].get( - lora_alpha_name) if adapter_updated else None + for transformer_name, transformer_lora_layers in self.lora_layers.items( + ): + for name, layer in transformer_lora_layers.items(): + lora_A_name = name + ".lora_A" + lora_B_name = name + ".lora_B" + lora_alpha_name = name + ".lora_alpha" + if lora_A_name in self.lora_adapters[lora_nickname]\ + and lora_B_name in self.lora_adapters[lora_nickname]: + # Get alpha value for this layer (defaults to None if not present) + lora_A = self.lora_adapters[lora_nickname][lora_A_name] + lora_B = self.lora_adapters[lora_nickname][lora_B_name] + # Simple lookup - alpha stored with same naming scheme as lora_A/lora_B + alpha = self.lora_adapters[lora_nickname].get( + lora_alpha_name) if adapter_updated else None - layer.set_lora_weights( - lora_A, - lora_B, - lora_alpha=alpha, - training_mode=self.fastvideo_args.training_mode, - lora_path=lora_path) - adapted_count += 1 - else: - if rank == 0: - logger.warning( - "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.", - lora_path, name) - layer.disable_lora = True + layer.set_lora_weights( + lora_A, + lora_B, + lora_alpha=alpha, + training_mode=self.fastvideo_args.training_mode, + lora_path=lora_path) + adapted_count += 1 + else: + if rank == 0: + logger.warning( + "LoRA adapter %s does not contain the weights for layer %s. LoRA will not be applied to it.", + lora_path, name) + layer.disable_lora = True logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank, lora_path, adapted_count) def merge_lora_weights(self) -> None: - for name, layer in self.lora_layers.items(): - layer.merge_lora_weights() + for transformer_name, transformer_lora_layers in self.lora_layers.items( + ): + for name, layer in transformer_lora_layers.items(): + layer.merge_lora_weights() def unmerge_lora_weights(self) -> None: - for name, layer in self.lora_layers.items(): - layer.unmerge_lora_weights() + for transformer_name, transformer_lora_layers in self.lora_layers.items( + ): + for name, layer in transformer_lora_layers.items(): + layer.unmerge_lora_weights() diff --git a/fastvideo/tests/training/VSA/h200_reference_wandb_summary_VSA.json b/fastvideo/tests/training/VSA/h200_reference_wandb_summary_VSA.json new file mode 100644 index 000000000..34aa73c67 --- /dev/null +++ b/fastvideo/tests/training/VSA/h200_reference_wandb_summary_VSA.json @@ -0,0 +1,18 @@ +{ + "step_time": 0.8274769484996796, + "avg_step_time": 1.4095150619745254, + "timing/optimizer_step": 0.02941332757472992, + "grad_norm": 0.24477368593215942, + "timing/forward_backward": 0.693624597042799, + "_timestamp": 1765590605.595408, + "timing/prepare_dit_inputs": 0.009604107588529587, + "learning_rate": 1e-05, + "train_loss": 0.30774877965450287, + "_step": 5, + "timing/normalize_input": 0.0003825277090072632, + "vsa_sparsity": 0.05, + "timing/clip_grad_norm": 0.025074530392885208, + "_runtime": 90.44591733, + "timing/get_next_batch": 0.0055067278444767, + "timing/reduce_loss": 0.00021916255354881287 +} diff --git a/fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json b/fastvideo/tests/training/VSA/l40s_reference_wandb_summary_VSA.json similarity index 100% rename from fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json rename to fastvideo/tests/training/VSA/l40s_reference_wandb_summary_VSA.json diff --git a/fastvideo/tests/training/VSA/test_training_loss_VSA.py b/fastvideo/tests/training/VSA/test_training_loss_VSA.py index 2f393b1fc..e0abe0907 100644 --- a/fastvideo/tests/training/VSA/test_training_loss_VSA.py +++ b/fastvideo/tests/training/VSA/test_training_loss_VSA.py @@ -4,6 +4,7 @@ import subprocess from pathlib import Path import json from huggingface_hub import snapshot_download +import torch # Import the training pipeline sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent)) @@ -12,7 +13,8 @@ from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.utils import FlexibleArgumentParser wandb_name = "test_training_loss_VSA" -reference_wandb_summary_file = "fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json" +l40s_reference_wandb_summary_file = "fastvideo/tests/training/VSA/reference_wandb_summary_VSA.json" +h200_reference_wandb_summary_file = "fastvideo/tests/training/VSA/h200_reference_wandb_summary_VSA.json" NUM_NODES = "1" NUM_GPUS_PER_NODE = "2" @@ -104,6 +106,14 @@ def test_distributed_training(): summary_file = 'data/wan_finetune_test_VSA/tracker/wandb/latest-run/files/wandb-summary.json' + device_name = torch.cuda.get_device_name() + if "L40S" in device_name: + reference_wandb_summary_file = l40s_reference_wandb_summary_file + elif "H200" in device_name: + reference_wandb_summary_file = h200_reference_wandb_summary_file + else: + raise ValueError(f"Unknown device: {device_name}") + reference_wandb_summary = json.load(open(reference_wandb_summary_file)) wandb_summary = json.load(open(summary_file)) diff --git a/fastvideo/tests/training/Vanilla/h200_reference_wandb_summary.json b/fastvideo/tests/training/Vanilla/h200_reference_wandb_summary.json new file mode 100644 index 000000000..ae7466b89 --- /dev/null +++ b/fastvideo/tests/training/Vanilla/h200_reference_wandb_summary.json @@ -0,0 +1,23 @@ +{ + "step_time": 1.171190332621336, + "vsa_sparsity": 0, + "_step": 5, + "avg_step_time": 1.2594494499266147, + "timing/forward_backward": 0.6414474211633205, + "timing/reduce_loss": 0.00025020167231559753, + "_runtime": 58.873311423, + "batch_size": 2, + "timing/prepare_dit_inputs": 0.011771254241466522, + "timing/normalize_input": 0.0005782917141914368, + "learning_rate": 1e-06, + "hidden_dim": 1536, + "train_loss": 0.13831664621829987, + "_timestamp": 1765589891.5887935, + "context_len": 512, + "grad_norm": 0.0643853023648262, + "ffn_dim": 8960, + "timing/optimizer_step": 0.02862849459052086, + "dit_seq_len": 6240, + "timing/clip_grad_norm": 0.02274264767765999, + "timing/get_next_batch": 0.40074965357780457 +} diff --git a/fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json b/fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json index 91908a005..e9b3fe87d 100644 --- a/fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json +++ b/fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json @@ -1 +1,11 @@ -{"train_loss":0.1021774671971798,"_step":5,"_wandb":{"runtime":35},"step_time":2.2575189135968685,"grad_norm":0.11582941561937332,"_runtime":35.701958502,"avg_step_time":2.4754387199878694,"_timestamp":1.7525528420185745e+09,"learning_rate":1e-06,"vsa_sparsity":0,"validation_videos_8_steps":{"videos":[{"caption":"A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","_type":"video-file","sha256":"ee83e83df073a648f89dcd288cccaed9af765e11fc35c77a6e0cb2ebaa1be5b0","size":475248,"path":"media/videos/validation_videos_8_steps_0_ee83e83df073a648f89d.mp4"},{"size":341490,"path":"media/videos/validation_videos_8_steps_0_d0b2758549d5c82845ca.mp4","caption":"A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","_type":"video-file","sha256":"d0b2758549d5c82845ca3c1ac0db6812877631495e77ac041acdc8ea31f0a5ee"},{"caption":"A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.","_type":"video-file","sha256":"08638381f1607d6ab10684772be38a58b18628fd402a208b25f6af96e765d454","size":436814,"path":"media/videos/validation_videos_8_steps_0_08638381f1607d6ab106.mp4"}],"captions":["A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.","A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.","A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table."],"_type":"videos","count":3}} \ No newline at end of file +{ + "train_loss": 0.1021774671971798, + "_step": 5, + "step_time": 2.2575189135968685, + "grad_norm": 0.11582941561937332, + "_runtime": 35.701958502, + "avg_step_time": 2.4754387199878694, + "_timestamp": 1752552842.0185745, + "learning_rate": 1e-06, + "vsa_sparsity": 0 +} diff --git a/fastvideo/tests/training/Vanilla/test_training_loss.py b/fastvideo/tests/training/Vanilla/test_training_loss.py index 8fb7826e4..0460268f0 100644 --- a/fastvideo/tests/training/Vanilla/test_training_loss.py +++ b/fastvideo/tests/training/Vanilla/test_training_loss.py @@ -18,6 +18,7 @@ from fastvideo.training.wan_training_pipeline import WanTrainingPipeline wandb_name = "test_training_loss" a40_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/a40_reference_wandb_summary.json" l40s_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/l40s_reference_wandb_summary.json" +h200_reference_wandb_summary_file = "fastvideo/tests/training/Vanilla/h200_reference_wandb_summary.json" NUM_NODES = "1" NUM_GPUS_PER_NODE = "2" @@ -126,6 +127,8 @@ def test_distributed_training(): reference_wandb_summary_file = a40_reference_wandb_summary_file elif "L40S" in device_name: reference_wandb_summary_file = l40s_reference_wandb_summary_file + elif "H200" in device_name: + reference_wandb_summary_file = h200_reference_wandb_summary_file else: raise ValueError(f"Unknown device: {device_name}") diff --git a/fastvideo/tests/training/lora/l40s_reference_lora_wandb_summary.json b/fastvideo/tests/training/lora/l40s_reference_lora_wandb_summary.json index 35fdeb535..1ee2b3739 100644 --- a/fastvideo/tests/training/lora/l40s_reference_lora_wandb_summary.json +++ b/fastvideo/tests/training/lora/l40s_reference_lora_wandb_summary.json @@ -8,37 +8,5 @@ "learning_rate": 0.00005, "step_time": 9.380061847000036, "train_loss": 0.1361890360713005, - "validation_videos_50_steps": { - "_type": "videos", - "captions": [ - "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", - "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.", - "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table." - ], - "count": 3, - "videos": [ - { - "_type": "video-file", - "caption": "A large metal cylinder is seen pressing down on a pile of Oreo cookies, flattening them as if they were under a hydraulic press.", - "path": "media/videos/validation_videos_50_steps_0_54dafffa2c2dea9fcb21.mp4", - "sha256": "54dafffa2c2dea9fcb21b734a6b4abb072b99a32e57603899d12b8eaf4324434", - "size": 611595 - }, - { - "_type": "video-file", - "caption": "A large metal cylinder is seen compressing colorful clay into a compact shape, demonstrating the power of a hydraulic press.", - "path": "media/videos/validation_videos_50_steps_0_6d82abfb6ab979f9df94.mp4", - "sha256": "6d82abfb6ab979f9df94e96fcd1ac344f40f03cb3f106a6b1d9c510dc17e8f91", - "size": 460476 - }, - { - "_type": "video-file", - "caption": "A large metal cylinder is seen pressing down on a pile of colorful candies, flattening them as if they were under a hydraulic press. The candies are crushed and broken into small pieces, creating a mess on the table.", - "path": "media/videos/validation_videos_50_steps_0_cedf3954f184052149c8.mp4", - "sha256": "cedf3954f184052149c83ac23776ad83f75cf911351fcb6f8e05854ba4d4c614", - "size": 738913 - } - ] - }, "vsa_sparsity": 0 - } \ No newline at end of file + } diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 4e78d9e7d..72a1cfeb3 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -7,7 +7,7 @@ import time from abc import abstractmethod from collections import deque from collections.abc import Iterator -from typing import Any +from typing import Any, cast import imageio import numpy as np @@ -58,6 +58,10 @@ class DistillationPipeline(TrainingPipeline): "transformer", "vae", ] + # used by lora training + trainable_transformer_names: list[str] = [ + "transformer", "fake_score_transformer" + ] validation_pipeline: ComposedPipelineBase train_dataloader: StatefulDataLoader train_loader_iter: Iterator[dict[str, Any]] @@ -72,6 +76,65 @@ class DistillationPipeline(TrainingPipeline): raise RuntimeError( "create_pipeline_stages should not be called for training pipeline") + def load_modules(self, + fastvideo_args: FastVideoArgs, + loaded_modules: dict[str, torch.nn.Module] | None = None): + modules = super().load_modules(fastvideo_args, loaded_modules) + training_args = cast(TrainingArgs, fastvideo_args) + + if training_args.real_score_model_path: + logger.info("Loading real score transformer from: %s", + training_args.real_score_model_path) + training_args.override_transformer_cls_name = "WanTransformer3DModel" + # TODO(will): can use deepcopy instead if the model is the same + self.real_score_transformer = self.load_module_from_path( + training_args.real_score_model_path, "transformer", + training_args) + modules["real_score_transformer"] = self.real_score_transformer + try: + self.real_score_transformer_2 = self.load_module_from_path( + training_args.real_score_model_path, "transformer_2", + training_args) + logger.info("Loaded real score transformer_2 for MoE support") + modules[ + "real_score_transformer_2"] = self.real_score_transformer_2 + except Exception: + logger.info( + "real score transformer_2 not found, using single transformer" + ) + self.real_score_transformer_2 = None + else: + raise ValueError( + "real_score_model_path is required for DMD distillation pipeline" + ) + + if training_args.fake_score_model_path: + logger.info("Loading fake score transformer from: %s", + training_args.fake_score_model_path) + training_args.override_transformer_cls_name = "WanTransformer3DModel" + self.fake_score_transformer = self.load_module_from_path( + training_args.fake_score_model_path, "transformer", + training_args) + modules["fake_score_transformer"] = self.fake_score_transformer + try: + self.fake_score_transformer_2 = self.load_module_from_path( + training_args.fake_score_model_path, "transformer_2", + training_args) + logger.info("Loaded fake score transformer_2 for MoE support") + modules[ + "fake_score_transformer_2"] = self.fake_score_transformer_2 + except Exception: + logger.info( + "fake score transformer_2 not found, using single transformer" + ) + self.fake_score_transformer_2 = None + else: + raise ValueError( + "fake_score_model_path is required for DMD distillation pipeline" + ) + + return modules + def initialize_training_pipeline(self, training_args: TrainingArgs): """Initialize the distillation training pipeline with multiple models.""" logger.info("Initializing distillation pipeline...") @@ -91,65 +154,13 @@ class DistillationPipeline(TrainingPipeline): else: self.boundary_timestep = None - if training_args.real_score_model_path: - logger.info("Loading real score transformer from: %s", - training_args.real_score_model_path) - training_args.override_transformer_cls_name = "WanTransformer3DModel" - self.real_score_transformer = self.load_module_from_path( - training_args.real_score_model_path, "transformer", - training_args) - try: - self.real_score_transformer_2 = self.load_module_from_path( - training_args.real_score_model_path, "transformer_2", - training_args) - logger.info("Loaded real score transformer_2 for MoE support") - except Exception: - logger.info( - "real score transformer_2 not found, using single transformer" - ) - self.real_score_transformer_2 = None - else: - self.real_score_transformer = self.get_module( - "real_score_transformer") - self.real_score_transformer_2 = self.get_module( - "real_score_transformer_2") - - if training_args.fake_score_model_path: - logger.info("Loading fake score transformer from: %s", - training_args.fake_score_model_path) - training_args.override_transformer_cls_name = "WanTransformer3DModel" - self.fake_score_transformer = self.load_module_from_path( - training_args.fake_score_model_path, "transformer", - training_args) - try: - self.fake_score_transformer_2 = self.load_module_from_path( - training_args.fake_score_model_path, "transformer_2", - training_args) - logger.info("Loaded fake score transformer_2 for MoE support") - except Exception: - logger.info( - "fake score transformer_2 not found, using single transformer" - ) - self.fake_score_transformer_2 = None - else: - self.fake_score_transformer = self.get_module( - "fake_score_transformer") - self.fake_score_transformer_2 = self.get_module( - "fake_score_transformer_2") - + # make sure the real score transformer is not trainable self.real_score_transformer.requires_grad_(False) self.real_score_transformer.eval() if self.real_score_transformer_2 is not None: self.real_score_transformer_2.requires_grad_(False) self.real_score_transformer_2.eval() - # Set training modes for fake score transformers (trainable) - self.fake_score_transformer.requires_grad_(True) - self.fake_score_transformer.train() - if self.fake_score_transformer_2 is not None: - self.fake_score_transformer_2.requires_grad_(True) - self.fake_score_transformer_2.train() - if training_args.enable_gradient_checkpointing_type is not None: self.fake_score_transformer = apply_activation_checkpointing( self.fake_score_transformer, @@ -349,22 +360,6 @@ class DistillationPipeline(TrainingPipeline): raise NotImplementedError( "Distillation pipelines must implement this method") - def _prepare_distillation(self, - training_batch: TrainingBatch) -> TrainingBatch: - """Prepare training environment for distillation.""" - self.transformer.requires_grad_(True) - self.transformer.train() - if self.transformer_2 is not None: - self.transformer_2.requires_grad_(True) - self.transformer_2.train() - self.fake_score_transformer.requires_grad_(True) - self.fake_score_transformer.train() - if self.fake_score_transformer_2 is not None: - self.fake_score_transformer_2.requires_grad_(True) - self.fake_score_transformer_2.train() - - return training_batch - def apply_ema_to_model(self, model): """Apply EMA weights to the model for validation or inference.""" if model is self.transformer and self.generator_ema is not None: @@ -948,8 +943,7 @@ class DistillationPipeline(TrainingPipeline): batches = [] # Collect N batches for gradient accumulation for _ in range(gradient_accumulation_steps): - batch = self._prepare_distillation(training_batch) - batch = self._get_next_batch(batch) + batch = self._get_next_batch(training_batch) batch = self._normalize_dit_input(batch) batch = self._prepare_dit_inputs(batch) batch = self._build_attention_metadata(batch) @@ -990,9 +984,6 @@ class DistillationPipeline(TrainingPipeline): # Only clip gradients for the model that is currently training self._clip_model_grad_norm_(batch_gen, self.transformer) - for param in self.transformer.parameters(): - # check if the gradient is not None and not zero - assert param.grad is not None and param.grad.abs().sum() > 0 self.optimizer.step() self.optimizer.zero_grad(set_to_none=True) @@ -1031,17 +1022,6 @@ class DistillationPipeline(TrainingPipeline): else: self._clip_model_grad_norm_(batch_fake, self.fake_score_transformer) - # Check gradients for fake score transformer - for param in self.fake_score_transformer.parameters(): - if param.requires_grad: - assert param.grad is not None and param.grad.abs().sum() > 0 - - # Check gradients for fake score transformer_2 if available - if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: - for param in self.fake_score_transformer_2.parameters(): - if param.requires_grad: - assert param.grad is not None and param.grad.abs().sum() > 0 - if self.train_fake_score_transformer_2 and self.fake_score_transformer_2 is not None: self.fake_score_optimizer_2.step() self.fake_score_lr_scheduler_2.step() diff --git a/fastvideo/training/self_forcing_distillation_pipeline.py b/fastvideo/training/self_forcing_distillation_pipeline.py index d6ab9d902..e23b97a8c 100644 --- a/fastvideo/training/self_forcing_distillation_pipeline.py +++ b/fastvideo/training/self_forcing_distillation_pipeline.py @@ -677,8 +677,7 @@ class SelfForcingDistillationPipeline(DistillationPipeline): batches = [] for _ in range(gradient_accumulation_steps): - batch = self._prepare_distillation(training_batch) - batch = self._get_next_batch(batch) + batch = self._get_next_batch(training_batch) batch = self._normalize_dit_input(batch) batch = self._prepare_dit_inputs(batch) batch = self._build_attention_metadata(batch) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 369ea0e67..2cb8bfafd 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -153,8 +153,6 @@ class TrainingPipeline(LoRAPipeline, ABC): ) if self.transformer_2 is not None: # Ensure transformer_2 has trainable parameters before creating optimizer - self.transformer_2.train() - self.transformer_2.requires_grad_(True) params_to_optimize_2 = self.transformer_2.parameters() params_to_optimize_2 = list( filter(lambda p: p.requires_grad, params_to_optimize_2)) @@ -238,29 +236,12 @@ class TrainingPipeline(LoRAPipeline, ABC): "Training pipelines must implement this method") def _prepare_training(self, training_batch: TrainingBatch) -> TrainingBatch: - self.transformer.train() self.optimizer.zero_grad() if self.transformer_2 is not None: - self.transformer_2.train() self.optimizer_2.zero_grad() training_batch.total_loss = 0.0 return training_batch - def _enable_training(self, model: torch.nn.Module, - optimizer: torch.optim.Optimizer) -> None: - """Enable training mode and gradients for the specified model.""" - for param in model.parameters(): - param.requires_grad = True - model.train() - optimizer.zero_grad() - - def _disable_training(self, model: torch.nn.Module, - optimizer: torch.optim.Optimizer) -> None: - """Disable training mode and gradients for the specified model.""" - for param in model.parameters(): - param.requires_grad = False - optimizer.zero_grad(set_to_none=True) - def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: with self.tracker.timed("timing/get_next_batch"): batch = next(self.train_loader_iter, None) # type: ignore @@ -311,15 +292,6 @@ class TrainingPipeline(LoRAPipeline, ABC): dtype=latents.dtype) timesteps = self._sample_timesteps(batch_size, latents.device) - # Enable training for the model that will be trained next and disable the other - if self.train_transformer_2: - self._enable_training(self.transformer_2, self.optimizer_2) - self._disable_training(self.transformer, self.optimizer) - else: - self._enable_training(self.transformer, self.optimizer) - if self.transformer_2 is not None: - self._disable_training(self.transformer_2, self.optimizer_2) - if self.training_args.sp_size > 1: # Make sure that the timesteps are the same across all sp processes. sp_group = get_sp_group()