720 lines
35 KiB
Python
720 lines
35 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""CPU contracts for PDDMethod and DMPDDMethod on fake packed models.
|
|
|
|
The fakes implement the ``PDDStudent`` contract on a two-modality packed
|
|
latent (video columns then audio columns) with analytic heads
|
|
``head_j = scale * (j + 1) * state`` so every rollout, loss, and carry
|
|
contract can be checked against closed-form references while gradients still
|
|
flow into one trainable parameter. Ports the behavioural tests of FastGen's
|
|
``tests/test_pddmodel.py`` and ``tests/test_pddmdmodel.py``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from fastvideo.layers.pdd import PDDModalitySchedule, pdd_fine_grid
|
|
from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
|
|
from fastvideo.train.methods.distribution_matching.dmpdd import DMPDDMethod
|
|
from fastvideo.train.methods.trajectory_matching.pdd import PDDMethod, parse_pdd_config
|
|
from fastvideo.train.methods.trajectory_matching.pdd_core import PDDModalityBlock, advance_block
|
|
from fastvideo.train.utils.training_config import (
|
|
DataConfig,
|
|
DistributedConfig,
|
|
OptimizerConfig,
|
|
TrainingConfig,
|
|
TrainingLoopConfig,
|
|
)
|
|
|
|
_VIDEO_NUMEL = 6
|
|
_AUDIO_NUMEL = 2
|
|
_NUMEL = _VIDEO_NUMEL + _AUDIO_NUMEL
|
|
_VIDEO = slice(0, _VIDEO_NUMEL)
|
|
_AUDIO = slice(_VIDEO_NUMEL, _NUMEL)
|
|
|
|
|
|
class _Scale(torch.nn.Module):
|
|
|
|
def __init__(self, value: float) -> None:
|
|
super().__init__()
|
|
self.scale = torch.nn.Parameter(torch.tensor(value))
|
|
|
|
|
|
class _FakePackedModel:
|
|
"""Packed two-modality fake implementing the PDD and DMD2 model contracts."""
|
|
|
|
device = torch.device("cpu")
|
|
pdd_model_dtype = torch.float32
|
|
num_train_timesteps = 1000
|
|
|
|
def __init__(self, *, trainable: bool, scale: float, pdd_steps: int | None = None) -> None:
|
|
self._trainable = trainable
|
|
self.transformer = _Scale(scale)
|
|
self.transformer.requires_grad_(trainable)
|
|
self.pdd_steps = pdd_steps
|
|
self.schedules = {
|
|
"video": PDDModalitySchedule(shift=12.0),
|
|
"audio": PDDModalitySchedule(shift=3.0),
|
|
}
|
|
self.head_calls: list[dict[str, Any]] = []
|
|
self.fused_calls: list[dict[str, Any]] = []
|
|
self.noise_calls: list[dict[str, Any]] = []
|
|
self.prepare_calls: list[str] = []
|
|
self.backward_calls: list[Any] = []
|
|
|
|
# -- ModelBase pieces the methods touch ---------------------------------------
|
|
|
|
def init_preprocessors(self, training_config: Any) -> None:
|
|
del training_config
|
|
|
|
def shift_and_clamp_timestep(self, timestep: torch.Tensor) -> torch.Tensor:
|
|
return timestep
|
|
|
|
def modality_slices(self) -> tuple[tuple[str, slice], ...]:
|
|
return (("video", _VIDEO), ("audio", _AUDIO))
|
|
|
|
def prepare_batch(self, raw_batch: dict[str, Any], *, generator: torch.Generator, latents_source: str) -> Any:
|
|
self.prepare_calls.append(latents_source)
|
|
if latents_source == "data":
|
|
latents = raw_batch["latent"].clone()
|
|
else:
|
|
latents = torch.zeros(1, _NUMEL)
|
|
return SimpleNamespace(
|
|
latents=latents,
|
|
timesteps=torch.tensor([0.0]),
|
|
attn_metadata=None,
|
|
attn_metadata_vsa="vsa-metadata",
|
|
dmd_latent_vis_dict={},
|
|
fake_score_latent_vis_dict={},
|
|
minimax_h3_dmd_layout=None,
|
|
prompt=raw_batch.get("prompt"),
|
|
)
|
|
|
|
def _sigmas(self, timestep: torch.Tensor) -> dict[str, torch.Tensor]:
|
|
base = (timestep.reshape(-1)[:1].to(torch.float64) / 1000.0).clamp(0.0, 0.999)
|
|
return {name: schedule.sigma(base) for name, schedule in self.schedules.items()}
|
|
|
|
def add_noise_for_batch(self, clean: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor,
|
|
batch: Any) -> torch.Tensor:
|
|
del batch
|
|
return self.add_noise(clean, noise, timestep)
|
|
|
|
def add_noise(self, clean: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
|
|
sigmas = self._sigmas(timestep)
|
|
out = clean.clone().to(torch.float64)
|
|
for name, columns in self.modality_slices():
|
|
sigma = sigmas[name].to(torch.float64)
|
|
out[:, columns] = (1.0 - sigma) * clean[:, columns].to(torch.float64) + sigma * noise[:, columns].to(
|
|
torch.float64)
|
|
return out.to(clean.dtype)
|
|
|
|
def backward(self, loss: torch.Tensor, ctx: Any, *, grad_accum_rounds: int) -> None:
|
|
self.backward_calls.append(ctx)
|
|
(loss / max(1, int(grad_accum_rounds))).backward()
|
|
|
|
# -- teacher / critic ---------------------------------------------------------
|
|
|
|
def predict_noise(self, state: torch.Tensor, timestep: torch.Tensor, batch: Any, *, conditional: bool,
|
|
cfg_uncond: dict[str, Any] | None, attn_kind: str) -> torch.Tensor:
|
|
self.noise_calls.append({
|
|
"timestep": timestep.detach().clone(),
|
|
"conditional": conditional,
|
|
"attn_kind": attn_kind,
|
|
"dtype": state.dtype,
|
|
})
|
|
# Like the H3 adapter, a forward rewrites the batch's clean-time field.
|
|
batch.timesteps = -timestep.detach().clone()
|
|
return self.transformer.scale * state + (0.25 if conditional else 0.0)
|
|
|
|
def predict_x0(self, state: torch.Tensor, timestep: torch.Tensor, batch: Any, *, conditional: bool,
|
|
cfg_uncond: dict[str, Any] | None, attn_kind: str) -> torch.Tensor:
|
|
del batch, cfg_uncond, attn_kind, timestep
|
|
return self.transformer.scale * state * (1.0 if conditional else 0.5)
|
|
|
|
# -- PDD student contract -------------------------------------------------------
|
|
|
|
def pdd_fine_grid(self, device: torch.device) -> torch.Tensor:
|
|
assert self.pdd_steps is not None
|
|
return pdd_fine_grid(self.pdd_steps, device=device)
|
|
|
|
def pdd_modality_schedules(self, batch: Any) -> tuple[PDDModalityBlock, ...]:
|
|
del batch
|
|
return tuple(PDDModalityBlock(name, columns, self.schedules[name]) for name, columns in self.modality_slices())
|
|
|
|
def pdd_timestep(self, base: torch.Tensor) -> torch.Tensor:
|
|
return base.to(torch.float64).reshape(-1) * 1000.0
|
|
|
|
def _heads(self, state: torch.Tensor) -> torch.Tensor:
|
|
assert self.pdd_steps is not None
|
|
factors = torch.arange(1, self.pdd_steps + 1, dtype=state.dtype)
|
|
return self.transformer.scale * factors[None, :, None] * state[:, None, :]
|
|
|
|
def predict_pdd_heads(self, state: torch.Tensor, timestep: torch.Tensor, batch: Any, *, conditional: bool,
|
|
cfg_uncond: dict[str, Any] | None, attn_kind: str,
|
|
head_window: tuple[int, int] | None = None) -> torch.Tensor:
|
|
del cfg_uncond
|
|
self.head_calls.append({
|
|
"timestep": timestep.detach().clone(),
|
|
"grad_enabled": torch.is_grad_enabled(),
|
|
"attn_kind": attn_kind,
|
|
"conditional": conditional,
|
|
"dtype": state.dtype,
|
|
"head_window": head_window,
|
|
})
|
|
batch.timesteps = timestep.detach().clone()
|
|
heads = self._heads(state)
|
|
if head_window is None:
|
|
return heads
|
|
return heads[:, head_window[0]:head_window[1]]
|
|
|
|
def predict_pdd_fused(self, state: torch.Tensor, timestep: torch.Tensor, start: int, end: int,
|
|
fine_grid: torch.Tensor, batch: Any, *, precision_decoding: torch.dtype, conditional: bool,
|
|
cfg_uncond: dict[str, Any] | None, attn_kind: str) -> torch.Tensor:
|
|
del timestep, cfg_uncond, conditional, attn_kind
|
|
self.fused_calls.append({"start": start, "end": end, "grad_enabled": torch.is_grad_enabled()})
|
|
heads = self._heads(state).to(precision_decoding)
|
|
out = torch.zeros_like(state, dtype=precision_decoding)
|
|
for block in self.pdd_modality_schedules(batch):
|
|
weights = block.schedule.integration_weights(fine_grid)[start:end]
|
|
alpha = (weights / weights.sum()).to(precision_decoding)
|
|
out[:, block.columns] = torch.einsum("n,bnk->bk", alpha, heads[:, start:end, block.columns])
|
|
return out.to(state.dtype)
|
|
|
|
|
|
def _training_config(*, grad_accum: int = 1, data_type: str = "text_only") -> TrainingConfig:
|
|
return TrainingConfig(
|
|
distributed=DistributedConfig(num_gpus=1, sp_size=1),
|
|
data=DataConfig(preprocessed_data_type=data_type, seed=0, train_batch_size=1),
|
|
optimizer=OptimizerConfig(learning_rate=1e-3, betas=(0.0, 0.999), lr_scheduler="constant"),
|
|
loop=TrainingLoopConfig(max_train_steps=8, gradient_accumulation_steps=grad_accum),
|
|
)
|
|
|
|
|
|
def _pdd_method_config(**overrides: Any) -> dict[str, Any]:
|
|
config: dict[str, Any] = {
|
|
"rollout_mode": "simulate",
|
|
"block_size_min": 2,
|
|
"block_size_max": 4,
|
|
"type_rk": "euler",
|
|
"precision_decoding": "fp32",
|
|
"cfg_uncond": {
|
|
"text": "zero"
|
|
},
|
|
}
|
|
config.update(overrides)
|
|
return config
|
|
|
|
|
|
def _dmd_method_config(**overrides: Any) -> dict[str, Any]:
|
|
config: dict[str, Any] = {
|
|
"rollout_mode": "simulate",
|
|
"generator_update_interval": 5,
|
|
"real_score_guidance_scale": 1.0,
|
|
"min_timestep_ratio": 0.02,
|
|
"max_timestep_ratio": 0.98,
|
|
"fake_score_loss_space": "x0",
|
|
"cfg_uncond": {
|
|
"text": "zero"
|
|
},
|
|
"fake_score_learning_rate": 1e-3,
|
|
"fake_score_betas": [0.0, 0.999],
|
|
"fake_score_lr_scheduler": "constant",
|
|
"block_size_min": 2,
|
|
"block_size_max": 8,
|
|
"type_rk": "euler",
|
|
"precision_decoding": "fp32",
|
|
"pdd_loss_weight": 1.0,
|
|
}
|
|
config.update(overrides)
|
|
return config
|
|
|
|
|
|
def _roles(*, pdd_steps: int = 8, critic: bool = True) -> dict[str, _FakePackedModel]:
|
|
roles = {
|
|
"student": _FakePackedModel(trainable=True, scale=1.0, pdd_steps=pdd_steps),
|
|
"teacher": _FakePackedModel(trainable=False, scale=0.5),
|
|
}
|
|
if critic:
|
|
roles["critic"] = _FakePackedModel(trainable=True, scale=0.25)
|
|
return roles
|
|
|
|
|
|
def _build(method_cls: type, method_config: dict[str, Any], *, grad_accum: int = 1, data_type: str = "text_only",
|
|
pdd_steps: int = 8, critic: bool = True, validation: dict[str, Any] | None = None) -> Any:
|
|
cfg = SimpleNamespace(
|
|
training=_training_config(grad_accum=grad_accum, data_type=data_type),
|
|
method=method_config,
|
|
validation=validation or {},
|
|
)
|
|
method = method_cls(cfg=cfg, role_models=_roles(pdd_steps=pdd_steps, critic=critic))
|
|
method.cuda_generator = torch.Generator(device="cpu").manual_seed(0)
|
|
return method
|
|
|
|
|
|
def _raw_batch(seed: int = 0) -> dict[str, Any]:
|
|
generator = torch.Generator().manual_seed(seed)
|
|
return {"latent": torch.randn(1, _NUMEL, generator=generator), "prompt": f"prompt-{seed}"}
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# Config parsing
|
|
# ----------------------------------------------------------------------------
|
|
|
|
|
|
def test_parse_pdd_config_defaults_and_validation() -> None:
|
|
config = parse_pdd_config({})
|
|
assert (config.block_size_min, config.block_size_max, config.type_rk) == (32, 128, "euler")
|
|
assert config.precision_decoding == torch.float32
|
|
assert config.teacher_guidance_scale == 1.0
|
|
assert config.pdd_loss_weight == 1.0
|
|
parsed = parse_pdd_config({
|
|
"block_size_min": 4,
|
|
"block_size_max": 8,
|
|
"type_rk": "midpoint",
|
|
"precision_decoding": "bf16",
|
|
"student_grid_indices": {
|
|
"3": [0, 2, 5, 8]
|
|
},
|
|
})
|
|
assert parsed.precision_decoding == torch.bfloat16
|
|
assert parsed.student_grid_indices == {3: [0, 2, 5, 8]}
|
|
with pytest.raises(ValueError, match="must not exceed"):
|
|
parse_pdd_config({"block_size_min": 8, "block_size_max": 4})
|
|
with pytest.raises(ValueError, match="type_rk must be one of"):
|
|
parse_pdd_config({"type_rk": "rk9"})
|
|
with pytest.raises(ValueError, match="precision_decoding must name a floating-point torch dtype"):
|
|
parse_pdd_config({"precision_decoding": "int8"})
|
|
assert parse_pdd_config({"precision_decoding": "torch.float64"}).precision_decoding == torch.float64
|
|
with pytest.raises(ValueError, match="pdd_loss_weight must be a non-negative"):
|
|
parse_pdd_config({"pdd_loss_weight": -1.0})
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# Standalone PDD
|
|
# ----------------------------------------------------------------------------
|
|
|
|
|
|
def test_pdd_data_latent_step_trains_the_student_from_aligned_block_starts() -> None:
|
|
method = _build(PDDMethod, _pdd_method_config(rollout_mode="data_latent"), data_type="t2va", critic=False)
|
|
student, teacher = method.student, method.teacher
|
|
|
|
loss_map, outputs, metrics = method.single_train_step(_raw_batch(), iteration=1)
|
|
|
|
assert student.prepare_calls == ["data"]
|
|
assert set(loss_map) == {"total_loss", "pdd_loss", "pdd_loss_video", "pdd_loss_audio"}
|
|
assert torch.isfinite(loss_map["total_loss"])
|
|
torch.testing.assert_close(loss_map["total_loss"], loss_map["pdd_loss_video"] + loss_map["pdd_loss_audio"])
|
|
assert metrics["pdd_block_start"] in {0.0, 2.0, 4.0, 6.0}
|
|
# One student forward at the model boundary and one Euler teacher forward.
|
|
assert len(student.head_calls) == 1
|
|
assert student.head_calls[0]["attn_kind"] == "vsa"
|
|
assert [call["attn_kind"] for call in teacher.noise_calls] == ["dense"]
|
|
assert method.latent_vis["generator_pred_video"].shape == (1, _NUMEL)
|
|
|
|
method.backward(loss_map, outputs)
|
|
assert student.transformer.scale.grad is not None
|
|
assert torch.isfinite(student.transformer.scale.grad)
|
|
assert teacher.transformer.scale.grad is None
|
|
assert method.get_optimizers(1) == [method._student_optimizer]
|
|
before = student.transformer.scale.detach().clone()
|
|
method.optimizers_schedulers_step(1)
|
|
assert student.transformer.scale.detach() != before
|
|
|
|
|
|
def test_pdd_loss_supervises_one_head_with_the_teacher_target(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
method = _build(PDDMethod, _pdd_method_config(rollout_mode="data_latent"), data_type="t2va", critic=False)
|
|
student = method.student
|
|
fine_grid = student.pdd_fine_grid(torch.device("cpu"))
|
|
batch = student.prepare_batch(_raw_batch(), generator=method.cuda_generator, latents_source="data")
|
|
x_n = batch.latents
|
|
step_idx = torch.tensor([2])
|
|
monkeypatch.setattr(
|
|
"fastvideo.train.methods.trajectory_matching.pdd.sample_supervision_index",
|
|
lambda *args, **kwargs: torch.tensor([4]),
|
|
)
|
|
pred = student.predict_pdd_heads(x_n, student.pdd_timestep(fine_grid[step_idx]), batch, conditional=True,
|
|
cfg_uncond=None, attn_kind="vsa")
|
|
|
|
head = method._compute_pdd_loss_from_prediction(x_n, step_idx, pred, fine_grid, batch)
|
|
|
|
assert head.k.tolist() == [4]
|
|
torch.testing.assert_close(head.p_k, pred[:, 4])
|
|
blocks = student.pdd_modality_schedules(batch)
|
|
expected_x_k = advance_block(x_n, pred.detach(), step_idx, head.k, fine_grid, blocks)
|
|
torch.testing.assert_close(head.x_k, expected_x_k)
|
|
assert not head.x_k.requires_grad
|
|
# Euler target: the teacher's native direction at x_k on the grid time.
|
|
expected_target = method.teacher.transformer.scale * head.x_k + 0.25
|
|
for name, columns in student.modality_slices():
|
|
expected = torch.mean((head.p_k[:, columns] - expected_target[:, columns])**2)
|
|
torch.testing.assert_close(head.per_modality[name], expected)
|
|
torch.testing.assert_close(head.total, head.per_modality["video"] + head.per_modality["audio"])
|
|
# A head window supervises the same head from the sliced prediction.
|
|
windowed = method._compute_pdd_loss_from_prediction(x_n, step_idx, pred[:, 2:6], fine_grid, batch,
|
|
block_end_idx=torch.tensor([6]), head_offset=2)
|
|
torch.testing.assert_close(windowed.total, head.total)
|
|
|
|
|
|
def test_pdd_carried_walk_advances_one_rung_per_call_and_clears_at_the_end() -> None:
|
|
# block_size_min=2 over 8 intervals: a 4-rung carried grid [0, 2, 4, 6, 8].
|
|
method = _build(PDDMethod, _pdd_method_config(), critic=False)
|
|
student = method.student
|
|
|
|
rungs, starts, head_counts, fused_counts = [], [], [], []
|
|
for call in range(6):
|
|
heads_before, fused_before = len(student.head_calls), len(student.fused_calls)
|
|
loss_map, outputs, metrics = method.single_train_step(_raw_batch(call), iteration=call + 1)
|
|
assert torch.isfinite(loss_map["total_loss"])
|
|
method.backward(loss_map, outputs)
|
|
rungs.append(metrics["rollout_step"])
|
|
starts.append(metrics["pdd_block_start"])
|
|
head_counts.append(len(student.head_calls) - heads_before)
|
|
fused_counts.append(len(student.fused_calls) - fused_before)
|
|
|
|
# Rank 0, one slot: the stream starts at rung 0, walks to rung 3, then
|
|
# restarts fresh from noise (no second pre-walk).
|
|
assert rungs == [0.0, 1.0, 2.0, 3.0, 0.0, 1.0]
|
|
assert starts == [0.0, 2.0, 4.0, 6.0, 0.0, 2.0]
|
|
assert head_counts == [1] * 6
|
|
# A single stagger group at offset 0 has nothing to pre-walk (FastGen
|
|
# skips it too); every call pays exactly one materialized forward.
|
|
assert fused_counts == [0] * 6
|
|
# A trajectory keeps the prompt it set out with.
|
|
assert method._pdd_carry.slots[0]["raw_batch"]["prompt"] == "prompt-4"
|
|
assert student.prepare_calls == ["zeros"] * 6
|
|
|
|
|
|
def test_pdd_second_slot_is_staggered_by_a_no_grad_pre_walk() -> None:
|
|
method = _build(PDDMethod, _pdd_method_config(), grad_accum=2, critic=False)
|
|
student = method.student
|
|
rungs = []
|
|
for call in range(4):
|
|
_, _, metrics = method.single_train_step(_raw_batch(call), iteration=call + 1)
|
|
rungs.append(metrics["rollout_step"])
|
|
# Slot 0 starts at rung 0 without a walk; slot 1 pre-walks the whole grid
|
|
# under no_grad and keeps the snapshot at rung 1.
|
|
assert rungs == [0.0, 1.0, 1.0, 2.0]
|
|
assert len(student.fused_calls) == 3
|
|
assert all(call["grad_enabled"] is False for call in student.fused_calls)
|
|
|
|
|
|
def test_pdd_carry_state_is_the_no_grad_block_advance_of_the_same_heads() -> None:
|
|
method = _build(PDDMethod, _pdd_method_config(), critic=False)
|
|
student = method.student
|
|
fine_grid = student.pdd_fine_grid(torch.device("cpu"))
|
|
method.single_train_step(_raw_batch(), iteration=1)
|
|
carried = method._pdd_carry.slots[0]
|
|
assert carried["rung"] == 1
|
|
state = carried["state"]
|
|
assert not state.requires_grad and state.grad_fn is None
|
|
# Replay the walk: fresh noise at u_0 then heads integrated over [0, 2).
|
|
batch = student.prepare_batch(_raw_batch(), generator=method.cuda_generator, latents_source="zeros")
|
|
generator = torch.Generator(device="cpu").manual_seed(0)
|
|
noise = torch.randn(1, _NUMEL, generator=generator)
|
|
fresh = student.add_noise(torch.zeros(1, _NUMEL), noise, student.pdd_timestep(fine_grid[:1]))
|
|
blocks = student.pdd_modality_schedules(batch)
|
|
with torch.no_grad():
|
|
expected = advance_block(fresh, student._heads(fresh), torch.tensor([0]), torch.tensor([2]), fine_grid, blocks)
|
|
torch.testing.assert_close(state, expected)
|
|
|
|
|
|
def test_pdd_method_configuration_guards() -> None:
|
|
with pytest.raises(ValueError, match="rollout_mode must be set"):
|
|
_build(PDDMethod, {"block_size_min": 2, "block_size_max": 4}, critic=False)
|
|
with pytest.raises(ValueError, match="requires method.rollout_mode='simulate'"):
|
|
_build(PDDMethod, _pdd_method_config(rollout_mode="data_latent"), data_type="text_only", critic=False)
|
|
with pytest.raises(ValueError, match="rollout_carry_slots must equal"):
|
|
_build(PDDMethod, _pdd_method_config(rollout_carry_slots=2), grad_accum=1, critic=False)
|
|
with pytest.raises(ValueError, match="exceeds the student's pdd_steps"):
|
|
_build(PDDMethod, _pdd_method_config(block_size_max=16), critic=False)
|
|
# Validation NFEs must not need blocks wider than the trained window.
|
|
_build(PDDMethod, _pdd_method_config(), critic=False, validation={"sampling_steps": [2, 4, 8]})
|
|
with pytest.raises(ValueError, match="wider than the trained method.block_size_max=4"):
|
|
_build(PDDMethod, _pdd_method_config(), critic=False, validation={"sampling_steps": [4, 1]})
|
|
with pytest.raises(ValueError, match="requires method.cfg_uncond"):
|
|
config = _pdd_method_config(teacher_guidance_scale=2.0)
|
|
config.pop("cfg_uncond")
|
|
_build(PDDMethod, config, critic=False)
|
|
with pytest.raises(ValueError, match="pdd_steps must be an int >= 2"):
|
|
cfg = SimpleNamespace(training=_training_config(), method=_pdd_method_config(), validation={})
|
|
PDDMethod(cfg=cfg, role_models={
|
|
"student": _FakePackedModel(trainable=True, scale=1.0),
|
|
"teacher": _FakePackedModel(trainable=False, scale=0.5),
|
|
})
|
|
with pytest.raises(TypeError, match="PDD-widened"):
|
|
narrow = _FakePackedModel(trainable=True, scale=1.0)
|
|
del narrow.pdd_steps
|
|
cfg = SimpleNamespace(training=_training_config(), method=_pdd_method_config(), validation={})
|
|
PDDMethod(cfg=cfg, role_models={
|
|
"student": narrow,
|
|
"teacher": _FakePackedModel(trainable=False, scale=0.5),
|
|
})
|
|
|
|
|
|
def test_pdd_teacher_guidance_runs_both_forwards_and_mixes_in_decoding_precision() -> None:
|
|
method = _build(PDDMethod, _pdd_method_config(teacher_guidance_scale=3.0), critic=False)
|
|
teacher = method.teacher
|
|
batch = method.student.prepare_batch(_raw_batch(), generator=method.cuda_generator, latents_source="zeros")
|
|
predict = method._pdd_teacher_predictor(batch)
|
|
state = torch.randn(1, _NUMEL)
|
|
out = predict(state, torch.tensor([0.5], dtype=torch.float64))
|
|
assert [call["conditional"] for call in teacher.noise_calls] == [True, False]
|
|
conditional = teacher.transformer.scale * state + 0.25
|
|
unconditional = teacher.transformer.scale * state
|
|
torch.testing.assert_close(out, unconditional + (conditional - unconditional) * 3.0)
|
|
assert out.dtype == torch.float32
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# DMPDD
|
|
# ----------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("rollout_mode", ["simulate", "data_latent"])
|
|
def test_dmpdd_matches_fastgen_clip_policy(rollout_mode: str) -> None:
|
|
method = _build(
|
|
DMPDDMethod,
|
|
_dmd_method_config(rollout_mode=rollout_mode, block_size_max=4, generator_update_interval=5),
|
|
data_type="t2va" if rollout_mode == "data_latent" else "text_only",
|
|
)
|
|
|
|
assert method.get_grad_clip_targets(5) == {"student": method.student.transformer}
|
|
assert method.get_grad_clip_targets(1) == {}
|
|
assert method.get_grad_sanitize_targets(5) == {"student": method.student.transformer}
|
|
assert method.get_grad_sanitize_targets(1) == {"critic": method.critic.transformer}
|
|
|
|
|
|
def test_dmpdd_rejects_negative_fake_score_weight_decay() -> None:
|
|
with pytest.raises(ValueError, match="fake_score_weight_decay must be non-negative"):
|
|
_build(DMPDDMethod, _dmd_method_config(fake_score_weight_decay=-0.01))
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# DMPDD
|
|
# ----------------------------------------------------------------------------
|
|
|
|
|
|
def test_dmpdd_student_step_uses_the_full_tail_and_carries_one_block() -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config(block_size_max=4))
|
|
student = method.student
|
|
fine_grid = student.pdd_fine_grid(torch.device("cpu"))
|
|
|
|
loss_map, outputs, metrics = method.single_train_step(_raw_batch(), iteration=5)
|
|
|
|
assert metrics["update_student"] == 1.0
|
|
assert metrics["rollout_step"] == 0.0
|
|
assert metrics["pdd_block_start"] == 0.0
|
|
assert {"pdd_loss_video", "pdd_loss_audio", "generator_loss_video", "generator_loss_audio"} <= metrics.keys()
|
|
torch.testing.assert_close(loss_map["generator_loss"], loss_map["dmd_loss"] + loss_map["pdd_loss"])
|
|
# One materialized student forward with gradients; the tail estimate fed DMD.
|
|
assert len(student.head_calls) == 1 and student.head_calls[0]["grad_enabled"]
|
|
generator_pred = method.latent_vis["generator_pred_video"]
|
|
batch = student.prepare_batch(_raw_batch(), generator=method.cuda_generator, latents_source="zeros")
|
|
blocks = student.pdd_modality_schedules(batch)
|
|
generator = torch.Generator(device="cpu").manual_seed(0)
|
|
noise = torch.randn(1, _NUMEL, generator=generator)
|
|
fresh = student.add_noise(torch.zeros(1, _NUMEL), noise, student.pdd_timestep(fine_grid[:1]))
|
|
with torch.no_grad():
|
|
heads = student._heads(fresh)
|
|
expected_tail = advance_block(fresh, heads, torch.tensor([0]), torch.tensor([8]), fine_grid, blocks)
|
|
expected_carry = advance_block(fresh, heads, torch.tensor([0]), torch.tensor([2]), fine_grid, blocks)
|
|
torch.testing.assert_close(generator_pred, expected_tail)
|
|
carried = method._pdd_carry.slots[0]
|
|
assert carried["rung"] == 1
|
|
torch.testing.assert_close(carried["state"], expected_carry)
|
|
|
|
method.backward(loss_map, outputs)
|
|
assert student.transformer.scale.grad is not None
|
|
assert method.critic.transformer.scale.grad is None
|
|
|
|
|
|
def test_dmpdd_critic_step_skips_pdd_loss_but_still_advances_the_carry() -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config(block_size_max=4, generator_update_interval=5))
|
|
student = method.student
|
|
|
|
loss_map, outputs, metrics = method.single_train_step(_raw_batch(), iteration=1)
|
|
|
|
assert metrics["update_student"] == 0.0
|
|
assert "pdd_loss" not in loss_map
|
|
assert loss_map["fake_score_loss"].item() > 0.0
|
|
assert loss_map["generator_loss"].item() == 0.0
|
|
assert not student.head_calls[0]["grad_enabled"]
|
|
assert method.teacher.noise_calls == []
|
|
assert method._pdd_carry.slots[0]["rung"] == 1
|
|
method.backward(loss_map, outputs)
|
|
assert method.critic.transformer.scale.grad is not None
|
|
assert student.transformer.scale.grad is None
|
|
|
|
|
|
def test_dmpdd_carry_walks_the_highest_nfe_grid_and_clears_after_the_last_rung() -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config(block_size_max=4))
|
|
rungs = []
|
|
for call in range(5):
|
|
_, _, metrics = method.single_train_step(_raw_batch(call), iteration=call + 1)
|
|
rungs.append(metrics["rollout_step"])
|
|
assert rungs == [0.0, 1.0, 2.0, 3.0, 0.0]
|
|
assert method._pdd_carry.slots[0]["rung"] == 1
|
|
|
|
|
|
def test_student_export_contract_samples_the_balanced_highest_nfe_partition() -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config(block_size_max=4))
|
|
requested = []
|
|
method.student.inference_contract = lambda *, pdd_step_indices: requested.append(pdd_step_indices) or {"ok": 1}
|
|
assert method.inference_checkpoint_contract("student") == {"ok": 1}
|
|
# pdd_steps=8 with block_size_min=2: four fused blocks, as validation samples four steps.
|
|
assert requested == [[0, 2, 4, 6, 8]]
|
|
assert method.inference_checkpoint_contract("critic") is None
|
|
|
|
|
|
def test_dmpdd_data_latent_mode_carries_nothing() -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config(rollout_mode="data_latent", block_size_max=4), data_type="t2va")
|
|
assert method._pdd_carry is None
|
|
loss_map, outputs, metrics = method.single_train_step(_raw_batch(), iteration=5)
|
|
assert method.student.prepare_calls == ["data"]
|
|
assert metrics["pdd_block_start"] in {0.0, 2.0, 4.0, 6.0}
|
|
assert torch.isfinite(loss_map["total_loss"])
|
|
method.backward(loss_map, outputs)
|
|
assert method.student.transformer.scale.grad is not None
|
|
|
|
|
|
def test_dmpdd_configuration_guards() -> None:
|
|
with pytest.raises(ValueError, match="Carried PDD block 6 exceeds"):
|
|
_build(
|
|
DMPDDMethod,
|
|
_dmd_method_config(block_size_min=3, block_size_max=3, student_grid_indices={3: [0, 1, 2, 8]}),
|
|
)
|
|
with pytest.raises(ValueError, match="generator_update_interval must be at least 2"):
|
|
# DMD2 reads 1 as "update both roles"; DMPDD would never train its critic.
|
|
_build(DMPDDMethod, _dmd_method_config(block_size_max=4, generator_update_interval=1))
|
|
with pytest.raises(ValueError, match="Carried DMPDDMethod cannot cover every student rung"):
|
|
# Four carried rungs with a student interval of 4 leave three phases untrained.
|
|
_build(DMPDDMethod, _dmd_method_config(block_size_max=4, generator_update_interval=4))
|
|
with pytest.raises(ValueError, match="rollout_carry_slots must equal"):
|
|
_build(DMPDDMethod, _dmd_method_config(block_size_max=4, rollout_carry_slots=2))
|
|
with pytest.raises(ValueError, match="rollout_carry must be false"):
|
|
_build(DMPDDMethod, _dmd_method_config(block_size_max=4, rollout_carry=True))
|
|
with pytest.raises(ValueError, match="remove method.dmd_denoising_steps"):
|
|
_build(DMPDDMethod, _dmd_method_config(block_size_max=4, dmd_denoising_steps=[999]))
|
|
|
|
|
|
def test_dmpdd_updates_exactly_one_role_and_scheduler_per_step() -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config(rollout_mode="data_latent", generator_update_interval=5),
|
|
data_type="t2va")
|
|
for iteration in range(1, 6):
|
|
for optimizer in method._optimizer_dict.values():
|
|
optimizer.zero_grad(set_to_none=True)
|
|
expected = "student" if iteration == 5 else "critic"
|
|
inactive = "critic" if expected == "student" else "student"
|
|
assert method.get_optimizers(iteration) == [method._optimizer_dict[expected]]
|
|
assert method.get_lr_schedulers(iteration) == [method._lr_scheduler_dict[expected]]
|
|
parameters = {role: model.transformer.scale.detach().clone() for role, model in method._role_models.items()}
|
|
loss_map, outputs, _ = method.single_train_step(_raw_batch(iteration), iteration)
|
|
method.backward(loss_map, outputs, grad_accum_rounds=2)
|
|
assert getattr(method, expected).transformer.scale.grad is not None
|
|
assert getattr(method, inactive).transformer.scale.grad is None
|
|
method.optimizers_schedulers_step(iteration)
|
|
assert not torch.equal(getattr(method, expected).transformer.scale, parameters[expected])
|
|
assert torch.equal(getattr(method, inactive).transformer.scale, parameters[inactive])
|
|
assert method._critic_lr_scheduler.last_epoch == 4
|
|
assert method._student_lr_scheduler.last_epoch == 1
|
|
# Teacher forwards mutate batch times after the student. Backward uses
|
|
# the context captured by that student forward, never the teacher's time.
|
|
context = method.student.backward_calls[0]
|
|
assert torch.equal(context[0], method.student.head_calls[-1]["timestep"])
|
|
assert context[1] == "vsa-metadata"
|
|
|
|
|
|
def test_public_dmd2_keeps_critic_update_on_student_steps(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
method = _build(DMD2Method, _dmd_method_config(rollout_mode="data_latent", generator_update_interval=5),
|
|
data_type="t2va")
|
|
|
|
def rollout(batch, *, with_grad):
|
|
with torch.set_grad_enabled(with_grad):
|
|
return method.student.transformer.scale * batch.latents
|
|
|
|
monkeypatch.setattr(method, "_student_rollout", rollout)
|
|
loss_map, outputs, _ = method.single_train_step(_raw_batch(), 5)
|
|
assert outputs["_fv_backward"]["student_ctx"] is not None
|
|
assert outputs["_fv_backward"]["critic_ctx"] is not None
|
|
method.backward(loss_map, outputs)
|
|
assert method.student.transformer.scale.grad is not None
|
|
assert method.critic.transformer.scale.grad is not None
|
|
assert method.get_optimizers(5) == [method._critic_optimizer, method._student_optimizer]
|
|
assert method.get_lr_schedulers(5) == [method._critic_lr_scheduler, method._student_lr_scheduler]
|
|
assert set(method.get_grad_clip_targets(5)) == {"student", "critic"}
|
|
|
|
|
|
def test_resume_seeds_both_alternating_roles_and_preserves_existing_state() -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config())
|
|
# Iteration zero selects only the student for DMPDD. A resumed checkpoint
|
|
# still needs slots for the critic's saved Adam state.
|
|
method.seed_optimizer_state_for_resume()
|
|
for role in ("student", "critic"):
|
|
optimizer = method._optimizer_dict[role]
|
|
parameter = getattr(method, role).transformer.scale
|
|
state = optimizer.state[parameter]
|
|
assert set(state) == {"step", "exp_avg", "exp_avg_sq"}
|
|
assert state["step"].device.type == "cpu"
|
|
state["exp_avg"].fill_(3.0)
|
|
method.seed_optimizer_state_for_resume()
|
|
for role in ("student", "critic"):
|
|
parameter = getattr(method, role).transformer.scale
|
|
assert method._optimizer_dict[role].state[parameter]["exp_avg"].item() == 3.0
|
|
|
|
|
|
def test_fsdp_boundary_uses_active_unclipped_critic(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
method = _build(DMPDDMethod, _dmd_method_config(generator_update_interval=5))
|
|
calls = []
|
|
for role in ("student", "critic"):
|
|
model = getattr(method, role)
|
|
model.supports_gradient_accumulation_no_sync = True
|
|
monkeypatch.setattr(model.transformer, "set_is_last_backward",
|
|
lambda flag, role=role: calls.append((role, "last", flag)), raising=False)
|
|
monkeypatch.setattr(model.transformer, "set_requires_gradient_sync",
|
|
lambda flag, *, recurse, role=role: calls.append((role, "sync", flag, recurse)), raising=False)
|
|
method.set_is_last_backward(1, False)
|
|
method.set_requires_gradient_sync(1, False)
|
|
assert calls == [("critic", "last", False), ("critic", "sync", False, True), ("critic", "last", False)]
|
|
assert method.get_grad_clip_targets(1) == {}
|
|
assert method.supports_gradient_accumulation_no_sync(1)
|
|
method.critic.supports_gradient_accumulation_no_sync = False
|
|
assert not method.supports_gradient_accumulation_no_sync(1)
|
|
assert method.supports_gradient_accumulation_no_sync(5)
|
|
|
|
|
|
def test_public_ref2va_example_constructs_and_trains_toy_roles(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
|
|
from fastvideo.train.utils.config import load_run_config
|
|
|
|
# Resolve the placeholder model path locally; exercise real config parsing
|
|
# without fetching weights or a Hub manifest.
|
|
monkeypatch.setattr("fastvideo.registry.get_pipeline_config_cls_from_name", lambda _: MiniMaxH3PipelineConfig)
|
|
repository = Path(__file__).resolve().parents[4]
|
|
cfg = load_run_config(str(repository / "examples/train/configs/distribution_matching/minimax_h3/dmpdd_ref2va.yaml"))
|
|
method = DMPDDMethod(cfg=cfg, role_models=_roles(pdd_steps=32))
|
|
method.cuda_generator = torch.Generator().manual_seed(43)
|
|
assert method._rollout_mode == "data_latent" and method._pdd_carry is None
|
|
assert method._fake_score_guidance_scale == 2.0
|
|
assert method._fake_score_cfg_dropout_rate == 0.1
|
|
assert method.pdd_config.pdd_loss_weight == 0.3
|
|
assert method.pdd_config.teacher_guidance_scale == 2.0
|
|
assert method.pdd_config.block_size_min == method.pdd_config.block_size_max == 4
|
|
assert method._score_timestep_continuous
|
|
assert method._fake_score_loss_space == {"__default__": "x0"}
|
|
for iteration in (1, 5):
|
|
losses, outputs, metrics = method.single_train_step(_raw_batch(iteration), iteration)
|
|
method.backward(losses, outputs)
|
|
assert torch.isfinite(losses["total_loss"])
|
|
assert metrics["pdd_block_start"] in set(map(float, range(0, 32, 4)))
|
|
assert method.student.prepare_calls[-1] == "data"
|
|
if iteration == 5:
|
|
assert metrics["dmd_real_uncond_forwards"] == 1
|
|
assert metrics["dmd_fake_uncond_forwards"] == 1
|
|
assert metrics["pdd_teacher_uncond_forwards"] == 1
|
|
torch.testing.assert_close(losses["generator_loss"], losses["dmd_loss"] + 0.3 * losses["pdd_loss"])
|