Zhang PeiyuanandClaude Opus 5.5
c520878fd2
Adds opt-in regional torch.compile to the YAML-driven modular Wan trainer (training.model.enable_torch_compile / torch_compile_kwargs) with FA4 and Wan modulation forward/backward as opaque custom ops, compiled repeated DiT blocks after FSDP setup, and activation checkpointing applied before FSDP for trainable Wan transformers. Merged with main's #1716 (WanModel resolves the checkpointing type in __init__; the PR's pre-FSDP wrap uses it and the post-load wrap is gone) and #1757 (pre_fsdp_model_transform runs before this PR's pre_fsdp_transform). Maintainer review (Swarm order #776): AdaLN skip names are canonicalized through checkpoint-wrapper prefixes, test_inference_regional_compile.py joins the unit lane, the Wan modulation custom-op test moved to the transformer lane, regional_compile is classified in the schema inventory, and a 1/2-GPU training smoke covers fully_shard(CheckpointWrapper) in eager and regional-compile modes. Validated on GB200: the FA4 CuTe forward/backward parity and fullgraph-backward tests and all four smoke cases pass. Also exempts MMAudio recipes, whose student builds its pipeline config at init, from the fine-tuning recipe construction check that #1648 left failing on main. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Wan activation-checkpoint and regional-compile setup tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from fastvideo.configs.pipelines.base import PipelineConfig
|
|
from fastvideo.train.models.wan import wan
|
|
from fastvideo.train.models.wan.wan import WanModel
|
|
from fastvideo.train.utils.training_config import (
|
|
ModelTrainingConfig,
|
|
TrainingConfig,
|
|
)
|
|
|
|
|
|
def test_wan_passes_activation_checkpointing_to_pre_fsdp_load(monkeypatch) -> None:
|
|
training_config = TrainingConfig(
|
|
model=ModelTrainingConfig(enable_gradient_checkpointing_type="full"),
|
|
pipeline_config=PipelineConfig(),
|
|
)
|
|
captured: dict = {}
|
|
transformer = torch.nn.Linear(1, 1)
|
|
|
|
def _fake_load_module_from_path(**kwargs):
|
|
captured.update(kwargs)
|
|
return transformer
|
|
|
|
monkeypatch.setattr(wan, "load_module_from_path", _fake_load_module_from_path)
|
|
monkeypatch.setattr(wan, "apply_trainable", lambda module, trainable: module)
|
|
monkeypatch.setattr(WanModel, "_enable_lora_if_configured", lambda self, module: False)
|
|
|
|
model = WanModel.__new__(WanModel)
|
|
# WanModel.__init__ resolves the role type against training.model and
|
|
# passes the result; _load_transformer does not fall back on its own.
|
|
result = model._load_transformer(
|
|
init_from="fake/model",
|
|
trainable=True,
|
|
disable_custom_init_weights=False,
|
|
enable_gradient_checkpointing_type="full",
|
|
training_config=training_config,
|
|
)
|
|
|
|
assert result is transformer
|
|
transform = captured["pre_fsdp_transform"]
|
|
assert transform is not None
|
|
assert transform.keywords == {"checkpointing_type": "full"}
|
|
|
|
|
|
def test_wan_skips_activation_checkpointing_for_frozen_role(monkeypatch) -> None:
|
|
training_config = TrainingConfig(
|
|
model=ModelTrainingConfig(enable_gradient_checkpointing_type="full"),
|
|
pipeline_config=PipelineConfig(),
|
|
)
|
|
captured: dict = {}
|
|
transformer = torch.nn.Linear(1, 1)
|
|
|
|
def _fake_load_module_from_path(**kwargs):
|
|
captured.update(kwargs)
|
|
return transformer
|
|
|
|
monkeypatch.setattr(wan, "load_module_from_path", _fake_load_module_from_path)
|
|
monkeypatch.setattr(wan, "apply_trainable", lambda module, trainable: module)
|
|
monkeypatch.setattr(WanModel, "_enable_lora_if_configured", lambda self, module: False)
|
|
|
|
model = WanModel.__new__(WanModel)
|
|
model._load_transformer(
|
|
init_from="fake/model",
|
|
trainable=False,
|
|
disable_custom_init_weights=True,
|
|
enable_gradient_checkpointing_type="full",
|
|
training_config=training_config,
|
|
)
|
|
|
|
assert captured["pre_fsdp_transform"] is None
|