diff --git a/fastvideo/tests/train/models/test_load_ltx2.py b/fastvideo/tests/train/models/test_load_ltx2.py index f7066a45e..1eb331ac4 100644 --- a/fastvideo/tests/train/models/test_load_ltx2.py +++ b/fastvideo/tests/train/models/test_load_ltx2.py @@ -29,7 +29,10 @@ from pathlib import Path import pytest import torch -from fastvideo.forward_context import set_forward_context +from fastvideo.forward_context import ( + get_forward_context, + set_forward_context, +) from fastvideo.pipelines import ForwardBatch from fastvideo.train.models.ltx2 import LTX2Model from fastvideo.train.utils.config import load_run_config @@ -47,6 +50,42 @@ _LTX2_TEXT_LEN = 1024 _MIN_GPU_MEMORY_GB = 60 +class _TinyLTX2Transformer(torch.nn.Module): + """Small stand-in that preserves the LTX-2 wrapper contract.""" + + def __init__(self, *, is_ltx2_3: bool) -> None: + super().__init__() + self.video_weight = torch.nn.Parameter(torch.tensor(1.0)) + self.audio_weight = torch.nn.Parameter(torch.tensor(1.0)) + self.a2v_weight = torch.nn.Parameter(torch.tensor(1.0)) + self.v2a_weight = torch.nn.Parameter(torch.tensor(1.0)) + self.av_ca_weight = torch.nn.Parameter(torch.tensor(1.0)) + self.config = type("Config", (), { + "arch_config": type("Arch", (), { + "caption_proj_before_connector": is_ltx2_3, + "cross_attention_dim": 4096, + "caption_channels": 3840, + })(), + })() + self.forward_fps: float | None = None + self.forward_timestep: torch.Tensor | None = None + self.forward_text_dim: int | None = None + + def forward( + self, + *, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + **_kwargs, + ) -> torch.Tensor: + self.forward_fps = get_forward_context().forward_batch.fps + self.forward_timestep = timestep + self.forward_text_dim = encoder_hidden_states.shape[-1] + sigma = timestep[:, 0].view(-1, 1, 1, 1, 1) + return hidden_states.float() - 2.0 * sigma + + def _gpu_too_small() -> bool: if not torch.cuda.is_available(): return True @@ -54,6 +93,120 @@ def _gpu_too_small() -> bool: return total < _MIN_GPU_MEMORY_GB * 1024**3 +def test_ltx2_rejects_audio_training_before_loading( + monkeypatch: pytest.MonkeyPatch) -> None: + cfg = load_run_config(str(_FIXTURE_DIR / _CASES["ltx2"][0])) + + def fail_if_loaded(*_args, **_kwargs): + raise AssertionError("transformer loading must not start") + + monkeypatch.setattr(LTX2Model, "_load_transformer", fail_if_loaded) + with pytest.raises(NotImplementedError, match="train_audio=True"): + LTX2Model( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + train_audio=True, + ) + + +@pytest.mark.parametrize("case", _CASES.keys()) +def test_ltx2_wrapper_contract_runs_on_cpu( + case: str, monkeypatch: pytest.MonkeyPatch) -> None: + fixture_name, text_dim = _CASES[case] + cfg = load_run_config(str(_FIXTURE_DIR / fixture_name)) + transformer = _TinyLTX2Transformer(is_ltx2_3=case == "ltx2_3") + monkeypatch.setattr( + LTX2Model, + "_load_transformer", + lambda *_args, **_kwargs: transformer, + ) + monkeypatch.setattr( + "fastvideo.train.models.base.get_local_torch_device", + lambda: torch.device("cpu"), + ) + + model = LTX2Model( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=True, + timestep_uniform_prob=0.0, + ) + + assert transformer.video_weight.requires_grad + for name, param in transformer.named_parameters(): + if any(pattern in name for pattern in ("audio", "a2v", "v2a", "av_ca")): + assert not param.requires_grad, f"{name} was not frozen" + + model._check_text_embedding_dim(torch.empty(1, 1, text_dim)) + with pytest.raises(ValueError, match="text_embedding width"): + model._check_text_embedding_dim(torch.empty(1, 1, text_dim + 1)) + + raw_latents = torch.randn(1, 128, 2, 2, 2) + batch = model.prepare_batch( + { + "text_embedding": torch.randn(1, 4, text_dim), + "text_attention_mask": torch.ones(1, 4), + "vae_latent": raw_latents, + "info_list": [{"fps": 12.0}], + }, + generator=torch.Generator(device="cpu").manual_seed(0), + ) + + assert batch.sigmas is not None + assert batch.timesteps is not None + assert batch.noisy_model_input is not None + assert torch.all((batch.sigmas >= 0.0) & (batch.sigmas <= 1.0)) + assert torch.allclose(batch.timesteps, batch.sigmas.flatten() * 1000.0) + expected_noisy = ((1.0 - batch.sigmas) * raw_latents.bfloat16().float() + + batch.sigmas * batch.noise.float()).bfloat16() + assert torch.equal(batch.noisy_model_input, expected_noisy) + + outer_batch = ForwardBatch(data_type="video", fps=30) + with set_forward_context( + current_timestep=torch.tensor([-1.0]), + attn_metadata=None, + forward_batch=outer_batch, + ): + velocity = model.predict_noise( + batch.noisy_model_input.permute(0, 2, 1, 3, 4), + batch.timesteps, + batch, + conditional=True, + ) + assert get_forward_context().forward_batch is outer_batch + + assert velocity.shape == (1, 2, 128, 2, 2) + assert velocity.device.type == "cpu" + assert velocity.dtype == torch.float32 + assert torch.allclose(velocity, torch.full_like(velocity, 2.0), atol=1e-3) + assert transformer.forward_fps == 12.0 + assert transformer.forward_text_dim == text_dim + assert transformer.forward_timestep is not None + assert transformer.forward_timestep.shape == (1, 8) + assert torch.allclose( + transformer.forward_timestep[:, 0], + batch.timesteps / 1000.0, + ) + + backward_fps: list[float] = [] + transformer.video_weight.register_hook( + lambda grad: backward_fps.append( + get_forward_context().forward_batch.fps) or grad) + with set_forward_context( + current_timestep=torch.tensor([-1.0]), + attn_metadata=None, + forward_batch=outer_batch, + ): + model.backward( + transformer.video_weight.square(), + (batch.timesteps, None), + grad_accum_rounds=2, + ) + assert get_forward_context().forward_batch is outer_batch + assert backward_fps == [12.0] + assert transformer.video_weight.grad.item() == pytest.approx(1.0) + + @pytest.mark.usefixtures("distributed_setup") @pytest.mark.parametrize("case", _CASES.keys()) def test_ltx2_model_loads_and_forwards(case: str): diff --git a/fastvideo/train/models/ltx2/ltx2.py b/fastvideo/train/models/ltx2/ltx2.py index 0e18e5e4e..645efeae3 100644 --- a/fastvideo/train/models/ltx2/ltx2.py +++ b/fastvideo/train/models/ltx2/ltx2.py @@ -22,7 +22,7 @@ Subclasses WanModel but replaces the pieces where LTX-2 differs: it to the data fps or validation (which uses the preset fps) would see different temporal frequencies than training - the ~5.8B audio / cross-modal (a2v, v2a, av_ca) parameters are - frozen by default: video-only forwards never give them + frozen: video-only forwards never give them gradients, and freezing keeps them out of DCP checkpoints and optimizer state """ @@ -97,6 +97,11 @@ class LTX2Model(WanModel): train_audio: bool = False, timestep_uniform_prob: float = 0.1, ) -> None: + if train_audio: + raise NotImplementedError("LTX2Model only supports video-only training; " + "train_audio=True requires audio batch, forward, " + "and loss plumbing.") + cfg_rate = float(getattr(training_config.data, "training_cfg_rate", 0.0) or 0.0) if cfg_rate > 0.0: raise NotImplementedError("LTX2Model only supports training_cfg_rate=0. CFG dropout " @@ -106,7 +111,6 @@ class LTX2Model(WanModel): "training.data.training_cfg_rate: 0.0.") self._timestep_uniform_prob = float(timestep_uniform_prob) - self._train_audio = bool(train_audio) self._rope_fps: float = _DEFAULT_ROPE_FPS self._rope_forward_batch: ForwardBatch | None = None @@ -124,7 +128,7 @@ class LTX2Model(WanModel): lora=lora, ) - if trainable and not self._train_audio: + if trainable: self._freeze_audio_parameters() # No negative-prompt cache: cfg_rate is forced to 0 above and