From eb0f1318f3646e4b9dd95516cf8c083ec31fe2c2 Mon Sep 17 00:00:00 2001 From: Jinzhe Pan <48981407+Eigensystem@users.noreply.github.com> Date: Tue, 1 Jul 2025 04:24:29 +0800 Subject: [PATCH] [Feat] activation checkpointing (#584) --- fastvideo/v1/fastvideo_args.py | 6 ++ .../v1/training/activation_checkpoint.py | 91 +++++++++++++++++++ fastvideo/v1/training/training_pipeline.py | 26 ++++-- 3 files changed, 114 insertions(+), 9 deletions(-) create mode 100644 fastvideo/v1/training/activation_checkpoint.py diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index 3cdeecbe2..737f72acb 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -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") diff --git a/fastvideo/v1/training/activation_checkpoint.py b/fastvideo/v1/training/activation_checkpoint.py new file mode 100644 index 000000000..db6d9ee9c --- /dev/null +++ b/fastvideo/v1/training/activation_checkpoint.py @@ -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) diff --git a/fastvideo/v1/training/training_pipeline.py b/fastvideo/v1/training/training_pipeline.py index 32e9d51b3..8b174746f 100644 --- a/fastvideo/v1/training/training_pipeline.py +++ b/fastvideo/v1/training/training_pipeline.py @@ -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()