adding gradient checkpointing (with errors)

docs/source/inference/inference_quick_start.md

@@ -79,7 +79,6 @@ def main():
     sampling_param = SamplingParam.from_pretrained(model_name)
     sampling_param.image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
     sampling_param.num_frames = 107
-    sampling_param.image_strength = 0.8  # How much to preserve the original image (0-1)
 
     # Generate video based on the image
     prompt = "A photograph coming to life with gentle movement"

fastvideo/v1/inference_engine.py

@@ -1,207 +0,0 @@
-# type: ignore
-
-# SPDX-License-Identifier: Apache-2.0
-"""
-Inference module for diffusion models.
-
-This module provides classes and functions for running inference with diffusion models.
-"""
-
-import time
-from typing import Any, Dict
-
-import torch
-
-from fastvideo.v1.fastvideo_args import FastVideoArgs
-from fastvideo.v1.logger import init_logger
-from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
-                                    build_pipeline)
-# TODO(will): remove, check if this is hunyuan specific
-from fastvideo.v1.utils import align_to
-
-logger = init_logger(__name__)
-
-
-class InferenceEngine:
-    """
-    Engine for running inference with diffusion models.
-    """
-
-    def __init__(
-        self,
-        pipeline: ComposedPipelineBase,
-        fastvideo_args: FastVideoArgs,
-    ):
-        """
-        Initialize the inference engine.
-        
-        Args:
-            pipeline: The pipeline to use for inference.
-            fastvideo_args: The inference arguments.
-            default_negative_prompt: The default negative prompt to use.
-        """
-        self.pipeline = pipeline
-        self.fastvideo_args = fastvideo_args
-
-    @classmethod
-    def create_engine(
-        cls,
-        fastvideo_args: FastVideoArgs,
-    ) -> "InferenceEngine":
-        """
-        Create an inference engine with the specified arguments.
-        
-        Args:
-            fastvideo_args: The inference arguments.
-            model_loader_cls: The model loader class to use. If None, it will be
-                determined from the model type.
-            pipeline_type: The type of pipeline to create. If None, it will be
-                determined from the model type.
-                
-        Returns:
-            The created inference engine.
-            
-        Raises:
-            ValueError: If the model type is not recognized or if the pipeline type
-                is not recognized.
-        """
-
-        logger.info("Building pipeline...")
-
-        # TODO(will): I don't really like this api.
-        # it should be something closer to pipeline_cls.from_pretrained(...)
-        # this way for training we can just do pipeline_cls.from_pretrained(
-        # checkpoint_path) and have it handle everything.
-        # TODO(Peiyuan): Then maybe we should only pass in model path and device, not the entire inference args?
-        pipeline = build_pipeline(fastvideo_args)
-        logger.info("Pipeline Ready")
-
-        # Create the inference engine
-        return cls(pipeline, fastvideo_args)
-
-    def run(
-        self,
-        prompt: str,
-        fastvideo_args: FastVideoArgs,
-    ) -> Dict[str, Any]:
-        """
-        Run inference with the pipeline.
-        
-        Args:
-            prompt: The prompt to use for generation.
-            negative_prompt: The negative prompt to use. If None, the default will be used.
-            seed: The random seed to use. If None, a random seed will be used.
-            **kwargs: Additional arguments to pass to the pipeline.
-            
-        Returns:
-            A dictionary containing the generated videos and metadata.
-        """
-        out_dict: Dict[str, Any] = dict()
-
-        num_videos_per_prompt = fastvideo_args.num_videos
-        seed = fastvideo_args.seed
-        height = fastvideo_args.height
-        width = fastvideo_args.width
-        video_length = fastvideo_args.num_frames
-        negative_prompt = fastvideo_args.neg_prompt
-        infer_steps = fastvideo_args.num_inference_steps
-        guidance_scale = fastvideo_args.guidance_scale
-        flow_shift = fastvideo_args.flow_shift
-        embedded_guidance_scale = fastvideo_args.embedded_cfg_scale
-        image_path = fastvideo_args.image_path
-
-        # ========================================================================
-        # Arguments: target_width, target_height, target_video_length
-        # ========================================================================
-        if width <= 0 or height <= 0 or video_length <= 0:
-            raise ValueError(
-                f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={video_length}"
-            )
-        if (video_length - 1) % 4 != 0:
-            raise ValueError(
-                f"`video_length-1` must be a multiple of 4, got {video_length}")
-
-        target_height = align_to(height, 16)
-        target_width = align_to(width, 16)
-        target_video_length = video_length
-
-        out_dict["size"] = (target_height, target_width, target_video_length)
-
-        # ========================================================================
-        # Arguments: prompt, new_prompt, negative_prompt
-        # ========================================================================
-        if not isinstance(prompt, str):
-            raise TypeError(
-                f"`prompt` must be a string, but got {type(prompt)}")
-        prompt = prompt.strip()
-
-        # negative prompt
-        if negative_prompt is not None:
-            negative_prompt = negative_prompt.strip()
-
-        # TODO(PY): move to hunyuan stage
-        latents_size = [(video_length - 1) // 4 + 1, height // 8, width // 8]
-        n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
-
-        # ========================================================================
-        # Print infer args
-        # ========================================================================
-        debug_str = f"""
-                        height: {target_height}
-                         width: {target_width}
-                  video_length: {target_video_length}
-                        prompt: {prompt}
-                    neg_prompt: {negative_prompt}
-                          seed: {seed}
-                   infer_steps: {infer_steps}
-         num_videos_per_prompt: {num_videos_per_prompt}
-                guidance_scale: {guidance_scale}
-                      n_tokens: {n_tokens}
-                    flow_shift: {flow_shift}
-       embedded_guidance_scale: {embedded_guidance_scale}"""
-        logger.info(debug_str)
-        # return
-        # sp_group = get_sp_group()
-        # local_rank = sp_group.rank
-        device = torch.device(fastvideo_args.device_str)
-        batch = ForwardBatch(
-            image_path=image_path,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_videos_per_prompt=num_videos_per_prompt,
-            height=fastvideo_args.height,
-            width=fastvideo_args.width,
-            num_frames=fastvideo_args.num_frames,
-            num_inference_steps=fastvideo_args.num_inference_steps,
-            guidance_scale=fastvideo_args.guidance_scale,
-            # generator=generator,
-            eta=0.0,
-            n_tokens=n_tokens,
-            data_type="video" if fastvideo_args.num_frames > 1 else "image",
-            device=device,
-            extra={},  # Any additional parameters
-        )
-
-        print('===============================================')
-        print(batch)
-        print('===============================================')
-        print('===============================================')
-        print(fastvideo_args)
-
-        # ========================================================================
-        # Pipeline inference
-        # ========================================================================
-        start_time = time.perf_counter()
-        samples = self.pipeline.forward(
-            batch=batch,
-            fastvideo_args=fastvideo_args,
-        ).output
-        # TODO(will): fix and move to hunyuan stage
-        # out_dict["seeds"] = batch.seeds
-        out_dict["samples"] = samples
-        out_dict["prompts"] = prompt
-
-        gen_time = time.perf_counter() - start_time
-        logger.info("Success, time: %s", gen_time)
-
-        return out_dict

fastvideo/v1/sample/v1_fastvideo_inference.py

@@ -1,77 +0,0 @@
-# type: ignore
-
-# SPDX-License-Identifier: Apache-2.0
-
-import os
-import sys
-
-import imageio
-import numpy as np
-import torch
-import torchvision
-from einops import rearrange
-
-from fastvideo.v1.distributed import (init_distributed_environment,
-                                      initialize_model_parallel)
-from fastvideo.v1.fastvideo_args import FastVideoArgs, prepare_fastvideo_args
-# Fix the import path
-from fastvideo.v1.inference_engine import InferenceEngine
-
-
-def initialize_distributed_and_parallelism(fastvideo_args: FastVideoArgs):
-    local_rank = int(os.environ.get("LOCAL_RANK", 0))
-    rank = int(os.environ.get("RANK", 0))
-    world_size = int(os.environ.get("WORLD_SIZE", 1))
-    torch.cuda.set_device(local_rank)
-    init_distributed_environment(world_size=world_size,
-                                 rank=rank,
-                                 local_rank=local_rank)
-    device_str = f"cuda:{local_rank}"
-    fastvideo_args.device_str = device_str
-    fastvideo_args.device = torch.device(device_str)
-    assert fastvideo_args.sp_size is not None
-    assert fastvideo_args.tp_size is not None
-    initialize_model_parallel(
-        sequence_model_parallel_size=fastvideo_args.sp_size,
-        tensor_model_parallel_size=fastvideo_args.tp_size,
-    )
-
-
-def main(fastvideo_args: FastVideoArgs):
-    initialize_distributed_and_parallelism(fastvideo_args)
-    engine = InferenceEngine.create_engine(fastvideo_args, )
-
-    if fastvideo_args.prompt_path is not None:
-        with open(fastvideo_args.prompt_path) as f:
-            prompts = [line.strip() for line in f.readlines()]
-    else:
-        if fastvideo_args.prompt is None:
-            raise ValueError("prompt or prompt_path is required")
-        prompts = [fastvideo_args.prompt]
-
-    # Process each prompt
-    for prompt in prompts:
-        outputs = engine.run(
-            prompt=prompt,
-            fastvideo_args=fastvideo_args,
-        )
-
-        # Process outputs
-        videos = rearrange(outputs["samples"], "b c t h w -> t b c h w")
-        frames = []
-        for x in videos:
-            x = torchvision.utils.make_grid(x, nrow=6)
-            x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
-            frames.append((x * 255).numpy().astype(np.uint8))
-
-        # Save video
-        os.makedirs(os.path.dirname(fastvideo_args.output_path), exist_ok=True)
-        imageio.mimsave(os.path.join(fastvideo_args.output_path,
-                                     f"{prompt[:100]}.mp4"),
-                        frames,
-                        fps=fastvideo_args.fps)
-
-
-if __name__ == "__main__":
-    fastvideo_args = prepare_fastvideo_args(sys.argv[1:])
-    main(fastvideo_args)

fastvideo/v1/training/training_pipeline.py

@@ -20,7 +20,8 @@ from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines import ComposedPipelineBase
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.training.training_utils import (
-    compute_density_for_timestep_sampling, get_sigmas, normalize_dit_input)
+    compute_density_for_timestep_sampling, get_sigmas, normalize_dit_input,
+    apply_activation_checkpointing)
 
 import wandb  # isort: skip
 
@@ -52,6 +53,16 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
         self.transformer = self.get_module("transformer")
         assert self.transformer is not None
 
+        # if training_args.gradient_checkpointing:
+        if True:
+            # with set_forward_context(current_timestep=0,
+            #                 attn_metadata=None):
+                logger.info(f"Enabling activation-checkpointing ")
+                self.transformer = apply_activation_checkpointing(
+                    self.transformer,
+                    checkpointing_type='full',
+                )
+                
         self.transformer.requires_grad_(True)
         self.transformer.train()
 

fastvideo/v1/training/training_utils.py

@@ -8,8 +8,13 @@ import torch.distributed as dist
 from torch.distributed.fsdp import FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import StateDictType
-
+from fastvideo.v1.forward_context import get_forward_context, set_forward_context
+import contextlib
 from fastvideo.v1.logger import init_logger
+import collections
+from enum import Enum
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper
+from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts
 
 logger = init_logger(__name__)
 
@@ -323,3 +328,109 @@ def _has_foreach_support(tensors: List[torch.Tensor],
                          device: torch.device) -> bool:
     return _device_has_foreach_support(device) and all(
         t is None or type(t) in [torch.Tensor] for t in tensors)
+
+DIFFUSERS_TRANSFORMER_BLOCK_NAMES = [
+    "transformer_blocks",
+    "single_transformer_blocks",
+    "temporal_transformer_blocks",
+    "blocks",
+    "layers",
+]
+
+def debug_tensor_creation():
+    original_new = torch.Tensor.__new__
+    
+    def new_with_debug(cls, *args, **kwargs):
+        tensor = original_new(cls, *args, **kwargs)
+        if tensor.dtype == torch.float32:
+            import traceback
+            print(f"FP32 tensor created during recompute: {tensor.shape}")
+            traceback.print_stack(limit=5)
+        return tensor
+    
+    torch.Tensor.__new__ = new_with_debug
+    return lambda: setattr(torch.Tensor, '__new__', original_new)
+
+def _ctx_fn():
+    saved = get_forward_context()          # ← runs while ctx is alive
+        
+    @contextlib.contextmanager
+    def _cm():
+        with torch.amp.autocast('cuda', enabled=True, dtype=torch.bfloat16):
+            if saved is None:
+                yield
+            else:
+                # current_timestep = saved.current_timestep
+                # if hasattr(current_timestep, 'dtype') and current_timestep.dtype == torch.float32:
+                #     current_timestep = current_timestep.to(torch.bfloat16)
+                with set_forward_context(
+                    current_timestep=saved.current_timestep.to(torch.bfloat16),
+                    attn_metadata=saved.attn_metadata,
+                    forward_batch=saved.forward_batch,
+                ):
+                    yield
+
+    return _cm(), _cm()
+
+class CheckpointType(str, Enum):
+    FULL = "full"
+    OPS = "ops"
+    BLOCK_SKIP = "block_skip"
+
+
+_SELECTIVE_ACTIVATION_CHECKPOINTING_OPS = {
+    torch.ops.aten.mm.default,
+    torch.ops.aten._scaled_dot_product_efficient_attention.default,
+    torch.ops.aten._scaled_dot_product_flash_attention.default,
+    torch.ops._c10d_functional.reduce_scatter_tensor.default,
+}
+
+
+def apply_activation_checkpointing(
+    module: torch.nn.Module, checkpointing_type: str = CheckpointType.FULL, n_layer: int = 1
+) -> torch.nn.Module:
+    if checkpointing_type == CheckpointType.FULL:
+        module = _apply_activation_checkpointing_blocks(module)
+    elif checkpointing_type == CheckpointType.OPS:
+        module = _apply_activation_checkpointing_ops(module, _SELECTIVE_ACTIVATION_CHECKPOINTING_OPS)
+    elif checkpointing_type == CheckpointType.BLOCK_SKIP:
+        module = _apply_activation_checkpointing_blocks(module, n_layer)
+    else:
+        raise ValueError(
+            f"Checkpointing type '{checkpointing_type}' not supported. Supported types are {CheckpointType.__members__.keys()}"
+        )
+    return module
+
+
+def _apply_activation_checkpointing_blocks(module: torch.nn.Module, n_layer: int = None) -> torch.nn.Module:
+    for transformer_block_name in DIFFUSERS_TRANSFORMER_BLOCK_NAMES:
+        blocks: torch.nn.Module = getattr(module, transformer_block_name, None)
+        if blocks is None:
+            continue
+        print("here")
+        for index, (layer_id, block) in enumerate(blocks.named_children()):
+            if n_layer is None or index % n_layer == 0:
+                block = checkpoint_wrapper(block, preserve_rng_state=False, context_fn=_ctx_fn)
+                blocks.register_module(layer_id, block)
+    return module
+
+
+def _apply_activation_checkpointing_ops(module: torch.nn.Module, ops) -> torch.nn.Module:
+
+    def _get_custom_policy(meta):
+        def _custom_policy(ctx, func, *args, **kwargs):
+            mode = "recompute" if ctx.is_recompute else "forward"
+            mm_count_key = f"{mode}_mm_count"
+            if func == torch.ops.aten.mm.default:
+                meta[mm_count_key] += 1
+            # Saves output of all compute ops, except every second mm
+            to_save = func in ops and not (func == torch.ops.aten.mm.default and meta[mm_count_key] % 2 == 0)
+            return CheckpointPolicy.MUST_SAVE if to_save else CheckpointPolicy.PREFER_RECOMPUTE
+
+        return _custom_policy
+
+    def selective_checkpointing_context_fn():
+        meta = collections.defaultdict(int)
+        return create_selective_checkpoint_contexts(_get_custom_policy(meta))
+
+    return checkpoint_wrapper(module, context_fn=selective_checkpointing_context_fn, preserve_rng_state=False)

fastvideo/v1/utils.py

@@ -10,6 +10,7 @@ import json
 import math
 import os
 import signal
+import socket
 import sys
 import tempfile
 import traceback
@@ -23,6 +24,7 @@ import filelock
 import torch
 import yaml
 from huggingface_hub import snapshot_download
+from remote_pdb import RemotePdb
 
 import fastvideo.v1.envs as envs
 from fastvideo.v1.logger import init_logger
@@ -645,3 +647,11 @@ class TypeBasedDispatcher:
             if isinstance(obj, ty):
                 return fn(obj)
         raise ValueError(f"Invalid object: {obj}")
+
+
+def remote_breakpoint() -> None:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        s.bind(("localhost", 0))  # Let the OS pick an ephemeral port.
+        port = s.getsockname()[1]
+        RemotePdb(host="localhost", port=port).wait_for_client()

pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
     "gradio>=5.22.0", "moviepy==1.0.3", "flask",
     "flask_restful", "aiohttp", "huggingface_hub", "cloudpickle",
     # System & Monitoring Tools
-    "gpustat", "watch",
+    "gpustat", "watch", "remote-pdb",
 
     # Kernel & Packaging
     "wheel",