From e4a70747bc608deff4138796e5de4ee6860aeb33 Mon Sep 17 00:00:00 2001 From: RandNMR73 Date: Thu, 9 Apr 2026 07:16:50 +0000 Subject: [PATCH] fix --- fastvideo/training/training_pipeline.py | 47 +------------------------ fastvideo/training/training_utils.py | 34 ++++++++++++++++++ 2 files changed, 35 insertions(+), 46 deletions(-) diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 371f161a7..b3ad66784 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -6,7 +6,6 @@ import os import shutil import tempfile import time -import types from abc import ABC, abstractmethod from collections import deque from collections.abc import Iterator @@ -45,7 +44,7 @@ from fastvideo.training.trackers import (DummyTracker, TrackerType, initialize_t from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases, compute_density_for_timestep_sampling, count_trainable, get_scheduler, get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint, - traverse_swap_module) + swap_fp4_linear, traverse_swap_module) from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict) try: @@ -58,50 +57,6 @@ except Exception: logger = init_logger(__name__) -def swap_fp4_linear(obj: Any, obj_path: str) -> int: - """ - Replace torch.nn.Linear and ReplicatedLinear layers with LinearFWD4BWD16, - preserving the input/output dimensions and bias setting. - Returns number of swaps performed. - """ - from fastvideo.layers.fp4linear import fp4_linear_forward - from fastvideo.layers.linear import ReplicatedLinear - - swaps_performed = 0 - - # Find all Linear-like attributes in the object - for attr_name in dir(obj): - if attr_name.startswith('_'): - continue - - try: - attr_value = getattr(obj, attr_name) - except Exception: - continue - - # Check if it's a linear layer we want to replace - should_replace = False - in_features = None - out_features = None - # if isinstance(attr_value, torch.nn.Linear): - # should_replace = True - # in_features = attr_value.in_features - # out_features = attr_value.out_features - # bias = attr_value.bias is not None - - if isinstance(attr_value, ReplicatedLinear): - should_replace = True - in_features = attr_value.input_size - out_features = attr_value.output_size - - if should_replace and in_features is not None and out_features is not None: - layer = getattr(obj, attr_name) # Replace the attribute - layer.forward = types.MethodType(fp4_linear_forward, layer) - swaps_performed += 1 - - return swaps_performed - - class TrainingPipeline(LoRAPipeline, ABC): """ A pipeline for training a model. All training pipelines should inherit from this class. diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 0047e3754..2610ee747 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -25,6 +25,40 @@ logger = init_logger(__name__) _HAS_ERRORED_CLIP_GRAD_NORM_WHILE_HANDLING_FAILING_DTENSOR_CASES = False +def swap_fp4_linear(obj: Any, obj_path: str) -> int: + """ + Swap supported linear layers to use the FP4 forward path in-place. + + Returns the number of layers updated on ``obj``. + """ + from fastvideo.layers.fp4linear import fp4_linear_forward + from fastvideo.layers.linear import ReplicatedLinear + + del obj_path # Unused today, but kept for traverse_swap_module compatibility. + + swaps_performed = 0 + + for attr_name in dir(obj): + if attr_name.startswith("_"): + continue + + try: + attr_value = getattr(obj, attr_name) + except Exception: + continue + + should_replace = False + if isinstance(attr_value, ReplicatedLinear): + should_replace = True + + if should_replace: + layer = getattr(obj, attr_name) + layer.forward = types.MethodType(fp4_linear_forward, layer) + swaps_performed += 1 + + return swaps_performed + + def gather_state_dict_on_cpu_rank0( model, device: torch.device | None = None,