Merge branch 'main' of github.com:hao-ai-lab/FastVideo

2025-07-01 02:34:09 +00:00
parent 9c135a3845 eb0f1318f3
commit 3973469f70
3 changed files with 114 additions and 9 deletions
@@ -414,6 +414,7 @@ class TrainingArgs(FastVideoArgs):
    lr_warmup_steps: int = 0
    max_grad_norm: float = 0.0
    gradient_checkpointing: bool = False
+    gradient_checkpointing_type: str = "full"
    selective_checkpointing: float = 0.0
    allow_tf32: bool = False
    mixed_precision: str = ""
@@ -615,6 +616,11 @@ class TrainingArgs(FastVideoArgs):
        parser.add_argument("--gradient-checkpointing",
                            action=StoreBoolean,
                            help="Whether to use gradient checkpointing")
+        parser.add_argument("--gradient-checkpointing-type",
+                            type=str,
+                            choices=["full", "ops", "block_skip"],
+                            default="full",
+                            help="Gradient checkpointing type")
        parser.add_argument("--selective-checkpointing",
                            type=float,
                            help="Selective checkpointing threshold")
@@ -0,0 +1,91 @@
+import collections
+from enum import Enum
+from typing import Optional
+
+import torch
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
+    checkpoint_wrapper)
+
+TRANSFORMER_BLOCK_NAMES = [
+    "blocks",
+    "double_blocks",
+    "single_blocks",
+    "transformer_blocks",
+    "temporal_transformer_blocks",
+    "transformer_double_blocks",
+    "transformer_single_blocks",
+]
+
+
+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: Optional[int] = None) -> torch.nn.Module:
+    for transformer_block_name in TRANSFORMER_BLOCK_NAMES:
+        blocks: torch.nn.Module = getattr(module, transformer_block_name, None)
+        if blocks is None:
+            continue
+        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)
+                blocks.register_module(layer_id, block)
+    return module
+
+
+def _apply_activation_checkpointing_ops(module: torch.nn.Module,
+                                        ops) -> torch.nn.Module:
+    from torch.utils.checkpoint import (CheckpointPolicy,
+                                        create_selective_checkpoint_contexts)
+
+    def _get_custom_policy(meta: dict[str, int]) -> CheckpointPolicy:
+
+        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: dict[str, int] = 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)
@@ -31,6 +31,8 @@ from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
                                    TrainingBatch)
+from fastvideo.v1.training.activation_checkpoint import (
+    apply_activation_checkpointing)
 from fastvideo.v1.training.training_utils import (
    clip_grad_norm_while_handling_failing_dtensor_cases,
    compute_density_for_timestep_sampling, get_sigmas, load_checkpoint,
@@ -83,6 +85,11 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
        self.transformer.requires_grad_(True)
        self.transformer.train()

+        if training_args.gradient_checkpointing:
+            self.transformer = apply_activation_checkpointing(
+                self.transformer,
+                checkpointing_type=training_args.gradient_checkpointing_type)
+
        noise_scheduler = self.modules["scheduler"]
        params_to_optimize = self.transformer.parameters()
        params_to_optimize = list(
@@ -309,17 +316,18 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                current_timestep=training_batch.current_timestep,
                attn_metadata=training_batch.attn_metadata):
            model_pred = self.transformer(**input_kwargs)
-        if self.training_args.precondition_outputs:
-            model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
-        target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents
+            if self.training_args.precondition_outputs:
+                model_pred = training_batch.noisy_model_input - model_pred * training_batch.sigmas
+            target = training_batch.latents if self.training_args.precondition_outputs else training_batch.noise - training_batch.latents

-        # make sure no implicit broadcasting happens
-        assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
-        loss = (torch.mean((model_pred.float() - target.float())**2) /
-                self.training_args.gradient_accumulation_steps)
+            # make sure no implicit broadcasting happens
+            assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
+            loss = (torch.mean((model_pred.float() - target.float())**2) /
+                    self.training_args.gradient_accumulation_steps)
+
+            loss.backward()
+            avg_loss = loss.detach().clone()

-        loss.backward()
-        avg_loss = loss.detach().clone()
        # logger.info(f"rank: {self.rank}, avg_loss: {avg_loss.item()}",
        #             local_main_process_only=False)
        world_group = get_world_group()