[bugfix]: validate LTX-2 train wrapper in CI

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):

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