From f72ced76112e92aa8b1620729303a0b70bf14bac Mon Sep 17 00:00:00 2001 From: Ohm-Rishabh <119029179+Ohm-Rishabh@users.noreply.github.com> Date: Sat, 3 Jan 2026 21:46:00 -0800 Subject: [PATCH] Layer offloading (#966) --- fastvideo/fastvideo_args.py | 19 ++ fastvideo/models/dits/wanvideo.py | 17 +- fastvideo/models/layerwise_offload.py | 224 ++++++++++++++++++++ fastvideo/models/loader/component_loader.py | 31 +++ fastvideo/pipelines/stages/denoising.py | 36 +++- 5 files changed, 320 insertions(+), 7 deletions(-) create mode 100644 fastvideo/models/layerwise_offload.py diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 7edf49028..5845449a9 100644 --- a/fastvideo/fastvideo_args.py +++ b/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( diff --git a/fastvideo/models/dits/wanvideo.py b/fastvideo/models/dits/wanvideo.py index fe3550e8a..e21ccca36 100644 --- a/fastvideo/models/dits/wanvideo.py +++ b/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: diff --git a/fastvideo/models/layerwise_offload.py b/fastvideo/models/layerwise_offload.py new file mode 100644 index 000000000..836d731ca --- /dev/null +++ b/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) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 2be4f32b2..c76ddc95b 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/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 diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index dd748c95d..7669b26df 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/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 \ No newline at end of file + return batch