[bugfix] LTX-2: train temporal RoPE at the clip's fps (#1892)

The legacy LTX-2 trainer did not give the clip fps to the DiT. The DiT then
kept temporal RoPE positions in frames during training, but inference uses
seconds (fps 24). The trainer now passes the fps of the stored latents (24 when
the value is missing).

The video transform stage resampled frames to train_fps but kept the source fps
in the batch. The saved clips then had the source fps, and the audio was cut to
the wrong length. The stage now sets the batch fps to train_fps.

The new regression test runs in the CI unit lane.

.buildkite/scripts/unit_test.sh

@@ -26,6 +26,7 @@ exec pytest \
   ./fastvideo/tests/worker/ \
   ./fastvideo/tests/training/test_runner.py \
   ./fastvideo/tests/training/test_trackers.py \
+  ./fastvideo/tests/training/test_ltx2_rope_fps.py \
   ./fastvideo/tests/inference/test_basic_fasth3_omniref_pdd.py \
   ./fastvideo/tests/inference/test_inference_regional_compile.py \
   ./fastvideo/tests/attention/ \

fastvideo/pipelines/preprocess/preprocess_stages.py

@@ -76,6 +76,8 @@ class VideoTransformStage(PipelineStage):
 
         video_pixel_values = video_pixel_values.float() / 255.0
         batch.latents = video_pixel_values
+        # The frames were resampled to train_fps, so describe the clip rather than the source video.
+        batch.fps = [self.train_fps] * len(batch.video_loader)
         batch.num_frames = [video_pixel_values.shape[2]] * len(batch.video_loader)
         batch.height = [video_pixel_values.shape[3]] * len(batch.video_loader)
         batch.width = [video_pixel_values.shape[4]] * len(batch.video_loader)

fastvideo/tests/training/test_ltx2_rope_fps.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The legacy LTX-2 trainer must hand the clip's fps to the DiT.
+
+The DiT divides temporal RoPE positions by ``forward_batch.fps`` from the
+forward context. Inference always sets it, so a training forward without it
+sees temporal positions in frames instead of seconds.
+"""
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fastvideo.configs.configs import VideoLoaderType
+from fastvideo.fastvideo_args import WorkloadType
+from fastvideo.forward_context import get_forward_context
+from fastvideo.pipelines.pipeline_batch_info import TrainingBatch
+from fastvideo.pipelines.preprocess.ltx2.ltx2_preprocess_pipelines import LTX2AudioEncodingStage
+from fastvideo.pipelines.preprocess.preprocess_stages import VideoTransformStage
+from fastvideo.training import ltx2_training_pipeline as ltx2_module
+from fastvideo.training.ltx2_training_pipeline import LTX2TrainingPipeline
+from fastvideo.training.trackers import DummyTracker
+from fastvideo.workflow.preprocess.preprocess_workflow_ltx2_t2v import LTX2PrecomputedSaver
+
+
+class _FpsRecordingDiT(torch.nn.Module):
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.ones(()))
+        self.seen_fps: list[float | None] = []
+
+    def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
+        forward_batch = get_forward_context().forward_batch
+        self.seen_fps.append(None if forward_batch is None else forward_batch.fps)
+        return hidden_states * self.scale
+
+
+def _pipeline(monkeypatch: pytest.MonkeyPatch) -> LTX2TrainingPipeline:
+    monkeypatch.setattr(ltx2_module, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(ltx2_module, "get_world_group",
+                        lambda: SimpleNamespace(all_reduce=lambda tensor, op=None: tensor))
+    pipeline = object.__new__(LTX2TrainingPipeline)
+    pipeline.transformer = _FpsRecordingDiT()
+    pipeline.tracker = DummyTracker()
+    pipeline.with_audio = False
+    pipeline.training_args = SimpleNamespace(num_latent_t=2,
+                                             ltx2_first_frame_conditioning_p=0.0,
+                                             gradient_accumulation_steps=1)
+    pipeline.noise_random_generator = torch.Generator("cpu").manual_seed(0)
+    pipeline.noise_gen_cuda = torch.Generator("cpu").manual_seed(0)
+    return pipeline
+
+
+def _pt_batch(fps: torch.Tensor | None) -> dict:
+    batch_size = 1 if fps is None else fps.numel()
+    latents = {"latents": torch.randn(batch_size, 4, 2, 2, 2)}
+    if fps is not None:
+        latents["fps"] = fps
+    return {
+        "latents": latents,
+        "conditions": {
+            "video_prompt_embeds": torch.randn(batch_size, 3, 8),
+            "audio_prompt_embeds": torch.randn(batch_size, 3, 8),
+            "prompt_attention_mask": torch.ones(batch_size, 3),
+        },
+    }
+
+
+def _train_step_fps(pipeline: LTX2TrainingPipeline, batch: dict) -> float | None:
+    training_batch = pipeline._get_next_batch_pt(batch, TrainingBatch())
+    training_batch = pipeline._prepare_dit_inputs(training_batch)
+    training_batch = pipeline._build_input_kwargs(training_batch)
+    training_batch.total_loss = 0.0
+    pipeline._transformer_forward_and_compute_loss(training_batch)
+    return pipeline.transformer.seen_fps[-1]
+
+
+def test_training_forward_sees_the_clip_fps(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = _pipeline(monkeypatch)
+    assert _train_step_fps(pipeline, _pt_batch(torch.tensor([24.0]))) == 24.0
+    assert _train_step_fps(pipeline, _pt_batch(torch.tensor([30.0]))) == 30.0
+
+
+def test_batch_without_fps_uses_the_preset_fps(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = _pipeline(monkeypatch)
+    assert _train_step_fps(pipeline, _pt_batch(None)) == 24.0
+
+
+def test_mixed_fps_batch_uses_the_first_sample(monkeypatch: pytest.MonkeyPatch) -> None:
+    pipeline = _pipeline(monkeypatch)
+    assert _train_step_fps(pipeline, _pt_batch(torch.tensor([25.0, 30.0]))) == 25.0
+
+
+class _Clip(str):
+    """A 60-frame video path whose frames read like a torchcodec decoder's."""
+
+    frames = torch.zeros(60, 3, 16, 16, dtype=torch.uint8)
+
+    def get_frames_at(self, indices) -> SimpleNamespace:
+        return SimpleNamespace(data=self.frames[list(indices)])
+
+
+def test_resampled_clip_keeps_train_fps_for_audio_and_saving(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
+    requested_audio_seconds: list[float] = []
+
+    def extract_audio(video_path: str, target_duration: float) -> None:
+        requested_audio_seconds.append(target_duration)
+
+    monkeypatch.setattr(LTX2AudioEncodingStage, "_extract_audio", staticmethod(extract_audio))
+    batch = SimpleNamespace(data_type="video",
+                            video_loader=[_Clip("clip.mp4")],
+                            video_file_name=["clip.mp4"],
+                            fps=[30.0],
+                            num_frames=[60],
+                            height=[16],
+                            width=[16],
+                            prompt_embeds=[torch.zeros(1, 3, 8)],
+                            prompt_attention_mask=[torch.ones(1, 3)],
+                            extra={})
+    args = SimpleNamespace(preprocess_config=SimpleNamespace(video_loader_type=VideoLoaderType.TORCHCODEC),
+                           workload_type=WorkloadType.T2V)
+
+    VideoTransformStage(train_fps=24, num_frames=9, max_height=16, max_width=16,
+                        do_temporal_sample=False).forward(batch, args)
+    LTX2AudioEncodingStage(torch.nn.Linear(1, 1), audio_processor=None, fallback_fps=24).forward(batch, args)
+    LTX2PrecomputedSaver(tmp_path).save_batch(batch)
+
+    # 9 frames sampled at 24 fps from a 30 fps source cover 9 / 24 s.
+    assert requested_audio_seconds == [pytest.approx(9 / 24)]
+    assert torch.load(tmp_path / "latents" / "clip.pt")["fps"] == 24.0
+    assert batch.num_frames == [9]
+    assert batch.fps == [24]

fastvideo/training/ltx2_training_pipeline.py

@@ -13,7 +13,7 @@ from fastvideo.forward_context import set_forward_context
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.ltx2 import VideoLatentShape
 from fastvideo.pipelines.basic.ltx2.ltx2_pipeline import LTX2Pipeline
-from fastvideo.pipelines.pipeline_batch_info import TrainingBatch
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch
 from fastvideo.training.training_pipeline import TrainingPipeline
 from fastvideo.training.trackers import (DummyTracker, TrackerType, Trackers, initialize_trackers)
 from fastvideo.training.training_utils import (clip_grad_norm_while_handling_failing_dtensor_cases, get_scheduler)
@@ -21,6 +21,9 @@ from fastvideo.utils import set_random_seed
 
 logger = init_logger(__name__)
 
+# Preset fps of the LTX-2 checkpoints, used when a batch carries no fps.
+_DEFAULT_ROPE_FPS = 24.0
+
 
 class LTX2TrainingPipeline(TrainingPipeline):
     """Training pipeline for LTX-2 text-to-video (optional audio)."""
@@ -37,6 +40,8 @@ class LTX2TrainingPipeline(TrainingPipeline):
     text_encoder: torch.nn.Module
     with_audio: bool = False
     tracker: TrackerType
+    # fps of the current batch; the DiT divides temporal RoPE positions by it.
+    _rope_fps: float = _DEFAULT_ROPE_FPS
 
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         # TODO (David): Change to port LTX2 scheduler into self.modules["scheduler"]
@@ -257,6 +262,7 @@ class LTX2TrainingPipeline(TrainingPipeline):
         training_batch.encoder_attention_mask = attention_mask.to(device, dtype=torch.bfloat16)
         training_batch.infos = []
         training_batch.raw_latent_shape = latents.shape
+        self._rope_fps = _DEFAULT_ROPE_FPS
         return training_batch
 
     def _get_next_batch_pt(
@@ -297,8 +303,20 @@ class LTX2TrainingPipeline(TrainingPipeline):
         else:
             training_batch.infos = []
         training_batch.raw_latent_shape = latents.shape
+        self._rope_fps = self._batch_fps(batch["latents"].get("fps"))
         return training_batch
 
+    @staticmethod
+    def _batch_fps(fps: torch.Tensor | list[float] | float | None) -> float:
+        if fps is None:
+            return _DEFAULT_ROPE_FPS
+        values = [float(v) for v in torch.as_tensor(fps).flatten().tolist()]
+        if not values or values[0] <= 0:
+            return _DEFAULT_ROPE_FPS
+        if any(v != values[0] for v in values):
+            logger.warning("Batch mixes fps %s; using %s for RoPE", values, values[0])
+        return values[0]
+
     def _normalize_dit_input(self, training_batch: TrainingBatch) -> TrainingBatch:
         return training_batch
 
@@ -432,8 +450,12 @@ class LTX2TrainingPipeline(TrainingPipeline):
         assert training_batch.noisy_model_input is not None
         assert training_batch.noise is not None
 
+        # Without fps in the forward context the DiT leaves temporal RoPE positions in
+        # frames, while inference divides them by the preset fps into seconds.
         with self.tracker.timed("timing/forward_backward"), set_forward_context(
-                current_timestep=training_batch.current_timestep, attn_metadata=training_batch.attn_metadata):
+                current_timestep=training_batch.current_timestep,
+                attn_metadata=training_batch.attn_metadata,
+                forward_batch=ForwardBatch(data_type="video", fps=self._rope_fps)):
             with torch.autocast("cuda", dtype=training_batch.latents.dtype), torch.autograd.set_detect_anomaly(True):
                 outputs = self.transformer(**input_kwargs)