diff --git a/docs/training/train_infra.md b/docs/training/train_infra.md index 766d261d27..7d070a6a07 100644 --- a/docs/training/train_infra.md +++ b/docs/training/train_infra.md @@ -51,7 +51,7 @@ torchrun --nproc_per_node=8 \ ## Config Format Every run is defined by a single YAML file with five top-level sections. -See `examples/train/example.yaml` for a fully-commented reference. +See `examples/train/configs/example.yaml` for a fully-commented reference. ### `models` — Role-based model instances @@ -80,9 +80,17 @@ Common model parameters: | `trainable` | `true` | Whether the model's parameters require gradients | | `disable_custom_init_weights` | `false` | Skip custom weight initialization (use for teacher/critic) | | `flow_shift` | `3.0` | Timestep shifting factor | -| `enable_gradient_checkpointing_type` | `null` | Gradient checkpointing (`"full"` or `null`) | +| `enable_gradient_checkpointing_type` | `null` | Gradient checkpointing (`"full"`, `"ops"`, or `null`) | | `attention_backend` | `null` | Optional role-local backend for Wan models (for example `ATTN_QAT_TRAIN`); overrides the process default only while this role's transformer is built | +`full` recomputes every transformer-block operation and is the +memory-conservative choice. `ops` retains outputs from supported, +dispatcher-visible fused attention operations. It can reduce recompute time at +the cost of higher activation memory, so use it only when the training shape has +verified memory headroom. Math SDPA, VMoBA, FA3, and `ATTN_QAT_TRAIN` do not +expose a retainable dispatcher boundary and therefore still use full attention +recomputation under `ops`. + Which roles are needed depends on the training method: | Method | Required roles | diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index 51dc03e66a..ec39f5f07d 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -147,7 +147,11 @@ training: mode_scale: 1.0 # default: 1.0 precondition_outputs: false # default: false moba_config: {} # default: {} - enable_gradient_checkpointing_type: full # default: null ("full" or null) + # default: null ("full", "ops", or null). "full" minimizes activation + # memory. "ops" retains supported dispatcher-visible attention outputs to + # reduce recompute time when memory headroom exists. Math SDPA, VMoBA, FA3, + # and ATTN_QAT_TRAIN still fully recompute attention under "ops". + enable_gradient_checkpointing_type: full # --- training top-level [TYPED] --- dit_precision: fp32 # default: "fp32" (master weight precision) diff --git a/fastvideo/tests/train/utils/test_activation_checkpoint.py b/fastvideo/tests/train/utils/test_activation_checkpoint.py new file mode 100644 index 0000000000..be904164bd --- /dev/null +++ b/fastvideo/tests/train/utils/test_activation_checkpoint.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import copy +import importlib + +import pytest +import torch +from torch import nn +from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper + +import fastvideo.train.utils.activation_checkpoint as activation_checkpoint +from fastvideo.train.utils.activation_checkpoint import apply_activation_checkpointing + +_TEST_OP_NAME = "fastvideo_activation_checkpoint_test::expensive_op" +_TEST_OP_CALLS = 0 +_KNOWN_TRAINING_BLOCK_SPARSE_ATTENTION_OP_NAMES = { + "fastvideo_kernel::block_sparse_attn_sm90", + "fastvideo_kernel::block_sparse_attn_sm100a", + "fastvideo_kernel::block_sparse_attn_triton", +} + + +@torch.library.custom_op(_TEST_OP_NAME, mutates_args=()) +def _expensive_op(value: torch.Tensor) -> torch.Tensor: + global _TEST_OP_CALLS + _TEST_OP_CALLS += 1 + return torch.sin(value) + + +@_expensive_op.register_fake +def _expensive_op_fake(value: torch.Tensor) -> torch.Tensor: + return torch.empty_like(value) + + +def _setup_expensive_op_context(ctx, inputs, output) -> None: + del output + ctx.save_for_backward(inputs[0]) + + +def _backward_expensive_op(ctx, grad_output: torch.Tensor) -> torch.Tensor: + (value,) = ctx.saved_tensors + return grad_output * torch.cos(value) + + +_expensive_op.register_autograd(_backward_expensive_op, setup_context=_setup_expensive_op_context) + + +class _ToyBlock(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.projection = nn.Linear(4, 4) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + return value + _expensive_op(self.projection(value)) + + +class _ToyTransformer(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.blocks = nn.ModuleList([_ToyBlock() for _ in range(4)]) + + def forward(self, value: torch.Tensor) -> torch.Tensor: + for block in self.blocks: + value = block(value) + return value + + +def _run_toy_model( + state_dict: dict[str, torch.Tensor], + checkpointing_type: str | None, +) -> tuple[int, torch.Tensor, torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: + global _TEST_OP_CALLS + _TEST_OP_CALLS = 0 + + model = _ToyTransformer() + model.load_state_dict(copy.deepcopy(state_dict)) + if checkpointing_type is not None: + apply_activation_checkpointing(model, checkpointing_type) + + value = torch.linspace(-0.5, 0.5, 12, dtype=torch.float64).reshape(3, 4).requires_grad_() + model.to(dtype=torch.float64) + output = model(value) + loss = output.square().mean() + loss.backward() + parameter_grads = { + name.replace("._checkpoint_wrapped_module", ""): parameter.grad.detach().clone() + for name, parameter in model.named_parameters() + } + return _TEST_OP_CALLS, output.detach(), loss.detach(), value.grad.detach().clone(), parameter_grads + + +def _toy_state_dict() -> dict[str, torch.Tensor]: + torch.manual_seed(17) + return _ToyTransformer().state_dict() + + +@pytest.mark.parametrize( + ("checkpointing_type", "retain_test_op", "expected_calls"), + [ + (None, True, 4), + ("full", True, 8), + ("ops", True, 4), + ("ops", False, 8), + ], +) +def test_checkpoint_policy_controls_expensive_op_recomputation( + monkeypatch: pytest.MonkeyPatch, + checkpointing_type: str | None, + retain_test_op: bool, + expected_calls: int, +) -> None: + op_names = activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES + if retain_test_op: + op_names = op_names | {_TEST_OP_NAME} + monkeypatch.setattr(activation_checkpoint, "_SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES", op_names) + + calls, *_ = _run_toy_model(_toy_state_dict(), checkpointing_type) + + assert calls == expected_calls + + +def test_checkpoint_modes_preserve_outputs_and_gradients(monkeypatch: pytest.MonkeyPatch) -> None: + op_names = activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES | {_TEST_OP_NAME} + monkeypatch.setattr(activation_checkpoint, "_SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES", op_names) + state_dict = _toy_state_dict() + + baseline = _run_toy_model(state_dict, None) + for checkpointing_type in ("full", "ops"): + result = _run_toy_model(state_dict, checkpointing_type) + torch.testing.assert_close(result[1], baseline[1], rtol=0, atol=0) + torch.testing.assert_close(result[2], baseline[2], rtol=0, atol=0) + torch.testing.assert_close(result[3], baseline[3], rtol=0, atol=0) + assert result[4].keys() == baseline[4].keys() + for name, baseline_grad in baseline[4].items(): + torch.testing.assert_close(result[4][name], baseline_grad, rtol=0, atol=0) + + +@pytest.mark.parametrize("checkpointing_type", ["full", "ops"]) +def test_checkpointing_wraps_each_block_not_transformer_root(checkpointing_type: str) -> None: + model = _ToyTransformer() + + result = apply_activation_checkpointing(model, checkpointing_type) + + assert result is model + assert not isinstance(model, CheckpointWrapper) + assert all(isinstance(block, CheckpointWrapper) for block in model.blocks) + + +@pytest.mark.parametrize("checkpointing_type", ["full", "ops"]) +def test_checkpointing_rejects_transformers_without_known_block_lists(checkpointing_type: str) -> None: + with pytest.raises(ValueError, match="Activation checkpointing is not applied successfully"): + apply_activation_checkpointing(nn.Linear(4, 4), checkpointing_type) + + +def test_known_training_block_sparse_attention_ops_are_retained() -> None: + missing_op_names = (_KNOWN_TRAINING_BLOCK_SPARSE_ATTENTION_OP_NAMES + - activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES) + + assert not missing_op_names, ( + f"Known training block-sparse attention ops missing from checkpoint save policy: {missing_op_names}") + + +def test_all_available_training_block_sparse_attention_ops_are_retained() -> None: + try: + block_sparse_attention = importlib.import_module("fastvideo_kernel.block_sparse_attn") + except ImportError as exc: + pytest.skip(f"fastvideo_kernel is unavailable on this platform: {exc}") + training_op_names = { + candidate._qualname + for candidate in vars(block_sparse_attention).values() + if getattr(candidate, "_qualname", "").startswith("fastvideo_kernel::block_sparse_attn_") + and getattr(candidate, "_backward_fn", None) is not None + } + + assert training_op_names + missing_op_names = training_op_names - activation_checkpoint._SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES + assert not missing_op_names, ( + f"Available training block-sparse attention ops missing from checkpoint save policy: {missing_op_names}") + + +@pytest.mark.parametrize( + "module_name", + [ + "fastvideo.train.models.kandinsky5.kandinsky5", + "fastvideo.train.models.ltx2.ltx2", + "fastvideo.train.models.minimax_h3.minimax_h3", + "fastvideo.train.models.wan.wan", + ], +) +def test_modular_model_plugins_use_modular_activation_checkpointing(module_name: str) -> None: + model_module = importlib.import_module(module_name) + + assert model_module.apply_activation_checkpointing is apply_activation_checkpointing + + +@pytest.mark.parametrize( + ("role_checkpointing_type", "fallback_checkpointing_type", "trainable", "expected_type", "expected_safe"), + [ + ("ops", "full", True, "ops", True), + (None, "full", True, "full", True), + (None, None, True, None, False), + ("ops", None, False, "ops", False), + ], +) +def test_wan_causal_checkpoint_safe_cache_uses_resolved_role_config( + monkeypatch: pytest.MonkeyPatch, + role_checkpointing_type: str | None, + fallback_checkpointing_type: str | None, + trainable: bool, + expected_type: str | None, + expected_safe: bool, +) -> None: + from fastvideo.train.models.wan.wan import WanModel + from fastvideo.train.models.wan.wan_causal import WanCausalModel + from fastvideo.train.utils.training_config import TrainingConfig + + training_config = TrainingConfig() + training_config.model.enable_gradient_checkpointing_type = fallback_checkpointing_type + monkeypatch.setattr(WanModel, "_load_transformer", lambda self, **kwargs: nn.Module()) + + model = WanCausalModel( + init_from="unused-by-test", + training_config=training_config, + trainable=trainable, + enable_gradient_checkpointing_type=role_checkpointing_type, + ) + + assert model._resolved_gradient_checkpointing_type == expected_type + assert model._should_use_checkpoint_safe_kv_cache() is expected_safe diff --git a/fastvideo/train/models/kandinsky5/kandinsky5.py b/fastvideo/train/models/kandinsky5/kandinsky5.py index 79fd37c461..a93943a44a 100644 --- a/fastvideo/train/models/kandinsky5/kandinsky5.py +++ b/fastvideo/train/models/kandinsky5/kandinsky5.py @@ -43,8 +43,6 @@ from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( FlowMatchEulerDiscreteScheduler, ) from fastvideo.pipelines import TrainingBatch -from fastvideo.training.activation_checkpoint import ( - apply_activation_checkpointing, ) from fastvideo.training.training_utils import ( compute_density_for_timestep_sampling, get_sigmas, @@ -53,6 +51,8 @@ ) from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.activation_checkpoint import ( + apply_activation_checkpointing, ) from fastvideo.train.utils.module_state import ( apply_trainable, ) from fastvideo.train.utils.moduleloader import ( diff --git a/fastvideo/train/models/ltx2/ltx2.py b/fastvideo/train/models/ltx2/ltx2.py index 9f891019e3..0deecd698e 100644 --- a/fastvideo/train/models/ltx2/ltx2.py +++ b/fastvideo/train/models/ltx2/ltx2.py @@ -39,10 +39,10 @@ from fastvideo.models.dits.ltx2 import VideoLatentShape from fastvideo.pipelines import ForwardBatch, TrainingBatch from fastvideo.platforms import AttentionBackendEnum -from fastvideo.training.activation_checkpoint import ( - apply_activation_checkpointing, ) from fastvideo.train.models.wan.wan import WanModel +from fastvideo.train.utils.activation_checkpoint import ( + apply_activation_checkpointing, ) from fastvideo.train.utils.module_state import ( apply_trainable, ) from fastvideo.train.utils.moduleloader import ( diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 7737639b36..3ad8563f0e 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -19,8 +19,6 @@ FlowMatchEulerDiscreteScheduler, ) from fastvideo.pipelines import TrainingBatch from fastvideo.platforms import AttentionBackendEnum -from fastvideo.training.activation_checkpoint import ( - apply_activation_checkpointing, ) from fastvideo.training.training_utils import ( compute_density_for_timestep_sampling, get_sigmas, @@ -33,6 +31,8 @@ ) from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.activation_checkpoint import ( + apply_activation_checkpointing, ) from fastvideo.train.utils.module_state import ( apply_trainable, ) from fastvideo.train.utils.moduleloader import ( @@ -80,12 +80,17 @@ def __init__( attention_backend=attention_backend, ) self._init_from = str(init_from) + self._resolved_gradient_checkpointing_type = (enable_gradient_checkpointing_type or getattr( + getattr(training_config, "model", None), + "enable_gradient_checkpointing_type", + None, + )) self.transformer = self._load_transformer( init_from=self._init_from, trainable=self._trainable, disable_custom_init_weights=(disable_custom_init_weights), - enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type), + enable_gradient_checkpointing_type=(self._resolved_gradient_checkpointing_type), training_config=training_config, transformer_override_safetensor=(transformer_override_safetensor), attention_backend=self.attention_backend, diff --git a/fastvideo/train/models/wan/wan_causal.py b/fastvideo/train/models/wan/wan_causal.py index 1f8f500126..b08a71305c 100644 --- a/fastvideo/train/models/wan/wan_causal.py +++ b/fastvideo/train/models/wan/wan_causal.py @@ -426,9 +426,7 @@ def _initialize_kv_cache( return kv_cache def _should_use_checkpoint_safe_kv_cache(self, ) -> bool: - tc = getattr(self, "training_config", None) - checkpointing_type = tc.model.enable_gradient_checkpointing_type if tc is not None else None - return (bool(checkpointing_type) and bool(self._trainable)) + return (bool(self._resolved_gradient_checkpointing_type) and bool(self._trainable)) def _should_snapshot_streaming_cache(self, ) -> bool: return (self._should_use_checkpoint_safe_kv_cache()) diff --git a/fastvideo/train/utils/activation_checkpoint.py b/fastvideo/train/utils/activation_checkpoint.py index 60fe8a2e5b..75546f12e9 100644 --- a/fastvideo/train/utils/activation_checkpoint.py +++ b/fastvideo/train/utils/activation_checkpoint.py @@ -5,9 +5,7 @@ model plugins within one training package. """ -import collections from enum import Enum -from typing import Any import torch from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper @@ -35,13 +33,40 @@ class CheckpointType(str, Enum): 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, +# Names rather than the op objects: the fastvideo ops register only when their +# backend module is imported, so torch.ops.fastvideo... would raise here on any +# build that has not loaded that backend. +_SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES = { + "aten::_scaled_dot_product_flash_attention", + "aten::_scaled_dot_product_efficient_attention", + "aten::_scaled_dot_product_cudnn_attention", + "fastvideo::_flash_attn_default_forward", + "fastvideo::_flash_attn_cute_forward", + "fastvideo::_flash_attn_cute_varlen_forward", + "fastvideo::_flash_attn_cute_fp4_forward", + "fastvideo::_flash_attn_no_pad_forward", + "fastvideo::_flash_attn_varlen_qk_no_pad_forward", + # VSA dispatches block_sparse_attn; video_sparse_attn is its Python entry + # point, not an op, and naming that here would match nothing. + "fastvideo_kernel::block_sparse_attn_sm90", + "fastvideo_kernel::block_sparse_attn_sm100a", + "fastvideo_kernel::block_sparse_attn_triton", } +# No collective is listed. The replaced set named +# _c10d_functional::reduce_scatter_tensor, which is correct in torchtitan, where +# Megatron-style sequence parallelism reduce-scatters inside the forward. Ulysses +# redistributes with all-to-all instead, so FastVideo has no forward +# reduce-scatter to retain; FSDP2's runs in the post-backward hook, outside the +# region. The one collective that does appear inside a block is the parameter +# all-gather, and retaining that would keep every checkpointed block's unsharded +# weights resident at once, which is what FSDP exists to avoid. + +# Math SDPA is decomposed before this policy sees it. VMoBA, the FA3 training +# path, and ATTN_QAT_TRAIN go through torch.autograd.Function rather than the +# dispatcher. No policy entry can retain those attention outputs, so they get +# full recomputation. + def apply_activation_checkpointing( module: torch.nn.Module, @@ -52,10 +77,7 @@ def apply_activation_checkpointing( 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, - ) + module = _apply_activation_checkpointing_ops(module) elif checkpointing_type == CheckpointType.BLOCK_SKIP: module = _apply_activation_checkpointing_blocks(module, n_layer) else: @@ -86,37 +108,33 @@ def _apply_activation_checkpointing_blocks( return module -def _apply_activation_checkpointing_ops( - module: torch.nn.Module, - ops: set[Any], -) -> torch.nn.Module: - """Checkpoint a module while retaining selected operation outputs.""" +def _apply_activation_checkpointing_ops(module: torch.nn.Module) -> torch.nn.Module: + """Checkpoint every block while retaining selected operation outputs.""" from torch.utils.checkpoint import CheckpointPolicy, create_selective_checkpoint_contexts - def _get_custom_policy(meta: dict[str, int]): - """Build a policy that alternates matrix-multiply output retention.""" + def selective_checkpointing_context_fn(): + """Retain selected expensive operations during recomputation.""" def _custom_policy(ctx, func, *args, **kwargs): - """Retain selected expensive operations during recomputation.""" - 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 - # Retain compute outputs except every second matrix multiplication. - to_save = func in ops and not (func == torch.ops.aten.mm.default and meta[mm_count_key] % 2 == 0) + # OpOverload.name() is e.g. "aten::_scaled_dot_product_flash_attention". + to_save = func.name() in _SELECTIVE_ACTIVATION_CHECKPOINTING_OP_NAMES return CheckpointPolicy.MUST_SAVE if to_save else CheckpointPolicy.PREFER_RECOMPUTE - return _custom_policy + return create_selective_checkpoint_contexts(_custom_policy) - def selective_checkpointing_context_fn(): - """Create independent operation counters for one checkpointed call.""" - meta: dict[str, int] = collections.defaultdict(int) - return create_selective_checkpoint_contexts(_get_custom_policy(meta)) - - # Selective checkpointing wraps modules without stochastic masks that must - # replay during recomputation. - return checkpoint_wrapper( - module, - context_fn=selective_checkpointing_context_fn, - preserve_rng_state=False, - ) + applied = False + for transformer_block_name in _TRANSFORMER_BLOCK_NAMES: + blocks: torch.nn.Module | None = getattr(module, transformer_block_name, None) + if blocks is None: + continue + for layer_id, block in blocks.named_children(): + # Selective checkpointing wraps modules without stochastic masks that + # must replay during recomputation. + checkpointed_block = checkpoint_wrapper(block, + context_fn=selective_checkpointing_context_fn, + preserve_rng_state=False) + blocks.register_module(layer_id, checkpointed_block) + applied = True + if not applied: + raise ValueError("Activation checkpointing is not applied successfully") + return module