139 lines
5.4 KiB
Python
139 lines
5.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Per-method GPU smoke test: ``WanCausalModel`` + ``CausalConsistencyDistillationMethod``.
|
|
|
|
Mirrors ``test_wan_causal_dfsft.py``. Causal consistency distillation
|
|
bootstraps a consistency MSE between the student's ``x0`` at ``t`` and an EMA
|
|
copy of the student at ``t_next``, where ``t_next`` is produced online by a
|
|
single CFG Euler step of a frozen teacher (all under clean-history teacher
|
|
forcing). This test exercises the full step: finite loss, nonzero student
|
|
gradients, frozen teacher, and a post-step EMA update.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
|
|
os.environ.setdefault("MASTER_ADDR", "localhost")
|
|
os.environ.setdefault("MASTER_PORT", "29519")
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.train.methods.consistency_model.causal_cd import (
|
|
CausalConsistencyDistillationMethod, )
|
|
from fastvideo.train.models.wan import WanCausalModel
|
|
from fastvideo.train.utils.config import load_run_config
|
|
|
|
_FIXTURE = str(Path(__file__).resolve().parent.parent / "fixtures" / "wan_causal_t2v_causal_cd_min.yaml")
|
|
|
|
|
|
def _build_synthetic_batch(
|
|
device: torch.device,
|
|
dtype: torch.dtype,
|
|
) -> dict[str, torch.Tensor]:
|
|
batch_size = 1
|
|
return {
|
|
"text_embedding": torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
|
|
"text_attention_mask": torch.ones(batch_size, 16, device=device),
|
|
"vae_latent": torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
|
|
}
|
|
|
|
|
|
@pytest.mark.usefixtures("distributed_setup")
|
|
def test_wan_causal_cd_single_train_step(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
if not torch.cuda.is_available():
|
|
pytest.skip("requires CUDA")
|
|
|
|
cfg = load_run_config(_FIXTURE)
|
|
|
|
device = torch.device("cuda:0")
|
|
dtype = torch.bfloat16
|
|
|
|
monkeypatch.setattr(
|
|
"fastvideo.train.utils.dataloader."
|
|
"build_parquet_t2v_train_dataloader",
|
|
lambda *args, **kwargs: None,
|
|
)
|
|
|
|
student = WanCausalModel(
|
|
init_from=cfg.models["student"]["init_from"],
|
|
training_config=cfg.training,
|
|
trainable=True,
|
|
)
|
|
student.transformer = student.transformer.to(device=device, dtype=dtype)
|
|
|
|
teacher = WanCausalModel(
|
|
init_from=cfg.models["teacher"]["init_from"],
|
|
training_config=cfg.training,
|
|
trainable=False,
|
|
)
|
|
teacher.transformer = teacher.transformer.to(device=device, dtype=dtype)
|
|
|
|
ema = WanCausalModel(
|
|
init_from=cfg.models["ema"]["init_from"],
|
|
training_config=cfg.training,
|
|
trainable=False,
|
|
)
|
|
ema.transformer = ema.transformer.to(device=device, dtype=dtype)
|
|
|
|
method = CausalConsistencyDistillationMethod(
|
|
cfg=cfg,
|
|
role_models={
|
|
"student": student,
|
|
"teacher": teacher,
|
|
"ema": ema
|
|
},
|
|
)
|
|
method.on_train_start()
|
|
|
|
batch = _build_synthetic_batch(device, dtype)
|
|
loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
|
|
|
|
loss = loss_map["total_loss"]
|
|
assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
|
|
assert torch.isfinite(loss).item(), (f"total_loss is not finite: {loss.item()}")
|
|
|
|
method.backward(loss_map, outputs, grad_accum_rounds=1)
|
|
|
|
blocks = student.transformer.blocks
|
|
assert blocks is not None and len(blocks) > 0
|
|
layer0 = blocks[0]
|
|
trainable = [p for p in layer0.parameters() if p.requires_grad]
|
|
assert len(trainable) > 0, "student layer 0 has no trainable parameters"
|
|
for i, p in enumerate(trainable):
|
|
assert p.grad is not None, f"student layer 0 param[{i}] has None grad"
|
|
assert torch.isfinite(p.grad).all().item(), (f"student layer 0 param[{i}] grad contains NaN/Inf")
|
|
assert any(p.grad.detach().float().norm().item() > 0.0
|
|
for p in trainable), ("all student layer-0 grads are exactly zero; consistency loss "
|
|
"did not reach the first transformer block")
|
|
|
|
# Teacher must stay frozen.
|
|
assert all(not p.requires_grad for p in teacher.transformer.parameters()), ("teacher must be frozen for Causal-CD")
|
|
|
|
# The EMA model and student start from the same checkpoint, so the first
|
|
# parameter must match before any update. FSDP fully_shard params are
|
|
# DTensors; compare the local shards (torch.equal is unsupported on
|
|
# DTensor).
|
|
def _local(p: torch.Tensor) -> torch.Tensor:
|
|
return p.to_local() if hasattr(p, "to_local") else p
|
|
|
|
ema_param = next(ema.transformer.parameters())
|
|
student_param = next(student.transformer.parameters())
|
|
assert torch.equal(_local(ema_param),
|
|
_local(student_param)), ("EMA model should start identical to the student (same checkpoint)")
|
|
|
|
# The EMA update must move EMA toward the student. Apply a visibly large
|
|
# perturbation so the bf16 lerp is well above rounding noise (the real
|
|
# optimizer step at lr=2e-6 would be sub-ULP in bf16).
|
|
with torch.no_grad():
|
|
student_param.add_(1.0)
|
|
before = _local(ema_param).detach().float().clone()
|
|
method._update_ema()
|
|
after = _local(ema_param).detach().float()
|
|
assert not torch.equal(before, after), ("EMA weights did not move after _update_ema")
|
|
# EMA = decay*ema + (1-decay)*student moves ~ (1-decay) of the gap.
|
|
expected = before + (1.0 - method._ema_decay) * (_local(student_param).detach().float() - before)
|
|
assert torch.allclose(after, expected, atol=1e-2), ("EMA update did not follow the expected lerp")
|