[feat] Hooks API and layerwise offloading for all DiTs (#1006)

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

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)

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)]

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
 

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
 
 

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

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()

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}")

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")

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",
     )

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

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,

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 = {