[feat] Add LTX-2.3 fine-tuning support and overfit verification
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).
This commit is contained in:
Will Lin
committed by
SolitaryThinker
SolitaryThinker
parent
f96b8e2ad2
commit
86a360ebc7
@@ -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: {}
|
||||
@@ -8,7 +8,10 @@ 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.
|
||||
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
|
||||
@@ -20,6 +23,7 @@ Usage:
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
from typing import Any
|
||||
|
||||
import av
|
||||
@@ -105,6 +109,9 @@ def main() -> None:
|
||||
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)
|
||||
@@ -170,17 +177,16 @@ def main() -> None:
|
||||
print(f" Latent shape: {latent.shape}")
|
||||
|
||||
with torch.no_grad():
|
||||
pre_embeds, attention_mask = text_encoder.preprocess_text_embeddings(
|
||||
prompts=[preprocess_text(caption)],
|
||||
tokenizer=tokenizer,
|
||||
tokenizer_kwargs=tokenizer_kwargs,
|
||||
padding_side=encoder_config.arch_config.padding_side,
|
||||
# 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),
|
||||
)
|
||||
video_embeds, _audio_embeds, _post_mask = text_encoder.run_connectors(
|
||||
pre_embeds,
|
||||
attention_mask,
|
||||
)
|
||||
text_embedding = video_embeds.squeeze(0).float().cpu()
|
||||
text_embedding = encoder_out.last_hidden_state.squeeze(0).float().cpu()
|
||||
print(f" Text embedding shape: {text_embedding.shape}")
|
||||
|
||||
record = {
|
||||
|
||||
@@ -1,16 +1,17 @@
|
||||
"""E2E overfit test for LTX-2 fine-tuning on the modular trainer (fastvideo/train).
|
||||
"""E2E overfit tests for LTX-2 fine-tuning on the modular trainer (fastvideo/train).
|
||||
|
||||
Downloads the single-sample cats dataset, preprocesses it with the
|
||||
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.
|
||||
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/LTX2-Distilled-Diffusers checkpoint (set
|
||||
HF_HOME to a cache that contains it, or allow ~60GB of downloads).
|
||||
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.
|
||||
"""
|
||||
@@ -35,11 +36,23 @@ NUM_GPUS_TRAINING = "4"
|
||||
|
||||
DATA_DIR = "data"
|
||||
LOCAL_RAW_DATA_DIR = Path(DATA_DIR) / "cats"
|
||||
LOCAL_PREPROCESSED_DATA_DIR = Path(DATA_DIR) / "ltx2_overfit_preprocessed"
|
||||
LOCAL_OUTPUT_DIR = Path(DATA_DIR) / "outputs_ltx2_overfit"
|
||||
|
||||
PREPROCESSING_SCRIPT = "fastvideo/pipelines/preprocess/preprocess_ltx2_overfit.py"
|
||||
TRAINING_CONFIG = "examples/train/configs/overfit_ltx2_t2v.yaml"
|
||||
|
||||
_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():
|
||||
@@ -55,17 +68,18 @@ def download_data():
|
||||
f"Download appeared to succeed but {LOCAL_RAW_DATA_DIR} does not exist")
|
||||
|
||||
|
||||
def run_preprocessing():
|
||||
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(LOCAL_PREPROCESSED_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}")
|
||||
print(f"Running preprocessing: {cmd} (model={case['model']})")
|
||||
subprocess.run(cmd, check=True, env=env)
|
||||
|
||||
|
||||
def run_training():
|
||||
def run_training(case: dict):
|
||||
env = os.environ.copy()
|
||||
env.setdefault("WANDB_MODE", "offline")
|
||||
cmd = [
|
||||
@@ -73,18 +87,18 @@ def run_training():
|
||||
"--nnodes", "1",
|
||||
"--nproc_per_node", NUM_GPUS_TRAINING,
|
||||
"-m", "fastvideo.train.entrypoint.train",
|
||||
"--config", TRAINING_CONFIG,
|
||||
"--training.checkpoint.output_dir", str(LOCAL_OUTPUT_DIR),
|
||||
"--training.data.data_path", str(LOCAL_PREPROCESSED_DATA_DIR),
|
||||
"--config", case["config"],
|
||||
"--training.checkpoint.output_dir", str(case["out_dir"]),
|
||||
"--training.data.data_path", str(case["prep_dir"]),
|
||||
"--callbacks.validation.dataset_file",
|
||||
str(LOCAL_PREPROCESSED_DATA_DIR / "validation_prompts.json"),
|
||||
str(case["prep_dir"] / "validation_prompts.json"),
|
||||
]
|
||||
print(f"Running training: {cmd}")
|
||||
subprocess.run(cmd, check=True, env=env)
|
||||
|
||||
|
||||
def _validation_videos_by_step() -> dict[int, str]:
|
||||
pattern = os.path.join(str(LOCAL_OUTPUT_DIR), "validation_step_*_video_0.mp4")
|
||||
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))
|
||||
@@ -97,27 +111,29 @@ def _validation_videos_by_step() -> dict[int, str]:
|
||||
os.environ.get("FASTVIDEO_NIGHTLY") != "1",
|
||||
reason="nightly e2e overfit test; set FASTVIDEO_NIGHTLY=1 to run",
|
||||
)
|
||||
def test_e2e_ltx2_overfit_new_stack():
|
||||
@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()
|
||||
run_training()
|
||||
run_preprocessing(case)
|
||||
run_training(case)
|
||||
|
||||
reference_video = LOCAL_PREPROCESSED_DATA_DIR / "training_sample_0.mp4"
|
||||
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()
|
||||
assert videos, f"No validation videos found under {LOCAL_OUTPUT_DIR}"
|
||||
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} =====")
|
||||
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 = {"final_step": final_step, "final_mean_ssim": final_mean}
|
||||
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(
|
||||
@@ -136,4 +152,5 @@ def test_e2e_ltx2_overfit_new_stack():
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ.setdefault("FASTVIDEO_NIGHTLY", "1")
|
||||
test_e2e_ltx2_overfit_new_stack()
|
||||
for _case_id in _CASES:
|
||||
test_e2e_ltx2_overfit_new_stack(_case_id)
|
||||
|
||||
@@ -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: {}
|
||||
@@ -3,6 +3,9 @@
|
||||
"H200": 5.2763,
|
||||
"L40S": 4.1553
|
||||
},
|
||||
"test_ltx2_3_finetune": {
|
||||
"GB200": 2.2158
|
||||
},
|
||||
"test_ltx2_finetune": {
|
||||
"GB200": 0.3127
|
||||
},
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Per-method GPU smoke test: ``LTX2Model`` + ``FineTuneMethod``.
|
||||
|
||||
Mirrors ``test_wan_finetune.py`` for the LTX-2 plugin. LTX2-specific
|
||||
differences: the synthetic ``raw_batch`` carries a single
|
||||
post-connector Gemma embedding ([1024, 3840]) 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).
|
||||
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
|
||||
@@ -30,12 +32,14 @@ from .grad_norm_regression import (
|
||||
resolve_blocks,
|
||||
)
|
||||
|
||||
_FIXTURE = str(
|
||||
Path(__file__).resolve().parent.parent / "fixtures"
|
||||
/ "ltx2_t2v_finetune_min.yaml")
|
||||
_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"),
|
||||
}
|
||||
|
||||
# Post-connector Gemma embedding width / length.
|
||||
_LTX2_TEXT_DIM = 3840
|
||||
_LTX2_TEXT_LEN = 1024
|
||||
|
||||
_MIN_GPU_MEMORY_GB = 60
|
||||
@@ -51,6 +55,7 @@ def _gpu_too_small() -> bool:
|
||||
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``.
|
||||
|
||||
@@ -62,7 +67,7 @@ def _build_synthetic_batch(
|
||||
"text_embedding":
|
||||
torch.randn(batch_size,
|
||||
_LTX2_TEXT_LEN,
|
||||
_LTX2_TEXT_DIM,
|
||||
text_dim,
|
||||
device=device,
|
||||
dtype=dtype),
|
||||
"text_attention_mask":
|
||||
@@ -73,13 +78,15 @@ def _build_synthetic_batch(
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("distributed_setup")
|
||||
@pytest.mark.parametrize("case", _CASES.keys())
|
||||
def test_ltx2_finetune_single_train_step(
|
||||
monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
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)")
|
||||
|
||||
cfg = load_run_config(_FIXTURE)
|
||||
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
|
||||
@@ -107,7 +114,7 @@ def test_ltx2_finetune_single_train_step(
|
||||
)
|
||||
method.on_train_start()
|
||||
|
||||
batch = _build_synthetic_batch(device, dtype)
|
||||
batch = _build_synthetic_batch(device, dtype, text_dim)
|
||||
loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
|
||||
|
||||
loss = loss_map["total_loss"]
|
||||
@@ -149,4 +156,4 @@ def test_ltx2_finetune_single_train_step(
|
||||
|
||||
# Device-keyed grad-norm regression on top of the same harness.
|
||||
# Skips when the current GPU has no seeded reference.
|
||||
check_grad_norm_regression("test_ltx2_finetune", model.transformer.model)
|
||||
check_grad_norm_regression(ref_name, model.transformer.model)
|
||||
|
||||
@@ -1,16 +1,17 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""GPU loading + forward smoke test for ``LTX2Model``.
|
||||
|
||||
Loads the real FastVideo/LTX2-Distilled-Diffusers checkpoint (18.9B at
|
||||
bf16, ~38GB — skips on GPUs with less than 60GB memory) via
|
||||
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 ([1024, 3840]);
|
||||
it returns the denoised x0 prediction. This mirrors the kwargs in
|
||||
[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``.
|
||||
"""
|
||||
|
||||
@@ -33,12 +34,14 @@ from fastvideo.pipelines import ForwardBatch
|
||||
from fastvideo.train.models.ltx2 import LTX2Model
|
||||
from fastvideo.train.utils.config import load_run_config
|
||||
|
||||
_FIXTURE = str(
|
||||
Path(__file__).resolve().parent.parent / "fixtures"
|
||||
/ "ltx2_t2v_finetune_min.yaml")
|
||||
_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),
|
||||
}
|
||||
|
||||
# Post-connector Gemma embedding width / length.
|
||||
_LTX2_TEXT_DIM = 3840
|
||||
_LTX2_TEXT_LEN = 1024
|
||||
|
||||
_MIN_GPU_MEMORY_GB = 60
|
||||
@@ -52,12 +55,14 @@ def _gpu_too_small() -> bool:
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("distributed_setup")
|
||||
def test_ltx2_model_loads_and_forwards():
|
||||
@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)")
|
||||
|
||||
cfg = load_run_config(_FIXTURE)
|
||||
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,
|
||||
@@ -81,7 +86,7 @@ def test_ltx2_model_loads_and_forwards():
|
||||
hidden_states = torch.randn(b, c, t, h, w, device=device, dtype=dtype)
|
||||
encoder_hidden_states = torch.randn(b,
|
||||
_LTX2_TEXT_LEN,
|
||||
_LTX2_TEXT_DIM,
|
||||
text_dim,
|
||||
device=device,
|
||||
dtype=dtype)
|
||||
timestep = torch.full((b, tokens), 0.5, device=device, dtype=torch.float32)
|
||||
|
||||
@@ -272,6 +272,8 @@ class LTX2Model(WanModel):
|
||||
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
|
||||
@@ -366,6 +368,28 @@ class LTX2Model(WanModel):
|
||||
# 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:
|
||||
|
||||
Reference in New Issue
Block a user