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

Extend the LTX-2 training plugin to LTX-2.3 checkpoints (gated
attention, cross-attention AdaLN, 4096-d post-connector text
embeddings, no in-DiT caption projection):
- preprocess captions via the text encoder's forward() (the inference
  path) instead of preprocess_text_embeddings + run_connectors, which
  crashes on 2.3 (4096-d video features fed into the 2048-d audio
  connector); output verified byte-identical for 2.0
- fail fast on parquet/checkpoint text-width mismatches; the check
  reads the loaded transformer's config since the transformer loader
  deepcopies dit_config before applying checkpoint arch flags
- clear the map-style dataset cache when regenerating overfit parquet
- add overfit_ltx2_3_t2v.yaml and parametrize the smoke, load, and
  nightly e2e tests over ltx2/ltx2_3; seed the GB200 grad-norm
  reference for 2.3

Verified on 4x GB200: LTX-2.3 overfit 300 steps, train loss -> 0.0085,
validation mean MS-SSIM vs the training clip 0.127 (step 0) -> 0.957
(step 300); parametrized nightly e2e passes for both versions
(0.964 / 0.957).

examples/train/configs/overfit_ltx2_3_t2v.yaml

@@ -0,0 +1,95 @@
+# LTX-2.3 T2V overfitting test config.
+#
+# Overfits on a single short video (480x832, 81 frames @ 24fps) to
+# verify the LTX-2 training plugin works end-to-end on an LTX-2.3
+# checkpoint (gated attention, cross-attention AdaLN, 4096-d
+# post-connector text embeddings, no in-DiT caption projection).
+# Uses the distilled checkpoint (validation is 8 sampling steps).
+#
+# Preprocess data first (writes data/ltx2_3_overfit_preprocessed):
+#   CUDA_VISIBLE_DEVICES=0 \
+#   LTX2_OVERFIT_MODEL=FastVideo/LTX-2.3-Distilled-Diffusers \
+#   LTX2_OVERFIT_OUTPUT_DIR=data/ltx2_3_overfit_preprocessed \
+#   python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
+#
+# Run:
+#   NUM_GPUS=4 bash examples/train/run.sh examples/train/configs/overfit_ltx2_3_t2v.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX-2.3-Distilled-Diffusers
+    trainable: true
+    enable_gradient_checkpointing_type: full
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/ltx2_3_overfit_preprocessed
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    # LTX2Model requires 0.0: CFG dropout would zero post-connector
+    # embeddings, which is not the model's unconditional input.
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 11  # (81 - 1) / 8 + 1
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 300
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/ltx2_3_overfit
+    # A full training-state checkpoint is ~150GB for the 13B trainable
+    # video branch; disable saves for the overfit smoke run.
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 1
+
+  tracker:
+    trackers: [wandb]
+    project_name: fastvideo_ltx2
+    run_name: ltx2_3_overfit
+
+  model:
+    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
+    # so the default noise-minus-clean target reproduces the official
+    # unweighted masked-MSE (mask is all-ones for plain T2V).
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
+    dataset_file: data/ltx2_3_overfit_preprocessed/validation_prompts.json
+    every_steps: 50
+    sampling_steps: [8]
+    guidance_scale: 1.0
+    num_frames: 81
+
+# Required so the LTX2T2VConfig pipeline config is resolved from
+# init_from (without a `pipeline:` key the loader falls back to a
+# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
+pipeline: {}

examples/train/configs/overfit_ltx2_t2v.yaml

@@ -0,0 +1,90 @@
+# LTX-2 T2V overfitting test config.
+#
+# Overfits on a single short video (480x832, 81 frames @ 24fps) to
+# verify the LTX-2 training plugin works end-to-end. Uses the
+# distilled checkpoint (validation is 8 sampling steps, single pass).
+#
+# Preprocess data first (writes data/ltx2_overfit_preprocessed):
+#   CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
+#
+# Run:
+#   NUM_GPUS=4 bash examples/train/run.sh examples/train/configs/overfit_ltx2_t2v.yaml
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX2-Distilled-Diffusers
+    trainable: true
+    enable_gradient_checkpointing_type: full
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/ltx2_overfit_preprocessed
+    dataloader_num_workers: 0
+    train_batch_size: 1
+    # LTX2Model requires 0.0: CFG dropout would zero post-connector
+    # embeddings, which is not the model's unconditional input.
+    training_cfg_rate: 0.0
+    seed: 42
+    num_latent_t: 11  # (81 - 1) / 8 + 1
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 300
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/ltx2_overfit
+    # A full training-state checkpoint is ~150GB for the 13B trainable
+    # video branch; disable saves for the overfit smoke run.
+    training_state_checkpointing_steps: 0
+    checkpoints_total_limit: 1
+
+  tracker:
+    trackers: [wandb]
+    project_name: fastvideo_ltx2
+    run_name: ltx2_overfit
+
+  model:
+    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
+    # so the default noise-minus-clean target reproduces the official
+    # unweighted masked-MSE (mask is all-ones for plain T2V).
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
+    dataset_file: data/ltx2_overfit_preprocessed/validation_prompts.json
+    every_steps: 50
+    sampling_steps: [8]
+    guidance_scale: 1.0
+    num_frames: 81
+
+# Required so the LTX2T2VConfig pipeline config is resolved from
+# init_from (without a `pipeline:` key the loader falls back to a
+# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
+pipeline: {}

fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py

@@ -0,0 +1,256 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Preprocess LTX-2 overfit data into parquet format.
+
+Encodes videos with the LTX-2 causal video VAE and captions with the
+Gemma text encoder (feature extractor + embedding connector) into the
+t2v parquet schema expected by the training framework.
+
+The stored text embeddings are POST-connector: the connector replaces
+pad positions with learnable registers and returns an all-valid mask,
+so the parquet collate's ones/zeros mask stays semantically correct
+and training needs no text encoder at all. Captions are encoded via
+the encoder's forward() (the exact inference path), which handles both
+LTX-2.0 (shared 3840-d features) and LTX-2.3 (separate 4096-d video /
+2048-d audio feature extractors).
+
+Videos are resampled to TRAIN_FPS and the preprocessed clip is also
+saved as an mp4 next to the parquet so overfit tests can use it as
+the SSIM reference.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 python fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py
+"""
+
+import json
+import os
+import shutil
+from typing import Any
+
+import av
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+import torch
+
+from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
+from fastvideo.utils import maybe_download_model, verify_model_config_and_directory
+
+# --- Config ---
+NUM_FRAMES = 81  # 8k+1 for temporal compression ratio 8
+MAX_HEIGHT = 480  # divisible by 32 (spatial compression)
+MAX_WIDTH = 832
+TRAIN_FPS = 24.0  # matches the LTX-2 preset fps used at validation
+
+DATA_DIR = os.environ.get("LTX2_OVERFIT_DATA_DIR", "data/cats")
+CAPTION_JSON = os.environ.get("LTX2_OVERFIT_CAPTION_JSON", "videos2caption_1_sample.json")
+VIDEO_SUBDIR = os.environ.get("LTX2_OVERFIT_VIDEO_SUBDIR", "video")
+OUTPUT_DIR = os.environ.get("LTX2_OVERFIT_OUTPUT_DIR", "data/ltx2_overfit_preprocessed")
+MODEL_REPO = os.environ.get("LTX2_OVERFIT_MODEL", "FastVideo/LTX2-Distilled-Diffusers")
+# The train dataloader samples with drop_last=True across data-parallel
+# groups, so the dataset must hold at least num_sp_groups * batch_size
+# rows or every rank gets zero batches. Replicate the overfit sample so
+# a 4-GPU FSDP run still sees one batch per rank.
+NUM_COPIES = int(os.environ.get("LTX2_OVERFIT_NUM_COPIES", "4"))
+
+
+def _init_single_process_distributed() -> None:
+    """FastVideo component loaders expect an initialized distributed env."""
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29511")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    from fastvideo.distributed import (
+        maybe_init_distributed_environment_and_model_parallel, )
+    maybe_init_distributed_environment_and_model_parallel(1, 1)
+
+
+def load_video(path: str, num_frames: int, target_fps: float, height: int,
+               width: int) -> tuple[torch.Tensor, np.ndarray]:
+    """Load a video as [1, C, T, H, W] in [-1, 1], resampled to target_fps.
+
+    Also returns the uint8 RGB frames [T, H, W, C] for reference-video export.
+    """
+    with av.open(path) as container:
+        stream = container.streams.video[0]
+        native_fps = float(stream.average_rate or target_fps)
+        decoded = [torch.from_numpy(frame.to_ndarray(format="rgb24")) for frame in container.decode(video=0)]
+    if not decoded:
+        raise RuntimeError(f"Could not read any frames from {path}")
+    raw = torch.stack(decoded).permute(0, 3, 1, 2)  # [T, C, H, W] uint8
+    step = native_fps / target_fps
+    wanted = [min(int(round(i * step)), raw.shape[0] - 1) for i in range(num_frames)]
+
+    frames = raw[wanted].float()  # [T, C, H, W] in [0, 255]
+    src_h, src_w = frames.shape[2], frames.shape[3]
+    scale = max(height / src_h, width / src_w)
+    new_h, new_w = int(round(src_h * scale)), int(round(src_w * scale))
+    frames = torch.nn.functional.interpolate(frames, size=(new_h, new_w), mode="bilinear", antialias=True)
+    top = (new_h - height) // 2
+    left = (new_w - width) // 2
+    frames = frames[:, :, top:top + height, left:left + width]
+
+    frames_np = (frames.permute(0, 2, 3, 1).clamp(0, 255).round().to(torch.uint8).numpy())
+    video = frames / 127.5 - 1.0  # [0,255] -> [-1,1]
+    video = video.permute(1, 0, 2, 3).unsqueeze(0)  # [1,C,T,H,W]
+    return video, frames_np
+
+
+def main() -> None:
+    _init_single_process_distributed()
+
+    from fastvideo.fastvideo_args import FastVideoArgs
+    from fastvideo.models.loader.component_loader import (
+        PipelineComponentLoader, )
+    from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
+
+    device = torch.device("cuda:0")
+    model_path = maybe_download_model(MODEL_REPO)
+    model_index = verify_model_config_and_directory(model_path)
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    # The map-style dataset caches parquet file metadata; a stale cache
+    # next to a regenerated parquet can crash or serve old rows.
+    shutil.rmtree(os.path.join(OUTPUT_DIR, "map_style_cache"), ignore_errors=True)
+
+    with open(os.path.join(DATA_DIR, CAPTION_JSON)) as f:
+        caption_data = json.load(f)
+
+    pipeline_config = LTX2T2VConfig()
+    fastvideo_args = FastVideoArgs(
+        model_path=model_path,
+        pipeline_config=pipeline_config,
+        num_gpus=1,
+        tp_size=1,
+        sp_size=1,
+        hsdp_shard_dim=1,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+    )
+
+    def load_component(name: str) -> Any:
+        transformers_or_diffusers, _ = model_index[name]
+        return PipelineComponentLoader.load_module(
+            module_name=name,
+            component_model_path=os.path.join(model_path, name),
+            transformers_or_diffusers=transformers_or_diffusers,
+            fastvideo_args=fastvideo_args,
+        )
+
+    print("Loading LTX-2 VAE...")
+    vae = load_component("vae")
+    vae_dtype = next(vae.parameters()).dtype
+    print(f"VAE loaded ({sum(p.numel() for p in vae.parameters())/1e6:.0f}M, {vae_dtype})")
+
+    print("Loading Gemma text encoder + tokenizer...")
+    text_encoder = load_component("text_encoder")
+    tokenizer = load_component("tokenizer")
+    tokenizer.padding_side = "left"
+    if tokenizer.pad_token is None and tokenizer.eos_token is not None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    encoder_config = pipeline_config.text_encoder_configs[0]
+    tokenizer_kwargs = dict(encoder_config.tokenizer_kwargs)
+    if "max_length" not in tokenizer_kwargs:
+        tokenizer_kwargs["max_length"] = encoder_config.arch_config.text_len
+    preprocess_text = pipeline_config.preprocess_text_funcs[0]
+
+    # --- Process each video ---
+    records = []
+    for idx, item in enumerate(caption_data):
+        video_name = item["path"]
+        record_id = f"{idx:04d}_{video_name}"
+        caption = item["cap"][0] if isinstance(item["cap"], list) else item["cap"]
+        video_path = os.path.join(DATA_DIR, VIDEO_SUBDIR, video_name)
+
+        print(f"\nProcessing: {video_name}")
+        print(f"  Caption: {caption[:80]}...")
+
+        video, frames_np = load_video(video_path, NUM_FRAMES, TRAIN_FPS, MAX_HEIGHT, MAX_WIDTH)
+        video = video.to(device=device, dtype=vae_dtype)
+        print(f"  Video shape: {video.shape}")
+
+        with torch.no_grad():
+            # LTX-2 encode() returns a deterministic distribution whose
+            # mean is already per-channel normalized; store as-is.
+            latent = vae.encode(video).mean.squeeze(0).float().cpu()
+        print(f"  Latent shape: {latent.shape}")
+
+        with torch.no_grad():
+            # Encode through forward() — the inference text path. The
+            # two-step preprocess_text_embeddings + run_connectors route
+            # breaks on LTX-2.3, whose separate audio feature extractor
+            # is narrower than the video one.
+            text_inputs = tokenizer([preprocess_text(caption)], **tokenizer_kwargs)
+            encoder_out = text_encoder(
+                input_ids=text_inputs["input_ids"].to(device),
+                attention_mask=text_inputs["attention_mask"].to(device),
+            )
+            text_embedding = encoder_out.last_hidden_state.squeeze(0).float().cpu()
+        print(f"  Text embedding shape: {text_embedding.shape}")
+
+        record = {
+            "id": record_id,
+            "vae_latent_bytes": latent.numpy().tobytes(),
+            "vae_latent_shape": list(latent.shape),
+            "vae_latent_dtype": str(latent.dtype).replace("torch.", ""),
+            "text_embedding_bytes": (text_embedding.numpy().tobytes()),
+            "text_embedding_shape": list(text_embedding.shape),
+            "text_embedding_dtype": str(text_embedding.dtype).replace("torch.", ""),
+            "file_name": video_name,
+            "caption": caption,
+            "media_type": "video",
+            "width": MAX_WIDTH,
+            "height": MAX_HEIGHT,
+            "num_frames": NUM_FRAMES,
+            "duration_sec": NUM_FRAMES / TRAIN_FPS,
+            "fps": TRAIN_FPS,
+        }
+        records.append(record)
+
+        # Save the preprocessed clip so overfit tests can compare
+        # validation output against the memorization target.
+        import imageio
+        ref_path = os.path.join(OUTPUT_DIR, f"training_sample_{idx}.mp4")
+        with imageio.get_writer(ref_path, fps=TRAIN_FPS) as writer:
+            for frame in frames_np:
+                writer.append_data(frame)
+        print(f"  Wrote reference clip to {ref_path}")
+
+    # Clean up
+    del text_encoder, tokenizer, vae
+    torch.cuda.empty_cache()
+
+    # Write parquet (replicated NUM_COPIES times; see comment at top)
+    replicated = []
+    for copy_idx in range(max(1, NUM_COPIES)):
+        for r in records:
+            row = dict(r)
+            row["id"] = f"{r['id']}_copy{copy_idx}"
+            replicated.append(row)
+    table = pa.table(
+        {k: [r[k] for r in replicated]
+         for k in replicated[0]},
+        schema=pyarrow_schema_t2v,
+    )
+    output_path = os.path.join(OUTPUT_DIR, "data_00000.parquet")
+    pq.write_table(table, output_path)
+    print(f"\nWrote {len(replicated)} records "
+          f"({len(records)} unique x {max(1, NUM_COPIES)} copies) to {output_path}")
+
+    # Write validation prompts for the validation callback
+    val_prompts = {
+        "data": [{
+            "caption": (item["cap"][0] if isinstance(item["cap"], list) else item["cap"]),
+        } for item in caption_data],
+    }
+    val_path = os.path.join(OUTPUT_DIR, "validation_prompts.json")
+    with open(val_path, "w") as f:
+        json.dump(val_prompts, f, indent=2)
+    print(f"Wrote validation prompts to {val_path}")
+
+    print("\nDone! Use data_path: " + OUTPUT_DIR + " in training config.")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/tests/nightly/test_e2e_ltx2_overfit_new_stack.py

@@ -0,0 +1,156 @@
+"""E2E overfit tests for LTX-2 fine-tuning on the modular trainer (fastvideo/train).
+
+Parametrized over LTX-2.0 and LTX-2.3 distilled checkpoints. Each case
+downloads the single-sample cats dataset, preprocesses it with the
+matching LTX-2 VAE + Gemma text encoder into the t2v parquet format,
+trains LTX2Model with FineTuneMethod for 300 steps on 4 GPUs, and
+checks that the final validation video reproduces the training clip
+(MS-SSIM against the preprocessed ground-truth clip) better than the
+step-0 baseline.
+
+GPU assumptions: 4 x large-memory GPUs (developed on 4x GB200 192GB;
+the 18.9B-param DiT trains FSDP-sharded with fp32 master weights).
+Requires the FastVideo LTX-2 distilled checkpoints (set HF_HOME to a
+cache that contains them, or allow ~60GB of downloads per version).
+
+Guarded by FASTVIDEO_NIGHTLY=1 so the default test suite stays fast.
+"""
+
+import glob
+import json
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+from huggingface_hub import snapshot_download
+
+sys.path.append(str(Path(__file__).parent.parent.parent.parent.parent))
+
+from fastvideo.tests.utils import compute_video_ssim_torchvision  # noqa: E402
+
+NUM_GPUS_TRAINING = "4"
+
+DATA_DIR = "data"
+LOCAL_RAW_DATA_DIR = Path(DATA_DIR) / "cats"
+
+PREPROCESSING_SCRIPT = "fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py"
+
+_CASES = {
+    "ltx2": {
+        "model": "FastVideo/LTX2-Distilled-Diffusers",
+        "config": "examples/train/configs/overfit_ltx2_t2v.yaml",
+        "prep_dir": Path(DATA_DIR) / "ltx2_overfit_preprocessed",
+        "out_dir": Path(DATA_DIR) / "outputs_ltx2_overfit",
+    },
+    "ltx2_3": {
+        "model": "FastVideo/LTX-2.3-Distilled-Diffusers",
+        "config": "examples/train/configs/overfit_ltx2_3_t2v.yaml",
+        "prep_dir": Path(DATA_DIR) / "ltx2_3_overfit_preprocessed",
+        "out_dir": Path(DATA_DIR) / "outputs_ltx2_3_overfit",
+    },
+}
+
+
+def download_data():
+    os.makedirs(DATA_DIR, exist_ok=True)
+    print(f"Downloading raw dataset to {LOCAL_RAW_DATA_DIR}...")
+    snapshot_download(
+        repo_id="wlsaidhi/cats-overfit-merged",
+        local_dir=str(LOCAL_RAW_DATA_DIR),
+        repo_type="dataset",
+        token=os.environ.get("HF_TOKEN"),
+    )
+    assert LOCAL_RAW_DATA_DIR.exists(), (
+        f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
+
+
+def run_preprocessing(case: dict):
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = env.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]
+    env["LTX2_OVERFIT_DATA_DIR"] = str(LOCAL_RAW_DATA_DIR)
+    env["LTX2_OVERFIT_OUTPUT_DIR"] = str(case["prep_dir"])
+    env["LTX2_OVERFIT_MODEL"] = case["model"]
+    cmd = [sys.executable, PREPROCESSING_SCRIPT]
+    print(f"Running preprocessing: {cmd} (model={case['model']})")
+    subprocess.run(cmd, check=True, env=env)
+
+
+def run_training(case: dict):
+    env = os.environ.copy()
+    env.setdefault("WANDB_MODE", "offline")
+    cmd = [
+        "torchrun",
+        "--nnodes", "1",
+        "--nproc_per_node", NUM_GPUS_TRAINING,
+        "-m", "fastvideo.train.entrypoint.train",
+        "--config", case["config"],
+        "--training.checkpoint.output_dir", str(case["out_dir"]),
+        "--training.data.data_path", str(case["prep_dir"]),
+        "--callbacks.validation.dataset_file",
+        str(case["prep_dir"] / "validation_prompts.json"),
+    ]
+    print(f"Running training: {cmd}")
+    subprocess.run(cmd, check=True, env=env)
+
+
+def _validation_videos_by_step(out_dir: Path) -> dict[int, str]:
+    pattern = os.path.join(str(out_dir), "validation_step_*_video_0.mp4")
+    videos: dict[int, str] = {}
+    for path in glob.glob(pattern):
+        match = re.search(r"validation_step_(\d+)_", os.path.basename(path))
+        if match:
+            videos[int(match.group(1))] = path
+    return videos
+
+
+@pytest.mark.skipif(
+    os.environ.get("FASTVIDEO_NIGHTLY") != "1",
+    reason="nightly e2e overfit test; set FASTVIDEO_NIGHTLY=1 to run",
+)
+@pytest.mark.parametrize("case_id", _CASES.keys())
+def test_e2e_ltx2_overfit_new_stack(case_id: str):
+    case = _CASES[case_id]
+    download_data()
+    run_preprocessing(case)
+    run_training(case)
+
+    reference_video = case["prep_dir"] / "training_sample_0.mp4"
+    assert reference_video.exists(), (
+        f"Reference (preprocessed training clip) not found at {reference_video}")
+
+    videos = _validation_videos_by_step(case["out_dir"])
+    assert videos, f"No validation videos found under {case['out_dir']}"
+    final_step = max(videos)
+    final_video = videos[final_step]
+    print(f"Final validation video (step {final_step}): {final_video}")
+
+    final_mean, final_min, final_max = compute_video_ssim_torchvision(
+        str(reference_video), final_video, use_ms_ssim=True)
+    print(f"\n===== MS-SSIM vs training clip at step {final_step} ({case_id}) =====")
+    print(f"Mean: {final_mean:.4f}  Min: {final_min:.4f}  Max: {final_max:.4f}")
+
+    results = {"case": case_id, "final_step": final_step, "final_mean_ssim": final_mean}
+    baseline_step = min(videos)
+    if baseline_step != final_step:
+        base_mean, _, _ = compute_video_ssim_torchvision(
+            str(reference_video), videos[baseline_step], use_ms_ssim=True)
+        print(f"Baseline (step {baseline_step}) mean MS-SSIM: {base_mean:.4f}")
+        results["baseline_mean_ssim"] = base_mean
+        assert final_mean > base_mean, (
+            f"Overfitting did not improve similarity to the training clip: "
+            f"step {final_step} mean SSIM {final_mean:.4f} <= "
+            f"step {baseline_step} mean SSIM {base_mean:.4f}")
+    print(json.dumps(results))
+
+    assert final_mean > 0.5, (
+        f"Mean MS-SSIM vs the training clip is below 0.5: {final_mean:.4f}")
+
+
+if __name__ == "__main__":
+    os.environ.setdefault("FASTVIDEO_NIGHTLY", "1")
+    for _case_id in _CASES:
+        test_e2e_ltx2_overfit_new_stack(_case_id)

fastvideo/tests/train/fixtures/ltx2_3_t2v_finetune_min.yaml

@@ -0,0 +1,51 @@
+# Minimum config to run a single training step of FineTuneMethod on
+# LTX2Model with an LTX-2.3 checkpoint (gated attention, cross-attn
+# AdaLN, 4096-d post-connector text embeddings, no in-DiT caption
+# projection) for the per-method smoke test. 18.9B params, ~38GB in
+# bf16 — the test skips below 60GB GPU memory. Same constraints as
+# LTX-2.0: training_cfg_rate=0 and precondition_outputs=false.
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX-2.3-Distilled-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 3
+    num_height: 256
+    num_width: 256
+    num_frames: 17
+
+  optimizer:
+    learning_rate: 1.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1
+    gradient_accumulation_steps: 1
+
+  model:
+    precondition_outputs: false
+
+# Required so the LTX2T2VConfig pipeline config is resolved from
+# init_from (without a `pipeline:` key the loader falls back to a
+# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
+pipeline: {}

fastvideo/tests/train/fixtures/ltx2_t2v_finetune_min.yaml

@@ -0,0 +1,53 @@
+# Minimum config to run a single training step of FineTuneMethod on
+# LTX2Model for the per-method smoke test. Uses the real
+# FastVideo/LTX2-Distilled-Diffusers checkpoint (18.9B params, ~38GB
+# in bf16 — needs a large-memory GPU; the test skips below 60GB) with
+# tiny synthetic latents. LTX2 requires training_cfg_rate=0 (CFG
+# dropout would zero post-connector text embeddings, which is not the
+# model's unconditional input) and precondition_outputs=false (the
+# plugin converts the DiT's x0 output to velocity).
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    init_from: FastVideo/LTX2-Distilled-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  dit_precision: bf16
+
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+
+  data:
+    seed: 42
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    num_latent_t: 3
+    num_height: 256
+    num_width: 256
+    num_frames: 17
+
+  optimizer:
+    learning_rate: 1.0e-6
+    betas: [0.9, 0.999]
+    weight_decay: 0.01
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 1
+    gradient_accumulation_steps: 1
+
+  model:
+    precondition_outputs: false
+
+# Required so the LTX2T2VConfig pipeline config is resolved from
+# init_from (without a `pipeline:` key the loader falls back to a
+# generic PipelineConfig and the LTX-2 DiT cannot be constructed).
+pipeline: {}

fastvideo/tests/train/methods/grad_norm_refs.json

@@ -3,6 +3,12 @@
     "H200": 5.2763,
     "L40S": 4.1553
   },
+  "test_ltx2_3_finetune": {
+    "GB200": 2.2158
+  },
+  "test_ltx2_finetune": {
+    "GB200": 0.3127
+  },
   "test_matrixgame2_finetune": {
     "H200": 4.8247
   },

fastvideo/tests/train/methods/test_ltx2_finetune.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Per-method GPU smoke test: ``LTX2Model`` + ``FineTuneMethod``.
+
+Mirrors ``test_wan_finetune.py`` for the LTX-2 plugin, parametrized
+over LTX-2.0 and LTX-2.3 checkpoints. LTX2-specific differences: the
+synthetic ``raw_batch`` carries a single post-connector Gemma
+embedding (3840-d for 2.0, 4096-d for 2.3 which has no in-DiT caption
+projection) and 128-channel VAE latents; the DiT is 18.9B params, so
+the test skips on GPUs with less than 60GB memory (e.g. the L40S CI
+runner).
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29522")
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from fastvideo.train.methods.fine_tuning.finetune import (
+    FineTuneMethod, )
+from fastvideo.train.models.ltx2 import LTX2Model
+from fastvideo.train.utils.config import load_run_config
+
+from .grad_norm_regression import (
+    check_grad_norm_regression,
+    resolve_blocks,
+)
+
+_FIXTURE_DIR = Path(__file__).resolve().parent.parent / "fixtures"
+
+# id -> (fixture, post-connector text embedding width, grad-norm ref name)
+_CASES = {
+    "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840, "test_ltx2_finetune"),
+    "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096, "test_ltx2_3_finetune"),
+}
+
+_LTX2_TEXT_LEN = 1024
+
+_MIN_GPU_MEMORY_GB = 60
+
+
+def _gpu_too_small() -> bool:
+    if not torch.cuda.is_available():
+        return True
+    total = torch.cuda.get_device_properties(0).total_memory
+    return total < _MIN_GPU_MEMORY_GB * 1024**3
+
+
+def _build_synthetic_batch(
+    device: torch.device,
+    dtype: torch.dtype,
+    text_dim: int,
+) -> dict[str, torch.Tensor]:
+    """Tiny synthetic ``raw_batch`` matching ``LTX2Model.prepare_batch``.
+
+    LTX-2 latents are 128-channel (temporal 8x / spatial 32x
+    compression); text embeddings are post-connector Gemma features.
+    """
+    batch_size = 1
+    return {
+        "text_embedding":
+        torch.randn(batch_size,
+                    _LTX2_TEXT_LEN,
+                    text_dim,
+                    device=device,
+                    dtype=dtype),
+        "text_attention_mask":
+        torch.ones(batch_size, _LTX2_TEXT_LEN, device=device),
+        "vae_latent":
+        torch.randn(batch_size, 128, 3, 8, 8, device=device, dtype=dtype),
+    }
+
+
+@pytest.mark.usefixtures("distributed_setup")
+@pytest.mark.parametrize("case", _CASES.keys())
+def test_ltx2_finetune_single_train_step(
+        case: str, monkeypatch: pytest.MonkeyPatch) -> None:
+    if _gpu_too_small():
+        pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB "
+                    "memory (LTX-2 DiT is 18.9B params)")
+
+    fixture_name, text_dim, ref_name = _CASES[case]
+    cfg = load_run_config(str(_FIXTURE_DIR / fixture_name))
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+
+    # Feed a synthetic ``raw_batch`` straight into ``single_train_step``,
+    # so the parquet train dataloader built by ``init_preprocessors`` is
+    # never iterated. Stub it out so construction does not require a real
+    # ``training.data.data_path``.
+    monkeypatch.setattr(
+        "fastvideo.train.utils.dataloader."
+        "build_parquet_t2v_train_dataloader",
+        lambda *args, **kwargs: None,
+    )
+
+    model = LTX2Model(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=True,
+    )
+    model.transformer = model.transformer.to(device=device, dtype=dtype)
+
+    method = FineTuneMethod(
+        cfg=cfg,
+        role_models={"student": model},
+    )
+    method.on_train_start()
+
+    batch = _build_synthetic_batch(device, dtype, text_dim)
+    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)
+
+    # LTX-2 nests its transformer_blocks under the ``model`` submodule.
+    blocks = resolve_blocks(model.transformer.model)
+    assert blocks is not None and len(blocks) > 0, (
+        "transformer is expected to expose a non-empty block list")
+    layer0 = blocks[0]
+
+    trainable = [p for p in layer0.parameters() if p.requires_grad]
+    assert len(trainable) > 0, "layer 0 has no trainable parameters"
+
+    for i, p in enumerate(trainable):
+        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
+        assert torch.isfinite(p.grad).all().item(), (
+            f"layer 0 param[{i}] grad contains NaN/Inf")
+
+    any_nonzero = any(
+        p.grad.detach().float().norm().item() > 0.0 for p in trainable)
+    assert any_nonzero, (
+        "all layer-0 grads are exactly zero; backward did not "
+        "reach the first transformer block")
+
+    # Audio / cross-modal parameters must be frozen by default.
+    audio_trainable = [
+        name for name, param in model.transformer.named_parameters()
+        if param.requires_grad and any(
+            pattern in name for pattern in ("audio", "a2v", "v2a", "av_ca"))
+    ]
+    assert not audio_trainable, (
+        f"audio/cross-modal params unexpectedly trainable: "
+        f"{audio_trainable[:5]}")
+
+    # Device-keyed grad-norm regression on top of the same harness.
+    # Skips when the current GPU has no seeded reference.
+    check_grad_norm_regression(ref_name, model.transformer.model)

fastvideo/tests/train/models/test_load_ltx2.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GPU loading + forward smoke test for ``LTX2Model``.
+
+Loads the real LTX-2.0 / LTX-2.3 distilled checkpoints (18.9B at bf16,
+~38GB — skips on GPUs with less than 60GB memory) via
+``LTX2Model.__init__`` and runs one transformer forward pass on
+synthetic inputs. Catches loader or forward-signature regressions in
+``fastvideo.train.models.ltx2.LTX2Model`` and the underlying
+``LTX2Transformer3DModel``.
+
+LTX-2's transformer takes per-token sigma timesteps in [0, 1] shaped
+[B, tokens] and a post-connector Gemma text embedding (3840-d for 2.0;
+4096-d for 2.3, which has no in-DiT caption projection); it returns
+the denoised x0 prediction. This mirrors the kwargs in
+``LTX2Model._build_distill_input_kwargs``.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Required by the ``distributed_setup`` fixture pulled from
+# ``fastvideo/tests/conftest.py``.  Set before any fastvideo import.
+os.environ.setdefault("MASTER_ADDR", "localhost")
+os.environ.setdefault("MASTER_PORT", "29519")
+
+from pathlib import Path
+
+import pytest
+import torch
+
+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
+
+_FIXTURE_DIR = Path(__file__).resolve().parent.parent / "fixtures"
+
+# id -> (fixture, post-connector text embedding width)
+_CASES = {
+    "ltx2": ("ltx2_t2v_finetune_min.yaml", 3840),
+    "ltx2_3": ("ltx2_3_t2v_finetune_min.yaml", 4096),
+}
+
+_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
+    total = torch.cuda.get_device_properties(0).total_memory
+    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):
+    if _gpu_too_small():
+        pytest.skip(f"requires a CUDA GPU with >= {_MIN_GPU_MEMORY_GB}GB "
+                    "memory (LTX-2 DiT is 18.9B params)")
+
+    fixture_name, text_dim = _CASES[case]
+    cfg = load_run_config(str(_FIXTURE_DIR / fixture_name))
+    model = LTX2Model(
+        init_from=cfg.models["student"]["init_from"],
+        training_config=cfg.training,
+        trainable=False,
+    )
+
+    transformer = model.transformer
+    assert isinstance(transformer, torch.nn.Module)
+    assert sum(p.numel() for p in transformer.parameters()) > 0
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+    transformer = transformer.to(device=device, dtype=dtype).eval()
+
+    # LTX-2 transformer takes [B, 128, T, H, W] latents, a post-connector
+    # Gemma embedding, and PER-TOKEN sigmas in [0, 1] ([B, T*H*W] with
+    # patch size 1x1x1). Small spatial + few frames so this fits next to
+    # the 18.9B model.
+    b, c, t, h, w = 1, 128, 3, 8, 8
+    tokens = t * h * w
+    hidden_states = torch.randn(b, c, t, h, w, device=device, dtype=dtype)
+    encoder_hidden_states = torch.randn(b,
+                                        _LTX2_TEXT_LEN,
+                                        text_dim,
+                                        device=device,
+                                        dtype=dtype)
+    timestep = torch.full((b, tokens), 0.5, device=device, dtype=torch.float32)
+
+    with torch.no_grad(), torch.autocast(device.type, dtype=dtype), \
+            set_forward_context(
+                current_timestep=timestep * 1000.0,
+                attn_metadata=None,
+                forward_batch=ForwardBatch(data_type="video", fps=24.0),
+            ):
+        out = transformer(
+            hidden_states=hidden_states,
+            encoder_hidden_states=encoder_hidden_states,
+            timestep=timestep,
+            encoder_attention_mask=None,
+            return_dict=False,
+        )
+
+    assert torch.is_tensor(out), f"expected a tensor output, got {type(out)}"
+    assert out.shape == hidden_states.shape, (
+        f"denoised output shape {tuple(out.shape)} != input latent shape "
+        f"{tuple(hidden_states.shape)}")
+    assert torch.isfinite(out).all().item(), "forward output contains NaN/Inf"

fastvideo/train/models/ltx2/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2 model plugin package."""
+
+from fastvideo.train.models.ltx2.ltx2 import (
+    LTX2Model as LTX2Model, )

fastvideo/train/models/ltx2/ltx2.py

@@ -0,0 +1,562 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2 model plugin (per-role instance).
+
+Subclasses WanModel but replaces the pieces where LTX-2 differs:
+  - transformer class name: LTX2Transformer3DModel (blocks nested
+    under ``transformer.model``, so activation checkpointing must
+    target the inner module)
+  - latents come pre-normalized from the LTX-2 VAE encoder
+    (per-channel statistics applied inside ``encode``), so no
+    ``normalize_dit_input`` call
+  - sigma sampling follows the official LTX-2 trainer: stretched
+    shifted-logit-normal with a token-count-dependent shift
+    (0.95 @ 1024 tokens -> 2.05 @ 4096 tokens) and a 10% uniform
+    mixture, instead of scheduler-index sampling
+  - the DiT consumes PER-TOKEN sigmas in [0, 1] shaped [B, tokens]
+    (not integer 0-1000 timesteps) and returns the DENOISED x0
+    prediction; ``predict_noise`` converts it back to the
+    framework's velocity convention v = (x_t - x0) / sigma so the
+    default FineTune target ``noise - clean`` applies unchanged
+  - temporal RoPE coordinates are divided by fps read from
+    ``get_forward_context().forward_batch.fps``; training must set
+    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: video-only forwards never give them
+    gradients, and freezing keeps them out of DCP checkpoints and
+    optimizer state
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.dits.ltx2 import VideoLatentShape
+from fastvideo.pipelines import ForwardBatch, TrainingBatch
+from fastvideo.training.activation_checkpoint import (
+    apply_activation_checkpointing, )
+
+from fastvideo.train.models.wan.wan import WanModel
+from fastvideo.train.utils.module_state import (
+    apply_trainable, )
+from fastvideo.train.utils.moduleloader import (
+    load_module_from_path, )
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+    from fastvideo.train.utils.lora import LoraConfig
+
+logger = init_logger(__name__)
+
+# Parameters whose names match any of these substrings belong to the
+# audio branch or the audio<->video cross-modal machinery. They are
+# unused (no gradients) in video-only forwards.
+_AUDIO_PARAM_PATTERNS = ("audio", "a2v", "v2a", "av_ca")
+
+# Official LTX-2 trainer timestep sampler constants
+# (ltx_trainer/timestep_samplers.py: ShiftedLogitNormalTimestepSampler).
+_SIGMA_MIN_TOKENS = 1024.0
+_SIGMA_MAX_TOKENS = 4096.0
+_SIGMA_MIN_SHIFT = 0.95
+_SIGMA_MAX_SHIFT = 2.05
+_SIGMA_STD = 1.0
+_SIGMA_EPS = 1e-3
+# 0.5% / 99.9% normal percentiles used to stretch the logit-normal
+# samples so the sigma distribution covers the full [0, 1] range.
+_SIGMA_Z_LO = -2.5758
+_SIGMA_Z_HI = 3.0902
+
+# Preset fps for LTX-2 checkpoints; used when a batch carries no fps
+# metadata so training RoPE still matches validation inference.
+_DEFAULT_ROPE_FPS = 24.0
+
+
+class LTX2Model(WanModel):
+    """LTX-2 per-role model for the modular trainer."""
+
+    _transformer_cls_name: str = "LTX2Transformer3DModel"
+
+    def __init__(
+        self,
+        *,
+        init_from: str,
+        training_config: TrainingConfig,
+        trainable: bool = True,
+        disable_custom_init_weights: bool = False,
+        enable_gradient_checkpointing_type: str
+        | None = None,
+        transformer_override_safetensor: str
+        | None = None,
+        lora: LoraConfig | dict[str, Any] | None = None,
+        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 "
+                                      "zeroes the post-connector text embeddings, which is not "
+                                      "what LTX-2 inference uses as the unconditional input "
+                                      "(an empty prompt encoded through Gemma + connector). Set "
+                                      "training.data.training_cfg_rate: 0.0.")
+
+        self._timestep_uniform_prob = float(timestep_uniform_prob)
+        self._rope_fps: float = _DEFAULT_ROPE_FPS
+        self._rope_forward_batch: ForwardBatch | None = None
+
+        super().__init__(
+            init_from=init_from,
+            training_config=training_config,
+            trainable=trainable,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            # The LTX-2 sigma sampler below does not use the scheduler;
+            # keep an unshifted scheduler for the inherited distillation
+            # helpers (add_noise / num_train_timesteps).
+            flow_shift=1.0,
+            enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
+            transformer_override_safetensor=(transformer_override_safetensor),
+            lora=lora,
+        )
+
+        if trainable:
+            self._freeze_audio_parameters()
+
+        # No negative-prompt cache: cfg_rate is forced to 0 above and
+        # loading Gemma (~23GB) on every rank just for an unused
+        # negative embedding is wasteful.
+        self.set_requires_negative_conditioning(False)
+
+    # ------------------------------------------------------------------
+    # Loading
+    # ------------------------------------------------------------------
+
+    def _load_transformer(
+        self,
+        *,
+        init_from: str,
+        trainable: bool,
+        disable_custom_init_weights: bool,
+        enable_gradient_checkpointing_type: str | None,
+        training_config: TrainingConfig,
+        transformer_override_safetensor: str | None = None,
+    ) -> torch.nn.Module:
+        transformer = load_module_from_path(
+            model_path=init_from,
+            module_type="transformer",
+            training_config=training_config,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            override_transformer_cls_name=(self._transformer_cls_name),
+            transformer_override_safetensor=(transformer_override_safetensor),
+        )
+        ckpt_type = (enable_gradient_checkpointing_type or getattr(
+            getattr(training_config, "model", None),
+            "enable_gradient_checkpointing_type",
+            None,
+        ))
+        if trainable and ckpt_type:
+            # LTX-2 nests transformer_blocks under ``.model``; applying
+            # checkpointing at the wrapper level raises because no block
+            # list is found there.
+            transformer.model = apply_activation_checkpointing(
+                transformer.model,
+                checkpointing_type=ckpt_type,
+            )
+        if self._enable_lora_if_configured(transformer):
+            return transformer
+        transformer = apply_trainable(transformer, trainable=trainable)
+        return transformer
+
+    def _freeze_audio_parameters(self) -> None:
+        frozen_params = 0
+        for name, param in self.transformer.named_parameters():
+            if any(pattern in name for pattern in _AUDIO_PARAM_PATTERNS):
+                param.requires_grad_(False)
+                frozen_params += param.numel()
+        logger.info(
+            "LTX2Model: froze %.2fB audio/cross-modal parameters "
+            "(video-only training)",
+            frozen_params / 1e9,
+        )
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def ensure_negative_conditioning(self) -> None:
+        raise NotImplementedError("LTX2Model does not implement negative conditioning; "
+                                  "training_cfg_rate must stay 0.")
+
+    @torch.no_grad()
+    def decode_latents(
+        self,
+        latents_b_t_c_h_w: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.vae is None:
+            raise RuntimeError("LTX-2 VAE is not initialized")
+        latents = latents_b_t_c_h_w.permute(0, 2, 1, 3, 4)
+        # The LTX-2 decoder un-normalizes internally via its
+        # per-channel statistics buffers; no external denorm.
+        media = self.vae.decode(latents.to(next(self.vae.parameters()).dtype))
+        return (media.float() / 2 + 0.5).clamp(0, 1)
+
+    def _init_timestep_mechanics(self) -> None:
+        assert self.training_config is not None
+        tc = self.training_config
+        # LTX2T2VConfig carries no flow_shift (sigmas are computed
+        # inline by the pipeline); the inherited shift/clamp helpers
+        # are unused for fine-tuning, so default to identity shift.
+        flow_shift = tc.pipeline_config.flow_shift  # type: ignore[union-attr]
+        self.timestep_shift = (float(flow_shift) if flow_shift is not None else 1.0)
+        self.num_train_timestep = int(self.noise_scheduler.num_train_timesteps)
+        self.min_timestep = 0
+        self.max_timestep = self.num_train_timestep
+
+    # ------------------------------------------------------------------
+    # Runtime primitives
+    # ------------------------------------------------------------------
+
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        assert self.training_config is not None
+        tc = self.training_config
+
+        dtype = self._get_training_dtype()
+        device = self.device
+
+        training_batch = TrainingBatch()
+        encoder_hidden_states = raw_batch["text_embedding"]
+        encoder_attention_mask = raw_batch["text_attention_mask"]
+        infos = raw_batch.get("info_list")
+
+        if latents_source == "zeros":
+            batch_size = encoder_hidden_states.shape[0]
+            vae_config = (
+                tc.pipeline_config.vae_config.arch_config  # type: ignore[union-attr]
+            )
+            num_channels = getattr(
+                vae_config,
+                "z_dim",
+                getattr(vae_config, "latent_channels", 128),
+            )
+            spatial_compression_ratio = (vae_config.spatial_compression_ratio)
+            latent_height = (tc.data.num_height // spatial_compression_ratio)
+            latent_width = (tc.data.num_width // spatial_compression_ratio)
+            latents = torch.zeros(
+                batch_size,
+                num_channels,
+                tc.data.num_latent_t,
+                latent_height,
+                latent_width,
+                device=device,
+                dtype=dtype,
+            )
+        elif latents_source == "data":
+            if "vae_latent" not in raw_batch:
+                raise ValueError("vae_latent not found in batch "
+                                 "and latents_source='data'")
+            latents = raw_batch["vae_latent"]
+            latents = latents[:, :, :tc.data.num_latent_t]
+            latents = latents.to(device, dtype=dtype)
+        else:
+            raise ValueError(f"Unknown latents_source: "
+                             f"{latents_source!r}")
+
+        self._check_text_embedding_dim(encoder_hidden_states)
+
+        # LTX-2 VAE encode() already applies per-channel normalization
+        # (scaling_factor is 1.0); latents are used as stored.
+        training_batch.latents = latents
+        training_batch.encoder_hidden_states = (encoder_hidden_states.to(device, dtype=dtype))
+        training_batch.encoder_attention_mask = (encoder_attention_mask.to(device, dtype=dtype))
+        training_batch.infos = infos
+
+        self._update_rope_fps(infos)
+
+        training_batch = self._prepare_dit_inputs(training_batch, generator)
+        training_batch = self._build_attention_metadata(training_batch)
+        training_batch.attn_metadata_vsa = None
+
+        return training_batch
+
+    def predict_noise(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+        clean_x: torch.Tensor | None = None,
+        aug_t: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if clean_x is not None or aug_t is not None:
+            raise NotImplementedError("LTX2Model does not support teacher forcing inputs")
+        device_type = self.device.type
+        dtype = self._get_training_dtype()
+        if conditional:
+            text_dict = batch.conditional_dict
+            if text_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+        else:
+            text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
+
+        if attn_kind not in ("dense", "vsa"):
+            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
+        attn_metadata = batch.attn_metadata
+
+        assert batch.sigmas is not None
+        # finetune.py hands noisy_latents in (B, T, C, H, W); the LTX-2
+        # DiT expects (B, C, T, H, W).
+        noisy_bcthw = noisy_latents.permute(0, 2, 1, 3, 4).to(dtype)
+
+        with torch.autocast(device_type, dtype=dtype), set_forward_context(
+                current_timestep=batch.timesteps,
+                attn_metadata=attn_metadata,
+                forward_batch=self._make_rope_forward_batch(),
+        ):
+            input_kwargs = self._build_distill_input_kwargs(
+                noisy_bcthw,
+                timestep,
+                text_dict,
+            )
+            transformer = self._get_transformer(timestep)
+            denoised = transformer(**input_kwargs)
+
+        if isinstance(denoised, tuple):
+            denoised = denoised[0]
+
+        # The wrapper returns the denoised x0 prediction
+        # (sample - velocity * sigma). Convert back to velocity so the
+        # framework's default target ``noise - clean`` applies. Compute
+        # in fp32: sigma can be ~1e-3 and bf16 division would amplify
+        # quantization error.
+        sigmas = batch.sigmas.float()
+        velocity = (noisy_bcthw.float() - denoised.float()) / sigmas
+        return velocity.permute(0, 2, 1, 3, 4)
+
+    def backward(
+        self,
+        loss: torch.Tensor,
+        ctx: Any,
+        *,
+        grad_accum_rounds: int,
+    ) -> None:
+        timesteps, attn_metadata = ctx
+        # Re-enter the forward context with the same fps-carrying batch
+        # so activation-checkpoint recompute sees identical RoPE inputs.
+        with set_forward_context(
+                current_timestep=timesteps,
+                attn_metadata=attn_metadata,
+                forward_batch=self._make_rope_forward_batch(),
+        ):
+            (loss / max(1, int(grad_accum_rounds))).backward()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _check_text_embedding_dim(self, encoder_hidden_states: torch.Tensor) -> None:
+        """Fail fast when the parquet was preprocessed with a mismatched
+        text stack (LTX-2.0 stores 3840-d post-connector embeddings; 2.3
+        has no in-DiT caption projection and expects 4096-d)."""
+        # Read the LOADED transformer's config: the transformer loader
+        # deepcopies pipeline_config.dit_config before applying the
+        # checkpoint's arch flags, so training_config.pipeline_config
+        # never sees version flags like caption_proj_before_connector.
+        arch = self.transformer.config.arch_config
+        if bool(getattr(arch, "caption_proj_before_connector", False)):
+            expected_dim = int(arch.cross_attention_dim)
+        else:
+            expected_dim = int(arch.caption_channels)
+        actual_dim = int(encoder_hidden_states.shape[-1])
+        if actual_dim != expected_dim:
+            raise ValueError(f"text_embedding width {actual_dim} does not match the "
+                             f"checkpoint's expected text context dim {expected_dim}. "
+                             "The parquet was likely preprocessed with a different "
+                             "LTX-2 version (2.0 stores 3840-d, 2.3 stores 4096-d); "
+                             "re-run preprocess_ltx2_overfit.py with LTX2_OVERFIT_MODEL "
+                             "set to the same checkpoint as models.student.init_from.")
+
+    def _update_rope_fps(self, infos: list[dict[str, Any]] | None) -> None:
+        fps: float | None = None
+        if infos:
+            raw_fps = infos[0].get("fps")
+            if raw_fps:
+                fps = float(raw_fps)
+        if fps is None or fps <= 0:
+            fps = _DEFAULT_ROPE_FPS
+        if self._rope_forward_batch is None or fps != self._rope_fps:
+            self._rope_fps = fps
+            self._rope_forward_batch = ForwardBatch(data_type="video", fps=fps)
+
+    def _make_rope_forward_batch(self) -> ForwardBatch:
+        if self._rope_forward_batch is None:
+            self._rope_forward_batch = ForwardBatch(
+                data_type="video",
+                fps=self._rope_fps,
+            )
+        return self._rope_forward_batch
+
+    def _sample_ltx2_sigmas(
+        self,
+        batch_size: int,
+        token_count: int,
+        device: torch.device,
+        generator: torch.Generator,
+    ) -> torch.Tensor:
+        """Official LTX-2 stretched shifted-logit-normal sigma sampler.
+
+        mu is linearly interpolated from the patchified video token
+        count (unclamped, matching the official trainer), the
+        logit-normal samples are stretched to cover [0, 1] using the
+        0.5%/99.9% percentiles, and 10% of samples are replaced by
+        U(eps, 1).
+        """
+        slope = ((_SIGMA_MAX_SHIFT - _SIGMA_MIN_SHIFT) / (_SIGMA_MAX_TOKENS - _SIGMA_MIN_TOKENS))
+        mu = slope * float(token_count) + (_SIGMA_MIN_SHIFT - slope * _SIGMA_MIN_TOKENS)
+
+        normal = torch.randn(
+            (batch_size, ),
+            generator=generator,
+            device=device,
+            dtype=torch.float32,
+        ) * _SIGMA_STD + mu
+        sigmas = torch.sigmoid(normal)
+
+        lo = torch.sigmoid(torch.tensor(mu + _SIGMA_Z_LO * _SIGMA_STD, device=device))
+        hi = torch.sigmoid(torch.tensor(mu + _SIGMA_Z_HI * _SIGMA_STD, device=device))
+        raw = (sigmas - lo) / (hi - lo)
+        stretched = torch.where(raw >= _SIGMA_EPS, raw, 2 * _SIGMA_EPS - raw)
+        stretched = stretched.clamp(0.0, 1.0)
+
+        if self._timestep_uniform_prob > 0.0:
+            prob = torch.rand(
+                (batch_size, ),
+                generator=generator,
+                device=device,
+            )
+            uniform = torch.rand(
+                (batch_size, ),
+                generator=generator,
+                device=device,
+            ) * (1.0 - _SIGMA_EPS) + _SIGMA_EPS
+            stretched = torch.where(
+                prob > self._timestep_uniform_prob,
+                stretched,
+                uniform,
+            )
+        return stretched
+
+    def _prepare_dit_inputs(
+        self,
+        training_batch: TrainingBatch,
+        generator: torch.Generator,
+    ) -> TrainingBatch:
+        assert self.training_config is not None
+        tc = self.training_config
+        latents = training_batch.latents
+        assert isinstance(latents, torch.Tensor)
+        batch_size = latents.shape[0]
+
+        video_shape = VideoLatentShape.from_torch_shape(latents.shape)
+        patchifier = getattr(self.transformer, "patchifier", None)
+        if patchifier is not None:
+            token_count = int(patchifier.get_token_count(video_shape))
+        else:
+            token_count = int(latents.shape[2] * latents.shape[3] * latents.shape[4])
+        self._token_count = token_count
+
+        sigmas = self._sample_ltx2_sigmas(
+            batch_size,
+            token_count,
+            latents.device,
+            generator,
+        )
+        if int(tc.distributed.sp_size or 1) > 1:
+            self.sp_group.broadcast(sigmas, src=0)
+
+        noise = torch.randn(
+            latents.shape,
+            generator=generator,
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+        sigmas_expanded = sigmas.view(-1, 1, 1, 1, 1)
+        noisy_model_input = ((1.0 - sigmas_expanded) * latents.float() + sigmas_expanded * noise.float()).to(
+            latents.dtype)
+
+        training_batch.noisy_model_input = noisy_model_input
+        # Keep the framework-wide "timesteps ~ sigma * 1000" convention
+        # (used for logging and the forward context); sigmas stay fp32
+        # for the velocity conversion.
+        training_batch.timesteps = sigmas * 1000.0
+        training_batch.sigmas = sigmas_expanded
+        training_batch.noise = noise
+        training_batch.raw_latent_shape = latents.shape
+
+        training_batch.conditional_dict = {
+            "encoder_hidden_states": (training_batch.encoder_hidden_states),
+            "encoder_attention_mask": (training_batch.encoder_attention_mask),
+        }
+
+        training_batch.latents = (training_batch.latents.permute(0, 2, 1, 3, 4))
+        return training_batch
+
+    def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch:
+        if envs.FASTVIDEO_ATTENTION_BACKEND in ("VIDEO_SPARSE_ATTN", "VMOBA_ATTN"):
+            raise NotImplementedError("LTX2Model does not support VSA/VMOBA attention backends")
+        training_batch.attn_metadata = None
+        return training_batch
+
+    def _build_distill_input_kwargs(
+        self,
+        noise_input: torch.Tensor,
+        timestep: torch.Tensor,
+        text_dict: dict[str, torch.Tensor] | None,
+        clean_x: torch.Tensor | None = None,
+        aug_t: torch.Tensor | None = None,
+    ) -> dict[str, Any]:
+        if text_dict is None:
+            raise ValueError("text_dict cannot be None for LTX-2 forward")
+        if clean_x is not None or aug_t is not None:
+            raise NotImplementedError("LTX2Model does not support teacher forcing inputs")
+
+        # noise_input arrives already in (B, C, T, H, W).
+        batch_size = noise_input.shape[0]
+        token_count = getattr(self, "_token_count", None)
+        if token_count is None:
+            token_count = int(noise_input.shape[2] * noise_input.shape[3] * noise_input.shape[4])
+
+        # The DiT wants per-token sigmas in [0, 1] shaped [B, tokens]
+        # (i2v-style conditioning would zero conditioned tokens; plain
+        # T2V uses the same sigma everywhere). ``timestep`` follows the
+        # framework's sigma*1000 convention.
+        sigma = (timestep.to(torch.float32) / 1000.0).view(batch_size, 1)
+        per_token_timestep = sigma.expand(batch_size, token_count).contiguous()
+
+        return {
+            "hidden_states": noise_input,
+            "encoder_hidden_states": (text_dict["encoder_hidden_states"]),
+            # Post-connector embeddings are all-valid; the connector
+            # replaced pad positions with learnable registers.
+            "encoder_attention_mask": None,
+            "timestep": per_token_timestep,
+            "return_dict": False,
+        }