Layer offloading (#966)

fastvideo/fastvideo_args.py

@@ -133,6 +133,7 @@ class FastVideoArgs:
     # CPU offload parameters
     dit_cpu_offload: bool = True
     use_fsdp_inference: bool = True
+    dit_layerwise_offload: bool = False
     text_encoder_cpu_offload: bool = True
     image_encoder_cpu_offload: bool = True
     vae_cpu_offload: bool = True
@@ -421,6 +422,11 @@ class FastVideoArgs:
             help=
             "Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
         )
+        parser.add_argument(
+            "--dit-layerwise-offload",
+            action=StoreBoolean,
+            help="Enable layerwise CPU offload with async H2D prefetch overlap.",
+        )
         parser.add_argument(
             "--use-fsdp-inference",
             action=StoreBoolean,
@@ -603,6 +609,19 @@ class FastVideoArgs:
 
         if current_platform.is_mps():
             self.use_fsdp_inference = False
+            self.dit_layerwise_offload = False
+
+        if self.dit_layerwise_offload:
+            if self.use_fsdp_inference:
+                logger.warning(
+                    "dit_layerwise_offload is enabled, automatically disabling use_fsdp_inference."
+                )
+                self.use_fsdp_inference = False
+            if self.dit_cpu_offload:
+                logger.warning(
+                    "dit_layerwise_offload is enabled, automatically disabling dit_cpu_offload."
+                )
+                self.dit_cpu_offload = False
 
         # Validate mode and inference_mode consistency
         assert isinstance(

fastvideo/models/dits/wanvideo.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import math
+from contextlib import nullcontext
 from typing import Any
 
 import numpy as np
@@ -749,9 +750,19 @@ class WanTransformer3DModel(CachableDiT):
                         block, hidden_states, encoder_hidden_states,
                         timestep_proj, freqs_cis, attention_mask)
             else:
-                for block in self.blocks:
-                    hidden_states = block(hidden_states, encoder_hidden_states,
-                                          timestep_proj, freqs_cis, attention_mask)
+                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,
+                                              timestep_proj, freqs_cis, attention_mask)
             # if teacache is enabled, we need to cache the original hidden states
 
             if enable_teacache:

fastvideo/models/layerwise_offload.py

@@ -0,0 +1,224 @@
+# SPDX-License-Identifier: Apache-2.0
+# Inspired by SGLang's layerwise offload implementation:
+# https://github.com/sgl-project/sglang/pull/15511
+#
+# This implementation provides a lightweight layerwise CPU offload manager
+# with async H2D prefetch using a dedicated CUDA stream, following SGLang's design.
+
+import re
+from contextlib import contextmanager
+from typing import Dict, Set, Optional, Tuple
+
+import torch
+
+
+class LayerwiseOffloadManager:
+    """A lightweight layerwise CPU offload manager.
+
+    Offloads per-layer parameters/buffers from GPU to CPU, and supports async H2D
+    prefetch using a dedicated CUDA stream.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        *,
+        module_list_attr: str,
+        num_layers: int,
+        enabled: bool,
+        pin_cpu_memory: bool = True,
+        auto_initialize: bool = False,
+    ) -> None:
+        self.model = model
+        self.module_list_attr = module_list_attr
+        self.num_layers = int(num_layers)
+        self.pin_cpu_memory = bool(pin_cpu_memory)
+
+        self.enabled = bool(enabled and torch.cuda.is_available())
+        self.device = (
+            torch.device("cuda", torch.cuda.current_device()) if self.enabled else None
+        )
+        self.copy_stream = torch.cuda.Stream() if self.enabled else None
+
+        self._layer_name_re = re.compile(
+            rf"(^|\.){re.escape(module_list_attr)}\.(\d+)(\.|$)"
+        )
+
+        self._cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
+        self._cpu_dtypes: Dict[int, Dict[str, torch.dtype]] = {}
+
+        self._gpu_layers: Dict[int, Set[str]] = {}
+
+        self._named_parameters: Dict[str, torch.nn.Parameter] = {}
+        self._named_buffers: Dict[str, torch.Tensor] = {}
+
+        self._meta: Dict[str, Tuple[int, torch.dtype]] = {}
+
+        if auto_initialize:
+            self.initialize()
+
+    def _match_layer_idx(self, name: str) -> Optional[int]:
+        m = self._layer_name_re.search(name)
+        if not m:
+            return None
+        try:
+            return int(m.group(2))
+        except Exception:
+            return None
+
+    def _record_meta(self, name: str, t: torch.Tensor) -> None:
+        if name not in self._meta:
+            self._meta[name] = (int(t.ndim), t.dtype)
+
+    def _make_placeholder(self, name: str) -> torch.Tensor:
+        """Rank-preserving empty placeholder on GPU."""
+        assert self.device is not None
+        ndim, dtype = self._meta[name]
+        shape = (0,) if ndim <= 0 else (0,) * ndim
+        return torch.empty(shape, device=self.device, dtype=dtype)
+
+    def _get_target(self, name: str) -> torch.Tensor:
+        if name in self._named_parameters:
+            return self._named_parameters[name]
+        return self._named_buffers[name]
+
+    def _offload_tensor(self, name: str, tensor: torch.Tensor, layer_idx: int) -> None:
+        if layer_idx not in self._cpu_weights:
+            self._cpu_weights[layer_idx] = {}
+            self._cpu_dtypes[layer_idx] = {}
+
+        self._record_meta(name, tensor)
+
+        cpu_weight = tensor.detach().to("cpu")
+        if self.pin_cpu_memory:
+            cpu_weight = cpu_weight.pin_memory()
+
+        self._cpu_weights[layer_idx][name] = cpu_weight
+        self._cpu_dtypes[layer_idx][name] = tensor.dtype
+
+        if self.device is not None:
+            tensor.data = self._make_placeholder(name)
+
+    @torch.compiler.disable
+    def initialize(self) -> None:
+        """Offload all matched layer tensors to CPU and prefetch layer 0 (sync)."""
+        if not self.enabled:
+            return
+
+        self._named_parameters = dict(self.model.named_parameters())
+        self._named_buffers = dict(self.model.named_buffers())
+
+        for name, param in self._named_parameters.items():
+            layer_idx = self._match_layer_idx(name)
+            if layer_idx is None or layer_idx >= self.num_layers:
+                continue
+            self._offload_tensor(name, param, layer_idx)
+
+        for name, buf in self._named_buffers.items():
+            layer_idx = self._match_layer_idx(name)
+            if layer_idx is None or layer_idx >= self.num_layers:
+                continue
+            self._offload_tensor(name, buf, layer_idx)
+
+        self.prefetch_layer(0, non_blocking=False)
+        if self.copy_stream is not None:
+            torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+    @torch.compiler.disable
+    def prefetch_layer(self, layer_idx: int, non_blocking: bool = True) -> None:
+        """Prefetch a layer's tensors from CPU to GPU (async on copy_stream)."""
+        if not self.enabled or self.device is None or self.copy_stream is None:
+            return
+        if layer_idx < 0 or layer_idx >= self.num_layers:
+            return
+        if layer_idx in self._gpu_layers:
+            return
+        if layer_idx not in self._cpu_weights:
+            return
+
+        self.copy_stream.wait_stream(torch.cuda.current_stream())
+
+        param_names: Set[str] = set()
+        with torch.cuda.stream(self.copy_stream):
+            for name, cpu_weight in self._cpu_weights[layer_idx].items():
+                target = self._get_target(name)
+
+                gpu_weight = torch.empty(
+                    cpu_weight.shape,
+                    dtype=self._cpu_dtypes[layer_idx][name],
+                    device=self.device,
+                )
+                gpu_weight.copy_(cpu_weight, non_blocking=non_blocking)
+
+                target.data = gpu_weight
+                param_names.add(name)
+
+        self._gpu_layers[layer_idx] = param_names
+
+    @contextmanager
+    def layer_scope(
+        self,
+        *,
+        prefetch_layer_idx: Optional[int],
+        release_layer_idx: Optional[int],
+        non_blocking: bool = True,
+    ):
+        if self.enabled and release_layer_idx is not None:
+            cur = release_layer_idx
+            if (
+                cur not in self._gpu_layers
+                and cur in self._cpu_weights
+                and self.device is not None
+                and self.copy_stream is not None
+            ):
+                self.prefetch_layer(cur, non_blocking=False)
+                torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+        if self.enabled and prefetch_layer_idx is not None:
+            self.prefetch_layer(prefetch_layer_idx, non_blocking=non_blocking)
+
+        try:
+            yield
+        finally:
+            if self.enabled and self.copy_stream is not None:
+                torch.cuda.current_stream().wait_stream(self.copy_stream)
+            if self.enabled and release_layer_idx is not None:
+                self.release_layer(release_layer_idx)
+
+
+    @torch.compiler.disable
+    def release_layer(self, layer_idx: int) -> None:
+        """Release a layer's tensors back to placeholders (free VRAM)."""
+        if not self.enabled or self.device is None:
+            return
+
+        if layer_idx < 0:
+            return
+
+        param_names = self._gpu_layers.pop(layer_idx, None)
+        if not param_names:
+            return
+
+        for name in param_names:
+            target = self._get_target(name)
+            # Ensure meta exists even if something unexpected happened
+            self._record_meta(name, target)
+            target.data = self._make_placeholder(name)
+
+    @torch.compiler.disable
+    def release_all(self) -> None:
+        """Release all currently-resident layers back to placeholders."""
+        if not self.enabled or self.device is None:
+            return
+
+        if self.copy_stream is not None:
+            torch.cuda.current_stream().wait_stream(self.copy_stream)
+
+        for layer_idx in list(self._gpu_layers.keys()):
+            param_names = self._gpu_layers.pop(layer_idx, None)
+            if not param_names:
+                continue
+            for name in param_names:
+                target = self._get_target(name)
+                self._record_meta(name, target)
+                target.data = self._make_placeholder(name)

fastvideo/models/loader/component_loader.py

@@ -36,6 +36,7 @@ from fastvideo.models.loader.weight_utils import (
 )
 from fastvideo.models.registry import ModelRegistry
 from fastvideo.utils import PRECISION_TO_TYPE
+from fastvideo.models.layerwise_offload import LayerwiseOffloadManager
 
 logger = init_logger(__name__)
 
@@ -609,6 +610,36 @@ 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)
+
         return model
 
 

fastvideo/pipelines/stages/denoising.py

@@ -253,19 +253,37 @@ class DenoisingStage(PipelineStage):
 
                 if boundary_timestep is None or t >= boundary_timestep:
                     if (fastvideo_args.dit_cpu_offload
+                            and not fastvideo_args.dit_layerwise_offload
                             and self.transformer_2 is not None and next(
                                 self.transformer_2.parameters()).device.type
                             == 'cuda'):
                         self.transformer_2.to('cpu')
                     current_model = self.transformer
+                    if (fastvideo_args.dit_cpu_offload
+                            and not fastvideo_args.dit_layerwise_offload
+                            and not fastvideo_args.use_fsdp_inference
+                            and current_model is not None):
+                        transformer_device = next(
+                            current_model.parameters()).device.type
+                        if transformer_device == 'cpu':
+                            current_model.to(get_local_torch_device())
                     current_guidance_scale = batch.guidance_scale
                 else:
                     # low-noise stage in wan2.2
-                    if fastvideo_args.dit_cpu_offload and next(
-                            self.transformer.parameters(
-                            )).device.type == 'cuda':
+                    if (fastvideo_args.dit_cpu_offload
+                            and not fastvideo_args.dit_layerwise_offload
+                            and next(self.transformer.parameters()).device.type
+                            == 'cuda'):
                         self.transformer.to('cpu')
                     current_model = self.transformer_2
+                    if (fastvideo_args.dit_cpu_offload
+                            and not fastvideo_args.dit_layerwise_offload
+                            and not fastvideo_args.use_fsdp_inference
+                            and current_model is not None):
+                        transformer_2_device = next(
+                            current_model.parameters()).device.type
+                        if transformer_2_device == 'cpu':
+                            current_model.to(get_local_torch_device())
                     current_guidance_scale = batch.guidance_scale_2
                 assert current_model is not None, "current_model is None"
 
@@ -458,6 +476,16 @@ class DenoisingStage(PipelineStage):
         # Update batch with final latents
         batch.latents = latents
 
+        if fastvideo_args.dit_layerwise_offload:
+            mgr = getattr(self.transformer, "_layerwise_offload_manager", None)
+            if mgr is not None and getattr(mgr, "enabled", False):
+                mgr.release_all()
+            if self.transformer_2 is not None:
+                mgr2 = getattr(self.transformer_2, "_layerwise_offload_manager",
+                               None)
+                if mgr2 is not None and getattr(mgr2, "enabled", False):
+                    mgr2.release_all()
+
         # Save STA mask search results if needed
         if st_attn_available and self.attn_backend == SlidingTileAttentionBackend and fastvideo_args.STA_mode == STA_Mode.STA_SEARCHING:
             self.save_sta_search_results(batch)
@@ -1184,4 +1212,4 @@ class DmdDenoisingStage(DenoisingStage):
         # Update batch with final latents
         batch.latents = latents
 
-        return batch
+        return batch