diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 2a41e6d67..f435a9356 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -133,7 +133,7 @@ class FastVideoArgs: # CPU offload parameters dit_cpu_offload: bool = True use_fsdp_inference: bool = False - dit_layerwise_offload: bool = False + dit_layerwise_offload: bool = True text_encoder_cpu_offload: bool = True image_encoder_cpu_offload: bool = True vae_cpu_offload: bool = True diff --git a/fastvideo/hooks/__init__.py b/fastvideo/hooks/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/hooks/hooks.py b/fastvideo/hooks/hooks.py new file mode 100644 index 000000000..7cfd2c454 --- /dev/null +++ b/fastvideo/hooks/hooks.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 + +import functools +from typing import Any +from torch import nn + + +class ForwardHook: + """ + Base class for forward hooks. + Hooks are used in the way: + modified_args, modified_kwargs = hook.pre_forward(module, *args, **kwargs) + output = module.forward(*modified_args, **modified_kwargs) + modified_output = hook.post_forward(module, output) + """ + + @classmethod + def name(cls) -> str: + raise NotImplementedError + + def on_attach(self, module: nn.Module): # noqa: B027 + """Called once when the hook is attached to the module.""" + pass + + def on_detach(self, module: nn.Module): # noqa: B027 + """ + Called once when the hook is detached from the module. + Note: this function is not guaranteed to be called if the module is + deleted before the hook is detached. + """ + pass + + def pre_forward(self, module: nn.Module, *args, + **kwargs) -> tuple[tuple[Any, ...], dict[str, Any]]: + """Called before the module's forward method is executed.""" + return args, kwargs + + def post_forward(self, module: nn.Module, output: Any) -> Any: + """Called after the module's forward method is executed.""" + return output + + +class ModuleHookManager: + module_hook_attribute = "_hook_manager" + + def __init__(self, module: nn.Module): + self.module = module + self.forward_hooks: dict[str, ForwardHook] = {} + self.original_forward = module.forward + + @classmethod + def get_from(cls, module: nn.Module) -> "ModuleHookManager | None": + if hasattr(module, cls.module_hook_attribute): + return getattr(module, cls.module_hook_attribute) + return None + + @classmethod + def get_from_or_default(cls, module: nn.Module) -> "ModuleHookManager": + if not hasattr(module, cls.module_hook_attribute): + setattr(module, cls.module_hook_attribute, cls(module)) + + def forward_hook_wrapper(mod: nn.Module, *args, **kwargs): + manager: ModuleHookManager = getattr(mod, + cls.module_hook_attribute) + for hook in manager.forward_hooks.values(): + args, kwargs = hook.pre_forward(mod, *args, **kwargs) + output = manager.original_forward(*args, **kwargs) + for hook in reversed(manager.forward_hooks.values()): + output = hook.post_forward(mod, output) + return output + + module.forward = functools.partial(forward_hook_wrapper, module) + + return getattr(module, cls.module_hook_attribute) + + @staticmethod + def remove_from_manager(module: nn.Module) -> None: + if hasattr(module, ModuleHookManager.module_hook_attribute): + manager: ModuleHookManager = getattr( + module, ModuleHookManager.module_hook_attribute) + module.forward = manager.original_forward + delattr(module, ModuleHookManager.module_hook_attribute) + + def _check_manager_attached(self) -> None: + if not hasattr(self.module, self.module_hook_attribute): + raise ValueError("ModuleHookManager is not attached to the module.") + if getattr(self.module, self.module_hook_attribute) is not self: + raise ValueError( + "ModuleHookManager attached to the module is different.") + + def append_forward_hook(self, hook: ForwardHook): + self._check_manager_attached() + if hook.name() in self.forward_hooks: + raise ValueError( + f"Hook with name {hook.name()} is already registered.") + # after python 3.7, dicts maintain insertion order + self.forward_hooks[hook.name()] = hook + hook.on_attach(self.module) + + def replace_forward_hook(self, + hook_name: str, + new_hook: ForwardHook, + run_on_attach: bool = True): + self._check_manager_attached() + if hook_name not in self.forward_hooks: + raise ValueError(f"No hook with name {hook_name} found.") + old_hook = self.forward_hooks[hook_name] + if run_on_attach: + old_hook.on_detach(self.module) + self.forward_hooks[hook_name] = new_hook + new_hook.on_attach(self.module) + + def remove_forward_hook(self, hook_name: str, run_detach: bool = True): + self._check_manager_attached() + if hook_name not in self.forward_hooks: + raise ValueError(f"No hook with name {hook_name} found.") + if run_detach: + self.forward_hooks[hook_name].on_detach(self.module) + del self.forward_hooks[hook_name] + + def get_forward_hook(self, hook_name: str) -> ForwardHook | None: + return self.forward_hooks.get(hook_name, None) diff --git a/fastvideo/hooks/layerwise_offload.py b/fastvideo/hooks/layerwise_offload.py new file mode 100644 index 000000000..b796b3c31 --- /dev/null +++ b/fastvideo/hooks/layerwise_offload.py @@ -0,0 +1,164 @@ +from contextlib import contextmanager +from typing import Any +import torch +from torch import nn +from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +def _tensor_placeholder(tensor: torch.Tensor, + device: torch.device) -> torch.Tensor: + """Create a rank-preserving empty placeholder on the specified device.""" + shape = (0, ) if tensor.ndim <= 0 else (0, ) * tensor.ndim + return torch.empty(shape, device=device, dtype=tensor.dtype) + + +class LayerwiseOffloadState: + + def __init__( + self, + async_copy_stream: torch.cuda.Stream, + device: torch.device, + next_state: "LayerwiseOffloadState | None" = None, + ) -> None: + self.async_copy_stream = async_copy_stream + self.next_state = next_state + self.gpu_named_parameters: dict[str, torch.Tensor] = {} + self.cpu_named_parameters: dict[str, torch.Tensor] = {} + self.module_ref: nn.Module = None # type: ignore + self.device: torch.device = device + + def _will_offload(self, name: str) -> bool: + return True + + @torch.compiler.disable + def on_init(self, module: nn.Module): + self.module_ref = module + for name, param in self.module_ref.named_parameters(): + if self._will_offload(name): + self.cpu_named_parameters[name] = ( + param.data.detach().to("cpu").pin_memory()) + param.data = _tensor_placeholder(param.data, self.device) + + @torch.compiler.disable + def wait_and_replace_params(self): + torch.cuda.current_stream().wait_stream(self.async_copy_stream) + # now gpu_named_parameters are ready + for name, param in self.module_ref.named_parameters(): + if not self._will_offload(name): + continue + if name not in self.gpu_named_parameters: + # first load with blocking load + self.gpu_named_parameters[name] = self.cpu_named_parameters[ + name].to(self.device) + param.data = self.gpu_named_parameters[name] + + @torch.compiler.disable + def prefetch_params(self): + compute_stream = torch.cuda.current_stream() + with torch.cuda.stream(self.async_copy_stream): + for name, param in self.module_ref.named_parameters(): + if not self._will_offload(name): + continue + assert name not in self.gpu_named_parameters + gpu_param = self.cpu_named_parameters[name].to( + self.device, non_blocking=True) + gpu_param.record_stream( + compute_stream + ) # ensure tensor will not be freed until forward is completed + self.gpu_named_parameters[name] = gpu_param + + @torch.compiler.disable + def release_gpu_params(self): + for name, param in self.module_ref.named_parameters(): + if self._will_offload(name): + param.data = _tensor_placeholder(param.data, self.device) + del self.gpu_named_parameters[name] + assert len(self.gpu_named_parameters) == 0 + + +class LayerwiseOffloadHook(ForwardHook): + """A hook that enables layerwise CPU offloading during forward pass.""" + + def __init__(self, state: LayerwiseOffloadState) -> None: + self.state = state + + def on_attach(self, module: nn.Module): + self.state.on_init(module) # pyright: ignore + + def on_detach(self, module: nn.Module): + named_parameters = dict(module.named_parameters()) + for name, cpu_tensor in self.state.cpu_named_parameters.items(): + if name not in self.state.gpu_named_parameters: + if name in named_parameters: + named_parameters[name].data = cpu_tensor.to( + device=self.state.device) + else: + logger.warning( + "Parameter {} not found in module during detachment.", + name, + ) + + @classmethod + def name(cls) -> str: + return "LayerwiseOffloadHook" + + def pre_forward(self, module: nn.Module, *args, **kwargs): + self.state.wait_and_replace_params() # pyright: ignore + if self.state.next_state is not None: + self.state.next_state.prefetch_params() # pyright: ignore + return args, kwargs + + def post_forward(self, module: torch.nn.Module, output: Any): + self.state.release_gpu_params() # pyright: ignore + return output + + @contextmanager + def mutate_params_scope(self): + try: + # load params to GPU and keep them there + self.state.wait_and_replace_params() # pyright: ignore + yield + finally: + # instead of releasing, we should overwrite the original params since they have been modified + self.state.cpu_named_parameters.clear() + self.state.gpu_named_parameters.clear() + self.state.on_init(self.state.module_ref) # pyright: ignore + + +def enable_layerwise_offload(model: nn.Module, is_replace: bool = False): + if torch.cuda.is_available(): + device = torch.device("cuda", torch.cuda.current_device()) + else: + logger.warning( + "CUDA is not available. Layerwise offloading is disabled.") + return + state_list = [] + async_stream = torch.cuda.Stream() + for name, submodule in model.named_children(): + if isinstance(submodule, nn.ModuleList): + for idx, module_entry in enumerate(submodule): + state = LayerwiseOffloadState(async_copy_stream=async_stream, + device=device) + state_list.append(state) + hook_mgr = ModuleHookManager.get_from_or_default(module_entry) + hook = LayerwiseOffloadHook(state) + if is_replace: + existing_hook = hook_mgr.forward_hooks.get(hook.name()) + if existing_hook is not None: + hook_mgr.replace_forward_hook(hook.name(), hook) + else: + raise AssertionError( + f"Expect hook exists in {name} for replacement.") + else: + hook_mgr.append_forward_hook(hook) + break + if len(state_list) == 0: + raise ValueError( + "No nn.ModuleList found in the model for layerwise offloading.") + + # circular linking of states + for i in range(len(state_list)): + state_list[i].next_state = state_list[(i + 1) % len(state_list)] diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index e21ccca36..1c611f7b7 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/fastvideo/models/dits/wanvideo.py @@ -750,18 +750,8 @@ class WanTransformer3DModel(CachableDiT): block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, attention_mask) else: - offload_mgr = getattr(self, "_layerwise_offload_manager", None) - use_offload = offload_mgr is not None and getattr(offload_mgr, "enabled", False) - - for i, block in enumerate(self.blocks): - scope = offload_mgr.layer_scope( - prefetch_layer_idx=i + 1 if i + 1 < len(self.blocks) else None, - release_layer_idx=i, - non_blocking=True, - ) if use_offload else nullcontext() - - with scope: - hidden_states = block(hidden_states, encoder_hidden_states, + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, attention_mask) # if teacache is enabled, we need to cache the original hidden states diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 510d487ab..d61c574fa 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -35,7 +35,7 @@ from fastvideo.models.loader.weight_utils import ( ) from fastvideo.models.registry import ModelRegistry from fastvideo.utils import PRECISION_TO_TYPE, is_pin_memory_available -from fastvideo.models.layerwise_offload import LayerwiseOffloadManager +from fastvideo.hooks.layerwise_offload import enable_layerwise_offload logger = init_logger(__name__) @@ -280,7 +280,7 @@ class TextEncoderLoader(ComponentLoader): target_device: torch.device, fastvideo_args: FastVideoArgs, dtype: str = "fp16", - use_text_encoder_override: bool = False, # prevent subclasses from misusing + use_text_encoder_override: bool = False, # prevent subclasses from misusing ): use_cpu_offload = ( fastvideo_args.text_encoder_cpu_offload @@ -297,7 +297,10 @@ class TextEncoderLoader(ComponentLoader): ) # Set quantization config if specified - if use_text_encoder_override and fastvideo_args.override_text_encoder_quant is not None: + if ( + use_text_encoder_override + and fastvideo_args.override_text_encoder_quant is not None + ): if fastvideo_args.override_text_encoder_safetensors is None: raise ValueError( "override_text_encoder_quant is set but override_text_encoder_safetensors is None" @@ -314,7 +317,10 @@ class TextEncoderLoader(ComponentLoader): model: TextEncoder = model_cls(model_config) # type: ignore weights_to_load = {name for name, _ in model.named_parameters()} - if use_text_encoder_override and fastvideo_args.override_text_encoder_safetensors is not None: + if ( + use_text_encoder_override + and fastvideo_args.override_text_encoder_safetensors is not None + ): loaded_weights: set[str] = model.load_weights( safetensors_weights_iterator( [fastvideo_args.override_text_encoder_safetensors], @@ -342,7 +348,6 @@ class TextEncoderLoader(ComponentLoader): if use_cpu_offload: pin_cpu_memory = fastvideo_args.pin_cpu_memory and is_pin_memory_available() - pin_cpu_memory = False # Disable FSDP for MPS as it's not compatible if current_platform.is_mps(): logger.info( @@ -673,35 +678,8 @@ class TransformerLoader(ComponentLoader): model = model.eval() - if fastvideo_args.dit_layerwise_offload and hasattr(model, "blocks"): - # Check if this is a Wan model (only Wan models support layerwise offload) - is_wan_model = "Wan" in cls_name - if not is_wan_model: - logger.warning( - "Layerwise offload is currently only supported for Wan models. " - "Model class '%s' does not support layerwise offload. " - "Disabling layerwise offload for this model.", - cls_name - ) - else: - try: - num_layers = len(getattr(model, "blocks")) - except TypeError: - num_layers = None - if isinstance(num_layers, int) and num_layers > 0: - # Ensure model is on the correct device (CUDA) before initializing manager - # This ensures non-managed parameters (embeddings, final norms) are on GPU - model = model.to(get_local_torch_device()) - mgr = LayerwiseOffloadManager( - model, - module_list_attr="blocks", - num_layers=num_layers, - enabled=True, - pin_cpu_memory=fastvideo_args.pin_cpu_memory, - auto_initialize=True, - ) - setattr(model, "_layerwise_offload_manager", mgr) - + if fastvideo_args.inference_mode and fastvideo_args.dit_layerwise_offload: + enable_layerwise_offload(model) return model diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 53faabc52..97b18bcb6 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -112,7 +112,6 @@ def maybe_load_fsdp_model( if use_fsdp: pin_cpu_memory = pin_cpu_memory and is_pin_memory_available() - pin_cpu_memory = False world_size = hsdp_replicate_dim * hsdp_shard_dim if not training_mode and not fsdp_inference: hsdp_replicate_dim = world_size diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index d9fd4b424..d2dbdde50 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/fastvideo/pipelines/lora_pipeline.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 from collections import defaultdict from collections.abc import Hashable +from contextlib import nullcontext from typing import Any +from collections.abc import Generator import torch import torch.distributed as dist @@ -12,8 +14,13 @@ from torch.distributed.tensor import DTensor from fastvideo.distributed import get_local_torch_device from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs -from fastvideo.layers.lora.linear import (BaseLayerWithLoRA, get_lora_layer, - replace_submodule) +from fastvideo.hooks.hooks import ModuleHookManager +from fastvideo.hooks.layerwise_offload import LayerwiseOffloadHook +from fastvideo.layers.lora.linear import ( + BaseLayerWithLoRA, + get_lora_layer, + replace_submodule, +) from fastvideo.logger import init_logger from fastvideo.models.loader.utils import get_param_names_mapping from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase @@ -22,17 +29,89 @@ from fastvideo.utils import maybe_download_lora logger = init_logger(__name__) +def _get_hook_ctx(module: nn.Module | None): + if module is None: + return nullcontext() + hook_mgr = ModuleHookManager.get_from(module) + if hook_mgr is not None: + offload_hook = hook_mgr.forward_hooks.get(LayerwiseOffloadHook.name()) + if offload_hook is not None: + return offload_hook.mutate_params_scope() # type: ignore + return nullcontext() + + +def _named_module_by_prefix( + module: nn.Module, prefixes: list[str] +) -> list[tuple[str | None, list[tuple[str, nn.Module]]]]: + none_list: list[tuple[str, nn.Module]] = [] + prefix_list: list[tuple[str, list[tuple[str, nn.Module]]]] = [ + (prefix, []) for prefix in prefixes + ] + for name, submodule in module.named_modules(): + for cur_prefix, cur_list in prefix_list: + # we should exclude e.g. block.1 and block.12.attn + if name.startswith(cur_prefix + "."): + cur_list.append((name, submodule)) + break + else: + none_list.append((name, submodule)) + return prefix_list + [(None, none_list)] # type: ignore + + +class LoRAModelLayers: + + def __init__(self, block_list: list[tuple[str, nn.Module]]) -> None: + # block_name -> {layer_name -> layer} + self.block_to_lora_layers: dict[str, dict[str, BaseLayerWithLoRA]] = {} + # layer_name -> block_name + self.lora_layers_to_block: dict[str, str | None] = {} + self.other_lora_layers: dict[str, BaseLayerWithLoRA] = {} + self.block_mapping = dict(block_list) + + def add_lora_layer(self, block_name: str | None, layer_name: str, + layer: BaseLayerWithLoRA): + if block_name is None: + self.other_lora_layers[layer_name] = layer + self.lora_layers_to_block[layer_name] = None + else: + if block_name not in self.block_to_lora_layers: + self.block_to_lora_layers[block_name] = {} + self.block_to_lora_layers[block_name][layer_name] = layer + self.lora_layers_to_block[layer_name] = block_name + + def all_lora_layers( + self, ) -> Generator[tuple[str, BaseLayerWithLoRA], Any, None]: + for block_layers in self.block_to_lora_layers.values(): + for name, layer in block_layers.items(): + yield name, layer + for name, layer in self.other_lora_layers.items(): + yield name, layer + + def lora_layers_by_block( + self, + ) -> Generator[ + tuple[nn.Module | None, dict[str, BaseLayerWithLoRA]], + Any, + None, + ]: + for block_name, layers in self.block_to_lora_layers.items(): + yield self.block_mapping[block_name], layers + yield None, self.other_lora_layers + + class LoRAPipeline(ComposedPipelineBase): """ Pipeline that supports injecting LoRA adapters into the diffusion transformer. TODO: support training. """ + lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict( dict ) # 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, dict[str, BaseLayerWithLoRA]] = {} + # model_name -> layers + lora_layers: dict[str, LoRAModelLayers] = {} fastvideo_args: FastVideoArgs | TrainingArgs exclude_lora_layers: dict[str, list[str]] = {} device: torch.device = get_local_torch_device() @@ -48,10 +127,10 @@ class LoRAPipeline(ComposedPipelineBase): self.device = get_local_torch_device() # 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] + 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( @@ -59,19 +138,23 @@ class LoRAPipeline(ComposedPipelineBase): ) secondary_transformer_name = transformer_name + "_2" - if secondary_transformer_name in self.modules and self.modules[ - secondary_transformer_name] is not None: + 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()) + 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 + 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 @@ -83,20 +166,33 @@ class LoRAPipeline(ComposedPipelineBase): self.fastvideo_args.lora_alpha = self.fastvideo_args.lora_rank self.lora_rank = self.fastvideo_args.lora_rank # type: ignore self.lora_alpha = self.fastvideo_args.lora_alpha # type: ignore - logger.info("Using LoRA training with rank %d and alpha %d", - self.lora_rank, self.lora_alpha) + logger.info( + "Using LoRA training with rank %d and alpha %d", + self.lora_rank, + self.lora_alpha, + ) 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_gate_compress" + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "to_q", + "to_k", + "to_v", + "to_out", + "to_qkv", + "to_gate_compress", ] logger.info( "Using default lora_target_modules for all transformers: %s", - self.lora_target_modules) + 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.lora_target_modules, + ) self.convert_to_lora_layers() # Inference @@ -104,7 +200,8 @@ class LoRAPipeline(ComposedPipelineBase): self.convert_to_lora_layers() self.set_lora_adapter( self.lora_nickname, # type: ignore - self.lora_path) # type: ignore + self.lora_path, + ) # type: ignore def is_target_layer(self, module_name: str) -> bool: if self.lora_target_modules is None: @@ -114,9 +211,9 @@ class LoRAPipeline(ComposedPipelineBase): def set_trainable(self) -> None: - def set_lora_grads(lora_layers: dict[str, BaseLayerWithLoRA], + def set_lora_grads(lora_layers: LoRAModelLayers, device_mesh: DeviceMesh): - for name, layer in lora_layers.items(): + for name, layer in lora_layers.all_lora_layers(): layer.lora_A.requires_grad_(True) layer.lora_B.requires_grad_(True) layer.base_layer.requires_grad_(False) @@ -131,10 +228,15 @@ class LoRAPipeline(ComposedPipelineBase): super().set_trainable() return - device_mesh = init_device_mesh("cuda", (dist.get_world_size(), 1), - mesh_dim_names=["fake", "replicate"]) - for transformer_name, transformer_module in self.trainable_transformer_modules.items( - ): + device_mesh = init_device_mesh( + "cuda", + (dist.get_world_size(), 1), + mesh_dim_names=["fake", "replicate"], + ) + 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: @@ -151,32 +253,71 @@ class LoRAPipeline(ComposedPipelineBase): if self.lora_initialized: return self.lora_initialized = True - for transformer_name, transformer_module in self.trainable_transformer_modules.items( - ): + for ( + transformer_name, + transformer_module, + ) in self.trainable_transformer_modules.items(): converted_count = 0 + # init bookkeeping structures 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 + # get block list + block_list = [] + for name, submodule in transformer_module.named_children(): + if isinstance(submodule, nn.ModuleList): + block_list = [(f"{name}.{i}", m) + for i, m in enumerate(submodule)] break - if excluded: - continue + self.lora_layers[transformer_name] = LoRAModelLayers(block_list) + logger.info("Converting %s to LoRA Transformer", transformer_name) + # scan every module and convert to LoRA layer if applicable - 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[transformer_name][name] = layer - replace_submodule(transformer_module, name, layer) - converted_count += 1 + for block_name, block_modules in _named_module_by_prefix( + transformer_module, + list(self.lora_layers[transformer_name].block_mapping), + ): + if block_name is not None and ( + not self.fastvideo_args.training_mode + and self.fastvideo_args.dit_layerwise_offload): + scope_ctx = _get_hook_ctx( + self.lora_layers[transformer_name]. + block_mapping[block_name]) + else: + scope_ctx = nullcontext() + with scope_ctx: + for name, layer in block_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: + block_name_split = name.split(".", 2) + if len(block_name_split) > 2: + block_name = (block_name_split[0] + "." + + block_name_split[1]) + else: + block_name = None + if (block_name + not in self.lora_layers[transformer_name]. + block_mapping): + block_name = None + self.lora_layers[transformer_name].add_lora_layer( + block_name, name, layer) + replace_submodule(transformer_module, name, layer) + converted_count += 1 logger.info("Converted %d layers to LoRA layers", converted_count) def set_lora_adapter(self, @@ -209,7 +350,7 @@ class LoRAPipeline(ComposedPipelineBase): # Extract alpha values and weights in a single pass to_merge_params: defaultdict[Hashable, - dict[Any, Any]] = defaultdict(dict) + dict[Any, Any]] = (defaultdict(dict)) for name, weight in lora_state_dict.items(): # Extract weights (lora_A, lora_B, and lora_alpha) name = name.replace("diffusion_model.", "") @@ -223,13 +364,14 @@ class LoRAPipeline(ComposedPipelineBase): target_name, _, _ = param_names_mapping_fn(layer_name) # Store alpha alongside weights with same target_name base alpha_key = target_name + ".lora_alpha" - self.lora_adapters[lora_nickname][alpha_key] = weight.item( - ) if weight.numel() == 1 else float(weight.mean()) + self.lora_adapters[lora_nickname][alpha_key] = ( + weight.item() + if weight.numel() == 1 else float(weight.mean())) continue name, _, _ = lora_param_names_mapping_fn(name) - target_name, merge_index, num_params_to_merge = param_names_mapping_fn( - name) + target_name, merge_index, num_params_to_merge = ( + param_names_mapping_fn(name)) # for (in_dim, r) @ (r, out_dim), we only merge (r, out_dim * n) where n is the number of linear layers to fuse # see param mapping in HunyuanVideoArchConfig if merge_index is not None and "lora_B" in name: @@ -261,45 +403,84 @@ class LoRAPipeline(ComposedPipelineBase): # Merge the new adapter adapted_count = 0 - 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 - logger.info("Rank %d: LoRA adapter %s applied to %d layers", rank, - lora_path, adapted_count) + for ( + transformer_name, + transformer_lora_layers, + ) in self.lora_layers.items(): + for ( + module, + layers, + ) in transformer_lora_layers.lora_layers_by_block(): + with _get_hook_ctx(module): + for name, layer in 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) + try: + layer.set_lora_weights( + lora_A, + lora_B, + lora_alpha=alpha, + training_mode=self.fastvideo_args. + training_mode, + lora_path=lora_path, + ) + except Exception as e: + logger.error( + "Error setting LoRA weights for layer %s: %s", + name, + str(e), + ) + raise e + 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 transformer_name, transformer_lora_layers in self.lora_layers.items( - ): - for name, layer in transformer_lora_layers.items(): - layer.merge_lora_weights() + for ( + transformer_name, + transformer_lora_layers, + ) in self.lora_layers.items(): + for ( + module, + layers, + ) in transformer_lora_layers.lora_layers_by_block(): + with _get_hook_ctx(module): + for name, layer in layers.items(): + layer.merge_lora_weights() def unmerge_lora_weights(self) -> None: - for transformer_name, transformer_lora_layers in self.lora_layers.items( - ): - for name, layer in transformer_lora_layers.items(): - layer.unmerge_lora_weights() + for ( + transformer_name, + transformer_lora_layers, + ) in self.lora_layers.items(): + for ( + module, + layers, + ) in transformer_lora_layers.lora_layers_by_block(): + with _get_hook_ctx(module): + for name, layer in layers.items(): + layer.unmerge_lora_weights() diff --git a/fastvideo/tests/hooks/test_hooks_api.py b/fastvideo/tests/hooks/test_hooks_api.py new file mode 100644 index 000000000..5142f8449 --- /dev/null +++ b/fastvideo/tests/hooks/test_hooks_api.py @@ -0,0 +1,60 @@ +from fastvideo.hooks.hooks import ForwardHook, ModuleHookManager +from torch import nn +from typing import Any +import torch + + +class EventHook(ForwardHook): + def __init__(self, content: str, event_list: list[str]): + self.content = content + self.event_list = event_list + + def name(self) -> str: + return f"EventHook_{self.content}" + + def pre_forward(self, module: nn.Module, *args, **kwargs): + print( + f"[{self.content}] Pre-forward called with args[0].shape: {args[0].shape}" + ) + self.event_list.append(f"[pre]{self.content}") + return args, kwargs + + def post_forward(self, module: nn.Module, output: Any): + print( + f"[{self.content}] Post-forward called with outputs.shape: {output.shape}" + ) + self.event_list.append(f"[post]{self.content}") + return output + + +def test_hook_execution_order(): + """Test that hooks are executed in the correct order: LIFO for pre-hooks, FIFO for post-hooks.""" + # Create a simple model + model = nn.Linear(10, 20) + + # Create event list to track hook execution order + events = [] + + # Create and push hooks in order: A then B + + manager = ModuleHookManager.get_from_or_default(model) + + hook_a = EventHook("A", events) + hook_b = EventHook("B", events) + + manager.append_forward_hook(hook_a) + manager.append_forward_hook(hook_b) + + # Perform a forward pass + input_tensor = torch.randn(2, 10) + model(input_tensor) + + # Verify the execution order is [pre_a, pre_b, post_b, post_a] + # Pre-hooks should be FILO (First In Last Out): A then B + # Post-hooks should be LIFO (Last In First Out): B then A + expected_events = ["[pre]A", "[pre]B", "[post]B", "[post]A"] + + assert events == expected_events, ( + f"Expected {expected_events}, but got {events}" + ) + print(f"✓ Hook execution order test passed: {events}") diff --git a/fastvideo/tests/hooks/test_layerwise_offload.py b/fastvideo/tests/hooks/test_layerwise_offload.py new file mode 100644 index 000000000..3ad5d3144 --- /dev/null +++ b/fastvideo/tests/hooks/test_layerwise_offload.py @@ -0,0 +1,395 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +import torch.nn as nn + +from fastvideo.hooks.layerwise_offload import ( + LayerwiseOffloadHook, + enable_layerwise_offload, +) +from fastvideo.hooks.hooks import ModuleHookManager + + +class SimpleBlock(nn.Module): + """A simple block with linear layers for testing.""" + + def __init__(self, hidden_size: int, dtype: torch.dtype = torch.float32): + super().__init__() + self.linear1 = nn.Linear(hidden_size, hidden_size, dtype=dtype) + self.linear2 = nn.Linear(hidden_size, hidden_size, dtype=dtype) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.linear1(x) + x = torch.relu(x) + x = self.linear2(x) + return x + + +class SimpleModelWithModuleList(nn.Module): + """A simple model with ModuleList for testing layerwise offloading.""" + + def __init__( + self, + num_blocks: int = 4, + hidden_size: int = 128, + dtype: torch.dtype = torch.float32, + ): + super().__init__() + self.blocks = nn.ModuleList( + [SimpleBlock(hidden_size, dtype=dtype) for _ in range(num_blocks)] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + x = block(x) + return x + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_basic(): + """Test basic functionality of layerwise offloading.""" + device = torch.device("cuda") + hidden_size = 128 + batch_size = 2 + seq_len = 16 + num_blocks = 4 + + # Create model + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + + # Get reference output without offloading + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) + with torch.no_grad(): + reference_output = model(input_tensor.clone()) + + # Enable layerwise offloading + enable_layerwise_offload(model) + + # Verify parameters are offloaded to CPU + for block in model.blocks: + for param in block.parameters(): + # Parameters should be placeholder tensors (empty) + assert param.numel() == 0, ( + "Parameters should be offloaded (empty tensors)" + ) + + # Run forward pass with offloading + with torch.no_grad(): + offloaded_output = model(input_tensor.clone()) + + # Check output correctness + assert torch.allclose( + reference_output, offloaded_output, rtol=1e-4, atol=1e-5 + ), "Output with offloading should match reference output" + + print( + f" Layerwise offload basic test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_bf16(): + """Test layerwise offloading with bfloat16 precision.""" + device = torch.device("cuda") + hidden_size = 256 + batch_size = 1 + seq_len = 32 + num_blocks = 3 + + # Create model + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.bfloat16 + ).to(device) + + # Get reference output without offloading + input_tensor = torch.randn( + batch_size, seq_len, hidden_size, device=device, dtype=torch.bfloat16 + ) + with torch.no_grad(): + reference_output = model(input_tensor.clone()) + + # Enable layerwise offloading + enable_layerwise_offload(model) + + # Run forward pass with offloading + with torch.no_grad(): + offloaded_output = model(input_tensor.clone()) + + # Check output correctness (looser tolerance for bf16) + assert torch.allclose( + reference_output, offloaded_output, rtol=1e-2, atol=1e-3 + ), "Output with offloading should match reference output for bf16" + + print( + f" Layerwise offload bf16 test passed: max diff = {(reference_output - offloaded_output).abs().max().item()}" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_multiple_forward_passes(): + """Test that layerwise offloading works correctly across multiple forward passes.""" + device = torch.device("cuda") + hidden_size = 64 + batch_size = 2 + seq_len = 8 + num_blocks = 3 + num_iterations = 5 + + # Create model + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + + # Get reference outputs without offloading + torch.manual_seed(42) + reference_outputs = [] + for i in range(num_iterations): + input_tensor = torch.randn( + batch_size, seq_len, hidden_size, device=device + ) + with torch.no_grad(): + reference_outputs.append(model(input_tensor.clone())) + + # Enable layerwise offloading + enable_layerwise_offload(model) + + # Run multiple forward passes with offloading + torch.manual_seed(42) + for i in range(num_iterations): + input_tensor = torch.randn( + batch_size, seq_len, hidden_size, device=device + ) + with torch.no_grad(): + offloaded_output = model(input_tensor.clone()) + + # Check output correctness for each iteration + assert torch.allclose( + reference_outputs[i], offloaded_output, rtol=1e-4, atol=1e-5 + ), f"Output mismatch in iteration {i}" + + print( + f" Multiple forward passes test passed for {num_iterations} iterations" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_parameter_integrity(): + """Test that parameters are correctly restored during forward pass.""" + device = torch.device("cuda") + hidden_size = 128 + num_blocks = 2 + + # Create model + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + + # Store original parameter values per block + original_params_per_block = [] + for block in model.blocks: + block_params = {} + for name, param in block.named_parameters(): + block_params[name] = param.data.clone() + original_params_per_block.append(block_params) + + # Enable layerwise offloading + enable_layerwise_offload(model) + + # Create hook managers and verify they exist + for block_idx, block in enumerate(model.blocks): + manager = ModuleHookManager.get_from(block) + assert manager is not None, "Hook manager should be attached to blocks" + hook: LayerwiseOffloadHook | None = manager.get_forward_hook( + "LayerwiseOffloadHook" + ) + assert hook is not None, "LayerwiseOffloadHook should be registered" + + # Verify parameters are stored in CPU + state = hook.state + assert len(state.cpu_named_parameters) > 0, ( + "CPU parameters should be stored" + ) + + # Verify CPU parameters match original values for this block + original_params = original_params_per_block[block_idx] + for name, cpu_param in state.cpu_named_parameters.items(): + assert name in original_params, ( + f"CPU parameter {name} not found in original params for block {block_idx}" + ) + assert torch.allclose( + cpu_param.cpu(), original_params[name].cpu(), rtol=1e-5, atol=1e-6 + ), f"CPU parameter {name} should match original for block {block_idx}" + + print(" Parameter integrity test passed") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_no_modulelist_error(): + """Test that enabling offload on a model without ModuleList raises an error.""" + + class ModelWithoutModuleList(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(128, 128) + + def forward(self, x): + return self.linear(x) + + model = ModelWithoutModuleList().to("cuda") + + with pytest.raises( + ValueError, + match="No nn.ModuleList found in the model for layerwise offloading", + ): + enable_layerwise_offload(model) + + print(" No ModuleList error test passed") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_memory_reduction(): + """Test that layerwise offloading reduces GPU memory usage.""" + device = torch.device("cuda") + hidden_size = 512 + num_blocks = 8 + batch_size = 1 + seq_len = 64 + + # Create model + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + + # Measure initial GPU memory + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) + with torch.no_grad(): + _ = model(input_tensor) + + memory_without_offload = torch.cuda.max_memory_allocated() + + # Reset model + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + + # Enable offloading + enable_layerwise_offload(model) + + # Measure GPU memory with offloading + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) + with torch.no_grad(): + _ = model(input_tensor) + + memory_with_offload = torch.cuda.max_memory_allocated() + + # Memory with offload should be less (parameters are offloaded) + # Note: This is a weak check as memory usage depends on many factors + print( + f"Memory without offload: {memory_without_offload / 1024**2:.2f} MB, " + f"with offload: {memory_with_offload / 1024**2:.2f} MB" + ) + + # We expect some reduction but this test is more informational + assert memory_with_offload < memory_without_offload * 1.5, ( + "Memory usage should not increase significantly with offloading" + ) + + print(" Memory reduction test passed") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_gradient_disabled(): + """Test that layerwise offloading works correctly with gradients disabled.""" + device = torch.device("cuda") + hidden_size = 128 + batch_size = 2 + seq_len = 16 + num_blocks = 3 + + # Create model + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + model.eval() + + # Enable layerwise offloading + enable_layerwise_offload(model) + + input_tensor = torch.randn(batch_size, seq_len, hidden_size, device=device) + + # Forward pass should work without gradients + with torch.no_grad(): + output = model(input_tensor) + + assert output.requires_grad is False, "Output should not require gradients" + assert output.shape == ( + batch_size, + seq_len, + hidden_size, + ), "Output shape should match input" + + print(" Gradient disabled test passed") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_layerwise_offload_different_batch_sizes(): + """Test layerwise offloading with different batch sizes.""" + device = torch.device("cuda") + hidden_size = 128 + seq_len = 16 + num_blocks = 3 + + model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + + # Get reference model without offloading + reference_model = SimpleModelWithModuleList( + num_blocks=num_blocks, hidden_size=hidden_size, dtype=torch.float32 + ).to(device) + reference_model.load_state_dict(model.state_dict()) + + # Enable layerwise offloading + enable_layerwise_offload(model) + + # Test different batch sizes + for batch_size in [1, 2, 4, 8]: + input_tensor = torch.randn( + batch_size, seq_len, hidden_size, device=device + ) + + with torch.no_grad(): + reference_output = reference_model(input_tensor.clone()) + offloaded_output = model(input_tensor.clone()) + + assert torch.allclose( + reference_output, offloaded_output, rtol=1e-4, atol=1e-5 + ), f"Output mismatch for batch_size={batch_size}" + + print(" Different batch sizes test passed") + + +if __name__ == "__main__": + # Run tests manually for debugging + if torch.cuda.is_available(): + print("Running layerwise offloading tests...") + test_layerwise_offload_basic() + test_layerwise_offload_bf16() + test_layerwise_offload_multiple_forward_passes() + test_layerwise_offload_parameter_integrity() + test_layerwise_offload_no_modulelist_error() + test_layerwise_offload_memory_reduction() + test_layerwise_offload_gradient_disabled() + test_layerwise_offload_different_batch_sizes() + print("\n All tests passed!") + else: + print("CUDA not available, skipping tests") diff --git a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py index ad919db5b..ef6d05041 100644 --- a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py +++ b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py @@ -93,8 +93,11 @@ def test_merge_lora_weights(model_id): lora_nickname = lora_config["lora_nickname"] lora_path = lora_config["lora_path"] + # When layerwise offload is enabled, placeholder tensors cannot be compared directly. args = FastVideoArgs.from_kwargs( model_path=model_id, + dit_layerwise_offload=False, + use_fsdp_inference=True, dit_cpu_offload=True, dit_precision="bf16", ) diff --git a/fastvideo/tests/ssim/test_inference_similarity.py b/fastvideo/tests/ssim/test_inference_similarity.py index 720a6fbf8..ba6b7afb0 100644 --- a/fastvideo/tests/ssim/test_inference_similarity.py +++ b/fastvideo/tests/ssim/test_inference_similarity.py @@ -234,6 +234,7 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id): "tp_size": BASE_PARAMS["tp_size"], "use_fsdp_inference": True, "dit_cpu_offload": False, + "dit_layerwise_offload": False, } if BASE_PARAMS.get("vae_sp"): init_kwargs["vae_sp"] = True diff --git a/fastvideo/tests/ssim/test_matrixgame_similarity.py b/fastvideo/tests/ssim/test_matrixgame_similarity.py index 65ec83564..e8c394a8e 100644 --- a/fastvideo/tests/ssim/test_matrixgame_similarity.py +++ b/fastvideo/tests/ssim/test_matrixgame_similarity.py @@ -88,6 +88,7 @@ def test_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id): init_kwargs = { "num_gpus": BASE_PARAMS["num_gpus"], "use_fsdp_inference": True, + "dit_layerwise_offload": False, "dit_cpu_offload": False, "vae_cpu_offload": False, "text_encoder_cpu_offload": True, diff --git a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py index 67eb419b3..147c63cd8 100644 --- a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py +++ b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py @@ -218,6 +218,7 @@ def test_turbodiffusion_i2v_inference_similarity(prompt, model_id): # Keep both transformers in VRAM - avoids CPU RAM bottleneck "dit_cpu_offload": False, "use_fsdp_inference": True, + "dit_layerwise_offload": False, } generation_kwargs = {