[feat] Add end-to-end MMAudio V2A training, inference, and evaluation (#1648)

Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

examples/inference/eval/mmaudio/README.md

@@ -0,0 +1,237 @@
+# MMAudio dataset inference and evaluation
+
+This example evaluates a trained MMAudio checkpoint without importing the
+upstream MMAudio package. Generation uses FastVideo's native
+`MMAudioPipeline`; the default metrics use `fastvideo.eval`.
+
+## Relation to official MMAudio
+
+Official `batch_eval.py` launches one complete model replica per GPU with
+`torchrun`, uses a `DistributedSampler`, runs batched bf16 inference, and
+writes audio without composing videos. Metric calculation is a separate step
+through `av_bench.extract` and `av_bench.evaluate`.
+
+This FastVideo path keeps that high-level separation:
+
+1. one complete MMAudio pipeline replica per GPU;
+2. data-parallel dataset inference to WAV;
+3. a canonical JSONL manifest that binds id, source video, generated audio,
+   caption, seed, and reference-audio source;
+4. independent multi-GPU metric workers over the complete corpus.
+
+There are two intentional differences. Rank-strided indices avoid the padded
+duplicates that PyTorch `DistributedSampler(drop_last=False)` can introduce,
+and cached CLIP/Synchformer/text features skip video decoding and conditioning
+encoder work already completed during FastVideo preprocessing. The current
+cached-feature stage accepts one sample per pipeline call; four GPUs still
+process four samples concurrently.
+
+## Environment
+
+Use the same FastVideo environment as training:
+
+```bash
+cd /mnt/lustre/vlm-kai/FastVideo
+source .venv/bin/activate
+uv pip install --python .venv/bin/python hear21passt pyloudnorm
+```
+
+The four default native metrics are:
+
+| FastVideo metric | Protocol |
+| --- | --- |
+| `audio.frechet_distance` | corpus PaSST FAD / FD_PaSST |
+| `audio.kl_divergence` | paired PaSST KL, `KL(gt || pred)` |
+| `audio.clap_score` | text-audio CLAP similarity |
+| `audio.desync` | source-video/generated-audio Synchformer DeSync |
+
+FAD and KL port the av-benchmark mathematics. DeSync vendors the same
+Synchformer family. FastVideo CLAP uses the Hugging Face LAION CLAP checkpoint
+and is not byte-identical to the external package's CLAP implementation; use
+the optional official backend below when exact paper-table reproduction is
+required.
+
+## 1. Export the final EMA
+
+The training result is a transformer state dict. Inference also needs the
+frozen audio VAE, vocoder, scheduler, and (for raw-video inference) the frozen
+conditioning encoders. Export a complete FastVideo model once:
+
+```bash
+bash examples/inference/eval/mmaudio/export_final_ema.sh
+```
+
+Default input:
+
+```text
+outputs/mmaudio_small_44k_ddp_from_scratch/posthoc_ema/official_ddp/
+  mmaudio_ema_final_sigma_0p05_step_000300000.pth
+```
+
+Default output:
+
+```text
+converted_weights/mmaudio/small_44k_ema_300000/
+```
+
+Override `EMA_CHECKPOINT`, `OUTPUT_MODEL`, or `ASSET_ROOT` before the command
+for another run or variant.
+
+## 2. Multi-GPU inference
+
+The launcher defaults to the VGGSound test feature cache and four GPUs:
+
+```bash
+bash examples/inference/eval/mmaudio/run_inference_vggsound.sh
+```
+
+For a four-sample smoke test:
+
+```bash
+MAX_SAMPLES=4 \
+OUTPUT_DIR=outputs/mmaudio_small_44k_ema_300000_vggsound_smoke \
+bash examples/inference/eval/mmaudio/run_inference_vggsound.sh
+```
+
+For a background full run with a persistent log:
+
+```bash
+mkdir -p logs
+nohup bash examples/inference/eval/mmaudio/run_inference_vggsound.sh \
+  > logs/mmaudio_vggsound_inference.log 2>&1 &
+tail -f logs/mmaudio_vggsound_inference.log
+```
+
+The output directory contains:
+
+```text
+audio/<sample-id>.wav
+eval_manifest.jsonl
+failures.jsonl
+summary.json
+manifest_rank_*.jsonl
+failures_rank_*.jsonl
+```
+
+Existing WAV files are resumed rather than regenerated. Add `--overwrite` to
+the Python runner when intentional regeneration is required. Seeds are
+`base_seed + global_index`, so results do not change with GPU count.
+
+## 3. Native FastVideo evaluation
+
+Run FAD, KL, CLAP, and DeSync with four metric replicas:
+
+```bash
+bash examples/inference/eval/mmaudio/run_evaluate_vggsound.sh
+```
+
+The first run extracts ground-truth audio from each source MP4 into
+`<OUTPUT_DIR>/reference_audio` and caches metric weights. Later runs reuse
+both. Results are written to `fastvideo_eval_results.json`; per-sample metrics
+are under `samples`, while FAD is under `corpus`.
+
+The underlying generic command is:
+
+```bash
+fastvideo eval run \
+  --manifest outputs/mmaudio_small_44k_ema_300000_vggsound_test/eval_manifest.jsonl \
+  --metrics audio.frechet_distance,audio.kl_divergence,audio.clap_score,audio.desync \
+  --num-gpus 4 \
+  --extract-audio outputs/mmaudio_small_44k_ema_300000_vggsound_test/reference_audio \
+  --extract-workers 16 \
+  --output-format full \
+  --output outputs/mmaudio_small_44k_ema_300000_vggsound_test/fastvideo_eval_results.json
+```
+
+Do not evaluate FAD by invoking the evaluator once per file. FAD is a
+set-vs-set statistic; this CLI submits the complete manifest in one call and
+serializes `EvalResults.corpus`.
+
+## Exact external av-benchmark backend
+
+For strict comparison to official MMAudio numbers, FastVideo provides the
+dedicated `fastvideo eval v2a` command. It launches the exact
+`av_bench.extract` and `av_bench.evaluate` functions with a Python interpreter
+from an isolated environment. This keeps av-benchmark's CLAP, ImageBind, and
+PyTorch dependencies out of the main FastVideo environment.
+
+Create the isolated environment once:
+
+```bash
+cd /mnt/lustre/vlm-kai
+git clone https://github.com/hkchengrex/av-benchmark.git
+uv venv av-benchmark/.venv --python 3.12
+uv pip install --python av-benchmark/.venv/bin/python -e av-benchmark
+# PyTorch 2.9+ routes torchaudio decoding through TorchCodec. Transformers 5
+# removed an API used by the vendored Synchformer AST, while SciPy 1.17 removed
+# an argument used by the official FAD implementation. Keep these constraints
+# in this isolated environment; select the wheel matching its PyTorch CUDA build.
+UV_TORCH_BACKEND=cu130 uv pip install \
+  --python av-benchmark/.venv/bin/python \
+  torchcodec==0.15.0 'transformers<5' 'scipy<1.17'
+
+mkdir -p av-benchmark/weights
+curl -L https://huggingface.co/lukewys/laion_clap/resolve/main/music_speech_audioset_epoch_15_esc_89.98.pt \
+  -o av-benchmark/weights/music_speech_audioset_epoch_15_esc_89.98.pt
+curl -L https://github.com/hkchengrex/MMAudio/releases/download/v0.1/synchformer_state_dict.pth \
+  -o av-benchmark/weights/synchformer_state_dict.pth
+```
+
+Download only the official VGGSound ground-truth cache (about 16 GB):
+
+```bash
+cd /mnt/lustre/vlm-kai/FastVideo
+.venv/bin/hf download hkchengrex/MMAudio-precomputed-results \
+  --repo-type dataset \
+  --include "vggsound-test-eval-cache/*" \
+  --local-dir /mnt/lustre/vlm-kai/datasets/VGGSound/av_benchmark
+```
+
+Then run the dedicated launcher:
+
+```bash
+bash examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh
+```
+
+For a background run:
+
+```bash
+nohup bash examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh \
+  > logs/mmaudio_vggsound_av_benchmark.log 2>&1 &
+tail -f logs/mmaudio_vggsound_av_benchmark.log
+```
+
+The official VGGSound cache produces the complete paper-style metric set:
+
+```text
+FD-VGG, FD-PANN, FD-PASST
+KL-PANNS-softmax, KL-PASST-softmax
+ISC-PANNS-mean/std, ISC-PASST-mean/std
+IB-Score, DeSync
+```
+
+The official VGGSound cache does not contain CLAP text features, so the launcher
+defaults to `SKIP_CLAP=1`; this saves prediction feature passes without changing
+the metric set. Set `SKIP_CLAP=0` for a custom GT cache that includes CLAP text
+features. Prediction features are cached and reused; set `RECOMPUTE=1` to force
+a clean extraction.
+
+The launcher also defaults to `ALIGN_PREDICTION_KEYS=1`. The official
+VGGSound cache prepends underscores to some sanitized ids; FastVideo keeps the
+original ids. The backend changes only unique matches in prediction feature
+caches, leaves WAV filenames untouched, and records exact/remapped/ambiguous/
+unmatched counts under `prediction_key_alignment` in the result JSON.
+
+The launcher defaults to `NUM_WORKERS=0`: this avoids both the machine's 64 MB
+`/dev/shm` limit and forking workers after the official extractor initializes
+CUDA. Raise it only after validating the container and backend combination.
+
+## Evaluation protocol warning
+
+The current FastVideo test cache contains the filtered natural-language
+descriptions used in this training run. Official MMAudio's VGGSound
+`batch_eval.py` reads class-label captions from the original VGGSound CSV.
+FAD, KL, and DeSync formulas do not read captions directly, but changing the
+generation caption changes the generated waveform and can therefore change all
+scores. CLAP additionally consumes the caption during evaluation. Record the
+generation manifest, test split, and ground-truth cache with reported results.

examples/inference/eval/mmaudio/eval_mmaudio_dataset.py

@@ -0,0 +1,304 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Distributed MMAudio inference over FastVideo preprocessing caches.
+
+Launch with ``torchrun``. Each process owns one complete MMAudio pipeline and
+processes a non-padded, rank-strided subset of the dataset. The output manifest
+is directly consumable by ``fastvideo eval run --manifest ...``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from scipy.io import wavfile
+from torch.utils.data import DataLoader, Dataset, Sampler
+from tqdm import tqdm
+
+from fastvideo.dataset.mmaudio_feature_dataset import build_mmaudio_feature_dataset
+from fastvideo.distributed import cleanup_dist_env_and_memory
+from fastvideo.pipelines.basic.mmaudio import MMAudioPipeline
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+
+
+class _RankStrideSampler(Sampler[int]):
+    """Assign every global index exactly once without tail padding."""
+
+    def __init__(self, size: int, rank: int, world_size: int) -> None:
+        self.indices = range(rank, size, world_size)
+
+    def __iter__(self):
+        return iter(self.indices)
+
+    def __len__(self) -> int:
+        return len(self.indices)
+
+
+class _IndexedDataset(Dataset):
+
+    def __init__(self, dataset: Dataset) -> None:
+        self.dataset = dataset
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, index: int) -> tuple[int, dict[str, Any]]:
+        return index, self.dataset[index]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-path", type=Path, required=True)
+    parser.add_argument("--feature-root", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--duration-seconds", type=float, default=8.0)
+    parser.add_argument("--num-inference-steps", type=int, default=25)
+    parser.add_argument("--guidance-scale", type=float, default=4.5)
+    parser.add_argument("--seed", type=int, default=14159265)
+    parser.add_argument("--negative-prompt", default="")
+    parser.add_argument("--num-workers", type=int, default=2)
+    parser.add_argument("--max-samples", type=int, default=0)
+    parser.add_argument("--overwrite", action="store_true")
+    parser.add_argument("--compile", action="store_true")
+    return parser.parse_args()
+
+
+def _safe_name(value: str) -> str:
+    value = re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._")
+    return value or "sample"
+
+
+def _feature_shapes(transformer: torch.nn.Module) -> dict[str, int]:
+    arch = transformer.config.arch_config
+    return {
+        "latent_seq_len": int(transformer.latent_seq_len),
+        "latent_dim": int(transformer.latent_dim),
+        "clip_seq_len": int(transformer.clip_seq_len),
+        "clip_dim": int(arch.clip_dim),
+        "sync_seq_len": int(transformer.sync_seq_len),
+        "sync_dim": int(arch.sync_dim),
+        "text_seq_len": int(arch.text_seq_len),
+        "text_dim": int(arch.text_dim),
+    }
+
+
+class _IgnoredComponent:
+    """Placeholder for a required module the direct-feature path never touches.
+
+    ``MMAudioPipeline`` requires the CLIP text/vision and Synchformer encoders,
+    but its direct-feature conditioning path returns before dereferencing them.
+    Callers that pass ``video_path`` (or omit the cached features) must supply
+    the real modules instead.
+    """
+
+    def __repr__(self) -> str:
+        return "<ignored component>"
+
+
+def _build_pipeline(args: argparse.Namespace, world_size: int) -> MMAudioPipeline:
+    # Cached features make the three conditioning encoders unnecessary. The
+    # stages still exist, but consume direct tensors and never dereference
+    # these sentinels.
+    ignored_component = _IgnoredComponent()
+    return MMAudioPipeline.from_pretrained(
+        str(args.model_path),
+        inference_mode=True,
+        loaded_modules={
+            "text_encoder": ignored_component,
+            "tokenizer": ignored_component,
+            "image_encoder": ignored_component,
+            "image_encoder_2": ignored_component,
+        },
+        workload_type="v2a",
+        num_gpus=world_size,
+        tp_size=1,
+        sp_size=1,
+        hsdp_replicate_dim=1,
+        hsdp_shard_dim=1,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        text_encoder_cpu_offload=False,
+        image_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        enable_torch_compile=args.compile,
+    )
+
+
+def _forward_batch(
+    sample: dict[str, Any],
+    *,
+    duration_seconds: float,
+    num_inference_steps: int,
+    guidance_scale: float,
+    seed: int,
+    negative_prompt: str,
+) -> ForwardBatch:
+    return ForwardBatch(
+        data_type="video",
+        prompt="",
+        negative_prompt=negative_prompt,
+        audio_start_in_s=0.0,
+        audio_end_in_s=duration_seconds,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        seed=seed,
+        num_videos_per_prompt=1,
+        height=8,
+        width=8,
+        num_frames=1,
+        save_video=False,
+        return_frames=False,
+        extra={
+            "mmaudio_clip_features": sample["clip_features"].unsqueeze(0),
+            "mmaudio_sync_features": sample["sync_features"].unsqueeze(0),
+            "mmaudio_text_features": sample["text_features"].unsqueeze(0),
+        },
+    )
+
+
+def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temporary = path.with_suffix(path.suffix + ".tmp")
+    with temporary.open("w", encoding="utf-8") as handle:
+        for row in rows:
+            handle.write(json.dumps(row, sort_keys=True) + "\n")
+    temporary.replace(path)
+
+
+def _merge_rank_outputs(output_dir: Path, world_size: int) -> None:
+    rows: list[dict[str, Any]] = []
+    failures: list[dict[str, Any]] = []
+    for rank in range(world_size):
+        manifest = output_dir / f"manifest_rank_{rank:05d}.jsonl"
+        failure_path = output_dir / f"failures_rank_{rank:05d}.jsonl"
+        if manifest.is_file():
+            rows.extend(json.loads(line) for line in manifest.read_text(encoding="utf-8").splitlines() if line.strip())
+        if failure_path.is_file():
+            failures.extend(
+                json.loads(line) for line in failure_path.read_text(encoding="utf-8").splitlines() if line.strip())
+    rows.sort(key=lambda row: int(row["index"]))
+    failures.sort(key=lambda row: int(row["index"]))
+    ids = [str(row["id"]) for row in rows]
+    if len(ids) != len(set(ids)):
+        raise RuntimeError("Duplicate sample ids found while merging inference manifests")
+    _write_jsonl(output_dir / "eval_manifest.jsonl", rows)
+    _write_jsonl(output_dir / "failures.jsonl", failures)
+    summary = {
+        "num_succeeded": len(rows),
+        "num_failed": len(failures),
+        "world_size": world_size,
+    }
+    (output_dir / "summary.json").write_text(
+        json.dumps(summary, indent=2, sort_keys=True) + "\n",
+        encoding="utf-8",
+    )
+
+
+@torch.inference_mode()
+def main() -> None:
+    args = parse_args()
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    if not torch.cuda.is_available():
+        raise RuntimeError("MMAudio dataset inference requires CUDA")
+    torch.cuda.set_device(local_rank)
+
+    output_dir = args.output_dir.expanduser().resolve()
+    audio_dir = output_dir / "audio"
+    audio_dir.mkdir(parents=True, exist_ok=True)
+    pipeline = _build_pipeline(args, world_size)
+    transformer = pipeline.get_module("transformer")
+    dataset = build_mmaudio_feature_dataset(
+        args.feature_root,
+        feature_shapes=_feature_shapes(transformer),
+        include_metadata=True,
+    )
+    dataset_size = len(dataset)
+    if args.max_samples > 0:
+        dataset_size = min(dataset_size, args.max_samples)
+    sampler = _RankStrideSampler(dataset_size, rank, world_size)
+    loader = DataLoader(
+        _IndexedDataset(dataset),
+        batch_size=None,
+        sampler=sampler,
+        num_workers=max(0, args.num_workers),
+        pin_memory=True,
+        persistent_workers=args.num_workers > 0,
+    )
+
+    manifest_rows: list[dict[str, Any]] = []
+    failure_rows: list[dict[str, Any]] = []
+    progress = tqdm(
+        loader,
+        total=len(sampler),
+        desc=f"MMAudio inference rank {rank}",
+        position=local_rank,
+    )
+    for index, sample in progress:
+        index = int(index)
+        sample_id = _safe_name(str(sample.get("sample_id", index)))
+        source_path = str(sample.get("source_path", ""))
+        caption = str(sample.get("caption", ""))
+        output_path = audio_dir / f"{sample_id}.wav"
+        sample_seed = args.seed + index
+        try:
+            if not source_path or not Path(source_path).is_file():
+                raise FileNotFoundError(f"source video is missing for sample {sample_id}: {source_path}")
+            if args.overwrite or not output_path.is_file():
+                batch = _forward_batch(
+                    sample,
+                    duration_seconds=args.duration_seconds,
+                    num_inference_steps=args.num_inference_steps,
+                    guidance_scale=args.guidance_scale,
+                    seed=sample_seed,
+                    negative_prompt=args.negative_prompt,
+                )
+                result = pipeline.forward(batch, pipeline.fastvideo_args)
+                audio = result.extra.get("audio")
+                sample_rate = result.extra.get("audio_sample_rate")
+                if not isinstance(audio, np.ndarray) or not isinstance(sample_rate, int):
+                    raise RuntimeError("MMAudio pipeline did not return decoded audio")
+                temporary = output_path.with_suffix(".tmp.wav")
+                wavfile.write(temporary, sample_rate, audio.astype(np.float32))
+                temporary.replace(output_path)
+            manifest_rows.append({
+                "id": sample_id,
+                "index": index,
+                "video": str(Path(source_path).resolve()),
+                "audio": str(output_path.resolve()),
+                "reference_audio_source": str(Path(source_path).resolve()),
+                "text_prompt": caption,
+                "seed": sample_seed,
+            })
+        except Exception as error:  # noqa: BLE001 - isolate corrupt benchmark samples
+            failure_rows.append({
+                "id": sample_id,
+                "index": index,
+                "source": source_path,
+                "error_type": type(error).__name__,
+                "error": str(error),
+            })
+
+    _write_jsonl(output_dir / f"manifest_rank_{rank:05d}.jsonl", manifest_rows)
+    _write_jsonl(output_dir / f"failures_rank_{rank:05d}.jsonl", failure_rows)
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    if rank == 0:
+        _merge_rank_outputs(output_dir, world_size)
+        print(f"Merged eval manifest: {output_dir / 'eval_manifest.jsonl'}")
+        print(f"Inference summary: {output_dir / 'summary.json'}")
+    if torch.distributed.is_initialized():
+        torch.distributed.barrier()
+    pipeline.close()
+    cleanup_dist_env_and_memory()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/eval/mmaudio/evaluate_official_av_benchmark.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run the exact external av-benchmark API used by official MMAudio."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--audio-dir", type=Path, required=True)
+    parser.add_argument("--gt-cache", type=Path, required=True)
+    parser.add_argument("--prediction-cache", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--audio-length", type=float, default=8.0)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    try:
+        from av_bench.evaluate import evaluate
+        from av_bench.extract import extract
+    except ImportError as error:
+        raise ImportError("The exact official backend is optional. Install hkchengrex/"
+                          "av-benchmark in a separate environment before running this file.") from error
+
+    args.prediction_cache.mkdir(parents=True, exist_ok=True)
+    extract(
+        audio_path=args.audio_dir,
+        output_path=args.prediction_cache,
+        device=args.device,
+        batch_size=args.batch_size,
+        audio_length=args.audio_length,
+    )
+    metrics = evaluate(
+        gt_audio_cache=args.gt_cache,
+        pred_audio_cache=args.prediction_cache,
+    )
+    payload = {
+        "backend": "hkchengrex/av-benchmark",
+        "audio_dir": str(args.audio_dir.resolve()),
+        "gt_cache": str(args.gt_cache.resolve()),
+        "prediction_cache": str(args.prediction_cache.resolve()),
+        "metrics": metrics,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(
+        json.dumps(payload, indent=2, sort_keys=True) + "\n",
+        encoding="utf-8",
+    )
+    print(json.dumps(payload, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

examples/inference/eval/mmaudio/export_final_ema.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# One-time export of the synthesized PostHocEMA transformer plus the frozen
+# official conditioning/decoder components into a complete FastVideo model.
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
+cd "${ROOT_DIR}"
+
+EMA_CHECKPOINT="${EMA_CHECKPOINT:-outputs/mmaudio_small_44k_ddp_from_scratch/posthoc_ema/official_ddp/mmaudio_ema_final_sigma_0p05_step_000300000.pth}"
+OUTPUT_MODEL="${OUTPUT_MODEL:-converted_weights/mmaudio/small_44k_ema_300000}"
+ASSET_ROOT="${ASSET_ROOT:-official_weights/mmaudio}"
+
+if [[ -e "${OUTPUT_MODEL}/model_index.json" ]]; then
+  echo "Model already exported: ${OUTPUT_MODEL}"
+  exit 0
+fi
+
+source .venv/bin/activate
+python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \
+  --variant small_44k \
+  --transformer-checkpoint "${EMA_CHECKPOINT}" \
+  --audio-vae-checkpoint "${ASSET_ROOT}/raw/ext_weights/v1-44.pth" \
+  --synchformer-checkpoint "${ASSET_ROOT}/raw/ext_weights/synchformer_state_dict.pth" \
+  --dfn5b-dir "${ASSET_ROOT}/DFN5B-CLIP-ViT-H-14-384" \
+  --bigvgan-dir "${ASSET_ROOT}/bigvgan_v2_44khz_128band_512x" \
+  --output "${OUTPUT_MODEL}"
+
+echo "Exported FastVideo MMAudio model: ${OUTPUT_MODEL}"

examples/inference/eval/mmaudio/run_evaluate_vggsound.sh

@@ -0,0 +1,25 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
+cd "${ROOT_DIR}"
+source .venv/bin/activate
+
+OUTPUT_DIR="${OUTPUT_DIR:-outputs/mmaudio_small_44k_ema_300000_vggsound_test}"
+MANIFEST="${MANIFEST:-${OUTPUT_DIR}/eval_manifest.jsonl}"
+REFERENCE_AUDIO_CACHE="${REFERENCE_AUDIO_CACHE:-${OUTPUT_DIR}/reference_audio}"
+RESULTS="${RESULTS:-${OUTPUT_DIR}/fastvideo_eval_results.json}"
+NUM_GPUS="${NUM_GPUS:-4}"
+EXTRACT_WORKERS="${EXTRACT_WORKERS:-16}"
+METRICS="${METRICS:-audio.frechet_distance,audio.kl_divergence,audio.clap_score,audio.desync}"
+
+fastvideo eval run \
+  --manifest "${MANIFEST}" \
+  --metrics "${METRICS}" \
+  --num-gpus "${NUM_GPUS}" \
+  --extract-audio "${REFERENCE_AUDIO_CACHE}" \
+  --extract-workers "${EXTRACT_WORKERS}" \
+  --output-format full \
+  --output "${RESULTS}"
+
+echo "FastVideo evaluation results: ${RESULTS}"

examples/inference/eval/mmaudio/run_evaluate_vggsound_av_benchmark.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
+cd "${ROOT_DIR}"
+
+OUTPUT_DIR="${OUTPUT_DIR:-outputs/mmaudio_small_44k_ema_300000_vggsound_test}"
+AUDIO_DIR="${AUDIO_DIR:-${OUTPUT_DIR}/audio}"
+GT_CACHE="${GT_CACHE:-/mnt/lustre/vlm-kai/datasets/VGGSound/av_benchmark/vggsound-test-eval-cache}"
+PREDICTION_CACHE="${PREDICTION_CACHE:-${OUTPUT_DIR}/av_benchmark_cache}"
+RESULTS="${RESULTS:-${OUTPUT_DIR}/av_benchmark_results.json}"
+AV_BENCH_PYTHON="${AV_BENCH_PYTHON:-/mnt/lustre/vlm-kai/av-benchmark/.venv/bin/python}"
+BATCH_SIZE="${BATCH_SIZE:-32}"
+# The official extractor initializes CUDA models before constructing its
+# DataLoader. Zero avoids CUDA-after-fork hangs and this container's 64 MB shm.
+NUM_WORKERS="${NUM_WORKERS:-0}"
+DEVICE="${DEVICE:-cuda}"
+RECOMPUTE="${RECOMPUTE:-0}"
+SKIP_VIDEO_RELATED="${SKIP_VIDEO_RELATED:-0}"
+# The official VGGSound GT cache has no CLAP text features, so CLAP is not an
+# official VGGSound output metric. Skip its prediction passes by default.
+SKIP_CLAP="${SKIP_CLAP:-1}"
+# Official VGGSound precomputed cache sanitizes a subset of sample ids by
+# prefixing underscores. Align only unique matches inside prediction caches.
+ALIGN_PREDICTION_KEYS="${ALIGN_PREDICTION_KEYS:-1}"
+
+if [[ ! -x "${AV_BENCH_PYTHON}" ]]; then
+  echo "Missing isolated av-benchmark Python: ${AV_BENCH_PYTHON}" >&2
+  echo "Set AV_BENCH_PYTHON=/path/to/av-benchmark/.venv/bin/python" >&2
+  exit 2
+fi
+if [[ ! -d "${GT_CACHE}" ]]; then
+  echo "Missing official VGGSound GT cache: ${GT_CACHE}" >&2
+  echo "Set GT_CACHE=/path/to/vggsound-test-eval-cache" >&2
+  exit 2
+fi
+
+EXTRA_ARGS=()
+if [[ "${RECOMPUTE}" == "1" ]]; then
+  EXTRA_ARGS+=(--recompute)
+fi
+if [[ "${SKIP_VIDEO_RELATED}" == "1" ]]; then
+  EXTRA_ARGS+=(--skip-video-related)
+fi
+if [[ "${SKIP_CLAP}" == "1" ]]; then
+  EXTRA_ARGS+=(--skip-clap)
+fi
+if [[ "${ALIGN_PREDICTION_KEYS}" == "1" ]]; then
+  EXTRA_ARGS+=(--align-prediction-keys)
+fi
+
+.venv/bin/fastvideo eval v2a \
+  --backend av-benchmark \
+  --audio-dir "${AUDIO_DIR}" \
+  --gt-cache "${GT_CACHE}" \
+  --prediction-cache "${PREDICTION_CACHE}" \
+  --output "${RESULTS}" \
+  --python-executable "${AV_BENCH_PYTHON}" \
+  --device "${DEVICE}" \
+  --batch-size "${BATCH_SIZE}" \
+  --num-workers "${NUM_WORKERS}" \
+  --audio-length 8 \
+  "${EXTRA_ARGS[@]}"
+
+echo "Official av-benchmark results: ${RESULTS}"

examples/inference/eval/mmaudio/run_inference_vggsound.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)"
+cd "${ROOT_DIR}"
+source .venv/bin/activate
+
+MODEL_PATH="${MODEL_PATH:-converted_weights/mmaudio/small_44k_ema_300000}"
+FEATURE_ROOT="${FEATURE_ROOT:-/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/test}"
+OUTPUT_DIR="${OUTPUT_DIR:-outputs/mmaudio_small_44k_ema_300000_vggsound_test}"
+NUM_GPUS="${NUM_GPUS:-4}"
+NUM_WORKERS="${NUM_WORKERS:-2}"
+MAX_SAMPLES="${MAX_SAMPLES:-0}"
+COMPILE="${COMPILE:-1}"
+MASTER_PORT="${MASTER_PORT:-29513}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
+
+EXTRA_ARGS=()
+if [[ "${MAX_SAMPLES}" -gt 0 ]]; then
+  EXTRA_ARGS+=(--max-samples "${MAX_SAMPLES}")
+fi
+if [[ "${COMPILE}" == "1" ]]; then
+  EXTRA_ARGS+=(--compile)
+fi
+
+torchrun \
+  --standalone \
+  --nproc_per_node="${NUM_GPUS}" \
+  --master_port="${MASTER_PORT}" \
+  examples/inference/eval/mmaudio/eval_mmaudio_dataset.py \
+  --model-path "${MODEL_PATH}" \
+  --feature-root "${FEATURE_ROOT}" \
+  --output-dir "${OUTPUT_DIR}" \
+  --duration-seconds 8 \
+  --num-inference-steps 25 \
+  --guidance-scale 4.5 \
+  --seed 14159265 \
+  --num-workers "${NUM_WORKERS}" \
+  "${EXTRA_ARGS[@]}"
+
+echo "Generated audio: ${OUTPUT_DIR}/audio"
+echo "Evaluation manifest: ${OUTPUT_DIR}/eval_manifest.jsonl"

examples/train/configs/example.yaml

@@ -105,6 +105,7 @@ training:
   optimizer:
     learning_rate: 2.0e-6             # default: 0.0
     betas: [0.9, 0.999]              # default: [0.9, 0.999]
+    eps: 1.0e-8                       # default: 1.0e-8
     weight_decay: 0.01                # default: 0.0
     lr_scheduler: constant            # default: "constant"
     lr_warmup_steps: 0                # default: 0

examples/train/configs/fine_tuning/mmaudio/large_44k_from_scratch.yaml

@@ -0,0 +1,100 @@
+# Official MMAudio large_44k transformer training from random initialization.
+# Requires 44.1 kHz TensorDict features with audio latents [345, 40].
+
+models:
+  student:
+    _target_: fastvideo.train.models.mmaudio.MMAudioModel
+    variant: large_44k
+    from_scratch: true
+    # Official standalone empty_string.pth. No DiT weights are loaded.
+    empty_string_features_path: ""
+    # Optional .pt cache. When empty, stats are recomputed at startup.
+    latent_statistics_path: ""
+    latent_statistics_chunk_size: 32
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod
+
+training:
+  # Matching component-tree template used only by DCP-to-Diffusers export.
+  model_path: converted_weights/mmaudio/large_44k
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+    pin_cpu_memory: false
+
+  data:
+    # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh.
+    data_path: ""
+    preprocessed_data_type: mmaudio_features
+    dataloader_num_workers: 2
+    train_batch_size: 8
+    training_cfg_rate: 0.1
+    seed: 14159265
+
+  optimizer:
+    learning_rate: 1.0e-4
+    betas: [0.9, 0.95]
+    eps: 1.0e-6
+    weight_decay: 1.0e-6
+    lr_scheduler: multistep_with_warmup
+    lr_warmup_steps: 1000
+    lr_milestones: [240000, 270000]
+    lr_gamma: 0.1
+
+  loop:
+    max_train_steps: 300000
+    # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512.
+    gradient_accumulation_steps: 16
+
+  checkpoint:
+    output_dir: outputs/mmaudio_large_44k_from_scratch
+    training_state_checkpointing_steps: 10000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: ""
+
+  tracker:
+    project_name: fastvideo_mmaudio
+    run_name: mmaudio_large_44k_from_scratch
+
+  model:
+    logit_mean: 0.0
+    logit_std: 1.0
+    precondition_outputs: false
+
+  # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast.
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  posthoc_ema:
+    _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback
+    sigma_rels: [0.05, 0.1]
+    update_every: 1
+    checkpoint_every: 5000
+    start_iter: 0
+    default_output_sigma: 0.05
+    step_size_correction: true
+  validation:
+    _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback
+    data_path: ""
+    every_steps: 5000
+    max_batches: 0
+    batch_size: 8
+    num_data_workers: 2
+    run_at_start: false
+    use_ema: false
+    inference_every_steps: 20000
+    inference_model_path: converted_weights/mmaudio/large_44k_v2
+    # Global total across all ranks; fixed by validation seed and sampler.
+    inference_num_samples: 16
+    inference_num_steps: 25
+    inference_guidance_scale: 4.5
+    inference_save_video: true
+    inference_log_to_tracker: true

examples/train/configs/fine_tuning/mmaudio/medium_44k_from_scratch.yaml

@@ -0,0 +1,100 @@
+# Official MMAudio medium_44k transformer training from random initialization.
+# Requires 44.1 kHz TensorDict features with audio latents [345, 40].
+
+models:
+  student:
+    _target_: fastvideo.train.models.mmaudio.MMAudioModel
+    variant: medium_44k
+    from_scratch: true
+    # Official standalone empty_string.pth. No DiT weights are loaded.
+    empty_string_features_path: ""
+    # Optional .pt cache. When empty, stats are recomputed at startup.
+    latent_statistics_path: ""
+    latent_statistics_chunk_size: 32
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod
+
+training:
+  # Matching component-tree template used only by DCP-to-Diffusers export.
+  model_path: converted_weights/mmaudio/medium_44k
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+    pin_cpu_memory: false
+
+  data:
+    # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh.
+    data_path: ""
+    preprocessed_data_type: mmaudio_features
+    dataloader_num_workers: 2
+    train_batch_size: 8
+    training_cfg_rate: 0.1
+    seed: 14159265
+
+  optimizer:
+    learning_rate: 1.0e-4
+    betas: [0.9, 0.95]
+    eps: 1.0e-6
+    weight_decay: 1.0e-6
+    lr_scheduler: multistep_with_warmup
+    lr_warmup_steps: 1000
+    lr_milestones: [240000, 270000]
+    lr_gamma: 0.1
+
+  loop:
+    max_train_steps: 300000
+    # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512.
+    gradient_accumulation_steps: 16
+
+  checkpoint:
+    output_dir: outputs/mmaudio_medium_44k_from_scratch
+    training_state_checkpointing_steps: 10000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: ""
+
+  tracker:
+    project_name: fastvideo_mmaudio
+    run_name: mmaudio_medium_44k_from_scratch
+
+  model:
+    logit_mean: 0.0
+    logit_std: 1.0
+    precondition_outputs: false
+
+  # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast.
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  posthoc_ema:
+    _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback
+    sigma_rels: [0.05, 0.1]
+    update_every: 1
+    checkpoint_every: 5000
+    start_iter: 0
+    default_output_sigma: 0.05
+    step_size_correction: true
+  validation:
+    _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback
+    data_path: ""
+    every_steps: 5000
+    max_batches: 0
+    batch_size: 8
+    num_data_workers: 2
+    run_at_start: false
+    use_ema: false
+    inference_every_steps: 20000
+    inference_model_path: converted_weights/mmaudio/large_44k_v2
+    # Global total across all ranks; fixed by validation seed and sampler.
+    inference_num_samples: 16
+    inference_num_steps: 25
+    inference_guidance_scale: 4.5
+    inference_save_video: true
+    inference_log_to_tracker: true

examples/train/configs/fine_tuning/mmaudio/small_16k_from_scratch.yaml

@@ -0,0 +1,100 @@
+# Official MMAudio small_16k transformer training from random initialization.
+# Requires 16 kHz TensorDict features with audio latents [250, 20].
+
+models:
+  student:
+    _target_: fastvideo.train.models.mmaudio.MMAudioModel
+    variant: small_16k
+    from_scratch: true
+    # Official standalone empty_string.pth. No DiT weights are loaded.
+    empty_string_features_path: ""
+    # Optional .pt cache. When empty, stats are recomputed at startup.
+    latent_statistics_path: ""
+    latent_statistics_chunk_size: 32
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod
+
+training:
+  # Matching component-tree template used only by DCP-to-Diffusers export.
+  model_path: converted_weights/mmaudio/small_16k
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+    pin_cpu_memory: false
+
+  data:
+    # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh.
+    data_path: ""
+    preprocessed_data_type: mmaudio_features
+    dataloader_num_workers: 2
+    train_batch_size: 8
+    training_cfg_rate: 0.1
+    seed: 14159265
+
+  optimizer:
+    learning_rate: 1.0e-4
+    betas: [0.9, 0.95]
+    eps: 1.0e-6
+    weight_decay: 1.0e-6
+    lr_scheduler: multistep_with_warmup
+    lr_warmup_steps: 1000
+    lr_milestones: [240000, 270000]
+    lr_gamma: 0.1
+
+  loop:
+    max_train_steps: 300000
+    # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512.
+    gradient_accumulation_steps: 16
+
+  checkpoint:
+    output_dir: outputs/mmaudio_small_16k_from_scratch
+    training_state_checkpointing_steps: 10000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: ""
+
+  tracker:
+    project_name: fastvideo_mmaudio
+    run_name: mmaudio_small_16k_from_scratch
+
+  model:
+    logit_mean: 0.0
+    logit_std: 1.0
+    precondition_outputs: false
+
+  # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast.
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  posthoc_ema:
+    _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback
+    sigma_rels: [0.05, 0.1]
+    update_every: 1
+    checkpoint_every: 5000
+    start_iter: 0
+    default_output_sigma: 0.05
+    step_size_correction: true
+  validation:
+    _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback
+    data_path: ""
+    every_steps: 5000
+    max_batches: 0
+    batch_size: 8
+    num_data_workers: 2
+    run_at_start: false
+    use_ema: false
+    # Set this after preparing a complete 16 kHz inference component tree.
+    inference_every_steps: 0
+    inference_model_path: ""
+    inference_num_samples: 16
+    inference_num_steps: 25
+    inference_guidance_scale: 4.5
+    inference_save_video: true
+    inference_log_to_tracker: true

examples/train/configs/fine_tuning/mmaudio/small_44k.yaml

@@ -0,0 +1,69 @@
+# MMAudio v1 small_44k conditional flow-matching fine-tuning.
+#
+# This config intentionally leaves data_path empty. Point it at an official-
+# compatible TensorDict memmap feature cache before starting training.
+
+models:
+  student:
+    _target_: fastvideo.train.models.mmaudio.MMAudioModel
+    init_from: converted_weights/mmaudio/small_44k
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod
+
+training:
+  distributed:
+    num_gpus: 1
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 1
+    pin_cpu_memory: false
+
+  data:
+    # Fill this after the precomputed MMAudio feature cache is ready.
+    data_path: ""
+    preprocessed_data_type: mmaudio_features
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.1
+    seed: 14159265
+
+  # These match the official MMAudio optimizer defaults. FastVideo currently
+  # uses constant-with-warmup here; the original recipe adds LR drops at
+  # iterations 240k and 270k.
+  optimizer:
+    learning_rate: 1.0e-4
+    betas: [0.9, 0.95]
+    eps: 1.0e-6
+    weight_decay: 1.0e-6
+    lr_scheduler: constant_with_warmup
+    lr_warmup_steps: 1000
+
+  loop:
+    max_train_steps: 300000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/mmaudio_small_44k_finetune
+    training_state_checkpointing_steps: 10000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: fastvideo_mmaudio
+    run_name: mmaudio_small_44k_finetune
+
+  model:
+    logit_mean: 0.0
+    logit_std: 1.0
+    precondition_outputs: false
+
+  # Official MMAudio trains FP32 weights under BF16 autocast.
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0

examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml

@@ -0,0 +1,108 @@
+# Official MMAudio small_44k training with one full DDP replica per GPU.
+# This is independent from the existing FSDP/HSDP scratch config.
+
+models:
+  student:
+    _target_: fastvideo.train.models.mmaudio.MMAudioModel
+    variant: small_44k
+    from_scratch: true
+    empty_string_features_path: ""
+    latent_statistics_path: ""
+    latent_statistics_chunk_size: 32
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod
+
+training:
+  model_path: converted_weights/mmaudio/small_44k
+  distributed:
+    strategy: ddp
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    # DDP replicates the full model and does not use HSDP meshes.
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 1
+    pin_cpu_memory: false
+
+  data:
+    data_path: ""
+    preprocessed_data_type: mmaudio_features
+    dataloader_num_workers: 0
+    train_batch_size: 128
+    training_cfg_rate: 0.1
+    seed: 14159265
+
+  optimizer:
+    learning_rate: 1.0e-4
+    betas: [0.9, 0.95]
+    eps: 1.0e-6
+    weight_decay: 1.0e-6
+    lr_scheduler: multistep_with_warmup
+    lr_warmup_steps: 1000
+    lr_milestones: [240000, 270000]
+    lr_gamma: 0.1
+    # Matches the upstream Runner's AdamW(..., fused=True).
+    fused: true
+
+  loop:
+    max_train_steps: 300000
+    # 4 GPUs x 128 samples/GPU = official effective global batch 512.
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/mmaudio_small_44k_ddp_from_scratch
+    # Keep this divisible by the 5,000-step PostHocEMA interval so resume
+    # always has an exactly aligned pair of official EMA snapshots.
+    training_state_checkpointing_steps: 10000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: ""
+
+  tracker:
+    project_name: fastvideo_mmaudio
+    run_name: mmaudio_small_44k_ddp_from_scratch
+
+  model:
+    logit_mean: 0.0
+    logit_std: 1.0
+    precondition_outputs: false
+    # Matches official MMAudio's torch.compile(self.train_fn). Set false to
+    # use eager execution when debugging or on an unsupported compiler stack.
+    compile_train_fn: true
+    torch_compile_kwargs:
+      fullgraph: true
+
+  # Official MMAudio uses FP32 master weights with BF16 autocast forward.
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  posthoc_ema:
+    _target_: fastvideo.train.callbacks.ddp_posthoc_ema.DDPPostHocEMACallback
+    sigma_rels: [0.05, 0.1]
+    update_every: 1
+    checkpoint_every: 5000
+    start_iter: 0
+    default_output_sigma: 0.05
+    step_size_correction: true
+  validation:
+    _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback
+    data_path: ""
+    every_steps: 5000
+    max_batches: 0
+    batch_size: 8
+    num_data_workers: 0
+    run_at_start: false
+    # Official periodic validation samples the online model; EMA is
+    # synthesized after training.
+    use_ema: false
+    inference_every_steps: 20000
+    inference_model_path: converted_weights/mmaudio/large_44k_v2
+    inference_num_samples: 16
+    inference_num_steps: 25
+    inference_guidance_scale: 4.5
+    inference_save_video: true
+    inference_log_to_tracker: true

examples/train/configs/fine_tuning/mmaudio/small_44k_from_scratch.yaml

@@ -0,0 +1,100 @@
+# Official MMAudio small_44k transformer training from random initialization.
+# Requires 44.1 kHz TensorDict features with audio latents [345, 40].
+
+models:
+  student:
+    _target_: fastvideo.train.models.mmaudio.MMAudioModel
+    variant: small_44k
+    from_scratch: true
+    # Official standalone empty_string.pth. No DiT weights are loaded.
+    empty_string_features_path: ""
+    # Optional .pt cache. When empty, stats are recomputed at startup.
+    latent_statistics_path: ""
+    latent_statistics_chunk_size: 32
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.flow_matching.FlowMatchingFineTuneMethod
+
+training:
+  # Matching component-tree template used only by DCP-to-Diffusers export.
+  model_path: converted_weights/mmaudio/small_44k
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+    pin_cpu_memory: false
+
+  data:
+    # Set with --training.data.data_path or run_train_vggsound_from_scratch.sh.
+    data_path: ""
+    preprocessed_data_type: mmaudio_features
+    dataloader_num_workers: 2
+    train_batch_size: 8
+    training_cfg_rate: 0.1
+    seed: 14159265
+
+  optimizer:
+    learning_rate: 1.0e-4
+    betas: [0.9, 0.95]
+    eps: 1.0e-6
+    weight_decay: 1.0e-6
+    lr_scheduler: multistep_with_warmup
+    lr_warmup_steps: 1000
+    lr_milestones: [240000, 270000]
+    lr_gamma: 0.1
+
+  loop:
+    max_train_steps: 300000
+    # 4 GPUs x 8 samples/GPU x 16 rounds = effective global batch 512.
+    gradient_accumulation_steps: 16
+
+  checkpoint:
+    output_dir: outputs/mmaudio_small_44k_from_scratch
+    training_state_checkpointing_steps: 10000
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: ""
+
+  tracker:
+    project_name: fastvideo_mmaudio
+    run_name: mmaudio_small_44k_from_scratch
+
+  model:
+    logit_mean: 0.0
+    logit_std: 1.0
+    precondition_outputs: false
+
+  # FP32 master weights; MMAudioModel runs transformer forward under BF16 autocast.
+  dit_precision: fp32
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  posthoc_ema:
+    _target_: fastvideo.train.callbacks.posthoc_ema.PostHocEMACallback
+    sigma_rels: [0.05, 0.1]
+    update_every: 1
+    checkpoint_every: 5000
+    start_iter: 0
+    default_output_sigma: 0.05
+    step_size_correction: true
+  validation:
+    _target_: fastvideo.train.callbacks.mmaudio_validation.MMAudioValidationCallback
+    data_path: ""
+    every_steps: 5000
+    max_batches: 0
+    batch_size: 8
+    num_data_workers: 2
+    run_at_start: false
+    use_ema: false
+    inference_every_steps: 20000
+    inference_model_path: converted_weights/mmaudio/large_44k_v2
+    # Global total across all ranks; fixed by validation seed and sampler.
+    inference_num_samples: 16
+    inference_num_steps: 25
+    inference_guidance_scale: 4.5
+    inference_save_video: true
+    inference_log_to_tracker: true

examples/training/finetune/mmaudio/preprocess_vggsound.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+GPU_NUM="${GPU_NUM:-4}"
+MODEL_PATH="${MODEL_PATH:-converted_weights/mmaudio/preprocess_44k}"
+DATASET_PATH="${DATASET_PATH:?Set DATASET_PATH to the extracted VGGSound directory}"
+SPLIT="${SPLIT:-train}"
+CAPTION_PATH="${CAPTION_PATH:-${DATASET_PATH}/sets/filtered_caption/vgg-${SPLIT}-filtered-caption.tsv}"
+OUTPUT_DIR="${OUTPUT_DIR:-${DATASET_PATH}/mmaudio_features/${SPLIT}}"
+MASTER_PORT="${MASTER_PORT:-29513}"
+BATCH_SIZE="${BATCH_SIZE:-1}"
+DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-2}"
+SAMPLES_PER_FILE="${SAMPLES_PER_FILE:-256}"
+VIDEO_LOADER_TYPE="${VIDEO_LOADER_TYPE:-torchcodec}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
+
+if [[ ! -f "${CAPTION_PATH}" ]]; then
+  echo "Caption manifest does not exist: ${CAPTION_PATH}" >&2
+  exit 1
+fi
+
+MODEL_LOAD_ARGS=()
+if [[ "${VIDEO_LOADER_TYPE}" == "torio" ]]; then
+  # PyTorch 2.7 torio environments can have a broken NCCL tensor broadcast.
+  # Loading on CPU makes every rank read its own local weights, matching the
+  # official MMAudio extractor; the feature stage moves modules to its GPU.
+  MODEL_LOAD_ARGS+=(--image-encoder-cpu-offload true)
+  MODEL_LOAD_ARGS+=(--text-encoder-cpu-offload true)
+  MODEL_LOAD_ARGS+=(--vae-cpu-offload true)
+fi
+
+torchrun --standalone \
+  --nproc_per_node="${GPU_NUM}" \
+  --master_port="${MASTER_PORT}" \
+  -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+  --model-path "${MODEL_PATH}" \
+  --mode preprocess \
+  --workload-type v2a \
+  "${MODEL_LOAD_ARGS[@]}" \
+  --preprocess.dataset-type vggsound \
+  --preprocess.dataset-path "${DATASET_PATH}" \
+  --preprocess.dataset-metadata-path "${CAPTION_PATH}" \
+  --preprocess.dataset-split "${SPLIT}" \
+  --preprocess.dataset-output-dir "${OUTPUT_DIR}" \
+  --preprocess.video-loader-type "${VIDEO_LOADER_TYPE}" \
+  --preprocess.preprocess-video-batch-size "${BATCH_SIZE}" \
+  --preprocess.dataloader-num-workers "${DATALOADER_NUM_WORKERS}" \
+  --preprocess.samples-per-file "${SAMPLES_PER_FILE}"

examples/training/finetune/mmaudio/run_preprocess_vggsound.sh

@@ -0,0 +1,119 @@
+#!/usr/bin/env bash
+
+# One-command VGGSound feature extraction for native FastVideo MMAudio training.
+# Every setting can be overridden with an environment variable; the defaults
+# match the four-GPU GB200 machine used for the validated smoke tests.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../../.." && pwd)"
+
+if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then
+  cat <<'EOF'
+Usage:
+  ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh
+
+Optional environment variables:
+  SPLITS="train val test"  Dataset splits to process sequentially
+  GPU_NUM=4                Number of GPUs/ranks
+  CUDA_VISIBLE_DEVICES=0,1,2,3
+  BATCH_SIZE=16            Per-GPU batch size
+  DATALOADER_NUM_WORKERS=16 Per-rank workers (torio: background decode threads)
+  DATASET_PATH=...         Extracted VGGSound root
+  MODEL_PATH=...           Converted MMAudio preprocessing model
+  OUTPUT_ROOT=...          Root for split feature caches
+  VIDEO_LOADER_TYPE=torchcodec Media reader (torchcodec or torio)
+  SAMPLES_PER_FILE=256     Samples per TensorDict shard
+  MASTER_PORT=29513        First torchrun rendezvous port
+
+Examples:
+  SPLITS=train ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh
+  BATCH_SIZE=8 GPU_NUM=2 CUDA_VISIBLE_DEVICES=0,1 \
+    ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh
+EOF
+  exit 0
+fi
+
+if [[ $# -ne 0 ]]; then
+  echo "Unknown argument: $1 (use --help)" >&2
+  exit 2
+fi
+
+VENV_PATH="${VENV_PATH:-${REPO_ROOT}/.venv}"
+if [[ ! -f "${VENV_PATH}/bin/activate" ]]; then
+  echo "FastVideo virtual environment does not exist: ${VENV_PATH}" >&2
+  exit 1
+fi
+# shellcheck disable=SC1091
+source "${VENV_PATH}/bin/activate"
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"
+export PYTHONUNBUFFERED=1
+
+GPU_NUM="${GPU_NUM:-4}"
+BATCH_SIZE="${BATCH_SIZE:-16}"
+DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-16}"
+SAMPLES_PER_FILE="${SAMPLES_PER_FILE:-256}"
+MASTER_PORT="${MASTER_PORT:-29513}"
+SPLITS="${SPLITS:-train val test}"
+DATASET_PATH="${DATASET_PATH:?Set DATASET_PATH to the extracted VGGSound directory}"
+MODEL_PATH="${MODEL_PATH:-${REPO_ROOT}/converted_weights/mmaudio/preprocess_44k}"
+OUTPUT_ROOT="${OUTPUT_ROOT:-${DATASET_PATH}/mmaudio_features}"
+VIDEO_LOADER_TYPE="${VIDEO_LOADER_TYPE:-torchcodec}"
+
+if [[ ! -d "${DATASET_PATH}/videos" ]]; then
+  echo "VGGSound video directory does not exist: ${DATASET_PATH}/videos" >&2
+  exit 1
+fi
+if [[ ! -f "${MODEL_PATH}/model_index.json" ]]; then
+  echo "Converted preprocessing model does not exist: ${MODEL_PATH}" >&2
+  exit 1
+fi
+
+echo "FastVideo MMAudio VGGSound preprocessing"
+echo "  repository: ${REPO_ROOT}"
+echo "  dataset:    ${DATASET_PATH}"
+echo "  model:      ${MODEL_PATH}"
+echo "  output:     ${OUTPUT_ROOT}"
+echo "  splits:     ${SPLITS}"
+echo "  GPUs:       ${CUDA_VISIBLE_DEVICES} (${GPU_NUM} ranks)"
+echo "  batch:      ${BATCH_SIZE} per GPU"
+echo "  workers:    ${DATALOADER_NUM_WORKERS} per GPU"
+echo "  reader:     ${VIDEO_LOADER_TYPE}"
+
+cd "${REPO_ROOT}"
+
+port="${MASTER_PORT}"
+for split in ${SPLITS}; do
+  caption_path="${DATASET_PATH}/sets/filtered_caption/vgg-${split}-filtered-caption.tsv"
+  output_dir="${OUTPUT_ROOT}/${split}"
+  if [[ ! -f "${caption_path}" ]]; then
+    echo "Caption manifest does not exist: ${caption_path}" >&2
+    exit 1
+  fi
+
+  echo
+  echo "===== Starting ${split} split ====="
+  echo "Progress is reported below in batches; existing cached IDs are skipped."
+
+  GPU_NUM="${GPU_NUM}" \
+  BATCH_SIZE="${BATCH_SIZE}" \
+  DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS}" \
+  SAMPLES_PER_FILE="${SAMPLES_PER_FILE}" \
+  MASTER_PORT="${port}" \
+  MODEL_PATH="${MODEL_PATH}" \
+  DATASET_PATH="${DATASET_PATH}" \
+  CAPTION_PATH="${caption_path}" \
+  OUTPUT_DIR="${output_dir}" \
+  SPLIT="${split}" \
+  VIDEO_LOADER_TYPE="${VIDEO_LOADER_TYPE}" \
+    bash "${SCRIPT_DIR}/preprocess_vggsound.sh"
+
+  echo "===== Finished ${split}: ${output_dir} ====="
+  port="$((port + 1))"
+done
+
+echo
+echo "All requested VGGSound splits finished successfully."

examples/training/finetune/mmaudio/run_preprocess_vggsound_torio.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+
+# Run the FastVideo MMAudio preprocessing pipeline with the optional torio
+# reference media reader in an isolated environment.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../../../.." && pwd)"
+
+if [[ "${1:-}" == "--help" || "${1:-}" == "-h" ]]; then
+  exec "${SCRIPT_DIR}/run_preprocess_vggsound.sh" "$@"
+fi
+
+TORIO_VENV_PATH="${TORIO_VENV_PATH:-${REPO_ROOT}/../envs/fastvideo-mmaudio-torio}"
+DATASET_PATH="${DATASET_PATH:?Set DATASET_PATH to the extracted VGGSound directory}"
+OUTPUT_ROOT="${OUTPUT_ROOT:-${DATASET_PATH}/mmaudio_features_torio}"
+DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}"
+
+if [[ ! -x "${TORIO_VENV_PATH}/bin/torchrun" ]]; then
+  echo "FastVideo torio environment does not exist: ${TORIO_VENV_PATH}" >&2
+  exit 1
+fi
+
+export PYTHONNOUSERSITE=1
+
+VENV_PATH="${TORIO_VENV_PATH}" \
+DATASET_PATH="${DATASET_PATH}" \
+OUTPUT_ROOT="${OUTPUT_ROOT}" \
+DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS}" \
+VIDEO_LOADER_TYPE=torio \
+  exec "${SCRIPT_DIR}/run_preprocess_vggsound.sh" "$@"

examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh

@@ -0,0 +1,109 @@
+#!/usr/bin/env bash
+# Train MMAudio small_44k with the official full-replica DDP strategy.
+# This launcher is independent from run_train_vggsound_from_scratch.sh.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
+
+# ---------------------------------------------------------------------------
+# Edit the DDP run here. Do not launch while another job owns the same GPUs.
+# ---------------------------------------------------------------------------
+FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train"
+VALIDATION_FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/val"
+EMPTY_STRING_FEATURES="${REPO_ROOT}/official_weights/mmaudio/raw/ext_weights/empty_string.pth"
+EMPTY_STRING_URL="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/empty_string.pth"
+
+NUM_GPUS=4
+MASTER_PORT=29502
+PER_GPU_BATCH_SIZE=128
+GRAD_ACCUM_STEPS=1
+# /dev/shm is currently 64 MB, so subprocess workers can exhaust shared memory.
+DATALOADER_NUM_WORKERS=0
+VALIDATION_NUM_WORKERS=0
+
+LEARNING_RATE="1.0e-4"
+MAX_TRAIN_STEPS=300000
+COMPILE_TRAIN_FN=true
+CHECKPOINT_EVERY_STEPS=10000
+VALIDATION_EVERY_STEPS=5000
+VALIDATION_MAX_BATCHES=0
+INFERENCE_EVERY_STEPS=20000
+INFERENCE_NUM_SAMPLES=16
+INFERENCE_NUM_STEPS=25
+INFERENCE_SAVE_VIDEO=true
+INFERENCE_LOG_TO_TRACKER=true
+INFERENCE_MODEL_PATH="${REPO_ROOT}/converted_weights/mmaudio/large_44k_v2"
+
+OUTPUT_DIR="${REPO_ROOT}/outputs/mmaudio_small_44k_ddp_from_scratch"
+RESUME_FROM_CHECKPOINT="${OUTPUT_DIR}/checkpoint-20000"
+TRACKER="wandb"
+TRACKER_PROJECT="fastvideo_mmaudio"
+
+# These W&B settings are used only when resuming an existing checkpoint.
+# For a new run, set RESUME_FROM_CHECKPOINT=""; the variables are then
+# explicitly unset so training cannot accidentally append to the old page.
+WANDB_RESUME_RUN_ID="vabj9rjb"
+WANDB_RESUME_ENTITY="alan-wang-university-of-toronto"
+
+CONFIG="${REPO_ROOT}/examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml"
+LATENT_STATISTICS="${FEATURE_DIR}/latent_statistics_44k.pt"
+
+if [[ -x "${REPO_ROOT}/.venv/bin/python" ]]; then
+    export PATH="${REPO_ROOT}/.venv/bin:${PATH}"
+fi
+
+for required_dir in "${FEATURE_DIR}" "${VALIDATION_FEATURE_DIR}"; do
+    if [[ ! -d "${required_dir}" ]]; then
+        echo "MMAudio feature directory does not exist: ${required_dir}" >&2
+        exit 2
+    fi
+done
+
+if [[ ! -f "${EMPTY_STRING_FEATURES}" ]]; then
+    mkdir -p "$(dirname "${EMPTY_STRING_FEATURES}")"
+    echo "Downloading official empty-string encoding to ${EMPTY_STRING_FEATURES}"
+    curl --fail --location --continue-at - \
+        --output "${EMPTY_STRING_FEATURES}" \
+        "${EMPTY_STRING_URL}"
+fi
+
+if [[ -n "${RESUME_FROM_CHECKPOINT}" ]]; then
+    export WANDB_RUN_ID="${WANDB_RESUME_RUN_ID}"
+    export WANDB_RESUME="must"
+    export WANDB_ENTITY="${WANDB_RESUME_ENTITY}"
+else
+    unset WANDB_RUN_ID WANDB_RESUME WANDB_ENTITY
+fi
+
+export NUM_GPUS MASTER_PORT
+cd "${REPO_ROOT}"
+exec bash examples/train/run.sh "${CONFIG}" \
+    --models.student.empty_string_features_path "${EMPTY_STRING_FEATURES}" \
+    --models.student.latent_statistics_path "${LATENT_STATISTICS}" \
+    --training.data.data_path "${FEATURE_DIR}" \
+    --training.data.train_batch_size "${PER_GPU_BATCH_SIZE}" \
+    --training.data.dataloader_num_workers "${DATALOADER_NUM_WORKERS}" \
+    --training.optimizer.learning_rate "${LEARNING_RATE}" \
+    --training.loop.max_train_steps "${MAX_TRAIN_STEPS}" \
+    --training.loop.gradient_accumulation_steps "${GRAD_ACCUM_STEPS}" \
+    --training.model.compile_train_fn "${COMPILE_TRAIN_FN}" \
+    --training.distributed.num_gpus "${NUM_GPUS}" \
+    --training.checkpoint.output_dir "${OUTPUT_DIR}" \
+    --training.checkpoint.training_state_checkpointing_steps "${CHECKPOINT_EVERY_STEPS}" \
+    --training.checkpoint.resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}" \
+    --training.tracker.trackers "[${TRACKER}]" \
+    --training.tracker.project_name "${TRACKER_PROJECT}" \
+    --training.tracker.run_name "mmaudio_small_44k_ddp_from_scratch" \
+    --callbacks.validation.data_path "${VALIDATION_FEATURE_DIR}" \
+    --callbacks.validation.every_steps "${VALIDATION_EVERY_STEPS}" \
+    --callbacks.validation.max_batches "${VALIDATION_MAX_BATCHES}" \
+    --callbacks.validation.num_data_workers "${VALIDATION_NUM_WORKERS}" \
+    --callbacks.validation.inference_every_steps "${INFERENCE_EVERY_STEPS}" \
+    --callbacks.validation.inference_num_samples "${INFERENCE_NUM_SAMPLES}" \
+    --callbacks.validation.inference_num_steps "${INFERENCE_NUM_STEPS}" \
+    --callbacks.validation.inference_save_video "${INFERENCE_SAVE_VIDEO}" \
+    --callbacks.validation.inference_log_to_tracker "${INFERENCE_LOG_TO_TRACKER}" \
+    --callbacks.validation.inference_model_path "${INFERENCE_MODEL_PATH}" \
+    "$@"

examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh

@@ -0,0 +1,126 @@
+#!/usr/bin/env bash
+# Train an official MMAudio v1 architecture from random DiT initialization.
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)"
+
+# ---------------------------------------------------------------------------
+# Edit the training run here. No environment variables are required.
+# ---------------------------------------------------------------------------
+VARIANT="small_44k"
+
+FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train"
+VALIDATION_FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/val"
+EMPTY_STRING_FEATURES="${REPO_ROOT}/official_weights/mmaudio/raw/ext_weights/empty_string.pth"
+EMPTY_STRING_URL="https://github.com/hkchengrex/MMAudio/releases/download/v0.1/empty_string.pth"
+
+NUM_GPUS=4
+# Validated for the default small_44k model on 4 x GB200 (about 30 GB/GPU).
+# Re-profile before switching VARIANT to medium_44k or large_44k.
+PER_GPU_BATCH_SIZE=128
+GRAD_ACCUM_STEPS=1
+# This container has a 64 MB /dev/shm. Keep these at zero unless the job is
+# launched with a larger shared-memory mount.
+DATALOADER_NUM_WORKERS=0
+VALIDATION_NUM_WORKERS=0
+
+LEARNING_RATE="1.0e-4"
+MAX_TRAIN_STEPS=300000
+CHECKPOINT_EVERY_STEPS=10000
+VALIDATION_EVERY_STEPS=5000
+# 0 evaluates the complete validation split.
+VALIDATION_MAX_BATCHES=0
+INFERENCE_EVERY_STEPS=20000
+INFERENCE_NUM_SAMPLES=16
+INFERENCE_NUM_STEPS=25
+INFERENCE_SAVE_VIDEO=true
+INFERENCE_LOG_TO_TRACKER=true
+INFERENCE_MODEL_PATH="${REPO_ROOT}/converted_weights/mmaudio/large_44k_v2"
+
+OUTPUT_DIR="${REPO_ROOT}/outputs/mmaudio_${VARIANT}_from_scratch"
+RESUME_FROM_CHECKPOINT=""
+# The API key is stored by `wandb login` in the current user's credentials;
+# never put it in this tracked launcher.
+TRACKER="wandb"
+TRACKER_PROJECT="fastvideo_mmaudio"
+
+# Prefer the repository environment so the script also works from a fresh
+# non-activated shell. An already activated external environment remains the
+# fallback when .venv is absent.
+if [[ -x "${REPO_ROOT}/.venv/bin/python" ]]; then
+    export PATH="${REPO_ROOT}/.venv/bin:${PATH}"
+fi
+
+for required_dir in "${FEATURE_DIR}" "${VALIDATION_FEATURE_DIR}"; do
+    if [[ ! -d "${required_dir}" ]]; then
+        echo "MMAudio feature directory does not exist: ${required_dir}" >&2
+        exit 2
+    fi
+done
+
+# Official TRAINING.md publishes this small standalone CLIP encoding. It is
+# fixed conditioning data, not a pretrained MMAudio transformer checkpoint.
+if [[ ! -f "${EMPTY_STRING_FEATURES}" ]]; then
+    mkdir -p "$(dirname "${EMPTY_STRING_FEATURES}")"
+    echo "Downloading official empty-string encoding to ${EMPTY_STRING_FEATURES}"
+    curl --fail --location --continue-at - \
+        --output "${EMPTY_STRING_FEATURES}" \
+        "${EMPTY_STRING_URL}"
+fi
+
+case "${VARIANT}" in
+    small_16k)
+        CONFIG_NAME="small_16k_from_scratch.yaml"
+        STATS_NAME="latent_statistics_16k.pt"
+        INFERENCE_EVERY_STEPS=0
+        INFERENCE_MODEL_PATH=""
+        ;;
+    small_44k|medium_44k|large_44k)
+        CONFIG_NAME="${VARIANT}_from_scratch.yaml"
+        STATS_NAME="latent_statistics_44k.pt"
+        ;;
+    large_44k_v2)
+        echo "Official MMAudio training does not support large_44k_v2." >&2
+        exit 2
+        ;;
+    *)
+        echo "Unknown VARIANT=${VARIANT}" >&2
+        exit 2
+        ;;
+esac
+
+CONFIG="${REPO_ROOT}/examples/train/configs/fine_tuning/mmaudio/${CONFIG_NAME}"
+LATENT_STATISTICS="${FEATURE_DIR}/${STATS_NAME}"
+
+export NUM_GPUS
+cd "${REPO_ROOT}"
+exec bash examples/train/run.sh "${CONFIG}" \
+    --models.student.empty_string_features_path "${EMPTY_STRING_FEATURES}" \
+    --models.student.latent_statistics_path "${LATENT_STATISTICS}" \
+    --training.data.data_path "${FEATURE_DIR}" \
+    --training.data.train_batch_size "${PER_GPU_BATCH_SIZE}" \
+    --training.data.dataloader_num_workers "${DATALOADER_NUM_WORKERS}" \
+    --training.optimizer.learning_rate "${LEARNING_RATE}" \
+    --training.loop.max_train_steps "${MAX_TRAIN_STEPS}" \
+    --training.loop.gradient_accumulation_steps "${GRAD_ACCUM_STEPS}" \
+    --training.distributed.num_gpus "${NUM_GPUS}" \
+    --training.distributed.hsdp_shard_dim "${NUM_GPUS}" \
+    --training.checkpoint.output_dir "${OUTPUT_DIR}" \
+    --training.checkpoint.training_state_checkpointing_steps "${CHECKPOINT_EVERY_STEPS}" \
+    --training.checkpoint.resume_from_checkpoint "${RESUME_FROM_CHECKPOINT}" \
+    --training.tracker.trackers "[${TRACKER}]" \
+    --training.tracker.project_name "${TRACKER_PROJECT}" \
+    --training.tracker.run_name "mmaudio_${VARIANT}_from_scratch" \
+    --callbacks.validation.data_path "${VALIDATION_FEATURE_DIR}" \
+    --callbacks.validation.every_steps "${VALIDATION_EVERY_STEPS}" \
+    --callbacks.validation.max_batches "${VALIDATION_MAX_BATCHES}" \
+    --callbacks.validation.num_data_workers "${VALIDATION_NUM_WORKERS}" \
+    --callbacks.validation.inference_every_steps "${INFERENCE_EVERY_STEPS}" \
+    --callbacks.validation.inference_num_samples "${INFERENCE_NUM_SAMPLES}" \
+    --callbacks.validation.inference_num_steps "${INFERENCE_NUM_STEPS}" \
+    --callbacks.validation.inference_save_video "${INFERENCE_SAVE_VIDEO}" \
+    --callbacks.validation.inference_log_to_tracker "${INFERENCE_LOG_TO_TRACKER}" \
+    --callbacks.validation.inference_model_path "${INFERENCE_MODEL_PATH}" \
+    "$@"

fastvideo/configs/configs.py

@@ -15,6 +15,7 @@ class DatasetType(str, Enum):
     """
     HF = "hf"
     MERGED = "merged"
+    VGGSOUND = "vggsound"
 
     @classmethod
     def from_string(cls, value: str) -> "DatasetType":
@@ -37,6 +38,7 @@ class VideoLoaderType(str, Enum):
     """
     TORCHCODEC = "torchcodec"
     TORCHVISION = "torchvision"
+    TORIO = "torio"
 
     @classmethod
     def from_string(cls, value: str) -> "VideoLoaderType":
@@ -60,6 +62,8 @@ class PreprocessConfig:
     # Model and dataset configuration
     model_path: str = ""
     dataset_path: str = ""
+    dataset_metadata_path: str = ""
+    dataset_split: str = "train"
     dataset_type: DatasetType = DatasetType.HF
     dataset_output_dir: str = "./output"
 
@@ -104,6 +108,14 @@ class PreprocessConfig:
                                      type=str,
                                      default=PreprocessConfig.dataset_path,
                                      help="Path to the dataset directory for preprocessing")
+        preprocess_args.add_argument(f"--{prefix_with_dot}dataset-metadata-path",
+                                     type=str,
+                                     default=PreprocessConfig.dataset_metadata_path,
+                                     help="Optional path to the dataset metadata or caption manifest")
+        preprocess_args.add_argument(f"--{prefix_with_dot}dataset-split",
+                                     type=str,
+                                     default=PreprocessConfig.dataset_split,
+                                     help="Dataset split to preprocess (for example train, val, or test)")
         preprocess_args.add_argument(f"--{prefix_with_dot}dataset-type",
                                      type=str,
                                      choices=DatasetType.choices(),
@@ -207,6 +219,8 @@ class PreprocessConfig:
     def check_preprocess_config(self) -> None:
         if self.dataset_path == "":
             raise ValueError("dataset_path must be set for preprocess mode")
+        if not self.dataset_split.strip():
+            raise ValueError("dataset_split must not be empty")
         if self.samples_per_file <= 0:
             raise ValueError("samples_per_file must be greater than 0")
         if self.flush_frequency <= 0:

fastvideo/configs/models/audio/mmaudio_vae.py

@@ -8,11 +8,15 @@ from fastvideo.configs.models.base import ArchConfig, ModelConfig
 
 @dataclass
 class MMAudioVAEArchConfig(ArchConfig):
+    """MMAudio mel-spectrogram VAE architecture.
+
+    ``mode`` is the single source of truth: ``MMAudioVAE`` derives the mel
+    dimension, the latent dimension, and the hidden width from it, so they are
+    not duplicated here.
+    """
+
     architectures: list[str] = field(default_factory=lambda: ["MMAudioVAE"])
     mode: str = "44k"
-    data_dim: int = 128
-    embed_dim: int = 40
-    hidden_dim: int = 512
     need_encoder: bool = False
 
 

fastvideo/configs/models/dits/__init__.py

@@ -12,7 +12,14 @@ from fastvideo.configs.models.dits.longcat import LongCatVideoConfig
 from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
 from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig
 from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config
-from fastvideo.configs.models.dits.mmaudio import MMAudioArchConfig, MMAudioTransformerConfig
+from fastvideo.configs.models.dits.mmaudio import (
+    MMAUDIO_44K_TRAINING_VARIANTS,
+    MMAUDIO_TRAINING_VARIANTS,
+    MMAUDIO_VARIANT_ARCHITECTURES,
+    MMAudioArchConfig,
+    MMAudioTransformerConfig,
+    get_mmaudio_transformer_config,
+)
 from fastvideo.configs.models.dits.stable_audio import StableAudioConfig
 from fastvideo.models.wan.config import WanVideoConfig
 from fastvideo.configs.models.dits.zimage import ZImageDiTConfig
@@ -30,6 +37,7 @@ __all__ = [
     "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig",
     "Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig",
     "StableAudioConfig", "GlmImageDiTConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig",
-    "LingBotWorldFastVideoConfig",
-    "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig", "MMAudioTransformerConfig"
+    "LingBotWorldFastVideoConfig", "LingBotVideoConfig", "MiniMaxH3Config", "ZImageDiTConfig", "MMAudioArchConfig",
+    "MMAudioTransformerConfig", "MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS",
+    "MMAUDIO_44K_TRAINING_VARIANTS", "get_mmaudio_transformer_config"
 ]

fastvideo/configs/models/dits/mmaudio.py

@@ -2,11 +2,83 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from typing import TypedDict
 
 from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
 from fastvideo.platforms import AttentionBackendEnum
 
 
+class MMAudioVariantArchitecture(TypedDict):
+    latent_dim: int
+    hidden_dim: int
+    depth: int
+    fused_depth: int
+    num_heads: int
+    latent_seq_len: int
+    v2: bool
+
+
+MMAUDIO_VARIANT_ARCHITECTURES: dict[str, MMAudioVariantArchitecture] = {
+    "small_16k": {
+        "latent_dim": 20,
+        "hidden_dim": 448,
+        "depth": 12,
+        "fused_depth": 8,
+        "num_heads": 7,
+        "latent_seq_len": 250,
+        "v2": False,
+    },
+    "small_44k": {
+        "latent_dim": 40,
+        "hidden_dim": 448,
+        "depth": 12,
+        "fused_depth": 8,
+        "num_heads": 7,
+        "latent_seq_len": 345,
+        "v2": False,
+    },
+    "medium_44k": {
+        "latent_dim": 40,
+        "hidden_dim": 896,
+        "depth": 12,
+        "fused_depth": 8,
+        "num_heads": 14,
+        "latent_seq_len": 345,
+        "v2": False,
+    },
+    "large_44k": {
+        "latent_dim": 40,
+        "hidden_dim": 896,
+        "depth": 21,
+        "fused_depth": 14,
+        "num_heads": 14,
+        "latent_seq_len": 345,
+        "v2": False,
+    },
+    "large_44k_v2": {
+        "latent_dim": 40,
+        "hidden_dim": 896,
+        "depth": 21,
+        "fused_depth": 14,
+        "num_heads": 14,
+        "latent_seq_len": 345,
+        "v2": True,
+    },
+}
+
+MMAUDIO_TRAINING_VARIANTS = (
+    "small_16k",
+    "small_44k",
+    "medium_44k",
+    "large_44k",
+)
+MMAUDIO_44K_TRAINING_VARIANTS = (
+    "small_44k",
+    "medium_44k",
+    "large_44k",
+)
+
+
 def _is_mmaudio_transformer_block(name: str, module) -> bool:
     del module
     parts = name.split(".")
@@ -47,3 +119,13 @@ class MMAudioArchConfig(DiTArchConfig):
 class MMAudioTransformerConfig(DiTConfig):
     arch_config: DiTArchConfig = field(default_factory=MMAudioArchConfig)
     prefix: str = "MMAudio"
+
+
+def get_mmaudio_transformer_config(variant: str) -> MMAudioTransformerConfig:
+    """Build the exact official architecture config for ``variant``."""
+    try:
+        architecture = MMAUDIO_VARIANT_ARCHITECTURES[variant]
+    except KeyError as exc:
+        supported = ", ".join(MMAUDIO_VARIANT_ARCHITECTURES)
+        raise ValueError(f"Unknown MMAudio variant {variant!r}; expected one of: {supported}") from exc
+    return MMAudioTransformerConfig(arch_config=MMAudioArchConfig(**architecture), )

fastvideo/configs/pipelines/__init__.py

@@ -28,6 +28,6 @@ __all__ = [
     "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HeliosPipelineConfig", "HYWorldConfig",
     "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig",
     "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig",
-    "LingBotVideoT2VConfig",
-    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", "get_pipeline_config_cls_from_name"
+    "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig",
+    "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/mmaudio.py

@@ -6,8 +6,9 @@ from __future__ import annotations
 from dataclasses import dataclass, field
 
 from fastvideo.configs.models import DiTConfig, EncoderConfig, ModelConfig
-from fastvideo.configs.models.audio import BigVGANV2Config, MMAudioVAEConfig
-from fastvideo.configs.models.dits import MMAudioTransformerConfig
+from fastvideo.configs.models.audio import (BigVGANV2ArchConfig, BigVGANV2Config, MMAudioVAEArchConfig,
+                                            MMAudioVAEConfig)
+from fastvideo.configs.models.dits import MMAudioTransformerConfig, get_mmaudio_transformer_config
 from fastvideo.configs.models.encoders import (
     MMAudioDFNCLIPTextConfig,
     MMAudioDFNCLIPVisionConfig,
@@ -39,6 +40,10 @@ class MMAudioV2AConfig(PipelineConfig):
 
     audio_decoder_config: ModelConfig = field(default_factory=MMAudioVAEConfig)
     audio_decoder_precision: str = "bf16"
+    # Offline training preprocessing uses the VAE encoder in fp32, matching
+    # the published MMAudio feature extraction recipe. Inference remains
+    # decoder-only bf16 through ``audio_decoder_precision`` above.
+    audio_encoder_precision: str = "fp32"
     vocoder_config: ModelConfig = field(default_factory=BigVGANV2Config)
     vocoder_precision: str = "bf16"
 
@@ -46,6 +51,8 @@ class MMAudioV2AConfig(PipelineConfig):
     # supports other durations, although quality can drop far away from the
     # eight-second training duration.
     duration_s: float = 8.0
+    # ``None`` means no hard cap; set it to reject requests beyond the
+    # checkpoint's validated maximum.
     max_audio_duration_s: float | None = None
     sampling_rate: int = 44100
     spectrogram_frame_rate: int = 512
@@ -62,5 +69,53 @@ class MMAudioV2AConfig(PipelineConfig):
 
     num_inference_steps: int = 25
     guidance_scale: float = 4.5
+    # Inherited from ``PipelineConfig``; MMAudio decodes audio through its own
+    # stage, so the video-VAE paths that read these are not exercised.
     vae_tiling: bool = False
     vae_sp: bool = False
+
+
+@dataclass
+class MMAudioSmall16kV2AConfig(MMAudioV2AConfig):
+    """Official small 16 kHz MMAudio inference configuration."""
+
+    dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("small_16k"))
+    # ``MMAudioVAEArchConfig.mode`` is the single source of truth for the VAE
+    # shapes; ``MMAudioVAE`` re-derives them from it.
+    audio_decoder_config: ModelConfig = field(
+        default_factory=lambda: MMAudioVAEConfig(arch_config=MMAudioVAEArchConfig(mode="16k", )))
+    vocoder_config: ModelConfig = field(default_factory=lambda: BigVGANV2Config(arch_config=BigVGANV2ArchConfig(
+        sample_rate=16000,
+        num_mels=80,
+    )))
+    sampling_rate: int = 16000
+    spectrogram_frame_rate: int = 256
+
+
+@dataclass
+class MMAudioSmall44kV2AConfig(MMAudioV2AConfig):
+    dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("small_44k"))
+
+
+@dataclass
+class MMAudioMedium44kV2AConfig(MMAudioV2AConfig):
+    dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("medium_44k"))
+
+
+@dataclass
+class MMAudioLarge44kV2AConfig(MMAudioV2AConfig):
+    dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("large_44k"))
+
+
+@dataclass
+class MMAudioLarge44kV2PipelineConfig(MMAudioV2AConfig):
+    dit_config: DiTConfig = field(default_factory=lambda: get_mmaudio_transformer_config("large_44k_v2"))
+
+
+MMAUDIO_PIPELINE_CONFIGS = {
+    "small_16k": MMAudioSmall16kV2AConfig,
+    "small_44k": MMAudioSmall44kV2AConfig,
+    "medium_44k": MMAudioMedium44kV2AConfig,
+    "large_44k": MMAudioLarge44kV2AConfig,
+    "large_44k_v2": MMAudioLarge44kV2PipelineConfig,
+}

fastvideo/dataset/__init__.py

@@ -4,6 +4,9 @@ from torchvision.transforms import Lambda
 
 from fastvideo.dataset.parquet_dataset_map_style import (build_parquet_map_style_dataloader)
 from fastvideo.dataset.ltx2_precomputed_dataset import (build_ltx2_precomputed_dataloader, LTX2PrecomputedDataset)
+from fastvideo.dataset.mmaudio_feature_dataset import (MMAudioFeatureDataset, build_mmaudio_feature_dataloader)
+from fastvideo.dataset.v2a_feature_cache import V2AFeatureShardWriter
+from fastvideo.dataset.vggsound import VGGSoundDataset
 from fastvideo.dataset.preprocessing_datasets import VideoCaptionMergedDataset, TextDataset
 from fastvideo.dataset.transform import (CenterCropResizeVideo, Normalize255, TemporalRandomCrop)
 from fastvideo.dataset.validation_dataset import ValidationDataset
@@ -43,6 +46,10 @@ __all__ = [
     "build_parquet_map_style_dataloader",
     "build_ltx2_precomputed_dataloader",
     "LTX2PrecomputedDataset",
+    "MMAudioFeatureDataset",
+    "build_mmaudio_feature_dataloader",
+    "V2AFeatureShardWriter",
+    "VGGSoundDataset",
     "ValidationDataset",
     "VideoCaptionMergedDataset",
     "TextDataset",

fastvideo/dataset/mmaudio_feature_dataset.py

@@ -0,0 +1,327 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Memory-mapped precomputed feature dataset for audio generation training."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch.utils.data import ConcatDataset, Dataset
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+from fastvideo.dataset.parquet_dataset_map_style import DP_SP_BatchSampler
+from fastvideo.distributed import get_sp_world_size, get_world_rank, get_world_size
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class MMAudioFeatureDataset(Dataset):
+    """Read MMAudio-compatible TensorDict memmaps without upstream imports.
+
+    Required tensors are ``mean``, ``std``, and ``text_features``. Video
+    caches additionally contain ``clip_features`` and ``sync_features``.
+    Audio-only caches omit both video tensors; fixed-size placeholders and a
+    false ``video_exists`` flag are returned so video and audio datasets can be
+    mixed by the same dataloader.
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        *,
+        latent_seq_len: int,
+        latent_dim: int,
+        clip_seq_len: int,
+        clip_dim: int,
+        sync_seq_len: int,
+        sync_dim: int,
+        text_seq_len: int,
+        text_dim: int,
+        include_metadata: bool = False,
+    ) -> None:
+        super().__init__()
+        self.root = Path(root).expanduser().resolve()
+        if not self.root.is_dir():
+            raise FileNotFoundError(f"MMAudio feature cache does not exist: {self.root}")
+
+        try:
+            from tensordict import TensorDict
+        except ImportError as exc:
+            raise ImportError(
+                "MMAudio training caches require TensorDict. Install with "
+                "`uv pip install -e '.[mmaudio-train]'`."
+            ) from exc
+
+        self.tensors = TensorDict.load_memmap(self.root)
+        keys = set(self.tensors.keys())
+        required = {"mean", "std", "text_features"}
+        missing = sorted(required - keys)
+        if missing:
+            raise ValueError(f"MMAudio feature cache {self.root} is missing tensors: {missing}")
+
+        self.mean = self.tensors["mean"]
+        self.std = self.tensors["std"]
+        self.text_features = self.tensors["text_features"]
+        self.clip_features = self.tensors.get("clip_features")
+        self.sync_features = self.tensors.get("sync_features")
+        if (self.clip_features is None) != (self.sync_features is None):
+            raise ValueError(
+                "MMAudio feature cache must contain both clip_features and "
+                "sync_features, or neither for audio-only data."
+            )
+
+        self._length = int(self.mean.shape[0])
+        for name in required:
+            if int(self.tensors[name].shape[0]) != self._length:
+                raise ValueError(f"MMAudio cache tensor {name!r} has a mismatched sample count")
+        if self.clip_features is not None and int(self.clip_features.shape[0]) != self._length:
+            raise ValueError("MMAudio cache clip_features has a mismatched sample count")
+        if self.sync_features is not None and int(self.sync_features.shape[0]) != self._length:
+            raise ValueError("MMAudio cache sync_features has a mismatched sample count")
+
+        self._require_shape("mean", self.mean, (latent_seq_len, latent_dim))
+        self._require_shape("std", self.std, (latent_seq_len, latent_dim))
+        self._require_shape("text_features", self.text_features, (text_seq_len, text_dim))
+        if self.clip_features is not None:
+            self._require_shape("clip_features", self.clip_features, (clip_seq_len, clip_dim))
+            assert self.sync_features is not None
+            self._require_shape("sync_features", self.sync_features, (sync_seq_len, sync_dim))
+
+        feature_dtype = self.text_features.dtype
+        self.empty_clip = torch.zeros(clip_seq_len, clip_dim, dtype=feature_dtype)
+        self.empty_sync = torch.zeros(sync_seq_len, sync_dim, dtype=feature_dtype)
+        self.has_video = self.clip_features is not None
+        self.metadata: list[dict[str, Any]] | None = None
+        if include_metadata:
+            metadata_path = self.root / "samples.jsonl"
+            if metadata_path.is_file():
+                with metadata_path.open(encoding="utf-8") as handle:
+                    self.metadata = [json.loads(line) for line in handle if line.strip()]
+                if len(self.metadata) != self._length:
+                    raise ValueError(
+                        f"MMAudio cache {metadata_path} has {len(self.metadata)} "
+                        f"metadata rows for {self._length} feature samples"
+                    )
+            else:
+                logger.warning(
+                    "MMAudio cache %s has no samples.jsonl; validation "
+                    "audio cannot be associated with source videos",
+                    self.root,
+                )
+                self.metadata = [{} for _ in range(self._length)]
+        logger.info("Loaded %d MMAudio feature samples from %s", self._length, self.root)
+
+    @staticmethod
+    def _require_shape(
+        name: str,
+        tensor: torch.Tensor,
+        expected: tuple[int, ...],
+    ) -> None:
+        actual = tuple(tensor.shape[1:])
+        if actual != expected:
+            raise ValueError(
+                f"MMAudio cache tensor {name!r} must have per-sample shape "
+                f"{expected}, got {actual}"
+            )
+
+    def __len__(self) -> int:
+        return self._length
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        clip_features = self.empty_clip
+        sync_features = self.empty_sync
+        if self.has_video:
+            assert self.clip_features is not None and self.sync_features is not None
+            clip_features = self.clip_features[index]
+            sync_features = self.sync_features[index]
+        sample: dict[str, Any] = {
+            "audio_latent_mean": self.mean[index],
+            "audio_latent_std": self.std[index],
+            "clip_features": clip_features,
+            "sync_features": sync_features,
+            "text_features": self.text_features[index],
+            "video_exists": torch.tensor(self.has_video, dtype=torch.bool),
+            "text_exists": torch.tensor(True, dtype=torch.bool),
+        }
+        if self.metadata is not None:
+            metadata = self.metadata[index]
+            sample.update({
+                "sample_id": str(metadata.get("id", index)),
+                "caption": str(metadata.get("caption", "")),
+                "source_path": str(metadata.get("source", "")),
+            })
+        return sample
+
+
+def _data_specs(data_path: str | Sequence[str] | dict[str, int]) -> list[tuple[str, int]]:
+    if isinstance(data_path, dict):
+        return [
+            (str(path), int(repeat))
+            for path, repeat in data_path.items()
+            if int(repeat) > 0
+        ]
+    if isinstance(data_path, Sequence) and not isinstance(data_path, str):
+        return [(str(path), 1) for path in data_path if str(path).strip()]
+    return [(part.strip(), 1) for part in str(data_path).split(",") if part.strip()]
+
+
+def _expand_cache_root(root: str | Path) -> list[Path]:
+    """Resolve one direct TensorDict cache or a directory of cache shards."""
+    path = Path(root).expanduser().resolve()
+    if not path.is_dir():
+        raise FileNotFoundError(f"MMAudio feature cache does not exist: {path}")
+    if (path / "meta.json").is_file():
+        return [path]
+    shards = sorted(metadata.parent for metadata in path.rglob("meta.json"))
+    if not shards:
+        raise FileNotFoundError(
+            f"No TensorDict cache shards (meta.json) were found under {path}")
+    return shards
+
+
+def compute_mmaudio_latent_stats(
+    data_path: str | Sequence[str] | dict[str, int],
+    *,
+    latent_seq_len: int,
+    latent_dim: int,
+    chunk_size: int = 32,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute official MMAudio normalization stats from the first cache.
+
+    The reference trainer uses the posterior means from its first video
+    dataset and reduces over sample and sequence dimensions. This chunked
+    implementation preserves that contract without materializing the complete
+    VGGSound tensor in RAM.
+    """
+    specs = _data_specs(data_path)
+    if not specs:
+        raise ValueError(
+            "training.data.data_path is empty; latent statistics require a "
+            "video feature cache."
+        )
+    root = specs[0][0]
+    chunk_size = max(1, int(chunk_size))
+    total = 0
+    value_sum = torch.zeros(latent_dim, dtype=torch.float64)
+    square_sum = torch.zeros(latent_dim, dtype=torch.float64)
+
+    try:
+        from tensordict import TensorDict
+    except ImportError as exc:
+        raise ImportError(
+            "MMAudio latent statistics require TensorDict. Install with "
+            "`uv pip install -e '.[mmaudio-train]'`."
+        ) from exc
+
+    for shard_root in _expand_cache_root(root):
+        tensors = TensorDict.load_memmap(shard_root)
+        mean = tensors.get("mean")
+        if not isinstance(mean, torch.Tensor):
+            raise ValueError(
+                f"MMAudio cache {shard_root} is missing tensor 'mean'"
+            )
+        if tuple(mean.shape[1:]) != (latent_seq_len, latent_dim):
+            raise ValueError(
+                f"MMAudio cache {shard_root} mean must have per-sample shape "
+                f"{(latent_seq_len, latent_dim)}, got {tuple(mean.shape[1:])}"
+            )
+        for start in range(0, int(mean.shape[0]), chunk_size):
+            values = mean[start:start + chunk_size].to(torch.float64)
+            flattened = values.reshape(-1, latent_dim)
+            value_sum += flattened.sum(dim=0)
+            square_sum += flattened.square().sum(dim=0)
+            total += int(flattened.shape[0])
+
+    if total < 2:
+        raise ValueError(
+            f"MMAudio latent statistics need at least two values, got {total}"
+        )
+    latent_mean = value_sum / total
+    variance = (
+        square_sum - value_sum.square() / total
+    ).clamp_min(0.0) / (total - 1)
+    latent_std = variance.sqrt()
+    return (
+        latent_mean.to(torch.float32).reshape(1, 1, latent_dim),
+        latent_std.to(torch.float32).reshape(1, 1, latent_dim),
+    )
+
+
+def build_mmaudio_feature_dataset(
+    data_path: str | Sequence[str] | dict[str, int],
+    *,
+    feature_shapes: dict[str, int],
+    include_metadata: bool = False,
+) -> Dataset:
+    """Build one map-style dataset over MMAudio TensorDict cache shards.
+
+    Unlike :func:`build_mmaudio_feature_dataloader`, this helper does not add
+    a distributed sampler. Dataset-scale inference can therefore assign
+    indices explicitly (for example ``range(rank, len(dataset), world_size)``)
+    without ``DistributedSampler`` padding the tail with duplicate samples.
+    """
+    specs = _data_specs(data_path)
+    if not specs:
+        raise ValueError(
+            "data_path is empty. Set it to one or more precomputed MMAudio "
+            "TensorDict mmap directories."
+        )
+
+    datasets: list[Dataset] = []
+    for root, repeat in specs:
+        for shard_root in _expand_cache_root(root):
+            dataset = MMAudioFeatureDataset(
+                shard_root,
+                include_metadata=include_metadata,
+                **feature_shapes,
+            )
+            datasets.extend([dataset] * repeat)
+    return datasets[0] if len(datasets) == 1 else ConcatDataset(datasets)
+
+
+def build_mmaudio_feature_dataloader(
+    data_path: str | Sequence[str] | dict[str, int],
+    *,
+    batch_size: int,
+    num_data_workers: int,
+    seed: int,
+    pin_memory: bool,
+    feature_shapes: dict[str, int],
+    include_metadata: bool = False,
+) -> StatefulDataLoader:
+    """Build a distributed/stateful loader over one or more feature caches."""
+    combined = build_mmaudio_feature_dataset(
+        data_path,
+        feature_shapes=feature_shapes,
+        include_metadata=include_metadata,
+    )
+
+    sp_world_size = get_sp_world_size()
+    sampler = DP_SP_BatchSampler(
+        batch_size=int(batch_size),
+        dataset_size=len(combined),
+        num_sp_groups=get_world_size() // sp_world_size,
+        sp_world_size=sp_world_size,
+        global_rank=get_world_rank(),
+        drop_last=True,
+        seed=int(seed),
+    )
+    return StatefulDataLoader(
+        combined,
+        batch_sampler=sampler,
+        num_workers=int(num_data_workers),
+        pin_memory=bool(pin_memory),
+        persistent_workers=int(num_data_workers) > 0,
+    )
+
+
+__all__ = [
+    "MMAudioFeatureDataset",
+    "build_mmaudio_feature_dataset",
+    "build_mmaudio_feature_dataloader",
+    "compute_mmaudio_latent_stats",
+]

fastvideo/dataset/v2a_feature_cache.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared sharded feature-cache writer for video-to-audio models."""
+
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class V2AFeatureShardWriter:
+    """Write model-specific feature dictionaries as resumable TensorDict shards.
+
+    Pipelines publish a ``dict[str, Tensor]`` with a common batch dimension;
+    this writer handles buffering, per-rank shard names, metadata, and resume
+    discovery without knowing what any feature means.
+    """
+
+    def __init__(
+        self,
+        output_root: str | Path,
+        *,
+        rank: int,
+        samples_per_shard: int,
+    ) -> None:
+        if samples_per_shard <= 0:
+            raise ValueError("samples_per_shard must be positive")
+        self.output_root = Path(output_root).expanduser().resolve()
+        self.worker_root = self.output_root / f"worker_{rank:05d}"
+        self.worker_root.mkdir(parents=True, exist_ok=True)
+        self.samples_per_shard = int(samples_per_shard)
+        self._features: dict[str, list[torch.Tensor]] = defaultdict(list)
+        self._metadata: list[dict[str, Any]] = []
+        self.existing_ids = self._read_existing_ids(self.output_root)
+        existing_indices = []
+        for path in self.worker_root.glob("shard_*"):
+            try:
+                existing_indices.append(int(path.name.split("_")[-1]))
+            except ValueError:
+                continue
+        self._next_shard = max(existing_indices, default=-1) + 1
+
+    @staticmethod
+    def _read_existing_ids(root: Path) -> set[str]:
+        existing: set[str] = set()
+        if not root.is_dir():
+            return existing
+        for metadata_path in root.rglob("samples.jsonl"):
+            with metadata_path.open(encoding="utf-8") as handle:
+                for line in handle:
+                    if not line.strip():
+                        continue
+                    sample_id = json.loads(line).get("id")
+                    if sample_id:
+                        existing.add(str(sample_id))
+        return existing
+
+    def contains(self, sample_id: str) -> bool:
+        return sample_id in self.existing_ids
+
+    @property
+    def buffered_samples(self) -> int:
+        return len(self._metadata)
+
+    def append(
+        self,
+        features: dict[str, torch.Tensor],
+        metadata: list[dict[str, Any]],
+    ) -> None:
+        if not features:
+            return
+        batch_sizes = {int(tensor.shape[0]) for tensor in features.values()}
+        if len(batch_sizes) != 1:
+            raise ValueError(f"V2A feature batch dimensions do not match: {batch_sizes}")
+        batch_size = batch_sizes.pop()
+        if batch_size != len(metadata):
+            raise ValueError(f"V2A feature batch has {batch_size} samples but {len(metadata)} metadata rows")
+        if self._features and set(features) != set(self._features):
+            raise ValueError(
+                f"V2A feature keys changed within one shard: {sorted(self._features)} vs {sorted(features)}")
+
+        for index, row in enumerate(metadata):
+            sample_id = str(row["id"])
+            if sample_id in self.existing_ids:
+                continue
+            for name, tensor in features.items():
+                self._features[name].append(tensor[index].detach().cpu().contiguous())
+            self._metadata.append(dict(row))
+            self.existing_ids.add(sample_id)
+            if self.buffered_samples >= self.samples_per_shard:
+                self.flush()
+
+    def flush(self) -> Path | None:
+        if not self._metadata:
+            return None
+        try:
+            from tensordict import TensorDict
+        except ImportError as exc:
+            raise ImportError("V2A feature caches require TensorDict; install fastvideo[mmaudio-train].") from exc
+
+        shard_path = self.worker_root / f"shard_{self._next_shard:06d}"
+        if shard_path.exists():
+            raise FileExistsError(f"Refusing to overwrite existing V2A cache shard: {shard_path}")
+        stacked = {name: torch.stack(values) for name, values in self._features.items()}
+        TensorDict(stacked, batch_size=[len(self._metadata)]).memmap_(str(shard_path))
+        with (shard_path / "samples.jsonl").open("w", encoding="utf-8") as handle:
+            for row in self._metadata:
+                handle.write(json.dumps(row, ensure_ascii=False) + "\n")
+        logger.info("Wrote %d V2A feature samples to %s", len(self._metadata), shard_path)
+        self._features.clear()
+        self._metadata.clear()
+        self._next_shard += 1
+        return shard_path
+
+    def close(self) -> None:
+        self.flush()
+
+
+__all__ = ["V2AFeatureShardWriter"]

fastvideo/dataset/vggsound.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Raw VGGSound metadata adapter for V2A preprocessing."""
+
+from __future__ import annotations
+
+import csv
+import os
+from pathlib import Path
+
+from torch.utils.data import Dataset
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+class VGGSoundDataset(Dataset):
+    """Map VGGSound metadata rows to local MP4 paths and captions.
+
+    ``Loie/VGGSound`` stores clips as ``<youtube-id>_<start:06d>.mp4``.
+    The downloaded tar archives must be extracted before random-access GPU
+    preprocessing; repeatedly seeking inside gzip archives is prohibitively
+    expensive for a shuffled training dataset.
+    """
+
+    def __init__(
+        self,
+        root: str | Path,
+        *,
+        split: str = "train",
+        metadata_path: str | Path | None = None,
+    ) -> None:
+        super().__init__()
+        self.root = Path(root).expanduser().resolve()
+        metadata = (Path(metadata_path).expanduser().resolve() if metadata_path is not None else self.root /
+                    "vggsound.csv")
+        if not metadata.is_file():
+            raise FileNotFoundError(f"VGGSound metadata file does not exist: {metadata}")
+
+        candidates = (self.root / "videos", self.root / "video", self.root)
+        self.video_root = next((path for path in candidates if path.is_dir()), self.root)
+        self.samples: list[tuple[str, str, Path]] = []
+        if metadata.suffix.lower() == ".tsv":
+            self._read_caption_manifest(metadata, split=split)
+        else:
+            self._read_vggsound_csv(metadata, split=split)
+
+        if not self.samples:
+            raise ValueError(f"No VGGSound samples found for split {split!r} in {metadata}")
+
+    def _read_caption_manifest(self, metadata: Path, *, split: str) -> None:
+        seen_ids: set[str] = set()
+        missing_videos = 0
+        available_videos = {
+            Path(name).stem
+            for name in os.listdir(self.video_root)
+            if Path(name).suffix.lower() == ".mp4"
+        }
+        with metadata.open(newline="", encoding="utf-8-sig") as handle:
+            reader = csv.DictReader(handle, delimiter="\t")
+            required_columns = {"id", "label"}
+            columns = set(reader.fieldnames or ())
+            if not required_columns.issubset(columns):
+                raise ValueError(
+                    f"VGGSound caption manifest {metadata} must contain columns {sorted(required_columns)}, "
+                    f"got {sorted(columns)}")
+            for row_number, row in enumerate(reader, start=2):
+                sample_id = str(row.get("id") or "").strip()
+                caption = row.get("label")
+                if not sample_id:
+                    raise ValueError(f"VGGSound caption manifest has an empty id at row {row_number}")
+                if caption is None or not caption.strip():
+                    raise ValueError(f"VGGSound caption manifest has an empty label at row {row_number}")
+                if sample_id in seen_ids:
+                    raise ValueError(f"Duplicate VGGSound id {sample_id!r} at row {row_number}")
+                seen_ids.add(sample_id)
+                video_path = self.video_root / f"{sample_id}.mp4"
+                if sample_id not in available_videos:
+                    missing_videos += 1
+                    continue
+                self.samples.append((sample_id, caption, video_path))
+        logger.info(
+            "Loaded %d VGGSound %s captions from %s (%d videos missing)",
+            len(self.samples),
+            split,
+            metadata,
+            missing_videos,
+        )
+
+    def _read_vggsound_csv(self, metadata: Path, *, split: str) -> None:
+        with metadata.open(newline="", encoding="utf-8") as handle:
+            reader = csv.reader(handle)
+            for row_number, row in enumerate(reader, start=1):
+                if len(row) != 4:
+                    raise ValueError(f"VGGSound CSV row {row_number} must have four columns, got {row!r}")
+                youtube_id, start_seconds, caption, row_split = row
+                if row_split.strip().lower() != split.lower():
+                    continue
+                try:
+                    start = int(start_seconds)
+                except ValueError as exc:
+                    raise ValueError(f"Invalid VGGSound start time at row {row_number}: {start_seconds!r}") from exc
+                sample_id = f"{youtube_id}_{start:06d}"
+                self.samples.append((sample_id, caption, self.video_root / f"{sample_id}.mp4"))
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> dict[str, str]:
+        sample_id, caption, path = self.samples[index]
+        return {
+            "id": sample_id,
+            "caption": caption,
+            "video_path": str(path),
+        }
+
+
+__all__ = ["VGGSoundDataset"]

fastvideo/entrypoints/cli/eval.py

@@ -43,20 +43,30 @@ class EvalSubcommand(CLISubcommand):
             _cmd_list(args)
         elif action == "run":
             _cmd_run(args)
+        elif action == "v2a":
+            _cmd_v2a(args)
         else:
             # Re-print help if no action was given.
             self._parser.print_help()  # type: ignore[attr-defined]
 
     def validate(self, args: argparse.Namespace) -> None:
         action = getattr(args, "eval_action", None)
-        if action == "run" and not args.videos:
-            raise SystemExit("`fastvideo eval run` requires --videos")
+        if action == "v2a":
+            if args.backend != "av-benchmark":
+                raise SystemExit(f"Unsupported V2A evaluation backend: {args.backend}")
+            return
+        if action != "run":
+            return
+        if not args.manifest and not args.videos and not args.audios:
+            raise SystemExit("`fastvideo eval run` requires --manifest, --videos, or --audios")
+        if args.manifest and (args.videos or args.audios or args.reference or args.reference_audio):
+            raise SystemExit("--manifest cannot be combined with media path arguments")
 
     def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
         eval_parser = subparsers.add_parser(
             "eval",
-            help="Run video-gen evaluation metrics",
-            usage="fastvideo eval {list,run} [...]",
+            help="Run video and V2A evaluation metrics",
+            usage="fastvideo eval {list,run,v2a} [...]",
         )
         sub = eval_parser.add_subparsers(dest="eval_action", required=False)
 
@@ -66,27 +76,101 @@ class EvalSubcommand(CLISubcommand):
 
         # `eval run`
         run_p = sub.add_parser("run", help="Evaluate videos against one or more metrics")
+        run_p.add_argument(
+            "--manifest",
+            type=str,
+            default=None,
+            help="JSON/JSONL canonical sample manifest (recommended for benchmarks).",
+        )
         run_p.add_argument("--videos",
                            type=str,
                            nargs="+",
                            required=False,
                            help="Path, glob, or directory of generated videos.")
+        run_p.add_argument("--audios",
+                           type=str,
+                           nargs="+",
+                           required=False,
+                           help="Path, glob, or directory of generated audio files.")
         run_p.add_argument("--reference",
                            type=str,
                            default=None,
                            help="Path / glob / dir of reference videos (for paired metrics).")
+        run_p.add_argument(
+            "--reference-audio",
+            type=str,
+            nargs="+",
+            default=None,
+            help="Path, glob, or directory of reference audio files.",
+        )
         run_p.add_argument("--metrics", type=str, default="all", help="Comma-separated metric names, or 'all'.")
         run_p.add_argument("--device", type=str, default="cuda", help="Torch device (e.g. 'cuda', 'cuda:0', 'cpu').")
-        run_p.add_argument("--text-prompt",
-                           type=str,
-                           nargs="*",
-                           default=None,
-                           help="Prompt(s) for text-conditioned metrics. One per video.")
+        run_p.add_argument("--num-gpus", type=int, default=1, help="Number of metric replicas for dataset evaluation.")
+        run_p.add_argument("--compile", action="store_true", help="Compile supported metric models.")
+        run_p.add_argument(
+            "--skip-missing-deps",
+            action="store_true",
+            help="Skip explicitly requested metrics whose optional dependencies are missing.",
+        )
+        run_p.add_argument(
+            "--extract-audio",
+            nargs="?",
+            const=True,
+            default=False,
+            help="Extract missing audio tracks; optionally provide a persistent cache directory.",
+        )
+        run_p.add_argument("--extract-workers", type=int, default=4)
+        run_p.add_argument(
+            "--text-prompt",
+            type=str,
+            nargs="*",
+            default=None,
+            help="Prompt(s) for text-conditioned metrics. One per video.",
+        )
         run_p.add_argument("--fps", type=float, default=None, help="Frame-rate annotation passed to fps-aware metrics.")
         run_p.add_argument("--output",
                            type=str,
                            default=None,
                            help="Write results as JSON to this path (default: stdout).")
+        run_p.add_argument(
+            "--output-format",
+            choices=("legacy", "full"),
+            default="legacy",
+            help="'full' includes corpus metrics and sample identities.",
+        )
+
+        # `eval v2a`: isolated corpus-level backends for video-to-audio.
+        v2a_p = sub.add_parser("v2a", help="Evaluate V2A audio with a corpus benchmark backend")
+        v2a_p.add_argument("--backend", choices=("av-benchmark", ), default="av-benchmark")
+        v2a_p.add_argument("--audio-dir", required=True, help="Directory containing generated WAV/FLAC files.")
+        v2a_p.add_argument("--gt-cache", required=True, help="Official ground-truth feature cache directory.")
+        v2a_p.add_argument(
+            "--prediction-cache",
+            default=None,
+            help="Generated-audio feature cache (default: <audio-dir>/../av_benchmark_cache).",
+        )
+        v2a_p.add_argument(
+            "--output",
+            default=None,
+            help="Result JSON (default: <audio-dir>/../av_benchmark_results.json).",
+        )
+        v2a_p.add_argument(
+            "--python-executable",
+            default=None,
+            help="Python from an isolated av-benchmark environment.",
+        )
+        v2a_p.add_argument("--device", default="cuda")
+        v2a_p.add_argument("--batch-size", type=int, default=32)
+        v2a_p.add_argument("--num-workers", type=int, default=0)
+        v2a_p.add_argument("--audio-length", type=float, default=8.0)
+        v2a_p.add_argument("--recompute", action="store_true", help="Re-extract the prediction cache.")
+        v2a_p.add_argument("--skip-video-related", action="store_true", help="Skip ImageBind and DeSync.")
+        v2a_p.add_argument("--skip-clap", action="store_true", help="Skip LAION/MS CLAP extraction and scores.")
+        v2a_p.add_argument(
+            "--align-prediction-keys",
+            action="store_true",
+            help="Align uniquely underscore-prefixed prediction ids to official GT cache ids.",
+        )
 
         # Stash the parser so cmd() can re-print help on no-action.
         self._parser = eval_parser  # type: ignore[attr-defined]
@@ -95,6 +179,7 @@ class EvalSubcommand(CLISubcommand):
 
 def _cmd_list(args: argparse.Namespace) -> None:
     from fastvideo.eval import list_metrics
+
     names = list_metrics()
     if args.group:
         prefix = args.group.rstrip(".") + "."
@@ -107,38 +192,73 @@ def _cmd_list(args: argparse.Namespace) -> None:
 
 
 def _cmd_run(args: argparse.Namespace) -> None:
-    from fastvideo.eval import create_evaluator
-    from fastvideo.eval.io import load_video
+    from fastvideo.eval import create_evaluator, samples_from, samples_from_manifest
 
-    video_paths = _expand_paths(args.videos)
-    if not video_paths:
-        raise SystemExit(f"No videos matched: {args.videos}")
-    ref_paths = _expand_paths([args.reference]) if args.reference else None
+    if args.manifest:
+        samples = samples_from_manifest(
+            args.manifest,
+            extract_audio=args.extract_audio,
+            extract_workers=args.extract_workers,
+        )
+    else:
+        video_paths = _expand_paths(args.videos, _VIDEO_EXTS) if args.videos else None
+        audio_paths = _expand_paths(args.audios, _AUDIO_EXTS) if args.audios else None
+        if video_paths is not None and not video_paths:
+            raise SystemExit(f"No videos matched: {args.videos}")
+        if audio_paths is not None and not audio_paths:
+            raise SystemExit(f"No audio files matched: {args.audios}")
+        generated_count = len(video_paths if video_paths is not None else audio_paths or [])
+        ref_paths = (_fill_reference_paths(_expand_paths([args.reference], _VIDEO_EXTS), generated_count)
+                     if args.reference else None)
+        ref_audio_paths = (_fill_reference_paths(_expand_paths(args.reference_audio, _AUDIO_EXTS), generated_count)
+                           if args.reference_audio else None)
+        text_prompts = args.text_prompt
+        if text_prompts is not None and len(text_prompts) == 1:
+            text_prompts = text_prompts * generated_count
+        samples = samples_from(
+            video=video_paths,
+            audio=audio_paths,
+            reference=ref_paths,
+            reference_audio=ref_audio_paths,
+            text_prompts=text_prompts,
+            fps=args.fps,
+            extract_audio=args.extract_audio,
+            extract_workers=args.extract_workers,
+        )
 
     metrics_arg: list[str] | str = ("all" if args.metrics == "all" else
                                     [m.strip() for m in args.metrics.split(",") if m.strip()])
 
-    evaluator = create_evaluator(metrics=metrics_arg, device=args.device)
+    evaluator = create_evaluator(
+        metrics=metrics_arg,
+        device=args.device,
+        num_gpus=max(1, args.num_gpus),
+        compile=args.compile,
+        skip_missing_deps=args.skip_missing_deps,
+    )
+    logger.info("Evaluating %d samples on %d GPU worker(s)", len(samples), evaluator.num_gpus)
+    results = evaluator.evaluate(samples=samples)
+    corpus = _serialize_results(results.corpus)
+    output_value: dict | list
+    if args.output_format == "full":
+        sample_results = [{
+            **_sample_identity(sample, index),
+            "scores": _serialize_results(result),
+        } for index, (sample, result) in enumerate(zip(samples, results, strict=True))]
+        output_value = {
+            "samples": sample_results,
+            "corpus": corpus,
+            "num_samples": len(samples),
+        }
+    else:
+        if corpus:
+            logger.warning("Corpus results are omitted in legacy output format; pass --output-format full to keep them")
+        output_value = [{
+            **_legacy_sample_identity(sample),
+            "scores": _serialize_results(result),
+        } for sample, result in zip(samples, results, strict=True)]
 
-    all_results: list[dict] = []
-    for i, vp in enumerate(video_paths):
-        logger.info("Evaluating %s (%d/%d)", vp, i + 1, len(video_paths))
-        kwargs: dict = {"video": load_video(vp)}
-        if ref_paths is not None:
-            ref = ref_paths[i] if i < len(ref_paths) else ref_paths[0]
-            kwargs["reference"] = load_video(ref)
-        if args.text_prompt is not None:
-            kwargs["text_prompt"] = (args.text_prompt[i] if i < len(args.text_prompt) else args.text_prompt[0])
-        if args.fps is not None:
-            kwargs["fps"] = args.fps
-
-        results = evaluator.evaluate(**kwargs)
-        all_results.append({
-            "video": str(vp),
-            "scores": _serialize_results(results),
-        })
-
-    payload = json.dumps(all_results, indent=2, default=_jsonable)
+    payload = json.dumps(output_value, indent=2, default=_jsonable)
     if args.output:
         Path(args.output).write_text(payload)
         logger.info("Wrote results to %s", args.output)
@@ -146,13 +266,38 @@ def _cmd_run(args: argparse.Namespace) -> None:
         print(payload)
 
 
-def _expand_paths(patterns: list[str]) -> list[str]:
+def _cmd_v2a(args: argparse.Namespace) -> None:
+    """Run an isolated corpus-level backend for generated V2A audio."""
+    from fastvideo.eval.v2a import run_av_benchmark
+
+    output = run_av_benchmark(
+        audio_dir=args.audio_dir,
+        gt_cache=args.gt_cache,
+        prediction_cache=args.prediction_cache,
+        output=args.output,
+        python_executable=args.python_executable,
+        device=args.device,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        audio_length=args.audio_length,
+        recompute=args.recompute,
+        skip_video_related=args.skip_video_related,
+        skip_clap=args.skip_clap,
+        align_prediction_keys=args.align_prediction_keys,
+    )
+    logger.info("Official av-benchmark results: %s", output)
+
+
+_VIDEO_EXTS = (".mp4", ".avi", ".mov", ".mkv", ".gif", ".webm")
+_AUDIO_EXTS = (".wav", ".mp3", ".flac", ".m4a", ".ogg", ".opus")
+
+
+def _expand_paths(patterns: list[str], extensions: tuple[str, ...]) -> list[str]:
     out: list[str] = []
     for pat in patterns:
         p = Path(pat)
         if p.is_dir():
-            for ext in (".mp4", ".avi", ".mov", ".mkv", ".gif"):
-                out.extend(sorted(str(f) for f in p.iterdir() if f.suffix.lower() == ext))
+            out.extend(sorted(str(f) for f in p.iterdir() if f.suffix.lower() in extensions))
         elif any(c in pat for c in "*?["):
             out.extend(sorted(glob.glob(pat)))
         else:
@@ -167,6 +312,43 @@ def _expand_paths(patterns: list[str]) -> list[str]:
     return deduped
 
 
+def _fill_reference_paths(paths: list[str], count: int) -> list[str]:
+    if not paths:
+        return paths
+    if len(paths) >= count:
+        return paths
+    return paths + [paths[0]] * (count - len(paths))
+
+
+def _legacy_sample_identity(sample: dict) -> dict:
+    """Preserve the original ``fastvideo eval run`` output schema."""
+    from fastvideo.eval.types import Video
+
+    video = sample.get("video")
+    if isinstance(video, Video) and video.source is not None:
+        return {"video": str(video.source)}
+    audio = sample.get("audio")
+    if isinstance(audio, str | Path):
+        return {"audio": str(audio)}
+    return {}
+
+
+def _sample_identity(sample: dict, index: int) -> dict:
+    from fastvideo.eval.types import Video
+
+    identity: dict = {"index": index}
+    if sample.get("id") is not None:
+        identity["id"] = sample["id"]
+    for key in ("video", "audio", "reference", "reference_audio", "text_prompt"):
+        value = sample.get(key)
+        if isinstance(value, Video):
+            if value.source is not None:
+                identity[key] = str(value.source)
+        elif isinstance(value, str | Path):
+            identity[key] = str(value)
+    return identity
+
+
 def _serialize_results(results) -> dict | list:
     """Turn evaluator output (dict or list-of-dicts) into JSON-friendly form."""
     if isinstance(results, list):

fastvideo/eval/README.md

@@ -157,6 +157,13 @@ pairs. Per-sample attachments (`text_prompt(s)`, `fps`,
 `auxiliary_info`, `extras=`) attach by kwarg; `extract_audio=True` pulls
 audio tracks off video sources for audio metrics.
 
+Samples whose `video` is a path (or a `Video` without frames) are decoded in
+the worker that picks them up. When the caller did not attach `fps`, the pool
+probes the container's frame rate and attaches it as `sample["fps"]` so
+fps-aware metrics see real timing; an explicit `fps` (for example `--fps` on
+`fastvideo eval run`) always wins and skips the probe. Scores recorded before
+this probe existed can therefore shift unless `--fps` is passed.
+
 `evaluate` also accepts a pre-loaded `(T, C, H, W)` tensor or a path
 string under `video` / `reference`. Paths are decoded inside the worker
 that picks up the sample, so peak memory stays bounded by `num_gpus`
@@ -429,6 +436,55 @@ translation); rigid whole-frame motion (e.g. pure camera rotation) is harder. To
 sweep several baselines into a table, see
 `examples/inference/eval/eval_third_person_separation.py`.
 
+## Exact V2A corpus evaluation
+
+`fastvideo eval v2a` is a dedicated corpus-level entry point for
+video-to-audio benchmarks. Unlike `fastvideo eval run`, which composes native
+per-sample and set metrics through the registry, the V2A command can launch an
+upstream benchmark as one protocol-preserving unit.
+
+The `av-benchmark` backend calls the official `av_bench.extract` and
+`av_bench.evaluate` functions. Its dependency stack is intentionally not part
+of any standard FastVideo extra. Point the command at a Python interpreter from
+an isolated av-benchmark environment:
+
+```bash
+fastvideo eval v2a \
+  --backend av-benchmark \
+  --audio-dir outputs/v2a/audio \
+  --gt-cache /path/to/official/ground-truth-cache \
+  --prediction-cache outputs/v2a/av_benchmark_cache \
+  --output outputs/v2a/av_benchmark_results.json \
+  --python-executable /path/to/av-benchmark/.venv/bin/python \
+  --audio-length 8 \
+  --batch-size 32 \
+  --num-workers 0
+```
+
+The backend is model-independent: generated filenames only need to match the
+keys in the selected official cache. It can therefore evaluate MMAudio or any
+other V2A model. A complete prediction cache is reused automatically;
+`--recompute` forces extraction again.
+
+Some official dataset caches sanitize filenames. For such a cache,
+`--align-prediction-keys` atomically remaps only prediction keys having a
+unique match after leading underscores are ignored. The source audio is never
+renamed, ambiguous keys remain unmatched, and alignment counts are serialized
+in the result JSON.
+
+The official VGGSound cache reports:
+
+```text
+FD-VGG, FD-PANN, FD-PASST
+KL-PANNS-softmax, KL-PASST-softmax
+ISC-PANNS-mean/std, ISC-PASST-mean/std
+IB-Score, DeSync
+```
+
+This command runs in a subprocess by design. It does not import av-benchmark,
+alter CUDA/PyTorch packages, or register metrics in ordinary FastVideo
+generation and evaluation processes.
+
 ## Out of scope (follow-up PRs)
 
 - **MIND** metrics. Depend on a separate `vipe` upstream submodule.

fastvideo/eval/__init__.py

@@ -26,7 +26,7 @@ from fastvideo.eval.metrics.base import BaseMetric  # noqa: E402
 from fastvideo.eval.registry import register, list_metrics, get_metric  # noqa: E402
 from fastvideo.eval.api import evaluate  # noqa: E402
 from fastvideo.eval.evaluator import Evaluator, create_evaluator  # noqa: E402
-from fastvideo.eval.io.inputs import as_video, samples_from  # noqa: E402
+from fastvideo.eval.io.inputs import as_video, samples_from, samples_from_manifest  # noqa: E402
 
 # Trigger metric auto-discovery
 import fastvideo.eval.metrics  # noqa: F401, E402
@@ -46,4 +46,5 @@ __all__ = [
     "get_cache_dir",
     "as_video",
     "samples_from",
+    "samples_from_manifest",
 ]

fastvideo/eval/io/__init__.py

@@ -1,6 +1,6 @@
 from fastvideo.eval.io.audio import NoAudioStreamError, extract_audio_track
-from fastvideo.eval.io.inputs import as_video, samples_from
-from fastvideo.eval.io.paths import (build_eval_kwargs, default_filename, glob_videos, sanitize_prompt)
+from fastvideo.eval.io.inputs import as_video, samples_from, samples_from_manifest
+from fastvideo.eval.io.paths import build_eval_kwargs, default_filename, glob_videos, sanitize_prompt
 from fastvideo.eval.io.video import extract_frames, load_video
 
 __all__ = [
@@ -12,6 +12,7 @@ __all__ = [
     "build_eval_kwargs",
     "as_video",
     "samples_from",
+    "samples_from_manifest",
     "extract_audio_track",
     "NoAudioStreamError",
 ]

fastvideo/eval/io/inputs.py

@@ -23,6 +23,7 @@ Design rules:
   set samples for corpus-shaped metrics (FVD / FAD) without disturbing
   per-sample paired metrics (LPIPS / PSNR / SSIM / gt_optical_flow).
 """
+
 from __future__ import annotations
 
 import json
@@ -199,6 +200,68 @@ def samples_from(
     return samples
 
 
+def samples_from_manifest(
+    manifest: str | Path,
+    *,
+    extract_audio: bool | str | Path = False,
+    extract_workers: int = 4,
+) -> list[dict]:
+    """Load canonical eval samples from a JSON or JSONL manifest.
+
+    Each row may contain ``video``, ``audio``, ``reference``,
+    ``reference_audio``, ``reference_audio_source``, ``text_prompt`` and
+    arbitrary metric metadata. ``reference_audio_source`` is a media file
+    whose audio track is extracted when ``extract_audio`` is enabled. Paths are
+    resolved relative to the manifest directory. Video paths are wrapped as
+    :class:`Video` handles for lazy decoding; all other keys are preserved.
+    This is the recommended interface for benchmark generation, where matching
+    by sample id is safer than independently sorting folders.
+    """
+    path = Path(manifest).expanduser().resolve()
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    if path.suffix == ".jsonl":
+        rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+    elif path.suffix == ".json":
+        rows = json.loads(path.read_text(encoding="utf-8"))
+        if not isinstance(rows, list):
+            raise ValueError(f"{path}: expected a top-level JSON list")
+    else:
+        raise ValueError(f"{path}: eval manifest must end in .json or .jsonl")
+
+    samples: list[dict[str, Any]] = []
+    for index, row in enumerate(rows):
+        if not isinstance(row, dict):
+            raise ValueError(f"{path}: row {index + 1} must be an object, got {type(row).__name__}")
+        sample = dict(row)
+        for key in (
+                "video",
+                "reference",
+                "audio",
+                "reference_audio",
+                "reference_audio_source",
+        ):
+            value = sample.get(key)
+            if value is None:
+                continue
+            if not isinstance(value, str | Path):
+                raise ValueError(f"{path}: row {index + 1} key {key!r} must be a path")
+            media_path = Path(value).expanduser()
+            if not media_path.is_absolute():
+                media_path = path.parent / media_path
+            media_path = media_path.resolve()
+            sample[key] = as_video(media_path) if key in {"video", "reference"} else str(media_path)
+        if "video" not in sample and "audio" not in sample:
+            raise ValueError(f"{path}: row {index + 1} needs at least video or audio")
+        samples.append(sample)
+
+    if not samples:
+        raise ValueError(f"{path}: eval manifest is empty")
+    if extract_audio:
+        _attach_extracted_audio(samples, extract_audio, extract_workers)
+    return samples
+
+
 # ------------------------------------------------------------------
 # Internal helpers
 # ------------------------------------------------------------------
@@ -280,12 +343,18 @@ def _attach_extracted_audio(
     # Build the work list — (sample_idx, dest_audio_key, source_path)
     work: list[tuple[int, str, str]] = []
     for i, s in enumerate(samples):
-        for vkey, akey in (("video", "audio"), ("reference", "reference_audio")):
+        for vkey, akey in (
+            ("video", "audio"),
+            ("reference", "reference_audio"),
+            ("reference_audio_source", "reference_audio"),
+        ):
             if akey in s:
                 continue
             v = s.get(vkey)
             if isinstance(v, Video) and v.source is not None:
                 work.append((i, akey, str(v.source)))
+            elif vkey == "reference_audio_source" and isinstance(v, str | Path):
+                work.append((i, akey, str(v)))
     if not work:
         return
 

fastvideo/eval/io/video.py

@@ -40,6 +40,29 @@ def extract_frames(video: torch.Tensor, n_frames: int | None = None) -> torch.Te
     return video[indices]
 
 
+def probe_video_fps(source: str | Path) -> float | None:
+    """Read a path-backed video's native frame rate without decoding frames."""
+    try:
+        import av
+
+        with av.open(str(source)) as container:
+            if not container.streams.video:
+                return None
+            stream = container.streams.video[0]
+            rate = stream.average_rate or stream.guessed_rate
+            if rate is not None and float(rate) > 0:
+                return float(rate)
+    except (ImportError, OSError, ValueError):
+        pass
+    try:
+        from decord import VideoReader, cpu
+
+        rate = float(VideoReader(str(source), ctx=cpu(0)).get_avg_fps())
+        return rate if rate > 0 else None
+    except (ImportError, OSError, RuntimeError, ValueError):
+        return None
+
+
 # --- Internal helpers ---
 
 

fastvideo/eval/metrics/audio/frechet_distance/metric.py

@@ -47,7 +47,13 @@ def _frechet_distance(mu1: np.ndarray,
     sigma2 = sigma2 + eps * np.eye(sigma2.shape[0])
 
     diff = mu1 - mu2
-    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    covariance_product = sigma1.dot(sigma2)
+    try:
+        # SciPy < 1.18 returns ``(root, error_estimate)`` with disp=False.
+        covmean, _ = linalg.sqrtm(covariance_product, disp=False)
+    except TypeError:
+        # SciPy 1.18 removed ``disp`` and now returns only the matrix root.
+        covmean = linalg.sqrtm(covariance_product)
     if np.iscomplexobj(covmean):
         if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
             raise ValueError(f"imaginary component {np.max(np.abs(covmean.imag))}")
@@ -108,6 +114,7 @@ class FrechetAudioDistanceMetric(BaseMetric):
     def setup(self) -> None:
         if self._model is None:
             from hear21passt.base import get_basic_model
+
             model = get_basic_model(mode="all")
             model.eval()
             model = model.to(self.device)
@@ -127,13 +134,11 @@ class FrechetAudioDistanceMetric(BaseMetric):
             ref = ref.numpy()
         ref = np.asarray(ref)
         if ref.ndim != 2 or ref.shape[1] != PASST_EMBED_DIM:
-            raise ValueError(f"Expected a 2-D tensor of shape (N, {PASST_EMBED_DIM}) "
-                             f"at {path}; got shape {ref.shape}.")
+            raise ValueError(f"Expected a 2-D tensor of shape (N, {PASST_EMBED_DIM}) at {path}; got shape {ref.shape}.")
         finite = np.isfinite(ref).all(axis=1)
         ref = ref[finite]
         if ref.shape[0] < 2:
-            raise ValueError(f"Cached reference features at {path} have only "
-                             f"{ref.shape[0]} finite rows; need >= 2.")
+            raise ValueError(f"Cached reference features at {path} have only {ref.shape[0]} finite rows; need >= 2.")
         self._cached_ref_mu = ref.mean(axis=0)
         self._cached_ref_sigma = np.cov(ref, rowvar=False)
         self._n_cached_ref = int(ref.shape[0])
@@ -209,8 +214,7 @@ class FrechetAudioDistanceMetric(BaseMetric):
                 name=self.name,
                 score=None,
                 details={
-                    "skipped": f"FAD needs >=2 finite-embed gen samples "
-                    f"(got {n_gen} valid of {len(self._gen_buf)})",
+                    "skipped": f"FAD needs >=2 finite-embed gen samples (got {n_gen} valid of {len(self._gen_buf)})",
                     "n_gen": n_gen,
                     "n_ref": n_ref,
                     "n_gen_dropped_nonfinite": n_gen_dropped,

fastvideo/eval/pool.py

@@ -9,6 +9,7 @@ write back into a result list in input order.
 
 Pool sizing: ``max_size = prefetch_factor * num_workers``.
 """
+
 from __future__ import annotations
 
 import queue
@@ -144,11 +145,16 @@ class VideoPool:
         :func:`fastvideo.eval.io.as_video` / :func:`samples_from`) to
         make them decode-eligible.
         """
-        from fastvideo.eval.io.video import load_video
+        from fastvideo.eval.io.video import load_video, probe_video_fps
 
         out = dict(sample)
         for key, val in sample.items():
             if isinstance(val, Video) and val.frames is None and val.source is not None:
+                if key == "video" and "fps" not in out:
+                    if val.fps is None:
+                        val.fps = probe_video_fps(val.source)
+                    if val.fps is not None:
+                        out["fps"] = float(val.fps)
                 val.frames = load_video(val.source)
                 out[key] = val
         return out

fastvideo/eval/v2a/__init__.py

@@ -0,0 +1,5 @@
+"""Video-to-audio corpus evaluation backends."""
+
+from fastvideo.eval.v2a.av_benchmark import run_av_benchmark
+
+__all__ = ["run_av_benchmark"]

fastvideo/eval/v2a/_av_benchmark_runner.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Standalone process runner for the official av-benchmark API.
+
+This file is executed directly by an interpreter from an isolated environment.
+Keep it free of ``fastvideo`` imports so the benchmark environment does not
+need FastVideo's runtime dependency stack.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from importlib import metadata
+from pathlib import Path
+from typing import Any
+from collections.abc import Callable
+
+_BASE_CACHE_FILES = (
+    "pann_features.pth",
+    "vggish_features.pth",
+    "passt_features_embed.pth",
+    "passt_logits.pth",
+)
+_VIDEO_CACHE_FILES = (
+    "imagebind_audio.pth",
+    "synchformer_audio.pth",
+)
+_CLAP_CACHE_FILES = (
+    "clap_laion_audio.pth",
+    "clap_ms_audio.pth",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--audio-dir", type=Path, required=True)
+    parser.add_argument("--gt-cache", type=Path, required=True)
+    parser.add_argument("--prediction-cache", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--num-workers", type=int, default=0)
+    parser.add_argument("--audio-length", type=float, default=8.0)
+    parser.add_argument("--recompute", action="store_true")
+    parser.add_argument("--skip-video-related", action="store_true")
+    parser.add_argument("--skip-clap", action="store_true")
+    parser.add_argument("--align-prediction-keys", action="store_true")
+    return parser.parse_args()
+
+
+def _load_backend() -> tuple[Callable[..., Any], Callable[..., Any]]:
+    try:
+        from av_bench.evaluate import evaluate
+        from av_bench.extract import extract
+    except ImportError as error:
+        raise ImportError("Official av_bench is not installed in this Python environment. "
+                          "Install https://github.com/hkchengrex/av-benchmark and pass its "
+                          "Python through --python-executable.") from error
+    return extract, evaluate
+
+
+def _required_prediction_files(*, skip_video_related: bool, skip_clap: bool) -> tuple[str, ...]:
+    required = list(_BASE_CACHE_FILES)
+    if not skip_video_related:
+        required.extend(_VIDEO_CACHE_FILES)
+    if not skip_clap:
+        required.extend(_CLAP_CACHE_FILES)
+    return tuple(required)
+
+
+def _validate_gt_cache(gt_cache: Path) -> None:
+    missing = [name for name in _BASE_CACHE_FILES if not (gt_cache / name).is_file()]
+    if missing:
+        raise FileNotFoundError(f"Ground-truth cache {gt_cache} is missing required files: {missing}")
+
+
+def _jsonable(value: Any) -> Any:
+    if isinstance(value, Path):
+        return str(value)
+    if hasattr(value, "detach"):
+        value = value.detach().cpu()
+    if hasattr(value, "item"):
+        try:
+            return value.item()
+        except ValueError:
+            pass
+    if hasattr(value, "tolist"):
+        return value.tolist()
+    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
+
+
+def _align_prediction_cache_keys(
+    *,
+    gt_cache: Path,
+    prediction_cache: Path,
+    required_files: tuple[str, ...],
+) -> dict[str, int]:
+    """Align filename-sanitized prediction keys to unique GT cache keys.
+
+    Official VGGSound caches prefix one or more underscores to some stems.
+    Generated audio keeps the original stem, so paired metrics otherwise drop
+    those samples silently. Only unique leading-underscore matches are changed.
+    """
+    import torch
+
+    gt_logits = torch.load(gt_cache / "passt_logits.pth", map_location="cpu", weights_only=True)
+    pred_logits = torch.load(prediction_cache / "passt_logits.pth", map_location="cpu", weights_only=True)
+    if not isinstance(gt_logits, dict) or not isinstance(pred_logits, dict):
+        raise TypeError("PaSST caches must be dictionaries keyed by sample id")
+
+    gt_keys = {str(key) for key in gt_logits}
+    pred_keys = {str(key) for key in pred_logits}
+    candidates_by_stem: dict[str, list[str]] = {}
+    for key in gt_keys:
+        candidates_by_stem.setdefault(key.lstrip("_"), []).append(key)
+
+    occupied = pred_keys & gt_keys
+    mapping: dict[str, str] = {}
+    ambiguous = 0
+    for key in sorted(pred_keys - gt_keys):
+        candidates = [
+            candidate for candidate in candidates_by_stem.get(key.lstrip("_"), []) if candidate not in occupied
+        ]
+        if len(candidates) == 1:
+            mapping[key] = candidates[0]
+            occupied.add(candidates[0])
+        elif len(candidates) > 1:
+            ambiguous += 1
+
+    if mapping:
+        for filename in required_files:
+            path = prediction_cache / filename
+            values = torch.load(path, map_location="cpu", weights_only=True)
+            if not isinstance(values, dict):
+                raise TypeError(f"Prediction cache must be a dictionary: {path}")
+            aligned = {mapping.get(str(key), str(key)): value for key, value in values.items()}
+            if len(aligned) != len(values):
+                raise RuntimeError(f"Prediction key alignment created a collision in {path}")
+            temporary = path.with_suffix(path.suffix + ".tmp")
+            torch.save(aligned, temporary)
+            temporary.replace(path)
+
+    return {
+        "exact_before_alignment": len(pred_keys & gt_keys),
+        "remapped": len(mapping),
+        "ambiguous": ambiguous,
+        "unmatched_after_alignment": len(pred_keys) - len(pred_keys & gt_keys) - len(mapping),
+    }
+
+
+def run_benchmark(
+    *,
+    audio_dir: Path,
+    gt_cache: Path,
+    prediction_cache: Path,
+    output: Path,
+    device: str,
+    batch_size: int,
+    num_workers: int,
+    audio_length: float,
+    recompute: bool,
+    skip_video_related: bool,
+    skip_clap: bool,
+    align_prediction_keys: bool = False,
+) -> dict[str, Any]:
+    audio_dir = audio_dir.expanduser().resolve()
+    gt_cache = gt_cache.expanduser().resolve()
+    prediction_cache = prediction_cache.expanduser().resolve()
+    output = output.expanduser().resolve()
+    if not audio_dir.is_dir():
+        raise FileNotFoundError(f"Generated audio directory does not exist: {audio_dir}")
+    if not gt_cache.is_dir():
+        raise FileNotFoundError(f"Ground-truth cache directory does not exist: {gt_cache}")
+    if batch_size <= 0:
+        raise ValueError(f"batch_size must be positive, got {batch_size}")
+    if num_workers < 0:
+        raise ValueError(f"num_workers must be non-negative, got {num_workers}")
+    if audio_length <= 0:
+        raise ValueError(f"audio_length must be positive, got {audio_length}")
+
+    audio_files = sorted((*audio_dir.glob("*.wav"), *audio_dir.glob("*.flac")))
+    if not audio_files:
+        raise FileNotFoundError(f"No WAV or FLAC files found in {audio_dir}")
+    _validate_gt_cache(gt_cache)
+    extract, evaluate = _load_backend()
+
+    required = _required_prediction_files(
+        skip_video_related=skip_video_related,
+        skip_clap=skip_clap,
+    )
+    cache_complete = all((prediction_cache / name).is_file() for name in required)
+    extraction_ran = recompute or not cache_complete
+    if extraction_ran:
+        prediction_cache.mkdir(parents=True, exist_ok=True)
+        extract(
+            audio_path=audio_dir,
+            output_path=prediction_cache,
+            device=device,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            audio_length=audio_length,
+            skip_video_related=skip_video_related,
+            skip_clap=skip_clap,
+        )
+
+    key_alignment = None
+    if align_prediction_keys:
+        key_alignment = _align_prediction_cache_keys(
+            gt_cache=gt_cache,
+            prediction_cache=prediction_cache,
+            required_files=required,
+        )
+
+    metrics = evaluate(
+        gt_audio_cache=gt_cache,
+        pred_audio_cache=prediction_cache,
+        skip_video_related=skip_video_related,
+        skip_clap=skip_clap,
+    )
+    try:
+        backend_version = metadata.version("av_bench")
+    except metadata.PackageNotFoundError:
+        backend_version = "editable-or-unknown"
+    payload = {
+        "backend": "hkchengrex/av-benchmark",
+        "backend_version": backend_version,
+        "audio_dir": str(audio_dir),
+        "gt_cache": str(gt_cache),
+        "prediction_cache": str(prediction_cache),
+        "num_generated_audio": len(audio_files),
+        "audio_length": audio_length,
+        "extraction_ran": extraction_ran,
+        "skip_video_related": skip_video_related,
+        "skip_clap": skip_clap,
+        "prediction_key_alignment": key_alignment,
+        "metrics": metrics,
+    }
+    output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = output.with_suffix(output.suffix + ".tmp")
+    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True, default=_jsonable) + "\n", encoding="utf-8")
+    temporary.replace(output)
+    return payload
+
+
+def main() -> None:
+    args = parse_args()
+    payload = run_benchmark(
+        audio_dir=args.audio_dir,
+        gt_cache=args.gt_cache,
+        prediction_cache=args.prediction_cache,
+        output=args.output,
+        device=args.device,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        audio_length=args.audio_length,
+        recompute=args.recompute,
+        skip_video_related=args.skip_video_related,
+        skip_clap=args.skip_clap,
+        align_prediction_keys=args.align_prediction_keys,
+    )
+    print(json.dumps(payload, indent=2, sort_keys=True, default=_jsonable))
+
+
+if __name__ == "__main__":
+    main()

fastvideo/eval/v2a/av_benchmark.py

@@ -0,0 +1,97 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Isolated launcher for the official ``hkchengrex/av-benchmark`` backend.
+
+The benchmark has a large dependency stack that can conflict with FastVideo's
+runtime. This module intentionally does not import ``av_bench``. It launches a
+small runner with a caller-selected Python interpreter, so exact official
+metrics can live in a separate virtual environment.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def run_av_benchmark(
+    *,
+    audio_dir: str | Path,
+    gt_cache: str | Path,
+    prediction_cache: str | Path | None = None,
+    output: str | Path | None = None,
+    python_executable: str | Path | None = None,
+    device: str = "cuda",
+    batch_size: int = 32,
+    num_workers: int = 0,
+    audio_length: float = 8.0,
+    recompute: bool = False,
+    skip_video_related: bool = False,
+    skip_clap: bool = False,
+    align_prediction_keys: bool = False,
+) -> Path:
+    """Extract prediction features and evaluate with official av-benchmark.
+
+    ``python_executable`` should point to an isolated environment containing
+    the official ``av_bench`` package. Omitting it uses the current process,
+    which is useful only when that package is already installed.
+    """
+    audio_path = Path(audio_dir).expanduser().resolve()
+    gt_path = Path(gt_cache).expanduser().resolve()
+    if not audio_path.is_dir():
+        raise FileNotFoundError(f"Generated audio directory does not exist: {audio_path}")
+    if not gt_path.is_dir():
+        raise FileNotFoundError(f"Ground-truth cache directory does not exist: {gt_path}")
+    if batch_size <= 0:
+        raise ValueError(f"batch_size must be positive, got {batch_size}")
+    if num_workers < 0:
+        raise ValueError(f"num_workers must be non-negative, got {num_workers}")
+    if audio_length <= 0:
+        raise ValueError(f"audio_length must be positive, got {audio_length}")
+
+    root = audio_path.parent
+    prediction_path = (Path(prediction_cache).expanduser().resolve() if prediction_cache is not None else root /
+                       "av_benchmark_cache")
+    output_path = (Path(output).expanduser().resolve() if output is not None else root / "av_benchmark_results.json")
+    runner = Path(__file__).with_name("_av_benchmark_runner.py")
+    python = str(Path(python_executable).expanduser()) if python_executable is not None else sys.executable
+
+    command = [
+        python,
+        str(runner),
+        "--audio-dir",
+        str(audio_path),
+        "--gt-cache",
+        str(gt_path),
+        "--prediction-cache",
+        str(prediction_path),
+        "--output",
+        str(output_path),
+        "--device",
+        device,
+        "--batch-size",
+        str(batch_size),
+        "--num-workers",
+        str(num_workers),
+        "--audio-length",
+        str(audio_length),
+    ]
+    if recompute:
+        command.append("--recompute")
+    if skip_video_related:
+        command.append("--skip-video-related")
+    if skip_clap:
+        command.append("--skip-clap")
+    if align_prediction_keys:
+        command.append("--align-prediction-keys")
+
+    logger.info("Launching official av-benchmark with isolated Python: %s", python)
+    logger.info("Prediction cache: %s", prediction_path)
+    subprocess.run(command, check=True)
+    if not output_path.is_file():
+        raise RuntimeError(f"av-benchmark completed without writing expected output: {output_path}")
+    return output_path

fastvideo/models/audio/__init__.py

@@ -5,5 +5,15 @@ from fastvideo.models.audio.ltx2_audio_vae import (
     LTX2AudioEncoder,
     LTX2Vocoder,
 )
+from fastvideo.models.audio.mmaudio_processing import (
+    MMAudioMelConverter,
+    build_mmaudio_mel_converter,
+)
 
-__all__ = ["LTX2AudioEncoder", "LTX2AudioDecoder", "LTX2Vocoder"]
+__all__ = [
+    "LTX2AudioEncoder",
+    "LTX2AudioDecoder",
+    "LTX2Vocoder",
+    "MMAudioMelConverter",
+    "build_mmaudio_mel_converter",
+]

fastvideo/models/audio/mmaudio_processing.py

@@ -0,0 +1,103 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Audio preprocessing used by native MMAudio training feature extraction."""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import torch
+import torch.nn as nn
+
+
+class MMAudioMelConverter(nn.Module):
+    """Convert waveforms to MMAudio's log-mel representation.
+
+    The FFT, padding, mel filter, and logarithm match the published MMAudio
+    preprocessing contract. Keeping this component in FastVideo avoids a
+    runtime import from the upstream training repository.
+    """
+
+    def __init__(
+        self,
+        *,
+        sampling_rate: int,
+        n_fft: int,
+        num_mels: int,
+        hop_size: int,
+        win_size: int,
+        fmin: float,
+        fmax: float,
+        log_base: Literal["e", "10"],
+    ) -> None:
+        super().__init__()
+        from librosa.filters import mel as librosa_mel_fn
+
+        mel = librosa_mel_fn(
+            sr=sampling_rate,
+            n_fft=n_fft,
+            n_mels=num_mels,
+            fmin=fmin,
+            fmax=fmax,
+        )
+        self.n_fft = n_fft
+        self.hop_size = hop_size
+        self.win_size = win_size
+        self.log_base = log_base
+        self.register_buffer("mel_basis", torch.from_numpy(mel).float())
+        self.register_buffer("hann_window", torch.hann_window(win_size))
+
+    def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor:
+        if waveform.ndim != 2:
+            raise ValueError(f"MMAudio waveforms must have shape [batch, samples], got {tuple(waveform.shape)}")
+        waveform = waveform.clamp(min=-1.0, max=1.0).to(self.mel_basis.device)
+        padding = int((self.n_fft - self.hop_size) / 2)
+        waveform = torch.nn.functional.pad(
+            waveform.unsqueeze(1),
+            [padding, padding],
+            mode="reflect",
+        ).squeeze(1)
+        spectrum = torch.stft(
+            waveform,
+            self.n_fft,
+            hop_length=self.hop_size,
+            win_length=self.win_size,
+            window=self.hann_window,
+            center=center,
+            pad_mode="reflect",
+            normalized=False,
+            onesided=True,
+            return_complex=True,
+        )
+        magnitude = torch.sqrt(torch.view_as_real(spectrum).pow(2).sum(-1) + 1e-9)
+        mel = torch.matmul(self.mel_basis, magnitude)
+        mel = torch.clamp(mel, min=1e-5)
+        return torch.log10(mel) if self.log_base == "10" else torch.log(mel)
+
+
+def build_mmaudio_mel_converter(mode: Literal["16k", "44k"] = "44k", ) -> MMAudioMelConverter:
+    if mode == "16k":
+        return MMAudioMelConverter(
+            sampling_rate=16_000,
+            n_fft=1024,
+            num_mels=80,
+            hop_size=256,
+            win_size=1024,
+            fmin=0,
+            fmax=8_000,
+            log_base="10",
+        )
+    if mode == "44k":
+        return MMAudioMelConverter(
+            sampling_rate=44_100,
+            n_fft=2048,
+            num_mels=128,
+            hop_size=512,
+            win_size=2048,
+            fmin=0,
+            fmax=44_100 / 2,
+            log_base="e",
+        )
+    raise ValueError(f"Unknown MMAudio mel mode: {mode}")
+
+
+__all__ = ["MMAudioMelConverter", "build_mmaudio_mel_converter"]

fastvideo/models/dits/mmaudio.py

@@ -10,6 +10,7 @@ separately verified optimization.
 
 from __future__ import annotations
 
+import math
 from dataclasses import dataclass
 from typing import Any
 
@@ -282,7 +283,15 @@ class MMAudioTransformer(BaseDiT):
     param_names_mapping = _DEFAULT_CONFIG.arch_config.param_names_mapping
     reverse_param_names_mapping = _DEFAULT_CONFIG.arch_config.reverse_param_names_mapping
 
-    def __init__(self, config: MMAudioTransformerConfig, hf_config: dict[str, Any], **kwargs) -> None:
+    def __init__(
+        self,
+        config: MMAudioTransformerConfig,
+        hf_config: dict[str, Any],
+        latent_mean: torch.Tensor | None = None,
+        latent_std: torch.Tensor | None = None,
+        empty_string_feat: torch.Tensor | None = None,
+        **kwargs,
+    ) -> None:
         del kwargs
         super().__init__(config=config, hf_config=hf_config)
         arch = config.arch_config
@@ -348,9 +357,57 @@ class MMAudioTransformer(BaseDiT):
             ]
         )
 
-        self.latent_mean = nn.Parameter(torch.full((1, 1, arch.latent_dim), float("nan")), requires_grad=False)
-        self.latent_std = nn.Parameter(torch.full((1, 1, arch.latent_dim), float("nan")), requires_grad=False)
-        self.empty_string_feat = nn.Parameter(torch.zeros((arch.text_seq_len, arch.text_dim)), requires_grad=False)
+        def fixed_tensor(
+            value: torch.Tensor | None,
+            shape: tuple[int, ...],
+            *,
+            fill_value: float,
+            name: str,
+        ) -> torch.Tensor:
+            if value is None:
+                return torch.full(
+                    shape,
+                    fill_value,
+                    device=self.sync_pos_emb.device,
+                    dtype=self.sync_pos_emb.dtype,
+                )
+            if value.numel() != math.prod(shape):
+                raise ValueError(
+                    f"MMAudio {name} must have {math.prod(shape)} values, "
+                    f"got {value.numel()}."
+                )
+            return value.to(
+                device=self.sync_pos_emb.device,
+                dtype=self.sync_pos_emb.dtype,
+            ).reshape(shape).clone()
+
+        self.latent_mean = nn.Parameter(
+            fixed_tensor(
+                latent_mean,
+                (1, 1, arch.latent_dim),
+                fill_value=float("nan"),
+                name="latent_mean",
+            ),
+            requires_grad=False,
+        )
+        self.latent_std = nn.Parameter(
+            fixed_tensor(
+                latent_std,
+                (1, 1, arch.latent_dim),
+                fill_value=float("nan"),
+                name="latent_std",
+            ),
+            requires_grad=False,
+        )
+        self.empty_string_feat = nn.Parameter(
+            fixed_tensor(
+                empty_string_feat,
+                (arch.text_seq_len, arch.text_dim),
+                fill_value=0.0,
+                name="empty_string_feat",
+            ),
+            requires_grad=False,
+        )
         self.empty_clip_feat = nn.Parameter(torch.zeros(1, arch.clip_dim), requires_grad=True)
         self.empty_sync_feat = nn.Parameter(torch.zeros(1, arch.sync_dim), requires_grad=True)
 
@@ -421,7 +478,12 @@ class MMAudioTransformer(BaseDiT):
         self._latent_seq_len = latent_seq_len
         self._clip_seq_len = clip_seq_len
         self._sync_seq_len = sync_seq_len
-        self.initialize_rotations()
+        # This method is also called by inference pipeline stages. Persistent
+        # buffers created under inference_mode become inference tensors and
+        # cannot later be saved for backward if a live transformer resumes
+        # training. Build the derived buffers as ordinary no-grad tensors.
+        with torch.inference_mode(False), torch.no_grad():
+            self.initialize_rotations()
 
     def initialize_weights(self) -> None:
 

fastvideo/models/loader/component_loader.py

@@ -1066,7 +1066,15 @@ class AudioDecoderLoader(ComponentLoader):
             audio_vae.load_state_dict(loaded, strict=True)
             return audio_vae.eval()
 
-        precision = getattr(fastvideo_args.pipeline_config, "audio_decoder_precision", "bf16")
+        # A preprocessing component may carry the MMAudio encoder in addition
+        # to the inference decoder. Feature extraction is defined in fp32 by
+        # the published recipe; do not quantize it to decoder inference dtype.
+        needs_encoder = bool(config.get("need_encoder", False))
+        precision_name = (
+            "audio_encoder_precision" if needs_encoder else
+            "audio_decoder_precision")
+        precision = getattr(fastvideo_args.pipeline_config, precision_name,
+                            "fp32" if needs_encoder else "bf16")
         # MMAudio normalizes its magnitude-preserving convolution weights in
         # fp32 and only then casts the whole feature utility module to bf16.
         # Constructing/loading directly in bf16 quantizes the unnormalized

fastvideo/models/loader/fsdp_load.py

@@ -494,6 +494,75 @@ def maybe_load_fsdp_model(
     return model
 
 
+def build_fsdp_model_from_scratch(
+    model_cls: type[nn.Module],
+    init_params: dict[str, Any],
+    device: torch.device,
+    hsdp_replicate_dim: int,
+    hsdp_shard_dim: int,
+    default_dtype: torch.dtype,
+    param_dtype: torch.dtype,
+    reduce_dtype: torch.dtype,
+    *,
+    seed: int,
+    cpu_offload: bool = False,
+    output_dtype: torch.dtype | None = None,
+    pin_cpu_memory: bool = True,
+) -> torch.nn.Module:
+    """Deterministically initialize and FSDP-shard a model without weights."""
+
+    mp_policy = MixedPrecisionPolicy(
+        param_dtype,
+        reduce_dtype,
+        output_dtype,
+        cast_forward_inputs=False,
+    )
+    set_mixed_precision_policy(
+        param_dtype=param_dtype,
+        reduce_dtype=reduce_dtype,
+        output_dtype=output_dtype,
+        mp_policy=mp_policy,
+    )
+
+    fork_devices: list[int] = []
+    if device.type == "cuda":
+        fork_devices = [
+            device.index if device.index is not None else torch.cuda.current_device()
+        ]
+    with torch.random.fork_rng(devices=fork_devices):
+        torch.manual_seed(int(seed))
+        with set_default_dtype(default_dtype), torch.device(device):
+            model = model_cls(**init_params)
+
+    from fastvideo.platforms import current_platform
+
+    if not current_platform.is_mps():
+        pin_cpu_memory = pin_cpu_memory and is_pin_memory_available()
+        if current_platform.is_npu():
+            device_mesh = init_device_mesh(
+                "npu",
+                mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
+                mesh_dim_names=("replicate", "shard"),
+            )
+        else:
+            device_mesh = init_device_mesh(
+                "cuda",
+                mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
+                mesh_dim_names=("replicate", "shard"),
+            )
+        shard_model(
+            model,
+            cpu_offload=cpu_offload,
+            reshard_after_forward=True,
+            mp_policy=mp_policy,
+            mesh=device_mesh,
+            fsdp_shard_conditions=model._fsdp_shard_conditions,
+            pin_cpu_memory=pin_cpu_memory,
+        )
+
+    for parameter in model.parameters():
+        parameter.requires_grad = False
+    return model
 def _regional_compile_unsupported_reason(
     init_params: dict[str, Any],
     *,

fastvideo/pipelines/basic/mmaudio/README.md

@@ -1,8 +1,9 @@
-# MMAudio Video-to-Audio Inference
+# MMAudio Video-to-Audio
 
 This directory contains FastVideo's native `MMAudioPipeline` implementation for
-video-to-audio (V2A) and text-to-audio (T2A) generation. The first supported
-checkpoint is `large_44k_v2`, which produces mono audio at 44.1 kHz.
+video-to-audio (V2A) and text-to-audio (T2A) generation. All five published
+MMAudio model architectures are registered independently so their sequence
+length, latent width, sample rate, VAE, and vocoder cannot be mixed silently.
 
 The production pipeline is implemented with FastVideo components and does not
 import the upstream `mmaudio` Python package. It loads the transformer, DFN5B
@@ -10,20 +11,41 @@ CLIP text/vision encoders, Synchformer visual encoder, audio VAE, BigVGAN-v2,
 tokenizer, and scheduler from a standard FastVideo/Diffusers-style checkpoint
 directory.
 
+## Official model variants
+
+| Variant | Audio | Hidden / layers / heads | Inference | Official training recipe |
+| --- | ---: | --- | --- | --- |
+| `small_16k` | 16 kHz | 448 / 12 / 7 | implemented | supported |
+| `small_44k` | 44.1 kHz | 448 / 12 / 7 | implemented | supported |
+| `medium_44k` | 44.1 kHz | 896 / 12 / 14 | implemented | supported |
+| `large_44k` | 44.1 kHz | 896 / 21 / 14 | implemented | supported |
+| `large_44k_v2` | 44.1 kHz | 896 / 21 / 14 | implemented and validated | not supported upstream |
+
+`small_16k` uses latent shape `[250, 20]`, the 16 kHz VAE, and the original
+80-mel BigVGAN. The 44.1 kHz variants use latent shape `[345, 40]`, the 44 kHz
+VAE, and 128-mel BigVGAN-v2. `_v2` changes transformer activations and timestep
+conditioning but not the large model dimensions.
+
 ## Status
 
-- `large_44k_v2` V2A inference: supported
+- Five variant-specific presets, pipeline configs, registry entries, and
+  converter choices: supported
 - T2A pipeline routing: supported; the provided example currently targets V2A
-- Output: mono WAV, 44.1 kHz
+- Output: mono WAV at 16 kHz or 44.1 kHz, according to the checkpoint
 - Default inference duration: 8 seconds
 - Variable-duration inference: supported
 - Single-GPU inference with FastVideo's default offloading: supported
+- `small_16k`, `small_44k`, `medium_44k`, and `large_44k` from-scratch training:
+  supported through FastVideo's modular trainer and official-compatible features
 - Source-video/audio muxing: not yet part of the pipeline output
-- Training and small/medium/16 kHz variants: not included in this port
+- `large_44k_v2` training: intentionally rejected because the upstream recipe
+  does not support training the v2 checkpoint
 
 The native pipeline has passed exact official-vs-FastVideo real-weight parity
 for a 25-step two-second waveform and a real ten-second V2A inference smoke
-test.
+test with `large_44k_v2`. The other four variants have architecture/config
+coverage; real-weight waveform parity still requires converting their official
+checkpoints.
 
 ## Requirements
 
@@ -49,6 +71,12 @@ source .venv/bin/activate
 UV_TORCH_BACKEND=cu126 uv pip install -e .
 ```
 
+Training additionally needs TensorDict for the memory-mapped feature cache:
+
+```bash
+UV_TORCH_BACKEND=cu126 uv pip install -e '.[mmaudio-train]'
+```
+
 Use `UV_TORCH_BACKEND=cu130` instead on CUDA 13. Conda is not required.
 
 ## Checkpoint Layout
@@ -149,6 +177,352 @@ python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \
 
 The converted checkpoint is approximately 9 GB.
 
+## Train MMAudio v1
+
+Training uses the native `fastvideo/train/` stack. FastVideo owns the training
+loop, HSDP/FSDP wrapping, optimizer and LR scheduler, gradient accumulation,
+gradient clipping, tracking, and distributed checkpoints. The reusable
+`FlowMatchingFineTuneMethod` implements shape-agnostic velocity supervision;
+`MMAudioModel` supplies MMAudio's audio posterior sampling, latent
+normalization, independent video/text CFG dropout, and multimodal forward.
+
+The first smoke-test target should be `small_44k`. The upstream documentation
+does not support training checkpoints ending in `_v2`, so the inference
+checkpoint `large_44k_v2` must not be used here.
+
+### Convert a trainable v1 transformer
+
+Download the official v1 checkpoint:
+
+```bash
+mkdir -p official_weights/mmaudio/raw/weights
+curl -L --continue-at - \
+  https://huggingface.co/hkchengrex/MMAudio/resolve/main/weights/mmaudio_small_44k.pth \
+  -o official_weights/mmaudio/raw/weights/mmaudio_small_44k.pth
+```
+
+Training only loads the transformer, so it does not need another copy of the
+VAE, DFN5B, Synchformer, or vocoder components:
+
+```bash
+python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \
+  --variant small_44k \
+  --transformer-only \
+  --transformer-checkpoint official_weights/mmaudio/raw/weights/mmaudio_small_44k.pth \
+  --output converted_weights/mmaudio/small_44k
+```
+
+Replace `small_44k` with `medium_44k` or `large_44k` in both arguments to train
+a larger v1 model. `small_16k` is also supported, but requires 16 kHz features.
+
+### Train the DiT from scratch
+
+From scratch means all trainable MMAudio transformer weights are randomly
+initialized. The VAE, DFN5B CLIP, and Synchformer remain pretrained feature
+extractors, as in the official recipe. The fixed CLIP embedding for an empty
+string is loaded either from upstream `empty_string.pth` or a converted official
+checkpoint, but is not trainable and does not initialize any DiT layer. To avoid
+downloading a DiT checkpoint for a pure scratch run:
+
+```bash
+curl -L --continue-at - \
+  https://github.com/hkchengrex/MMAudio/releases/download/v0.1/empty_string.pth \
+  -o official_weights/mmaudio/raw/ext_weights/empty_string.pth
+```
+
+Four configs mirror the official v1 architectures:
+
+```text
+examples/train/configs/fine_tuning/mmaudio/small_16k_from_scratch.yaml
+examples/train/configs/fine_tuning/mmaudio/small_44k_from_scratch.yaml
+examples/train/configs/fine_tuning/mmaudio/medium_44k_from_scratch.yaml
+examples/train/configs/fine_tuning/mmaudio/large_44k_from_scratch.yaml
+```
+
+They use AdamW (`1e-4`, betas `0.9/0.95`, epsilon `1e-6`, weight decay
+`1e-6`), 1,000 warmup steps, MultiStepLR drops at 240k and 270k, gradient clip
+1.0, and 300k steps. The FSDP-aware `PostHocEMACallback` mirrors upstream
+MMAudio's `nitrous_ema` settings: sigma profiles `[0.05, 0.1]`, an update every
+step, snapshots every 5,000 steps, and final synthesis at sigma `0.05`. It is an
+additional callback; FastVideo's existing conventional EMA remains unchanged
+for other models.
+
+Edit dataset paths, variant, batch size, learning rate, validation intervals,
+output directory, resume path, and tracker at the top of the launcher. No
+shell environment variables or command-line overrides are required. The
+launcher also downloads the official standalone `empty_string.pth` if it is
+missing. Then run the current 44.1 kHz VGGSound cache on four GPUs:
+
+```bash
+bash examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh
+```
+
+### Optional full-replica DDP for `small_44k`
+
+MMAudio's upstream runner uses native PyTorch DDP with one complete model per
+GPU and `broadcast_buffers=False`. FastVideo exposes that behavior as the
+opt-in `training.distributed.strategy: ddp`; the default remains `fsdp`, so
+existing model configs and the launcher above are unchanged. The initial DDP
+plugin is intentionally limited to MMAudio from-scratch training.
+
+The independent four-GPU config and launcher are:
+
+```text
+examples/train/configs/fine_tuning/mmaudio/small_44k_ddp_from_scratch.yaml
+examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh
+```
+
+Run it only after the current GPU job has finished:
+
+```bash
+bash examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh
+```
+
+The script uses port `29502` and writes to
+`outputs/mmaudio_small_44k_ddp_from_scratch`, so it does not reuse the FSDP
+job's rendezvous or output directory. DDP can be faster for this small model
+when every GPU can hold the full model, optimizer, activations, and rank 0's
+two EMA copies. The DDP config also enables upstream's fused AdamW. It uses
+more per-GPU model memory than FSDP.
+
+The DDP config also enables `training.model.compile_train_fn: true`.
+Following upstream MMAudio, compilation happens after DDP wrapping. For
+PyTorch 2.12 compatibility, FastVideo keeps explicit-Generator posterior,
+flow-time, prior, and CFG-mask sampling in eager code, applies masks with
+broadcasted `torch.where`, and compiles the inner transformer's tensor-only
+forward with `fullgraph: true`. The outer DDP reducer/control flow and the
+lightweight MSE loss remain eager. This preserves the official RNG order while
+making any graph break inside the expensive transformer graph a hard error
+instead of a silent fallback.
+
+Set `COMPILE_TRAIN_FN=false` near the top of the DDP launcher to disable it.
+The typed default is `false`, and the option is currently rejected for FSDP,
+so existing FastVideo pipelines are unaffected. Validation loss reuses the
+compiled transformer forward; periodic generation keeps its separate eager
+`guided_flow` path. Optional `torch.compile` arguments can be placed under
+`training.model.torch_compile_kwargs`.
+
+For parity with upstream MMAudio, rank 0 runs the exact
+`nitrous_ema.PostHocEMA` implementation with sigma profiles `[0.05, 0.1]`,
+updates every optimizer step, and snapshots every 5,000 steps. FastVideo's
+distributed checkpoint interval is 10,000 steps so every resumable training
+checkpoint has a matching EMA snapshot. Periodic validation/inference still
+uses online weights, and the final sigma-0.05 synthesis is saved under
+`posthoc_ema/official_ddp/`.
+
+Set `VARIANT` to `medium_44k` or `large_44k` without changing the feature
+cache. The launcher computes official latent normalization statistics once,
+saves `latent_statistics_44k.pt` inside the feature directory, and reuses it
+across those three variants. `small_16k` needs a separately preprocessed 16 kHz
+cache because its latent shape is `[250, 20]`. `large_44k_v2` is rejected before
+weights or data are loaded because the published training implementation does
+not support v2.
+
+### Validation and training-time inference
+
+`MMAudioValidationCallback` evaluates the cached validation split every 5,000
+steps by default. It follows the official validation loss construction:
+posterior latent sampling, logit-normal flow time, prior noise, and independent
+video/text CFG masks. The validation RNG is reset on every pass so losses are
+comparable between checkpoints. Set `VALIDATION_MAX_BATCHES=0` in the launcher
+to use the full split, or a positive value for a faster subset.
+
+Every 20,000 steps the same callback runs the native FastVideo MMAudio pipeline
+with the live training transformer and cached CLIP, Synchformer, and text
+features. It generates a fixed global set of 16 validation outputs (four per
+rank on a four-GPU job) spanning the validation batch's source rows, not the
+complete validation split. The validation seed,
+sampler, and per-sample noise seeds stay fixed across training steps.
+`INFERENCE_MODEL_PATH` supplies only the frozen scheduler, audio VAE, and
+vocoder needed to decode a waveform; its transformer weights and feature
+encoders are not loaded. Samples are written to:
+
+```text
+outputs/mmaudio_<variant>_from_scratch/validation_audio/step_<step>/
+```
+
+Both a standalone WAV and an MP4 containing the original validation video plus
+generated audio are saved. MP4 composition mirrors official MMAudio demo's
+`reencode_with_audio`: decode source frames through the requested endpoint,
+re-encode H.264 at the source's guessed frame rate and 10 Mbps in `yuv420p`,
+then feed the mono generated waveform to the AAC encoder at the model sample
+rate. It uses PyAV and does not import Torio. As in the official 44.1 kHz path,
+the encoded container can be slightly longer than 8 seconds because 353,280
+audio samples and AAC frame padding do not end at exactly 8.0 seconds. If
+composition or source-video loading fails, the WAV is retained and training
+continues.
+
+Post-hoc EMA rank-local profiles and the final synthesized sigma-0.05 shards
+are written under `posthoc_ema/rank_<rank>/`. Periodic validation uses online
+weights by default, matching the official training loop; official MMAudio only
+synthesizes sigma-0.05 EMA weights after training. Set `use_ema: true` in the
+YAML only when intentionally evaluating an available post-hoc snapshot. The
+current callback reports flow-matching validation loss and generates media
+samples; dataset-level FAD/CLAP/synchronization benchmark metrics remain a
+separate evaluation job.
+
+Set `TRACKER="wandb"` near the top of the launcher after `wandb login` to log
+`total_loss`, `flow_matching_loss`, gradient norm, step time, validation loss,
+and the 16 composed validation MP4 files. Set
+`INFERENCE_LOG_TO_TRACKER=false` to keep media local while retaining scalar
+logging.
+
+Each YAML's `training.model_path` names a matching converted component tree.
+Training never reads its DiT weights; the path is retained so FastVideo can
+copy the component configs and replace the transformer when exporting a DCP
+checkpoint:
+
+```bash
+python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \
+  --variant small_44k \
+  --transformer-config-only \
+  --output converted_weights/mmaudio/small_44k
+```
+
+This export skeleton contains no weights and is sufficient for transformer-only
+checkpoint export. Then run:
+
+```bash
+python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+  --checkpoint outputs/mmaudio_small_44k_from_scratch/checkpoint-10000 \
+  --output-dir converted_weights/mmaudio/small_44k_trained \
+  --overwrite
+```
+
+For a directly runnable inference export, the template at
+`converted_weights/mmaudio/small_44k` must contain all pipeline components.
+A transformer-only template produces a transformer-only export, which can be
+combined with the shared 44 kHz encoders/VAE/vocoder component tree.
+
+### Preprocess VGGSound with FastVideo
+
+Feature extraction is a native FastVideo preprocessing workflow. The upstream
+MMAudio repository is useful as a numerical reference, but is not imported by
+the preprocessing or training command.
+
+The reference training loader imports the now-removed
+`torio.io.StreamingMediaDecoder`. Current FastVideo uses `torchaudio.load` for
+audio and PyAV timestamp sampling for the two video frame rates, while retaining
+the reference normalization, resampling, transforms, mel, and encoder contracts.
+
+First create a preprocessing-only component tree. It contains the 44.1 kHz
+VAE encoder, DFN5B CLIP text/vision encoders, Synchformer, and tokenizer; it
+does not duplicate the trainable DiT or inference vocoder:
+
+```bash
+python scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py \
+  --preprocessor-only \
+  --audio-vae-checkpoint official_weights/mmaudio/raw/ext_weights/v1-44.pth \
+  --synchformer-checkpoint official_weights/mmaudio/raw/ext_weights/synchformer_state_dict.pth \
+  --dfn5b-dir official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384 \
+  --output converted_weights/mmaudio/preprocess_44k
+```
+
+For the `Loie/VGGSound` release, extract the gzip shards once for random-access
+decoding. The archives contain a seven-component directory prefix:
+
+```bash
+mkdir -p /path/to/VGGSound/videos
+for shard in /path/to/VGGSound/vggsound_*.tar.gz; do
+  tar -xzf "$shard" -C /path/to/VGGSound/videos --strip-components=7
+done
+```
+
+Run the shared FastVideo preprocessing entrypoint. Each rank gets disjoint
+samples and writes resumable TensorDict shards. Existing sample IDs are skipped
+when the same output directory is resumed. A filtered caption manifest must
+have the same `id<TAB>label` schema used by MMAudio; `label` is the full text
+description encoded by DFN5B, not the original VGGSound class label.
+
+```bash
+torchrun --standalone --nproc_per_node=4 \
+  -m fastvideo.pipelines.preprocess.v1_preprocessing_new \
+  --model-path converted_weights/mmaudio/preprocess_44k \
+  --mode preprocess \
+  --workload-type v2a \
+  --preprocess.dataset-type vggsound \
+  --preprocess.dataset-path /path/to/VGGSound \
+  --preprocess.dataset-metadata-path /path/to/VGGSound/sets/filtered_caption/vgg-train-filtered-caption.tsv \
+  --preprocess.dataset-split train \
+  --preprocess.dataset-output-dir /path/to/VGGSound/mmaudio_features/train \
+  --preprocess.preprocess-video-batch-size 1 \
+  --preprocess.dataloader-num-workers 2 \
+  --preprocess.samples-per-file 256
+```
+
+The configurable launcher defaults to four GPUs and derives the manifest from
+`SPLIT`:
+
+```bash
+DATASET_PATH=/path/to/VGGSound SPLIT=train \
+  bash examples/training/finetune/mmaudio/preprocess_vggsound.sh
+```
+
+Use the same pipeline for every split by changing only `SPLIT`; write each one
+to its own cache directory:
+
+```bash
+for split in train val test; do
+  DATASET_PATH=/path/to/VGGSound SPLIT="$split" GPU_NUM=4 \
+    bash examples/training/finetune/mmaudio/preprocess_vggsound.sh
+done
+```
+
+MMAudio's split contract peak-normalizes train audio to 0.95 and leaves val/test
+audio amplitudes unchanged. Corrupt, silent, or short train inputs and invalid
+val/test inputs are skipped and recorded in `failures_rank_*.jsonl`.
+
+### Precomputed feature contract
+
+The online training step deliberately does not run the audio VAE, DFN5B, or
+Synchformer. Point `training.data.data_path` at either one TensorDict
+memory-mapped directory or the parent directory of FastVideo's feature shards:
+
+```text
+video + audio cache
+├── mean             [N, 345, 40]
+├── std              [N, 345, 40]
+├── clip_features    [N,  64, 1024]
+├── sync_features    [N, 192, 768]
+└── text_features    [N,  77, 1024]
+
+audio-only cache
+├── mean             [N, 345, 40]
+├── std              [N, 345, 40]
+└── text_features    [N,  77, 1024]
+```
+
+Video and audio-only caches can be mixed by using a mapping of cache paths to
+repeat counts in YAML. Missing video conditions are replaced by MMAudio's
+learned null-video tokens. The `text_features` values are normalized per-token
+DFN5B CLIP hidden states computed from the dataset captions. The cache loader
+has no runtime import dependency on the upstream `mmaudio` package.
+
+The example config intentionally contains an empty data path until a dataset is
+prepared:
+
+```text
+examples/train/configs/fine_tuning/mmaudio/small_44k.yaml
+```
+
+Start a single-GPU run while supplying the cache path on the command line:
+
+```bash
+NUM_GPUS=1 bash examples/train/run.sh \
+  examples/train/configs/fine_tuning/mmaudio/small_44k.yaml \
+  --training.data.data_path /path/to/VGGSound/mmaudio_features/train
+```
+
+Leaving the path empty produces an explicit error instead of silently starting
+with the wrong data. The pretrained single-GPU FSDP path and
+scratch-initialized four-GPU FSDP path have both been validated. Sequence and
+tensor parallelism are not implemented, so keep `sp_size: 1` and `tp_size: 1`. The
+example preserves the official learning rate, AdamW betas/epsilon, weight
+decay, 1,000-step warmup, CFG dropout, logit-normal timestep sampling, and
+300,000-step duration. Its per-GPU batch size is intentionally one for initial
+smoke testing and should be scaled with gradient accumulation or additional
+data-parallel GPUs.
+
 ## Run V2A Inference
 
 The runnable example is

fastvideo/pipelines/basic/mmaudio/mmaudio_pipeline.py

@@ -14,7 +14,16 @@ from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
 
 
 class MMAudioPipeline(ComposedPipelineBase):
-    """MMAudio large-44k-v2 V2A/T2A inference pipeline."""
+    """MMAudio large-44k-v2 V2A/T2A inference pipeline.
+
+    ``text_encoder``, ``tokenizer``, ``image_encoder`` and ``image_encoder_2``
+    are required modules, but the direct-feature path of the conditioning
+    stages returns before dereferencing them. Callers that pass cached
+    ``mmaudio_clip_features``/``mmaudio_sync_features`` (and
+    ``mmaudio_text_features``) may therefore inject a placeholder for those
+    four modules; callers that pass ``video_path`` or omit the cached features
+    must supply the real modules.
+    """
 
     _required_config_modules = [
         "transformer",

fastvideo/pipelines/basic/mmaudio/presets.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Published MMAudio large-44k-v2 inference defaults."""
+"""Published inference defaults for all official MMAudio variants."""
 
 from fastvideo.api.presets import InferencePreset, PresetStageSpec
 
@@ -26,15 +26,52 @@ _DEFAULTS = {
     "return_frames": False,
 }
 
-MMAUDIO_LARGE_44K_V2 = InferencePreset(
-    name="mmaudio_large_44k_v2",
-    version=1,
-    model_family="mmaudio",
-    description=("MMAudio large-44k-v2 video-to-audio generation with DFN5B CLIP, "
-                 "Synchformer, a 44.1 kHz audio VAE, and BigVGAN-v2."),
-    workload_type="v2a",
-    stage_schemas=(_DENOISE_STAGE, ),
-    defaults=dict(_DEFAULTS),
+
+def _mmaudio_preset(
+    variant: str,
+    description: str,
+) -> InferencePreset:
+    return InferencePreset(
+        name=f"mmaudio_{variant}",
+        version=1,
+        model_family="mmaudio",
+        description=description,
+        workload_type="v2a",
+        stage_schemas=(_DENOISE_STAGE, ),
+        defaults=dict(_DEFAULTS),
+    )
+
+
+MMAUDIO_SMALL_16K = _mmaudio_preset(
+    "small_16k",
+    ("MMAudio small 16 kHz video-to-audio generation with DFN5B CLIP, "
+     "Synchformer, the 16 kHz audio VAE, and the 16 kHz BigVGAN vocoder."),
+)
+MMAUDIO_SMALL_44K = _mmaudio_preset(
+    "small_44k",
+    ("MMAudio small 44.1 kHz video-to-audio generation with DFN5B CLIP, "
+     "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."),
+)
+MMAUDIO_MEDIUM_44K = _mmaudio_preset(
+    "medium_44k",
+    ("MMAudio medium 44.1 kHz video-to-audio generation with DFN5B CLIP, "
+     "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."),
+)
+MMAUDIO_LARGE_44K = _mmaudio_preset(
+    "large_44k",
+    ("MMAudio large 44.1 kHz video-to-audio generation with DFN5B CLIP, "
+     "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."),
+)
+MMAUDIO_LARGE_44K_V2 = _mmaudio_preset(
+    "large_44k_v2",
+    ("MMAudio large-44k-v2 video-to-audio generation with DFN5B CLIP, "
+     "Synchformer, the 44.1 kHz audio VAE, and BigVGAN-v2."),
 )
 
-ALL_PRESETS = (MMAUDIO_LARGE_44K_V2, )
+ALL_PRESETS = (
+    MMAUDIO_SMALL_16K,
+    MMAUDIO_SMALL_44K,
+    MMAUDIO_MEDIUM_44K,
+    MMAUDIO_LARGE_44K,
+    MMAUDIO_LARGE_44K_V2,
+)

fastvideo/pipelines/basic/mmaudio/stages.py

@@ -1,9 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Inference stages for MMAudio V2A/T2A.
+"""Inference stages and shared video preprocessing for MMAudio V2A/T2A.
 
 The video sampler intentionally mirrors ``mmaudio.data.av_utils.read_frames``:
 timestamps are sampled independently at 8 FPS and 25 FPS, and a decoded frame
 is repeated when the source FPS is lower than a requested sampling rate.
+
+Training preprocessing can instead select the FFmpeg ``fps`` filter sampler,
+matching MMAudio's original ``StreamingMediaDecoder`` training path.
 """
 
 from __future__ import annotations
@@ -68,6 +71,90 @@ def _read_frames_at_fps(
     return [np.stack(frames) for frames in outputs]
 
 
+def _read_frames_with_fps_filters(
+    video_path: str | Path,
+    frame_rates: tuple[float, ...],
+    *,
+    start_s: float,
+    end_s: float,
+) -> list[np.ndarray]:
+    """Decode RGB frames using one FFmpeg graph with an ``fps`` branch per rate.
+
+    MMAudio's original training loader configures
+    ``StreamingMediaDecoder.add_basic_video_stream(frame_rate=...)``, which
+    delegates resampling to FFmpeg's ``fps`` filter. A single split graph keeps
+    both output rates on the same decode and avoids sharing a VideoFrame between
+    multiple filter graphs.
+    """
+    import av
+
+    if not frame_rates or any(fps <= 0 for fps in frame_rates):
+        raise ValueError("MMAudio video frame rates must be positive")
+    if start_s < 0 or end_s <= start_s:
+        raise ValueError(f"Invalid MMAudio video interval [{start_s}, {end_s}] seconds")
+
+    expected_counts = [int((end_s - start_s) * fps) for fps in frame_rates]
+    if any(count <= 0 for count in expected_counts):
+        raise ValueError(f"MMAudio video interval [{start_s}, {end_s}] is too short for "
+                         f"frame rates {frame_rates}")
+    outputs: list[list[np.ndarray]] = [[] for _ in frame_rates]
+
+    with av.open(str(video_path)) as container:
+        stream = container.streams.video[0]
+        stream.thread_type = "AUTO"
+
+        graph = av.filter.Graph()
+        source = graph.add_buffer(template=stream)
+        trim = graph.add("trim", f"start={start_s}:end={end_s}")
+        source.link_to(trim)
+        setpts = graph.add("setpts", "PTS-STARTPTS")
+        trim.link_to(setpts)
+
+        if len(frame_rates) > 1:
+            split = graph.add("split", str(len(frame_rates)))
+            setpts.link_to(split)
+            branch_source = split
+        else:
+            branch_source = setpts
+
+        sinks = []
+        for index, fps in enumerate(frame_rates):
+            fps_filter = graph.add("fps", f"fps={fps}")
+            if len(frame_rates) > 1:
+                branch_source.link_to(fps_filter, index, 0)
+            else:
+                branch_source.link_to(fps_filter)
+            rgb_filter = graph.add("format", "pix_fmts=rgb24")
+            fps_filter.link_to(rgb_filter)
+            sink = graph.add("buffersink")
+            rgb_filter.link_to(sink)
+            sinks.append(sink)
+        graph.configure()
+
+        def drain_outputs() -> None:
+            for index, (sink, expected_count) in enumerate(zip(sinks, expected_counts, strict=True)):
+                while len(outputs[index]) < expected_count:
+                    try:
+                        output_frame = sink.pull()
+                    except (av.error.BlockingIOError, av.error.EOFError):
+                        break
+                    outputs[index].append(output_frame.to_ndarray(format="rgb24"))
+
+        for frame in container.decode(stream):
+            source.push(frame)
+            drain_outputs()
+            if all(len(frames) >= expected for frames, expected in zip(outputs, expected_counts, strict=True)):
+                break
+        else:
+            source.push(None)
+            drain_outputs()
+
+    if any(not frames for frames in outputs):
+        raise ValueError(f"Could not decode video frames from {video_path} in "
+                         f"[{start_s}, {end_s}] seconds.")
+    return [np.stack(frames) for frames in outputs]
+
+
 def preprocess_mmaudio_video(
     video_path: str | Path,
     *,
@@ -76,13 +163,19 @@ def preprocess_mmaudio_video(
     sync_fps: int = 25,
     clip_size: int = 384,
     sync_size: int = 224,
+    use_ffmpeg_fps_filter: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, float]:
     """Return official-format CLIP frames, sync frames, and effective duration.
 
     CLIP output is float32 ``[T,3,384,384]`` in ``[0,1]``. Synchformer
     output is float32 ``[T,3,224,224]`` in ``[-1,1]``.
+
+    The default timestamp sampler preserves inference parity. Training callers
+    should enable ``use_ffmpeg_fps_filter`` to match MMAudio's original
+    StreamingMediaDecoder/FFmpeg resampling behavior.
     """
-    clip_array, sync_array = _read_frames_at_fps(
+    frame_reader = (_read_frames_with_fps_filters if use_ffmpeg_fps_filter else _read_frames_at_fps)
+    clip_array, sync_array = frame_reader(
         video_path,
         (float(clip_fps), float(sync_fps)),
         start_s=0.0,
@@ -144,8 +237,10 @@ class MMAudioInputValidationStage(PipelineStage):
         if isinstance(batch.negative_prompt, list) and len(batch.negative_prompt) != 1:
             raise ValueError("MMAudio currently supports one negative prompt per request.")
 
-        direct_video = (batch.extra.get("mmaudio_clip_frames") is not None
-                        and batch.extra.get("mmaudio_sync_frames") is not None)
+        direct_video = ((batch.extra.get("mmaudio_clip_frames") is not None
+                         and batch.extra.get("mmaudio_sync_frames") is not None)
+                        or (batch.extra.get("mmaudio_clip_features") is not None
+                            and batch.extra.get("mmaudio_sync_features") is not None))
         if (fastvideo_args.workload_type is WorkloadType.V2A and batch.video_path is None and not direct_video):
             raise ValueError("MMAudio V2A requires `video_path` or preprocessed MMAudio frame tensors.")
 
@@ -187,6 +282,51 @@ class MMAudioVideoConditioningStage(PipelineStage):
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         pc = fastvideo_args.pipeline_config
         duration_s = float(batch.extra["mmaudio_duration_s"])
+        direct_clip_features = batch.extra.get("mmaudio_clip_features")
+        direct_sync_features = batch.extra.get("mmaudio_sync_features")
+        if (direct_clip_features is None) != (direct_sync_features is None):
+            raise ValueError("MMAudio direct conditioning requires both CLIP and sync features.")
+        if direct_clip_features is not None:
+            if direct_clip_features.ndim != 3 or direct_sync_features.ndim != 3:
+                raise ValueError("MMAudio direct CLIP/sync features must have shape [B,T,C].")
+            if direct_clip_features.shape[0] != 1 or direct_sync_features.shape[0] != 1:
+                raise ValueError("MMAudio direct features currently require batch size 1.")
+            latent_length, clip_length, sync_length = mmaudio_sequence_lengths(
+                duration_s,
+                pc,
+            )
+            expected_clip = (1, clip_length, pc.dit_config.arch_config.clip_dim)
+            expected_sync = (1, sync_length, pc.dit_config.arch_config.sync_dim)
+            if tuple(direct_clip_features.shape) != expected_clip:
+                raise ValueError(f"MMAudio direct CLIP features must have shape {expected_clip}, "
+                                 f"got {tuple(direct_clip_features.shape)}.")
+            if tuple(direct_sync_features.shape) != expected_sync:
+                raise ValueError(f"MMAudio direct sync features must have shape {expected_sync}, "
+                                 f"got {tuple(direct_sync_features.shape)}.")
+            self.transformer.update_seq_lengths(
+                latent_length,
+                clip_length,
+                sync_length,
+            )
+            device = get_local_torch_device()
+            dtype = next(self.transformer.parameters()).dtype
+            batch.extra["mmaudio_sequence_lengths"] = (
+                latent_length,
+                clip_length,
+                sync_length,
+            )
+            batch.extra["mmaudio_clip_features"] = direct_clip_features.to(
+                device=device,
+                dtype=dtype,
+                non_blocking=True,
+            )
+            batch.extra["mmaudio_sync_features"] = direct_sync_features.to(
+                device=device,
+                dtype=dtype,
+                non_blocking=True,
+            )
+            return batch
+
         clip_frames = batch.extra.get("mmaudio_clip_frames")
         sync_frames = batch.extra.get("mmaudio_sync_frames")
         use_video = batch.video_path is not None or (clip_frames is not None and sync_frames is not None)
@@ -291,42 +431,73 @@ class MMAudioTextConditioningStage(PipelineStage):
     def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
         device = get_local_torch_device()
         model_dtype = next(self.transformer.parameters()).dtype
-        self.text_encoder = self.text_encoder.to(device)
+        direct_text_features = batch.extra.get("mmaudio_text_features")
 
-        def encode(text: str | list[str] | None) -> torch.Tensor | None:
-            if text is None:
-                return None
-            texts = [text] if isinstance(text, str) else text
-            tokens = self.tokenizer(
-                texts,
-                padding="max_length",
-                truncation=True,
-                max_length=77,
-                return_tensors="pt",
+        if direct_text_features is not None:
+            if direct_text_features.ndim != 3 or direct_text_features.shape[0] != 1:
+                raise ValueError("MMAudio direct text features must have shape [1,T,C].")
+            expected_shape = (
+                1,
+                self.transformer.config.arch_config.text_seq_len,
+                self.transformer.config.arch_config.text_dim,
             )
-            # Hugging Face CLIPTokenizer pads with EOS by default, while
-            # OpenCLIP/MMAudio pads with token id zero.
-            input_ids = tokens.input_ids.masked_fill(tokens.attention_mask == 0, 0).to(device)
-            with set_forward_context(current_timestep=0, attn_metadata=None):
-                return self.text_encoder(input_ids).last_hidden_state.to(model_dtype)
+            if tuple(direct_text_features.shape) != expected_shape:
+                raise ValueError(f"MMAudio direct text features must have shape {expected_shape}, "
+                                 f"got {tuple(direct_text_features.shape)}.")
+            text_features = direct_text_features.to(
+                device=device,
+                dtype=model_dtype,
+                non_blocking=True,
+            )
+            negative_text_features = None
+        else:
+            self.text_encoder = self.text_encoder.to(device)
 
-        text_features = encode(batch.prompt)
-        if text_features is None:
-            text_features = self.transformer.get_empty_string_sequence(1)
-        negative_text_features = encode(batch.negative_prompt)
+            def encode(text: str | list[str] | None) -> torch.Tensor | None:
+                if text is None:
+                    return None
+                texts = [text] if isinstance(text, str) else text
+                tokens = self.tokenizer(
+                    texts,
+                    padding="max_length",
+                    truncation=True,
+                    max_length=77,
+                    return_tensors="pt",
+                )
+                # Hugging Face CLIPTokenizer pads with EOS by default, while
+                # OpenCLIP/MMAudio pads with token id zero.
+                input_ids = tokens.input_ids.masked_fill(tokens.attention_mask == 0, 0).to(device)
+                with set_forward_context(current_timestep=0, attn_metadata=None):
+                    return self.text_encoder(input_ids).last_hidden_state.to(model_dtype)
 
-        if fastvideo_args.text_encoder_cpu_offload:
-            self.text_encoder = self.text_encoder.to("cpu")
+            text_features = encode(batch.prompt)
+            if text_features is None:
+                text_features = self.transformer.get_empty_string_sequence(1)
+            negative_text_features = encode(batch.negative_prompt)
 
-        conditions = self.transformer.preprocess_conditions(
-            batch.extra["mmaudio_clip_features"],
-            batch.extra["mmaudio_sync_features"],
-            text_features,
-        )
-        empty_conditions = self.transformer.get_empty_conditions(
-            1,
-            negative_text_features=negative_text_features,
-        )
+            if fastvideo_args.text_encoder_cpu_offload:
+                self.text_encoder = self.text_encoder.to("cpu")
+
+        # A live training transformer can be FSDP2-sharded. Condition
+        # preprocessing reads positional-embedding parameters before the
+        # transformer's normal pre-forward hook would all-gather them.
+        unshard = getattr(self.transformer, "unshard", None)
+        if callable(unshard):
+            unshard()
+        with torch.autocast(
+                device_type=device.type,
+                dtype=torch.bfloat16,
+                enabled=device.type == "cuda",
+        ):
+            conditions = self.transformer.preprocess_conditions(
+                batch.extra["mmaudio_clip_features"],
+                batch.extra["mmaudio_sync_features"],
+                text_features,
+            )
+            empty_conditions = self.transformer.get_empty_conditions(
+                1,
+                negative_text_features=negative_text_features,
+            )
         batch.extra["mmaudio_conditions"] = conditions
         batch.extra["mmaudio_empty_conditions"] = empty_conditions
         return batch
@@ -390,13 +561,22 @@ class MMAudioDenoisingStage(PipelineStage):
             # Stop if interrupted
             if getattr(self, "interrupt", False):
                 break
-            flow = self.transformer.guided_flow(
-                timestep / self.scheduler.config.num_train_timesteps,
-                latents,
-                conditions,
-                empty_conditions,
-                float(batch.guidance_scale),
-            )
+            # Converted inference checkpoints carry bf16 parameters directly,
+            # while a live FSDP training model keeps fp32 master parameters and
+            # runs its forward pass under bf16 autocast. Use the same forward
+            # precision here so training-time validation can reuse this stage.
+            with torch.autocast(
+                    device_type=latents.device.type,
+                    dtype=torch.bfloat16,
+                    enabled=latents.device.type == "cuda",
+            ):
+                flow = self.transformer.guided_flow(
+                    timestep / self.scheduler.config.num_train_timesteps,
+                    latents,
+                    conditions,
+                    empty_conditions,
+                    float(batch.guidance_scale),
+                )
             # The shared FlowMatch scheduler supplies the exact inverted
             # forward-time sigma schedule. Its generic ``step`` deliberately
             # upcasts samples to fp32, however, while MMAudio's published

fastvideo/pipelines/pipeline_batch_info.py

@@ -353,6 +353,7 @@ class TrainingBatch:
     timesteps: torch.Tensor | None = None
     sigmas: torch.Tensor | None = None
     noise: torch.Tensor | None = None
+    training_target: torch.Tensor | None = None
 
     # MiniMax H3 reuses the packed row boundaries from batch preparation to
     # split the transformer's joint sequence back into video and audio outputs.

fastvideo/pipelines/pipeline_registry.py

@@ -19,6 +19,7 @@ logger = init_logger(__name__)
 _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = {
     WorkloadType.I2V: "PreprocessPipelineI2V",
     WorkloadType.T2V: "PreprocessPipelineT2V",
+    WorkloadType.V2A: "PreprocessPipelineV2A",
 }
 
 
@@ -57,9 +58,24 @@ class _PipelineRegistry:
         """Get supported pipelines for the given pipeline type."""
         return set(self.pipelines.get(pipeline_type.value, {}).keys())
 
-    def _load_preprocess_pipeline_cls(self, workload_type: WorkloadType) -> type[ComposedPipelineBase] | None:
-        pipeline_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[workload_type]
-        return self.pipelines.get(PipelineType.PREPROCESS.value, {}).get(pipeline_name)
+    def _load_preprocess_pipeline_cls(
+        self,
+        pipeline_name_in_config: str,
+        workload_type: WorkloadType,
+    ) -> type[ComposedPipelineBase] | None:
+        pipelines = self.pipelines.get(PipelineType.PREPROCESS.value, {})
+
+        # Existing T2V/I2V integrations retain their exact workload-level
+        # lookup. New V2A families may define different feature contracts and
+        # therefore use a family-specific preprocessing class.
+        if (workload_type is WorkloadType.V2A and pipeline_name_in_config.endswith("Pipeline")):
+            family_pipeline_name = (pipeline_name_in_config.removesuffix("Pipeline") + "PreprocessPipeline")
+            family_pipeline = pipelines.get(family_pipeline_name)
+            if family_pipeline is not None:
+                return family_pipeline
+
+        generic_name = _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME[workload_type]
+        return pipelines.get(generic_name)
 
     def _try_load_pipeline_cls(self, pipeline_name_in_config: str, pipeline_type: PipelineType,
                                workload_type: WorkloadType) -> type[ComposedPipelineBase] | type[LoRAPipeline] | None:
@@ -68,7 +84,10 @@ class _PipelineRegistry:
             return None
 
         if pipeline_type == PipelineType.PREPROCESS:
-            return self._load_preprocess_pipeline_cls(workload_type)
+            return self._load_preprocess_pipeline_cls(
+                pipeline_name_in_config,
+                workload_type,
+            )
         elif pipeline_type == PipelineType.BASIC or pipeline_type == PipelineType.TRAINING:
             return self.pipelines[pipeline_type.value].get(pipeline_name_in_config)
         else:

fastvideo/pipelines/preprocess/mmaudio/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.pipelines.preprocess.mmaudio.mmaudio_preprocess_pipeline import (
+    MMAudioPreprocessPipeline, )
+
+__all__ = ["MMAudioPreprocessPipeline"]

fastvideo/pipelines/preprocess/mmaudio/mmaudio_preprocess_pipeline.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Native FastVideo preprocessing pipeline for MMAudio training."""
+
+from __future__ import annotations
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.models.audio.mmaudio_processing import build_mmaudio_mel_converter
+from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.pipelines.preprocess.mmaudio.stages import MMAudioFeatureExtractionStage
+
+
+class MMAudioPreprocessPipeline(ComposedPipelineBase):
+    """Encode raw audio, video, and captions into MMAudio training features."""
+
+    _required_config_modules = [
+        "audio_vae",
+        "text_encoder",
+        "tokenizer",
+        "image_encoder",
+        "image_encoder_2",
+    ]
+
+    def __init__(self, model_path: str, fastvideo_args: FastVideoArgs, **kwargs) -> None:
+        # Official training features are extracted in fp32. These overrides
+        # are local to preprocess mode and do not change inference defaults.
+        config = fastvideo_args.pipeline_config
+        config.text_encoder_precisions = ("fp32", )
+        config.image_encoder_precisions = ("fp32", "fp32")
+        config.audio_encoder_precision = "fp32"
+        super().__init__(model_path, fastvideo_args, **kwargs)
+
+    def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None:
+        audio_vae = self.get_module("audio_vae")
+        if not hasattr(audio_vae, "encoder"):
+            raise ValueError("MMAudio preprocessing requires an audio_vae component converted "
+                             "with need_encoder=true. Use the converter's --preprocessor-only mode.")
+        if getattr(audio_vae, "mode", None) != "44k":
+            raise ValueError("The current MMAudio training integration requires the 44k audio VAE")
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        mel_converter = build_mmaudio_mel_converter("44k")
+        self.add_stage(
+            "feature_extraction_stage",
+            MMAudioFeatureExtractionStage(
+                audio_vae=self.get_module("audio_vae"),
+                mel_converter=mel_converter,
+                image_encoder=self.get_module("image_encoder"),
+                sync_encoder=self.get_module("image_encoder_2"),
+                text_encoder=self.get_module("text_encoder"),
+                tokenizer=self.get_module("tokenizer"),
+            ),
+        )
+
+
+EntryClass = MMAudioPreprocessPipeline

fastvideo/pipelines/preprocess/mmaudio/stages.py

@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio-specific feature stages for the shared FastVideo V2A workflow."""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Any, cast
+
+import torch
+import torchaudio
+
+from fastvideo.configs.configs import VideoLoaderType
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.mmaudio.stages import (
+    _CLIP_MEAN,
+    _CLIP_STD,
+    preprocess_mmaudio_video,
+)
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, PreprocessBatch
+from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import (PREPROCESSED_MEDIA_KEY, PREPROCESS_ERROR_KEY,
+                                                                       preprocess_mmaudio_media_with_torio)
+from fastvideo.pipelines.stages.base import PipelineStage
+from fastvideo.pipelines.stages.validators import VerificationResult
+
+logger = init_logger(__name__)
+
+
+class MMAudioFeatureExtractionStage(PipelineStage):
+    """Extract VAE posterior statistics plus CLIP/Synchformer features."""
+
+    def __init__(
+        self,
+        *,
+        audio_vae: torch.nn.Module,
+        mel_converter: torch.nn.Module,
+        image_encoder: torch.nn.Module,
+        sync_encoder: torch.nn.Module,
+        text_encoder: torch.nn.Module,
+        tokenizer: Any,
+    ) -> None:
+        super().__init__()
+        self.audio_vae = audio_vae.eval()
+        self.mel_converter = mel_converter.eval()
+        self.image_encoder = image_encoder.eval()
+        self.sync_encoder = sync_encoder.eval()
+        self.text_encoder = text_encoder.eval()
+        self.tokenizer = tokenizer
+
+    def verify_input(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    def verify_output(self, batch, fastvideo_args):
+        return VerificationResult()
+
+    @staticmethod
+    def _load_audio(
+        path: str | Path,
+        *,
+        target_sample_rate: int,
+        target_samples: int,
+        normalize_audio: bool,
+    ) -> torch.Tensor:
+        waveform, sample_rate = torchaudio.load(str(path))
+        if waveform.ndim != 2 or waveform.shape[0] == 0:
+            raise ValueError(f"Decoded audio has an invalid shape: {tuple(waveform.shape)}")
+        waveform = waveform.mean(dim=0)
+        if normalize_audio:
+            abs_max = waveform.abs().max()
+            if abs_max <= 1e-6:
+                raise ValueError("Decoded audio is silent")
+            waveform = waveform / abs_max * 0.95
+        if sample_rate != target_sample_rate:
+            waveform = torchaudio.functional.resample(
+                waveform,
+                sample_rate,
+                target_sample_rate,
+                lowpass_filter_width=64,
+                rolloff=0.9475937167399596,
+                resampling_method="sinc_interp_kaiser",
+                beta=14.769656459379492,
+            )
+        if waveform.shape[0] < target_samples:
+            raise ValueError(
+                f"Audio is too short: need {target_samples} samples at {target_sample_rate} Hz, got {waveform.shape[0]}"
+            )
+        return waveform[:target_samples].contiguous()
+
+    @torch.inference_mode()
+    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+        batch = cast(PreprocessBatch, batch)
+        if not isinstance(batch.prompt, list):
+            raise ValueError("MMAudio preprocessing expects a list of captions")
+        if len(batch.video_loader) != len(batch.prompt):
+            raise ValueError("MMAudio preprocessing paths and captions have different batch sizes")
+
+        pc = fastvideo_args.pipeline_config
+        duration_s = float(pc.duration_s)
+        sample_rate = int(pc.sampling_rate)
+        alignment = int(pc.spectrogram_frame_rate * pc.latent_downsample_rate)
+        target_samples = math.ceil(duration_s * sample_rate / alignment) * alignment
+        expected_clip_frames = int(duration_s * pc.clip_frame_rate)
+        expected_sync_frames = int(duration_s * pc.sync_frame_rate)
+        preprocess_config = fastvideo_args.preprocess_config
+        dataset_split = "train" if preprocess_config is None else preprocess_config.dataset_split.lower()
+        normalize_audio = dataset_split == "train"
+        use_torio = (preprocess_config is not None and preprocess_config.video_loader_type == VideoLoaderType.TORIO)
+
+        preprocessed_media = batch.extra.get(PREPROCESSED_MEDIA_KEY)
+        preprocess_errors = batch.extra.get(PREPROCESS_ERROR_KEY)
+        for name, values in (
+            (PREPROCESSED_MEDIA_KEY, preprocessed_media),
+            (PREPROCESS_ERROR_KEY, preprocess_errors),
+        ):
+            if values is not None and (not isinstance(values, list) or len(values) != len(batch.video_loader)):
+                raise ValueError(f"MMAudio {name} must align with the input batch")
+
+        valid_ids: list[str] = []
+        valid_paths: list[str] = []
+        valid_captions: list[str] = []
+        audio_batch: list[torch.Tensor] = []
+        clip_batch: list[torch.Tensor] = []
+        sync_batch: list[torch.Tensor] = []
+        failures: list[dict[str, str]] = []
+
+        for index, (sample_id, path,
+                    caption) in enumerate(zip(
+                        batch.video_file_name,
+                        batch.video_loader,
+                        batch.prompt,
+                        strict=True,
+                    )):
+            try:
+                path_str = str(path)
+                if not Path(path_str).is_file():
+                    raise FileNotFoundError(path_str)
+                preprocess_error = (None if preprocess_errors is None else preprocess_errors[index])
+                if preprocess_error is not None:
+                    raise ValueError(str(preprocess_error))
+                media = (None if preprocessed_media is None else preprocessed_media[index])
+                if media is not None:
+                    if not isinstance(media, dict):
+                        raise ValueError("Preprocessed MMAudio media must be a dictionary")
+                    audio = media["audio"]
+                    clip_frames = media["clip_frames"]
+                    sync_frames = media["sync_frames"]
+                    effective_duration = float(media["effective_duration"])
+                elif use_torio:
+                    audio, clip_frames, sync_frames, effective_duration = (preprocess_mmaudio_media_with_torio(
+                        path_str,
+                        duration_s=duration_s,
+                        target_sample_rate=sample_rate,
+                        target_samples=target_samples,
+                        normalize_audio=normalize_audio,
+                        clip_fps=pc.clip_frame_rate,
+                        sync_fps=pc.sync_frame_rate,
+                        clip_size=pc.clip_image_size,
+                        sync_size=pc.sync_image_size,
+                    ))
+                else:
+                    audio = self._load_audio(
+                        path_str,
+                        target_sample_rate=sample_rate,
+                        target_samples=target_samples,
+                        normalize_audio=normalize_audio,
+                    )
+                    clip_frames, sync_frames, effective_duration = preprocess_mmaudio_video(
+                        path_str,
+                        duration_s=duration_s,
+                        clip_fps=pc.clip_frame_rate,
+                        sync_fps=pc.sync_frame_rate,
+                        clip_size=pc.clip_image_size,
+                        sync_size=pc.sync_image_size,
+                        use_ffmpeg_fps_filter=True,
+                    )
+                if effective_duration < duration_s:
+                    raise ValueError(f"Video is too short: need {duration_s}s, got {effective_duration:.3f}s")
+                if clip_frames.shape[0] != expected_clip_frames:
+                    raise ValueError(f"Expected {expected_clip_frames} CLIP frames, got {clip_frames.shape[0]}")
+                if sync_frames.shape[0] != expected_sync_frames:
+                    raise ValueError(f"Expected {expected_sync_frames} sync frames, got {sync_frames.shape[0]}")
+            except Exception as exc:
+                logger.warning("Skipping MMAudio preprocessing sample %s: %s", sample_id, exc)
+                failures.append({"id": str(sample_id), "error": str(exc)})
+                continue
+            valid_ids.append(str(sample_id))
+            valid_paths.append(path_str)
+            valid_captions.append(str(caption))
+            audio_batch.append(audio)
+            clip_batch.append(clip_frames)
+            sync_batch.append(sync_frames)
+
+        batch.extra["failed_samples"] = failures
+        if not valid_ids:
+            batch.video_file_name = []
+            batch.video_loader = []
+            batch.prompt = []
+            batch.extra["precomputed_features"] = {}
+            batch.extra["precomputed_metadata"] = []
+            return batch
+
+        device = get_local_torch_device()
+        self.audio_vae = self.audio_vae.to(device=device, dtype=torch.float32)
+        self.mel_converter = self.mel_converter.to(device=device, dtype=torch.float32)
+        waveforms = torch.stack(audio_batch).to(device=device, dtype=torch.float32)
+        posterior = self.audio_vae.encode(self.mel_converter(waveforms))
+        latent_mean = posterior.mean.transpose(1, 2).detach().cpu()
+        latent_std = posterior.std.transpose(1, 2).detach().cpu()
+
+        self.image_encoder = self.image_encoder.to(device=device, dtype=torch.float32)
+        clip_video = torch.stack(clip_batch).to(device=device, dtype=torch.float32)
+        mean = torch.tensor(_CLIP_MEAN, device=device).view(1, 1, 3, 1, 1)
+        std = torch.tensor(_CLIP_STD, device=device).view(1, 1, 3, 1, 1)
+        clip_video = ((clip_video - mean) / std).flatten(0, 1)
+        clip_outputs: list[torch.Tensor] = []
+        chunk_size = int(pc.clip_batch_size_multiplier)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            for start in range(0, clip_video.shape[0], chunk_size):
+                encoded = self.image_encoder(clip_video[start:start + chunk_size]).last_hidden_state
+                clip_outputs.append(encoded)
+        clip_features = torch.cat(clip_outputs).reshape(len(valid_ids), expected_clip_frames, -1).detach().cpu()
+
+        self.sync_encoder = self.sync_encoder.to(device=device, dtype=torch.float32)
+        sync_video = torch.stack(sync_batch).to(device=device, dtype=torch.float32)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            sync_features = self.sync_encoder(sync_video).last_hidden_state.detach().cpu()
+
+        self.text_encoder = self.text_encoder.to(device=device, dtype=torch.float32)
+        tokens = self.tokenizer(
+            valid_captions,
+            padding="max_length",
+            truncation=True,
+            max_length=77,
+            return_tensors="pt",
+        )
+        input_ids = tokens.input_ids.masked_fill(tokens.attention_mask == 0, 0).to(device)
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            text_features = self.text_encoder(input_ids).last_hidden_state.detach().cpu()
+
+        batch.video_file_name = valid_ids
+        batch.video_loader = valid_paths
+        batch.prompt = valid_captions
+        batch.extra["precomputed_features"] = {
+            "mean": latent_mean,
+            "std": latent_std,
+            "clip_features": clip_features,
+            "sync_features": sync_features,
+            "text_features": text_features,
+        }
+        batch.extra["precomputed_metadata"] = [{
+            "id": sample_id,
+            "caption": caption,
+            "source": path,
+            "split": dataset_split,
+        } for sample_id, caption, path in zip(
+            valid_ids,
+            valid_captions,
+            valid_paths,
+            strict=True,
+        )]
+        return batch
+
+
+__all__ = ["MMAudioFeatureExtractionStage"]

fastvideo/pipelines/preprocess/mmaudio/torio_media_reader.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Optional torio media reader for MMAudio training preprocessing."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torchaudio
+from torchvision.transforms import v2
+
+PREPROCESSED_MEDIA_KEY = "preprocessed_media"
+PREPROCESS_ERROR_KEY = "preprocess_error"
+
+
+def preprocess_mmaudio_media_with_torio(
+    media_path: str | Path,
+    *,
+    duration_s: float,
+    target_sample_rate: int,
+    target_samples: int,
+    normalize_audio: bool,
+    clip_fps: int = 8,
+    sync_fps: int = 25,
+    clip_size: int = 384,
+    sync_size: int = 224,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]:
+    """Decode audio and dual-rate video with StreamingMediaDecoder.
+
+    The dependency is imported lazily so normal FastVideo environments do not
+    require the deprecated torio package. This backend is intended for offline
+    reference-compatible training preprocessing only.
+    """
+    try:
+        from torio.io import StreamingMediaDecoder
+    except ImportError as exc:
+        raise ImportError("MMAudio torio preprocessing requires the isolated torio "
+                          "environment. Use run_preprocess_vggsound_torio.sh.") from exc
+
+    expected_clip_frames = int(clip_fps * duration_s)
+    expected_sync_frames = int(sync_fps * duration_s)
+    reader = StreamingMediaDecoder(str(media_path))
+    reader.add_basic_video_stream(
+        frames_per_chunk=expected_clip_frames,
+        frame_rate=float(clip_fps),
+        format="rgb24",
+    )
+    reader.add_basic_video_stream(
+        frames_per_chunk=expected_sync_frames,
+        frame_rate=float(sync_fps),
+        format="rgb24",
+    )
+    reader.add_basic_audio_stream(frames_per_chunk=2**30)
+    reader.fill_buffer()
+    clip_frames, sync_frames, audio_frames = reader.pop_chunks()
+
+    if clip_frames is None:
+        raise ValueError("CLIP video stream returned no frames")
+    if clip_frames.shape[0] < expected_clip_frames:
+        raise ValueError(f"CLIP video is too short: expected {expected_clip_frames} frames, "
+                         f"got {clip_frames.shape[0]}")
+    if sync_frames is None:
+        raise ValueError("Synchformer video stream returned no frames")
+    if sync_frames.shape[0] < expected_sync_frames:
+        raise ValueError(f"Synchformer video is too short: expected {expected_sync_frames} "
+                         f"frames, got {sync_frames.shape[0]}")
+    if audio_frames is None or audio_frames.ndim != 2:
+        shape = None if audio_frames is None else tuple(audio_frames.shape)
+        raise ValueError(f"Decoded audio has an invalid shape: {shape}")
+
+    waveform = audio_frames.transpose(0, 1).mean(dim=0)
+    if normalize_audio:
+        abs_max = waveform.abs().max()
+        if abs_max <= 1e-6:
+            raise ValueError("Decoded audio is silent")
+        waveform = waveform / abs_max * 0.95
+
+    source_sample_rate = int(reader.get_out_stream_info(2).sample_rate)
+    if source_sample_rate != target_sample_rate:
+        waveform = torchaudio.transforms.Resample(
+            source_sample_rate,
+            target_sample_rate,
+            lowpass_filter_width=64,
+            rolloff=0.9475937167399596,
+            resampling_method="sinc_interp_kaiser",
+            beta=14.769656459379492,
+        )(waveform)
+    if waveform.shape[0] < target_samples:
+        raise ValueError(f"Audio is too short: need {target_samples} samples at "
+                         f"{target_sample_rate} Hz, got {waveform.shape[0]}")
+    waveform = waveform[:target_samples].contiguous()
+
+    clip_transform = v2.Compose([
+        v2.Resize(
+            (clip_size, clip_size),
+            interpolation=v2.InterpolationMode.BICUBIC,
+        ),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+    ])
+    sync_transform = v2.Compose([
+        v2.Resize(
+            sync_size,
+            interpolation=v2.InterpolationMode.BICUBIC,
+        ),
+        v2.CenterCrop(sync_size),
+        v2.ToImage(),
+        v2.ToDtype(torch.float32, scale=True),
+        v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+    clip_frames = clip_transform(clip_frames[:expected_clip_frames])
+    sync_frames = sync_transform(sync_frames[:expected_sync_frames])
+    effective_duration = min(
+        duration_s,
+        clip_frames.shape[0] / clip_fps,
+        sync_frames.shape[0] / sync_fps,
+    )
+    return waveform, clip_frames, sync_frames, effective_duration
+
+
+@dataclass(frozen=True)
+class MMAudioTorioRowPreprocessor:
+    """Decode one V2A metadata row on a background CPU thread."""
+
+    duration_s: float
+    target_sample_rate: int
+    target_samples: int
+    normalize_audio: bool
+    clip_fps: int = 8
+    sync_fps: int = 25
+    clip_size: int = 384
+    sync_size: int = 224
+
+    def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
+        output = dict(row)
+        media_path = Path(str(output["video_path"]))
+        try:
+            if not media_path.is_file():
+                raise FileNotFoundError(media_path)
+            audio, clip, sync, duration = preprocess_mmaudio_media_with_torio(
+                media_path,
+                duration_s=self.duration_s,
+                target_sample_rate=self.target_sample_rate,
+                target_samples=self.target_samples,
+                normalize_audio=self.normalize_audio,
+                clip_fps=self.clip_fps,
+                sync_fps=self.sync_fps,
+                clip_size=self.clip_size,
+                sync_size=self.sync_size,
+            )
+            output[PREPROCESSED_MEDIA_KEY] = {
+                "audio": audio,
+                "clip_frames": clip,
+                "sync_frames": sync,
+                "effective_duration": duration,
+            }
+        except Exception as exc:
+            output[PREPROCESS_ERROR_KEY] = f"{type(exc).__name__}: {exc}"
+        return output
+
+
+__all__ = [
+    "MMAudioTorioRowPreprocessor",
+    "PREPROCESSED_MEDIA_KEY",
+    "PREPROCESS_ERROR_KEY",
+    "preprocess_mmaudio_media_with_torio",
+]

fastvideo/registry.py

@@ -45,7 +45,14 @@ from fastvideo.configs.pipelines.flux_2 import (
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig
-from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig
+from fastvideo.configs.pipelines.mmaudio import (
+    MMAudioLarge44kV2PipelineConfig,
+    MMAudioLarge44kV2AConfig,
+    MMAudioMedium44kV2AConfig,
+    MMAudioSmall16kV2AConfig,
+    MMAudioSmall44kV2AConfig,
+    MMAudioV2AConfig,
+)
 from fastvideo.configs.pipelines.turbodiffusion import (
     TurboDiffusionI2V_A14B_Config,
     TurboDiffusionT2V_14B_Config,
@@ -291,6 +298,31 @@ def _get_config_info(
     return None
 
 
+def _mmaudio_variant_detector(expected: str) -> Callable[[str], bool]:
+
+    def detector(path: str) -> bool:
+        normalized = path.lower().replace("-", "_")
+        return expected in normalized and (expected == "large_44k_v2" or "large_44k_v2" not in normalized)
+
+    return detector
+
+
+def _unversioned_mmaudio_detector(*, mode: str) -> Callable[[str], bool]:
+    variants = ("small_16k", "small_44k", "medium_44k", "large_44k")
+
+    def detector(path: str) -> bool:
+        normalized = path.lower().replace("-", "_")
+        # This fallback is path-only. Matching ``mmaudiopipeline`` would also
+        # fire for every variant-specific model_index and create ambiguity.
+        if "mmaudio" not in normalized or "mmaudiopipeline" in normalized:
+            return False
+        if any(variant in normalized for variant in variants):
+            return False
+        return ("16k" in normalized) if mode == "16k" else ("16k" not in normalized)
+
+    return detector
+
+
 def _register_wan_configs(definitions: tuple[WanModelDefinition, ...]) -> None:
     for definition in definitions:
         register_configs(
@@ -316,17 +348,48 @@ def _register_configs() -> None:
         pipeline_cls_name="HeliosPyramidPipeline",
     )
 
-    # MMAudio large-44k-v2 (video/text-to-audio). The checkpoint is converted
-    # into standard per-component FastVideo/Diffusers-style directories by
-    # scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py.
+    # Official MMAudio video/text-to-audio variants. Converted directories keep
+    # the variant in their path, allowing the registry to select the matching
+    # 16 kHz or 44.1 kHz component configuration before any modules are loaded.
+    mmaudio_variants = (
+        ("small_16k", MMAudioSmall16kV2AConfig),
+        ("small_44k", MMAudioSmall44kV2AConfig),
+        ("medium_44k", MMAudioMedium44kV2AConfig),
+        ("large_44k_v2", MMAudioLarge44kV2PipelineConfig),
+        ("large_44k", MMAudioLarge44kV2AConfig),
+    )
+    for variant, pipeline_config_cls in mmaudio_variants:
+        register_configs(
+            sampling_param_cls=None,
+            pipeline_config_cls=pipeline_config_cls,
+            workload_types=(WorkloadType.V2A, WorkloadType.T2A),
+            hf_model_paths=[
+                f"FastVideo/MMAudio-{variant.replace('_', '-')}-Diffusers",
+            ],
+            model_detectors=[_mmaudio_variant_detector(variant)],
+            model_family="mmaudio",
+            default_preset=f"mmaudio_{variant}",
+            pipeline_cls_name="MMAudioPipeline",
+        )
+
+    # Preserve converted preprocessing trees and older local checkpoint names
+    # that predate explicit variant tags. These detectors intentionally do not
+    # inspect ``_class_name=MMAudioPipeline`` so they cannot overlap the five
+    # exact variant entries above.
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=MMAudioSmall16kV2AConfig,
+        workload_types=(WorkloadType.V2A, WorkloadType.T2A),
+        model_detectors=[_unversioned_mmaudio_detector(mode="16k")],
+        model_family="mmaudio",
+        default_preset="mmaudio_small_16k",
+        pipeline_cls_name="MMAudioPipeline",
+    )
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=MMAudioV2AConfig,
         workload_types=(WorkloadType.V2A, WorkloadType.T2A),
-        hf_model_paths=["FastVideo/MMAudio-large-44k-v2-Diffusers"],
-        model_detectors=[
-            lambda path: "mmaudio" in path.lower() or "mmaudiopipeline" in path.lower(),
-        ],
+        model_detectors=[_unversioned_mmaudio_detector(mode="44k")],
         model_family="mmaudio",
         default_preset="mmaudio_large_44k_v2",
         pipeline_cls_name="MMAudioPipeline",

fastvideo/tests/api/test_mmaudio.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Registry coverage for every official MMAudio model variant."""
+
+import json
+from pathlib import Path
+
+import pytest
+
+from fastvideo.configs.pipelines.mmaudio import (
+    MMAudioLarge44kV2PipelineConfig,
+    MMAudioLarge44kV2AConfig,
+    MMAudioMedium44kV2AConfig,
+    MMAudioSmall16kV2AConfig,
+    MMAudioSmall44kV2AConfig,
+    MMAudioV2AConfig,
+)
+from fastvideo.registry import get_default_preset, get_pipeline_config_cls_from_name
+
+
+@pytest.mark.parametrize(
+    ("variant", "config_cls"),
+    [
+        ("small-16k", MMAudioSmall16kV2AConfig),
+        ("small-44k", MMAudioSmall44kV2AConfig),
+        ("medium-44k", MMAudioMedium44kV2AConfig),
+        ("large-44k", MMAudioLarge44kV2AConfig),
+        ("large-44k-v2", MMAudioLarge44kV2PipelineConfig),
+    ],
+)
+def test_mmaudio_variant_registry(variant: str, config_cls: type) -> None:
+    model_path = f"FastVideo/MMAudio-{variant}-Diffusers"
+
+    assert get_pipeline_config_cls_from_name(model_path) is config_cls
+    assert get_default_preset(model_path) == f"mmaudio_{variant.replace('-', '_')}"
+
+
+@pytest.mark.parametrize(
+    ("directory_name", "config_cls", "preset"),
+    [
+        ("preprocess_44k", MMAudioV2AConfig, "mmaudio_large_44k_v2"),
+        ("preprocess_16k", MMAudioSmall16kV2AConfig, "mmaudio_small_16k"),
+    ],
+)
+def test_unversioned_mmaudio_preprocessor_paths_remain_compatible(
+    tmp_path: Path,
+    directory_name: str,
+    config_cls: type,
+    preset: str,
+) -> None:
+    model_path = tmp_path / "mmaudio" / directory_name
+    model_path.mkdir(parents=True)
+    (model_path / "model_index.json").write_text(
+        json.dumps({
+            "_class_name": "MMAudioPipeline",
+            "_diffusers_version": "0.36.0",
+            "_fastvideo_preprocessor_only": True,
+        }),
+        encoding="utf-8",
+    )
+
+    assert get_pipeline_config_cls_from_name(str(model_path)) is config_cls
+    assert get_default_preset(str(model_path)) == preset

fastvideo/tests/api/test_mmaudio_conversion.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Lightweight conversion coverage for official MMAudio variants."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+from pathlib import Path
+
+import torch
+
+
+def _converter_module():
+    path = Path(__file__).resolve().parents[3] / (
+        "scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py"
+    )
+    spec = importlib.util.spec_from_file_location("mmaudio_variant_converter", path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_small_16k_transformer_conversion_writes_variant_contract(
+    tmp_path: Path,
+) -> None:
+    converter = _converter_module()
+    checkpoint = tmp_path / "mmaudio_small_16k.pth"
+    output = tmp_path / "converted"
+    torch.save({"empty_string_feat": torch.zeros(77, 1024)}, checkpoint)
+
+    converter.convert(
+        argparse.Namespace(
+            variant="small_16k",
+            transformer_checkpoint=checkpoint,
+            audio_vae_checkpoint=None,
+            synchformer_checkpoint=None,
+            dfn5b_dir=None,
+            bigvgan_dir=None,
+            bigvgan_checkpoint=None,
+            transformer_only=True,
+            transformer_config_only=False,
+            preprocessor_only=False,
+            output=output,
+        )
+    )
+
+    transformer_config = json.loads(
+        (output / "transformer/config.json").read_text(encoding="utf-8")
+    )
+    model_index = json.loads(
+        (output / "model_index.json").read_text(encoding="utf-8")
+    )
+    assert transformer_config["latent_dim"] == 20
+    assert transformer_config["latent_seq_len"] == 250
+    assert transformer_config["hidden_dim"] == 448
+    assert transformer_config["v2"] is False
+    assert model_index["_fastvideo_mmaudio_variant"] == "small_16k"
+    assert model_index["_fastvideo_audio_sample_rate"] == 16000
+
+
+def test_converter_lists_all_five_official_variants() -> None:
+    converter = _converter_module()
+    assert tuple(converter.TRANSFORMER_VARIANTS) == (
+        "small_16k",
+        "small_44k",
+        "medium_44k",
+        "large_44k",
+        "large_44k_v2",
+    )
+    assert converter.BIGVGAN_16K_CONFIG["num_mels"] == 80
+    assert converter.BIGVGAN_16K_CONFIG["upsample_rates"] == [4, 4, 2, 2, 2, 2]
+
+
+def test_weight_free_transformer_export_skeleton(tmp_path: Path) -> None:
+    converter = _converter_module()
+    output = tmp_path / "skeleton"
+
+    converter.convert(
+        argparse.Namespace(
+            variant="medium_44k",
+            transformer_checkpoint=None,
+            audio_vae_checkpoint=None,
+            synchformer_checkpoint=None,
+            dfn5b_dir=None,
+            bigvgan_dir=None,
+            bigvgan_checkpoint=None,
+            transformer_only=False,
+            transformer_config_only=True,
+            preprocessor_only=False,
+            output=output,
+        )
+    )
+
+    config = json.loads(
+        (output / "transformer/config.json").read_text(encoding="utf-8")
+    )
+    assert config["hidden_dim"] == 896
+    assert config["depth"] == 12
+    assert not list((output / "transformer").glob("*.safetensors"))

fastvideo/tests/dataset/test_mmaudio_feature_dataset.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the MMAudio-compatible precomputed feature cache."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+
+from fastvideo.dataset.mmaudio_feature_dataset import (
+    MMAudioFeatureDataset,
+    _expand_cache_root,
+    compute_mmaudio_latent_stats,
+)
+
+tensordict = pytest.importorskip("tensordict")
+
+
+def _write_cache(path: Path, *, with_video: bool) -> None:
+    tensors = {
+        "mean": torch.randn(2, 4, 3),
+        "std": torch.rand(2, 4, 3),
+        "text_features": torch.randn(2, 5, 7),
+    }
+    if with_video:
+        tensors["clip_features"] = torch.randn(2, 6, 8)
+        tensors["sync_features"] = torch.randn(2, 9, 10)
+    tensordict.TensorDict(tensors, batch_size=[2]).memmap_(str(path))
+
+
+def _load(path: Path, *, include_metadata: bool = False) -> MMAudioFeatureDataset:
+    return MMAudioFeatureDataset(
+        path,
+        latent_seq_len=4,
+        latent_dim=3,
+        clip_seq_len=6,
+        clip_dim=8,
+        sync_seq_len=9,
+        sync_dim=10,
+        text_seq_len=5,
+        text_dim=7,
+        include_metadata=include_metadata,
+    )
+
+
+@pytest.mark.parametrize("with_video", [False, True])
+def test_mmaudio_feature_cache_contract(
+    tmp_path: Path,
+    with_video: bool,
+) -> None:
+    cache = tmp_path / "cache"
+    _write_cache(cache, with_video=with_video)
+
+    dataset = _load(cache)
+    sample = dataset[0]
+
+    assert len(dataset) == 2
+    assert sample["audio_latent_mean"].shape == (4, 3)
+    assert sample["audio_latent_std"].shape == (4, 3)
+    assert sample["clip_features"].shape == (6, 8)
+    assert sample["sync_features"].shape == (9, 10)
+    assert sample["text_features"].shape == (5, 7)
+    assert sample["video_exists"].item() is with_video
+    assert sample["text_exists"].item() is True
+    if not with_video:
+        assert torch.count_nonzero(sample["clip_features"]) == 0
+        assert torch.count_nonzero(sample["sync_features"]) == 0
+
+
+def test_mmaudio_feature_cache_rejects_wrong_shape(tmp_path: Path) -> None:
+    cache = tmp_path / "cache"
+    _write_cache(cache, with_video=True)
+
+    with pytest.raises(ValueError, match="mean.*per-sample shape"):
+        MMAudioFeatureDataset(
+            cache,
+            latent_seq_len=99,
+            latent_dim=3,
+            clip_seq_len=6,
+            clip_dim=8,
+            sync_seq_len=9,
+            sync_dim=10,
+            text_seq_len=5,
+            text_dim=7,
+        )
+
+
+def test_mmaudio_feature_cache_loads_validation_metadata(tmp_path: Path) -> None:
+    cache = tmp_path / "cache"
+    _write_cache(cache, with_video=True)
+    rows = [
+        {"id": "sample-a", "caption": "first", "source": "/video/a.mp4"},
+        {"id": "sample-b", "caption": "second", "source": "/video/b.mp4"},
+    ]
+    with (cache / "samples.jsonl").open("w", encoding="utf-8") as handle:
+        for row in rows:
+            handle.write(json.dumps(row) + "\n")
+
+    sample = _load(cache, include_metadata=True)[1]
+
+    assert sample["sample_id"] == "sample-b"
+    assert sample["caption"] == "second"
+    assert sample["source_path"] == "/video/b.mp4"
+
+
+def test_mmaudio_feature_cache_discovers_nested_shards(tmp_path: Path) -> None:
+    first = tmp_path / "worker_00000" / "shard_000000"
+    second = tmp_path / "worker_00001" / "shard_000000"
+    _write_cache(first, with_video=True)
+    _write_cache(second, with_video=True)
+
+    assert _expand_cache_root(tmp_path) == [first, second]
+
+
+def test_mmaudio_latent_stats_match_official_reduction(tmp_path: Path) -> None:
+    first = tmp_path / "worker_00000" / "shard_000000"
+    second = tmp_path / "worker_00001" / "shard_000000"
+    _write_cache(first, with_video=True)
+    _write_cache(second, with_video=True)
+
+    first_tensors = tensordict.TensorDict.load_memmap(first)["mean"]
+    second_tensors = tensordict.TensorDict.load_memmap(second)["mean"]
+    complete = torch.cat([first_tensors, second_tensors], dim=0)
+    expected_mean = complete.mean(dim=(0, 1), keepdim=True)
+    expected_std = complete.std(dim=(0, 1), keepdim=True)
+
+    actual_mean, actual_std = compute_mmaudio_latent_stats(
+        tmp_path,
+        latent_seq_len=4,
+        latent_dim=3,
+        chunk_size=1,
+    )
+
+    torch.testing.assert_close(actual_mean, expected_mean)
+    torch.testing.assert_close(actual_std, expected_std)

fastvideo/tests/dataset/test_v2a_preprocessing.py

@@ -0,0 +1,410 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for shared raw V2A preprocessing infrastructure."""
+
+from __future__ import annotations
+
+import json
+import math
+import sys
+import types
+from fractions import Fraction
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+from fastvideo.configs.configs import VideoLoaderType
+from fastvideo.dataset.mmaudio_feature_dataset import _expand_cache_root
+from fastvideo.dataset.v2a_feature_cache import V2AFeatureShardWriter
+from fastvideo.dataset.vggsound import VGGSoundDataset
+from fastvideo.fastvideo_args import WorkloadType
+from fastvideo.pipelines.basic.mmaudio import stages as basic_mmaudio_stages
+from fastvideo.pipelines.pipeline_registry import PipelineType, _PipelineRegistry
+from fastvideo.pipelines.preprocess.mmaudio import stages as mmaudio_stages
+from fastvideo.pipelines.preprocess.mmaudio import torio_media_reader
+from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import (
+    MMAudioTorioRowPreprocessor,
+    PREPROCESSED_MEDIA_KEY,
+    PREPROCESS_ERROR_KEY,
+    preprocess_mmaudio_media_with_torio,
+)
+from fastvideo.workflow.preprocess.preprocess_workflow_v2a import (
+    PreprocessWorkflowV2A, V2AForwardBatchBuilder)
+
+pytest.importorskip("tensordict")
+
+
+class _GenericT2VPreprocessPipeline:
+    pass
+
+
+class _FamilyPreprocessPipeline:
+    pass
+
+
+def _write_fractional_fps_video(
+    path: Path,
+    *,
+    duration_s: float,
+    frame_rate: Fraction = Fraction(30_000, 1_001),
+) -> None:
+    av = pytest.importorskip("av")
+    with av.open(str(path), mode="w") as container:
+        stream = container.add_stream("mpeg4", rate=frame_rate)
+        stream.width = 32
+        stream.height = 32
+        stream.pix_fmt = "yuv420p"
+        time_base = Fraction(frame_rate.denominator, frame_rate.numerator)
+        frame_count = math.ceil(duration_s * float(frame_rate))
+        for index in range(frame_count):
+            pixels = np.full((32, 32, 3), index % 255, dtype=np.uint8)
+            frame = av.VideoFrame.from_ndarray(pixels, format="rgb24")
+            frame.pts = index
+            frame.time_base = time_base
+            for packet in stream.encode(frame):
+                container.mux(packet)
+        for packet in stream.encode():
+            container.mux(packet)
+
+
+def test_vggsound_dataset_maps_hf_metadata_to_clip_name(tmp_path: Path) -> None:
+    (tmp_path / "videos").mkdir()
+    (tmp_path / "vggsound.csv").write_text(
+        "abc123,7,a train sound,train\nxyz789,42,a test sound,test\n",
+        encoding="utf-8",
+    )
+
+    dataset = VGGSoundDataset(tmp_path, split="train")
+
+    assert len(dataset) == 1
+    assert dataset[0] == {
+        "id": "abc123_000007",
+        "caption": "a train sound",
+        "video_path": str(tmp_path / "videos" / "abc123_000007.mp4"),
+    }
+
+
+def test_vggsound_dataset_reads_filtered_caption_manifest(tmp_path: Path) -> None:
+    videos = tmp_path / "videos"
+    videos.mkdir()
+    (videos / "video-one_000010.mp4").touch()
+    manifest = tmp_path / "vgg-val-filtered-caption.tsv"
+    manifest.write_text(
+        "id\tlabel\n"
+        "video-one_000010\tA detailed description of the audible event.\n"
+        "missing_000020\tThis row has no corresponding video.\n",
+        encoding="utf-8",
+    )
+
+    dataset = VGGSoundDataset(tmp_path, split="val", metadata_path=manifest)
+
+    assert len(dataset) == 1
+    assert dataset[0] == {
+        "id": "video-one_000010",
+        "caption": "A detailed description of the audible event.",
+        "video_path": str(videos / "video-one_000010.mp4"),
+    }
+
+
+def test_mmaudio_audio_normalization_matches_split_contract(monkeypatch: pytest.MonkeyPatch) -> None:
+    waveform = torch.tensor([[0.25, -0.5, 0.125, 0.0]], dtype=torch.float32)
+    monkeypatch.setattr(
+        mmaudio_stages.torchaudio,
+        "load",
+        lambda _path: (waveform.clone(), 44_100),
+    )
+
+    train_audio = mmaudio_stages.MMAudioFeatureExtractionStage._load_audio(
+        "sample.mp4",
+        target_sample_rate=44_100,
+        target_samples=4,
+        normalize_audio=True,
+    )
+    eval_audio = mmaudio_stages.MMAudioFeatureExtractionStage._load_audio(
+        "sample.mp4",
+        target_sample_rate=44_100,
+        target_samples=4,
+        normalize_audio=False,
+    )
+
+    torch.testing.assert_close(train_audio.abs().max(), torch.tensor(0.95))
+    torch.testing.assert_close(eval_audio, waveform[0])
+
+
+def test_mmaudio_torio_loader_is_an_explicit_config_choice() -> None:
+    assert VideoLoaderType.from_string("torio") is VideoLoaderType.TORIO
+
+
+def test_mmaudio_torio_reader_preserves_reference_media_contract(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls: list[tuple[str, int, float | None]] = []
+
+    class _StreamInfo:
+        sample_rate = 44_100
+
+    class _FakeStreamingMediaDecoder:
+        def __init__(self, path: str) -> None:
+            assert path == "sample.mp4"
+
+        def add_basic_video_stream(
+            self,
+            *,
+            frames_per_chunk: int,
+            frame_rate: float,
+            format: str,
+        ) -> None:
+            assert format == "rgb24"
+            calls.append(("video", frames_per_chunk, frame_rate))
+
+        def add_basic_audio_stream(self, *, frames_per_chunk: int) -> None:
+            calls.append(("audio", frames_per_chunk, None))
+
+        def fill_buffer(self) -> None:
+            pass
+
+        def pop_chunks(self):
+            clip = torch.arange(8 * 3 * 16 * 20, dtype=torch.uint8).reshape(8, 3, 16, 20)
+            sync = torch.arange(25 * 3 * 16 * 20, dtype=torch.uint8).reshape(25, 3, 16, 20)
+            audio = torch.stack([
+                torch.linspace(-0.5, 0.5, 44_100),
+                torch.linspace(-0.25, 0.25, 44_100),
+            ], dim=1)
+            return clip, sync, audio
+
+        def get_out_stream_info(self, index: int) -> _StreamInfo:
+            assert index == 2
+            return _StreamInfo()
+
+    torio_module = types.ModuleType("torio")
+    torio_module.__path__ = []
+    io_module = types.ModuleType("torio.io")
+    io_module.StreamingMediaDecoder = _FakeStreamingMediaDecoder
+    torio_module.io = io_module
+    monkeypatch.setitem(sys.modules, "torio", torio_module)
+    monkeypatch.setitem(sys.modules, "torio.io", io_module)
+
+    audio, clip, sync, duration = preprocess_mmaudio_media_with_torio(
+        "sample.mp4",
+        duration_s=1.0,
+        target_sample_rate=44_100,
+        target_samples=44_100,
+        normalize_audio=True,
+        clip_size=32,
+        sync_size=24,
+    )
+
+    assert calls == [
+        ("video", 8, 8.0),
+        ("video", 25, 25.0),
+        ("audio", 2**30, None),
+    ]
+    assert audio.shape == (44_100,)
+    torch.testing.assert_close(audio.abs().max(), torch.tensor(0.95))
+    assert clip.shape == (8, 3, 32, 32)
+    assert sync.shape == (25, 3, 24, 24)
+    assert duration == 1.0
+
+
+def test_mmaudio_torio_row_preprocessor_returns_worker_ready_media(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    video_path = tmp_path / "sample.mp4"
+    video_path.touch()
+    expected = (
+        torch.ones(16),
+        torch.ones(8, 3, 4, 4),
+        torch.ones(25, 3, 4, 4),
+        1.0,
+    )
+    monkeypatch.setattr(
+        torio_media_reader,
+        "preprocess_mmaudio_media_with_torio",
+        lambda *_args, **_kwargs: expected,
+    )
+    preprocessor = MMAudioTorioRowPreprocessor(
+        duration_s=1.0,
+        target_sample_rate=16,
+        target_samples=16,
+        normalize_audio=True,
+        clip_size=4,
+        sync_size=4,
+    )
+
+    output = preprocessor({
+        "id": "sample",
+        "caption": "sound",
+        "video_path": str(video_path),
+    })
+
+    assert PREPROCESS_ERROR_KEY not in output
+    media = output[PREPROCESSED_MEDIA_KEY]
+    assert media["audio"] is expected[0]
+    assert media["clip_frames"] is expected[1]
+    assert media["sync_frames"] is expected[2]
+    assert media["effective_duration"] == 1.0
+
+
+def test_v2a_batch_builder_preserves_preprocessed_media_alignment() -> None:
+    media = {
+        "audio": torch.ones(1),
+        "clip_frames": torch.ones(1),
+        "sync_frames": torch.ones(1),
+        "effective_duration": 1.0,
+    }
+    rows = [{
+        "id": "good",
+        "caption": "good sound",
+        "video_path": "good.mp4",
+        PREPROCESSED_MEDIA_KEY: media,
+    }, {
+        "id": "bad",
+        "caption": "bad sound",
+        "video_path": "bad.mp4",
+        PREPROCESS_ERROR_KEY: "decode failed",
+    }]
+
+    batch = V2AForwardBatchBuilder(seed=123)(rows)
+
+    assert batch.extra[PREPROCESSED_MEDIA_KEY] == [media, None]
+    assert batch.extra[PREPROCESS_ERROR_KEY] == [None, "decode failed"]
+
+
+def test_v2a_torio_thread_prefetch_preserves_batch_order() -> None:
+    workflow = object.__new__(PreprocessWorkflowV2A)
+    workflow.training_dataloader = [
+        [{"id": "0"}, {"id": "1"}],
+        [{"id": "2"}, {"id": "3"}],
+        [{"id": "4"}],
+    ]
+    workflow.torio_decode_workers = 2
+    workflow.torio_row_preprocessor = lambda row: {
+        **row,
+        "decoded": True,
+    }
+
+    batches = list(workflow._iter_preprocessed_rows())
+
+    assert [[row["id"] for row in batch] for batch in batches] == [
+        ["0", "1"],
+        ["2", "3"],
+        ["4"],
+    ]
+    assert all(row["decoded"] for batch in batches for row in batch)
+
+
+def test_mmaudio_ffmpeg_sampler_recovers_fractional_fps_boundary(
+    tmp_path: Path,
+) -> None:
+    video_path = tmp_path / "fractional-fps.mp4"
+    _write_fractional_fps_video(video_path, duration_s=10.0)
+
+    clip_frames, sync_frames = basic_mmaudio_stages._read_frames_with_fps_filters(
+        video_path,
+        (8.0, 25.0),
+        start_s=0.0,
+        end_s=8.0,
+    )
+
+    assert clip_frames.shape == (64, 32, 32, 3)
+    assert sync_frames.shape == (200, 32, 32, 3)
+
+
+def test_mmaudio_ffmpeg_sampler_does_not_pad_short_video(
+    tmp_path: Path,
+) -> None:
+    video_path = tmp_path / "short.mp4"
+    _write_fractional_fps_video(video_path, duration_s=7.0)
+
+    clip_frames, sync_frames = basic_mmaudio_stages._read_frames_with_fps_filters(
+        video_path,
+        (8.0, 25.0),
+        start_s=0.0,
+        end_s=8.0,
+    )
+
+    assert clip_frames.shape[0] < 64
+    assert sync_frames.shape[0] < 200
+
+
+def test_mmaudio_inference_keeps_timestamp_sampler_by_default(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls: list[str] = []
+    sampled_frames = [
+        np.zeros((8, 8, 8, 3), dtype=np.uint8),
+        np.zeros((25, 8, 8, 3), dtype=np.uint8),
+    ]
+
+    def timestamp_sampler(*_args, **_kwargs):
+        calls.append("timestamp")
+        return sampled_frames
+
+    def ffmpeg_sampler(*_args, **_kwargs):
+        calls.append("ffmpeg")
+        return sampled_frames
+
+    monkeypatch.setattr(
+        basic_mmaudio_stages,
+        "_read_frames_at_fps",
+        timestamp_sampler,
+    )
+    monkeypatch.setattr(
+        basic_mmaudio_stages,
+        "_read_frames_with_fps_filters",
+        ffmpeg_sampler,
+    )
+
+    basic_mmaudio_stages.preprocess_mmaudio_video(
+        "unused.mp4",
+        duration_s=1.0,
+        clip_size=8,
+        sync_size=8,
+    )
+    basic_mmaudio_stages.preprocess_mmaudio_video(
+        "unused.mp4",
+        duration_s=1.0,
+        clip_size=8,
+        sync_size=8,
+        use_ffmpeg_fps_filter=True,
+    )
+
+    assert calls == ["timestamp", "ffmpeg"]
+
+
+def test_v2a_writer_creates_discoverable_resumable_shards(tmp_path: Path) -> None:
+    writer = V2AFeatureShardWriter(tmp_path, rank=0, samples_per_shard=1)
+    features = {
+        "mean": torch.randn(1, 4, 3),
+        "std": torch.rand(1, 4, 3),
+        "clip_features": torch.randn(1, 6, 8),
+        "sync_features": torch.randn(1, 9, 10),
+        "text_features": torch.randn(1, 5, 7),
+    }
+    writer.append(features, [{"id": "sample-1", "caption": "sound"}])
+    writer.close()
+
+    shards = _expand_cache_root(tmp_path)
+    assert shards == [tmp_path / "worker_00000" / "shard_000000"]
+    metadata = json.loads((shards[0] / "samples.jsonl").read_text().strip())
+    assert metadata["id"] == "sample-1"
+
+    resumed = V2AFeatureShardWriter(tmp_path, rank=0, samples_per_shard=1)
+    assert resumed.contains("sample-1")
+    resumed.append(features, [{"id": "sample-1", "caption": "duplicate"}])
+    resumed.close()
+    assert len(_expand_cache_root(tmp_path)) == 1
+
+
+def test_family_preprocessing_is_enabled_only_for_v2a() -> None:
+    registry = _PipelineRegistry({
+        PipelineType.PREPROCESS.value: {
+            "PreprocessPipelineT2V": _GenericT2VPreprocessPipeline,
+            "ExamplePreprocessPipeline": _FamilyPreprocessPipeline,
+        }
+    })
+
+    assert registry._load_preprocess_pipeline_cls("ExamplePipeline", WorkloadType.V2A) is _FamilyPreprocessPipeline
+    assert registry._load_preprocess_pipeline_cls("ExamplePipeline", WorkloadType.T2V) is _GenericT2VPreprocessPipeline

fastvideo/tests/eval/test_audio_fad_math.py

@@ -0,0 +1,17 @@
+import numpy as np
+import pytest
+
+from fastvideo.eval.metrics.audio.frechet_distance.metric import (
+    _frechet_distance, )
+
+
+def test_frechet_distance_supports_current_scipy_sqrtm_api():
+    mean_a = np.array([0.0, 1.0])
+    mean_b = np.array([1.0, 3.0])
+    covariance = np.array([[2.0, 0.25], [0.25, 1.0]])
+
+    same = _frechet_distance(mean_a, covariance, mean_a, covariance)
+    shifted = _frechet_distance(mean_a, covariance, mean_b, covariance)
+
+    assert same == pytest.approx(0.0, abs=1e-10)
+    assert shifted == pytest.approx(5.0, abs=1e-10)

fastvideo/tests/eval/test_inputs.py

@@ -0,0 +1,97 @@
+import json
+
+import pytest
+import torch
+
+from fastvideo.eval import Video, samples_from_manifest
+from fastvideo.eval.pool import VideoPool
+
+
+def test_samples_from_manifest_resolves_media_paths(tmp_path):
+    media = tmp_path / "media"
+    media.mkdir()
+    video = media / "sample.mp4"
+    audio = media / "sample.wav"
+    video.touch()
+    audio.touch()
+    manifest = tmp_path / "samples.jsonl"
+    manifest.write_text(
+        json.dumps({
+            "id": "sample",
+            "video": "media/sample.mp4",
+            "audio": "media/sample.wav",
+            "text_prompt": "A bell rings.",
+        }) + "\n",
+        encoding="utf-8",
+    )
+
+    samples = samples_from_manifest(manifest)
+
+    assert len(samples) == 1
+    assert samples[0]["id"] == "sample"
+    assert isinstance(samples[0]["video"], Video)
+    assert samples[0]["video"].source == str(video)
+    assert samples[0]["audio"] == str(audio)
+    assert samples[0]["text_prompt"] == "A bell rings."
+
+
+def test_samples_from_manifest_rejects_rows_without_generated_media(tmp_path):
+    manifest = tmp_path / "samples.jsonl"
+    manifest.write_text(
+        json.dumps({
+            "id": "sample",
+            "text_prompt": "A bell rings."
+        }) + "\n",
+        encoding="utf-8",
+    )
+
+    with pytest.raises(ValueError, match="needs at least video or audio"):
+        samples_from_manifest(manifest)
+
+
+def test_samples_from_manifest_rejects_empty_manifest(tmp_path):
+    manifest = tmp_path / "samples.jsonl"
+    manifest.write_text("", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="manifest is empty"):
+        samples_from_manifest(manifest)
+
+
+def test_video_pool_keeps_explicit_fps_over_probe(monkeypatch):
+    frames = torch.zeros(2, 3, 4, 4)
+    probes = []
+
+    def _probe(source):
+        probes.append(source)
+        return 29.97
+
+    monkeypatch.setattr("fastvideo.eval.io.video.probe_video_fps", _probe)
+    monkeypatch.setattr(
+        "fastvideo.eval.io.video.load_video",
+        lambda _source: frames,
+    )
+    video = Video(source="sample.mp4")
+
+    decoded = VideoPool([])._decode({"video": video, "fps": 8.0})
+
+    assert decoded["fps"] == pytest.approx(8.0)
+    assert probes == []
+    assert decoded["video"].frames is frames
+
+
+def test_video_pool_propagates_path_video_fps(monkeypatch):
+    frames = torch.zeros(2, 3, 4, 4)
+    monkeypatch.setattr(
+        "fastvideo.eval.io.video.probe_video_fps",
+        lambda _source: 29.97,
+    )
+    monkeypatch.setattr(
+        "fastvideo.eval.io.video.load_video",
+        lambda _source: frames,
+    )
+    video = Video(source="sample.mp4")
+
+    decoded = VideoPool([])._decode({"video": video})
+
+    assert decoded["fps"] == pytest.approx(29.97)
+    assert decoded["video"].frames is frames

fastvideo/tests/eval/test_v2a_av_benchmark.py

@@ -0,0 +1,163 @@
+import json
+import subprocess
+from pathlib import Path
+
+from fastvideo.eval.v2a import run_av_benchmark
+from fastvideo.eval.v2a import _av_benchmark_runner as runner
+
+
+def _make_gt_cache(path: Path) -> None:
+    path.mkdir()
+    for name in runner._BASE_CACHE_FILES:
+        (path / name).touch()
+
+
+def test_run_av_benchmark_uses_selected_python(monkeypatch, tmp_path):
+    audio_dir = tmp_path / "audio"
+    gt_cache = tmp_path / "gt"
+    audio_dir.mkdir()
+    gt_cache.mkdir()
+    selected_python = tmp_path / "av-benchmark-venv" / "python"
+    captured = {}
+
+    def fake_subprocess_run(command, check):
+        captured["command"] = command
+        captured["check"] = check
+        output = Path(command[command.index("--output") + 1])
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text("{}", encoding="utf-8")
+        return subprocess.CompletedProcess(command, 0)
+
+    monkeypatch.setattr(subprocess, "run", fake_subprocess_run)
+
+    output = run_av_benchmark(
+        audio_dir=audio_dir,
+        gt_cache=gt_cache,
+        python_executable=selected_python,
+        batch_size=8,
+        num_workers=3,
+        skip_clap=True,
+    )
+
+    command = captured["command"]
+    assert command[0] == str(selected_python)
+    assert command[1].endswith("_av_benchmark_runner.py")
+    assert command[command.index("--batch-size") + 1] == "8"
+    assert command[command.index("--num-workers") + 1] == "3"
+    assert "--skip-clap" in command
+    assert captured["check"] is True
+    assert output == audio_dir.parent / "av_benchmark_results.json"
+
+
+def test_official_runner_extracts_and_serializes_metrics(monkeypatch, tmp_path):
+    audio_dir = tmp_path / "audio"
+    gt_cache = tmp_path / "gt"
+    prediction_cache = tmp_path / "pred-cache"
+    output = tmp_path / "results.json"
+    audio_dir.mkdir()
+    (audio_dir / "sample.wav").touch()
+    _make_gt_cache(gt_cache)
+    calls = {}
+
+    def fake_extract(**kwargs):
+        calls["extract"] = kwargs
+        prediction_cache.mkdir(parents=True, exist_ok=True)
+        for name in runner._required_prediction_files(skip_video_related=False, skip_clap=True):
+            (prediction_cache / name).touch()
+
+    def fake_evaluate(**kwargs):
+        calls["evaluate"] = kwargs
+        return {
+            "FD-PASST": 12.5,
+            "KL-PASST-softmax": 0.25,
+            "DeSync": 0.4,
+        }
+
+    monkeypatch.setattr(runner, "_load_backend", lambda: (fake_extract, fake_evaluate))
+
+    payload = runner.run_benchmark(
+        audio_dir=audio_dir,
+        gt_cache=gt_cache,
+        prediction_cache=prediction_cache,
+        output=output,
+        device="cuda",
+        batch_size=4,
+        num_workers=2,
+        audio_length=8.0,
+        recompute=False,
+        skip_video_related=False,
+        skip_clap=True,
+    )
+
+    assert calls["extract"]["audio_path"] == audio_dir.resolve()
+    assert calls["extract"]["batch_size"] == 4
+    assert calls["extract"]["num_workers"] == 2
+    assert calls["evaluate"]["gt_audio_cache"] == gt_cache.resolve()
+    assert payload["num_generated_audio"] == 1
+    assert payload["extraction_ran"] is True
+    assert payload["metrics"]["FD-PASST"] == 12.5
+    assert json.loads(output.read_text(encoding="utf-8"))["metrics"]["DeSync"] == 0.4
+
+
+def test_official_runner_reuses_complete_prediction_cache(monkeypatch, tmp_path):
+    audio_dir = tmp_path / "audio"
+    gt_cache = tmp_path / "gt"
+    prediction_cache = tmp_path / "pred-cache"
+    output = tmp_path / "results.json"
+    audio_dir.mkdir()
+    prediction_cache.mkdir()
+    (audio_dir / "sample.flac").touch()
+    _make_gt_cache(gt_cache)
+    for name in runner._required_prediction_files(skip_video_related=True, skip_clap=True):
+        (prediction_cache / name).touch()
+
+    def fail_extract(**_kwargs):
+        raise AssertionError("complete prediction cache should be reused")
+
+    monkeypatch.setattr(runner, "_load_backend", lambda: (fail_extract, lambda **_kwargs: {"FD-PASST": 1.0}))
+
+    payload = runner.run_benchmark(
+        audio_dir=audio_dir,
+        gt_cache=gt_cache,
+        prediction_cache=prediction_cache,
+        output=output,
+        device="cuda",
+        batch_size=4,
+        num_workers=0,
+        audio_length=8.0,
+        recompute=False,
+        skip_video_related=True,
+        skip_clap=True,
+    )
+
+    assert payload["extraction_ran"] is False
+    assert payload["metrics"] == {"FD-PASST": 1.0}
+
+
+def test_prediction_keys_align_only_unique_leading_underscore_matches(tmp_path):
+    import torch
+
+    gt_cache = tmp_path / "gt"
+    prediction_cache = tmp_path / "pred"
+    gt_cache.mkdir()
+    prediction_cache.mkdir()
+    torch.save({"_123": torch.tensor(1), "normal": torch.tensor(2)}, gt_cache / "passt_logits.pth")
+    torch.save(
+        {"123": torch.tensor(3), "normal": torch.tensor(4), "unknown": torch.tensor(5)},
+        prediction_cache / "passt_logits.pth",
+    )
+
+    stats = runner._align_prediction_cache_keys(
+        gt_cache=gt_cache,
+        prediction_cache=prediction_cache,
+        required_files=("passt_logits.pth",),
+    )
+
+    aligned = torch.load(prediction_cache / "passt_logits.pth", weights_only=True)
+    assert set(aligned) == {"_123", "normal", "unknown"}
+    assert stats == {
+        "exact_before_alignment": 1,
+        "remapped": 1,
+        "ambiguous": 0,
+        "unmatched_after_alignment": 1,
+    }

fastvideo/tests/train/callbacks/test_ddp_posthoc_ema.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only lifecycle test for official MMAudio DDP PostHocEMA."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.train.callbacks.ddp_posthoc_ema import DDPPostHocEMACallback
+from fastvideo.train.utils.distributed_strategy import wrap_module_ddp
+
+
+def test_official_ddp_posthoc_ema_snapshots_and_resumes(
+    tmp_path: Path,
+) -> None:
+    pytest.importorskip("nitrous_ema")
+    rendezvous = tmp_path / "ema_rendezvous"
+    dist.init_process_group(
+        "gloo",
+        init_method=f"file://{rendezvous}",
+        rank=0,
+        world_size=1,
+    )
+    try:
+        transformer = wrap_module_ddp(
+            torch.nn.Linear(2, 2, bias=False),
+            device=torch.device("cpu"),
+            broadcast_buffers=False,
+        )
+        method = SimpleNamespace(
+            student=SimpleNamespace(
+                transformer=transformer,
+                device=torch.device("cpu"),
+            ),
+        )
+        callback = DDPPostHocEMACallback(
+            checkpoint_every=2,
+            checkpoint_folder=str(tmp_path / "snapshots"),
+        )
+        callback.training_config = SimpleNamespace(
+            checkpoint=SimpleNamespace(output_dir=str(tmp_path)),
+        )
+        callback.on_train_start(method)
+        callback.on_training_step_end(method, {}, iteration=1)
+        with torch.no_grad():
+            transformer.module.weight.add_(1)
+        callback.on_training_step_end(method, {}, iteration=2)
+
+        snapshot_root = tmp_path / "snapshots"
+        assert (snapshot_root / "0.2.pt").is_file()
+        assert (snapshot_root / "1.2.pt").is_file()
+
+        restored = DDPPostHocEMACallback(
+            checkpoint_every=2,
+            checkpoint_folder=str(snapshot_root),
+        )
+        restored.training_config = callback.training_config
+        restored.on_train_start(method)
+        restored.load_state_dict(callback.state_dict())
+        assert restored._ema is not None
+        assert int(restored._ema.step) == 2
+
+        callback.on_train_end(method, iteration=2)
+        final = list(snapshot_root.glob(
+            "mmaudio_ema_final_sigma_0p05_step_*.pth",
+        ))
+        assert len(final) == 1
+    finally:
+        dist.destroy_process_group()

fastvideo/tests/train/callbacks/test_mmaudio_validation.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.train.callbacks.mmaudio_validation import (
+    _global_inference_indices,
+    _inference_sample_index, )
+
+
+def test_mmaudio_inference_uses_fixed_global_budget() -> None:
+    assignments = [_global_inference_indices(16, 4, rank) for rank in range(4)]
+
+    assert all(len(rank_indices) == 4 for rank_indices in assignments)
+    assert sorted(index for rank_indices in assignments
+                  for index in rank_indices) == list(range(16))
+
+
+def test_mmaudio_inference_rows_cover_validation_batch() -> None:
+    available = 8
+    assignments = [_global_inference_indices(16, 4, rank) for rank in range(4)]
+
+    rows = {
+        _inference_sample_index(global_index, local_index, available)
+        for rank_indices in assignments
+        for local_index, global_index in enumerate(rank_indices)
+    }
+
+    assert rows == set(range(min(16, available)))
+
+
+def test_mmaudio_inference_pads_collective_calls() -> None:
+    assignments = [_global_inference_indices(16, 6, rank) for rank in range(6)]
+
+    assert all(len(rank_indices) == 3 for rank_indices in assignments)
+    valid = [
+        index for rank_indices in assignments for index in rank_indices
+        if index is not None
+    ]
+    assert sorted(valid) == list(range(16))
+    assert sum(index is None for rank_indices in assignments
+               for index in rank_indices) == 2

fastvideo/tests/train/callbacks/test_posthoc_ema.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+from fastvideo.train.callbacks.posthoc_ema import (
+    PostHocEMACallback,
+    sigma_rel_to_gamma,
+    solve_posthoc_weights,
+)
+
+
+class _FakeMethod:
+
+    def __init__(self) -> None:
+        self.student = SimpleNamespace(
+            transformer=torch.nn.Linear(2, 2, bias=False),
+        )
+
+
+def test_posthoc_math_is_finite() -> None:
+    gamma_005 = sigma_rel_to_gamma(0.05)
+    gamma_01 = sigma_rel_to_gamma(0.1)
+    weights = solve_posthoc_weights(
+        torch.tensor([5, 5, 10, 10]),
+        torch.tensor([gamma_005, gamma_01, gamma_005, gamma_01]),
+        10,
+        gamma_005,
+    )
+    assert weights.shape == (4, )
+    assert torch.isfinite(weights).all()
+
+
+def test_posthoc_update_matches_nitrous_ema_step_order(tmp_path: Path) -> None:
+    callback = PostHocEMACallback(
+        checkpoint_every=100,
+        checkpoint_folder=str(tmp_path),
+    )
+    callback.training_config = SimpleNamespace(
+        checkpoint=SimpleNamespace(output_dir=str(tmp_path)),
+    )
+    method = _FakeMethod()
+    callback.on_train_start(method)
+
+    first_weight = method.student.transformer.weight.detach().float().clone()
+    callback.on_training_step_end(method, {}, iteration=1)
+    torch.testing.assert_close(callback._ema_models[0].shadow["weight"], first_weight)
+
+    with torch.no_grad():
+        method.student.transformer.weight.add_(1)
+    second_weight = method.student.transformer.weight.detach().float().clone()
+    callback.on_training_step_end(method, {}, iteration=2)
+
+    gamma = sigma_rel_to_gamma(0.05)
+    # nitrous-ema increments KarrasEMA.step before reading beta.
+    official_beta = ((1 - 1 / 3)**(1 + gamma) *
+                     (1 - 1 / 4)**(1 + gamma))**0.5
+    expected = first_weight * official_beta + second_weight * (1 - official_beta)
+    torch.testing.assert_close(callback._ema_models[0].shadow["weight"], expected)
+
+
+def test_posthoc_callback_snapshots_and_synthesizes(tmp_path: Path) -> None:
+    callback = PostHocEMACallback(
+        checkpoint_every=2,
+        checkpoint_folder=str(tmp_path),
+    )
+    callback.training_config = SimpleNamespace(
+        checkpoint=SimpleNamespace(output_dir=str(tmp_path)),
+    )
+    method = _FakeMethod()
+    callback.on_train_start(method)
+
+    callback.on_training_step_end(method, {}, iteration=1)
+    with torch.no_grad():
+        method.student.transformer.weight.add_(1)
+    callback.on_training_step_end(method, {}, iteration=2)
+
+    assert len(list((tmp_path / "rank_00000").glob("*.pt"))) == 2
+    synthesized = callback.synthesize_local_shard()
+    assert synthesized is not None
+    assert set(synthesized) == {"weight"}
+    assert torch.isfinite(synthesized["weight"]).all()
+
+    state = callback.state_dict()
+    restored = PostHocEMACallback(
+        checkpoint_every=2,
+        checkpoint_folder=str(tmp_path),
+    )
+    restored.training_config = callback.training_config
+    restored.on_train_start(method)
+    restored.load_state_dict(state)
+    assert restored.state_dict()["calls"] == 2
+
+    callback.on_train_end(method, iteration=2)
+    final_paths = list((tmp_path / "rank_00000").glob(
+        "synthesized_sigma_0p05_step_*.pt",
+    ))
+    assert len(final_paths) == 1
+    final_state = torch.load(final_paths[0], weights_only=True)
+    assert set(final_state) == {"weight"}

fastvideo/tests/train/methods/test_mmaudio_flow_matching.py

@@ -0,0 +1,491 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU unit tests for the MMAudio flow-matching training adapter."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from fastvideo.configs.models.dits import MMAUDIO_VARIANT_ARCHITECTURES
+from fastvideo.configs.pipelines.mmaudio import MMAudioSmall44kV2AConfig
+from fastvideo.models.dits.mmaudio import MMAudioTransformer
+from fastvideo.train.methods.fine_tuning.flow_matching import (
+    FlowMatchingFineTuneMethod,
+)
+from fastvideo.train.models.mmaudio import MMAudioModel
+from fastvideo.train.models.mmaudio.mmaudio import _load_or_compute_latent_stats
+from fastvideo.train.utils.config import RunConfig
+from fastvideo.train.utils.training_config import (
+    DataConfig,
+    DistributedConfig,
+    ModelTrainingConfig,
+    OptimizerConfig,
+    TrainingConfig,
+    TrainingLoopConfig,
+)
+
+
+class _TinyMMAudioTransformer(nn.Module):
+    latent_seq_len = 4
+    latent_dim = 3
+    clip_seq_len = 6
+    sync_seq_len = 9
+
+    def __init__(self, *, v2: bool = False) -> None:
+        super().__init__()
+        self.v2 = v2
+        self.unshard_calls = 0
+        arch = SimpleNamespace(
+            clip_dim=8,
+            sync_dim=10,
+            text_seq_len=5,
+            text_dim=7,
+        )
+        self.config = SimpleNamespace(arch_config=arch)
+        self.proj = nn.Linear(self.latent_dim, self.latent_dim, bias=False)
+        self.latent_mean = nn.Parameter(
+            torch.tensor([[[0.5, -0.25, 1.0]]]),
+            requires_grad=False,
+        )
+        self.latent_std = nn.Parameter(
+            torch.tensor([[[2.0, 0.5, 4.0]]]),
+            requires_grad=False,
+        )
+        self.empty_string_feat = nn.Parameter(
+            torch.full((5, 7), 13.0),
+            requires_grad=False,
+        )
+        self.empty_clip_feat = nn.Parameter(torch.full((1, 8), 11.0))
+        self.empty_sync_feat = nn.Parameter(torch.full((1, 10), 12.0))
+
+    def unshard(self) -> None:
+        self.unshard_calls += 1
+
+    def normalize(self, value: torch.Tensor) -> torch.Tensor:
+        return value.sub(self.latent_mean).div(self.latent_std)
+
+    def get_empty_clip_sequence(self, batch_size: int) -> torch.Tensor:
+        return self.empty_clip_feat.unsqueeze(0).expand(
+            batch_size, self.clip_seq_len, -1
+        )
+
+    def get_empty_sync_sequence(self, batch_size: int) -> torch.Tensor:
+        return self.empty_sync_feat.unsqueeze(0).expand(
+            batch_size, self.sync_seq_len, -1
+        )
+
+    def get_empty_string_sequence(self, batch_size: int) -> torch.Tensor:
+        return self.empty_string_feat.unsqueeze(0).expand(batch_size, -1, -1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: dict[str, torch.Tensor],
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        del timestep
+        condition_dependency = (
+            encoder_hidden_states["clip_features"].sum()
+            + encoder_hidden_states["sync_features"].sum()
+        ) * 0.0
+        return self.proj(hidden_states) + condition_dependency
+
+
+def _training_config() -> TrainingConfig:
+    return TrainingConfig(
+        data=DataConfig(
+            data_path="",
+            preprocessed_data_type="mmaudio_features",
+            train_batch_size=2,
+            training_cfg_rate=0.4,
+            seed=123,
+        ),
+        optimizer=OptimizerConfig(
+            learning_rate=1e-4,
+            betas=(0.9, 0.95),
+            eps=1e-6,
+        ),
+        loop=TrainingLoopConfig(max_train_steps=2),
+    )
+
+
+def _raw_batch() -> dict[str, torch.Tensor]:
+    return {
+        "audio_latent_mean": (
+            torch.arange(24, dtype=torch.float32).reshape(2, 4, 3) / 10
+        ),
+        "audio_latent_std": torch.full((2, 4, 3), 0.2),
+        "clip_features": torch.zeros(2, 6, 8),
+        "sync_features": torch.zeros(2, 9, 10),
+        "text_features": torch.zeros(2, 5, 7),
+        "video_exists": torch.tensor([False, True]),
+        "text_exists": torch.tensor([True, False]),
+    }
+
+
+def test_prepare_batch_matches_official_rng_and_flow_contract() -> None:
+    config = _training_config()
+    transformer = _TinyMMAudioTransformer()
+    model = MMAudioModel(
+        init_from="unused",
+        training_config=config,
+        transformer=transformer,
+    )
+    raw = _raw_batch()
+    actual_generator = torch.Generator().manual_seed(7)
+    expected_generator = torch.Generator().manual_seed(7)
+
+    batch = model.prepare_batch(raw, generator=actual_generator)
+
+    posterior_noise = torch.empty_like(raw["audio_latent_mean"]).normal_(
+        generator=expected_generator
+    )
+    clean = raw["audio_latent_mean"] + raw["audio_latent_std"] * posterior_noise
+    clean = (clean - transformer.latent_mean) / transformer.latent_std
+    timestep = torch.sigmoid(
+        torch.randn(2, generator=expected_generator)
+        * config.model.logit_std
+        + config.model.logit_mean
+    )
+    prior = torch.empty_like(clean).normal_(generator=expected_generator)
+    noisy = (1 - timestep[:, None, None]) * prior + timestep[:, None, None] * clean
+    null_video = torch.rand(2, generator=expected_generator) < 0.4
+    null_text = torch.rand(2, generator=expected_generator) < 0.4
+
+    torch.testing.assert_close(batch.latents, clean)
+    torch.testing.assert_close(batch.timesteps, timestep)
+    torch.testing.assert_close(batch.noise, prior)
+    torch.testing.assert_close(batch.noisy_model_input, noisy)
+    torch.testing.assert_close(batch.training_target, clean - prior)
+    assert transformer.unshard_calls == 1
+
+    assert batch.conditional_dict is not None
+    expected_video_null = torch.tensor([True, False]) | null_video
+    expected_text_null = torch.tensor([False, True]) | null_text
+    clip = batch.conditional_dict["clip_features"]
+    sync = batch.conditional_dict["sync_features"]
+    text = batch.conditional_dict["text_features"]
+    torch.testing.assert_close(
+        clip[expected_video_null],
+        torch.full_like(clip[expected_video_null], 11.0),
+    )
+    torch.testing.assert_close(
+        sync[expected_video_null],
+        torch.full_like(sync[expected_video_null], 12.0),
+    )
+    torch.testing.assert_close(
+        text[expected_text_null],
+        torch.full_like(text[expected_text_null], 13.0),
+    )
+
+    assert transformer.latent_mean.requires_grad is False
+    assert transformer.latent_std.requires_grad is False
+    assert transformer.empty_string_feat.requires_grad is False
+    assert transformer.empty_clip_feat.requires_grad is True
+    assert transformer.empty_sync_feat.requires_grad is True
+
+
+def test_sequence_length_update_in_inference_mode_keeps_normal_buffers() -> None:
+    owner = SimpleNamespace()
+    inference_modes: list[bool] = []
+
+    def initialize_rotations() -> None:
+        inference_modes.append(torch.is_inference_mode_enabled())
+        owner.latent_rot = torch.ones(1)
+        owner.clip_rot = torch.ones(1)
+
+    owner.initialize_rotations = initialize_rotations
+
+    with torch.inference_mode():
+        MMAudioTransformer.update_seq_lengths(owner, 4, 6, 9)
+
+    assert owner._latent_seq_len == 4
+    assert owner._clip_seq_len == 6
+    assert owner._sync_seq_len == 9
+    assert inference_modes == [False]
+    assert not torch.is_inference(owner.latent_rot)
+    assert not torch.is_inference(owner.clip_rot)
+
+
+def test_flow_matching_method_runs_loss_and_backward(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    config = _training_config()
+    monkeypatch.setattr(
+        "fastvideo.dataset.mmaudio_feature_dataset."
+        "build_mmaudio_feature_dataloader",
+        lambda *args, **kwargs: None,
+    )
+    model = MMAudioModel(
+        init_from="unused",
+        training_config=config,
+        transformer=_TinyMMAudioTransformer(),
+    )
+    cfg = RunConfig(
+        models={},
+        method={
+            "_target_": (
+                "fastvideo.train.methods.fine_tuning.flow_matching."
+                "FlowMatchingFineTuneMethod"
+            )
+        },
+        training=config,
+        callbacks={},
+        raw={},
+    )
+    method = FlowMatchingFineTuneMethod(
+        cfg=cfg,
+        role_models={"student": model},
+    )
+    method.cuda_generator = torch.Generator().manual_seed(17)
+
+    eager_parameters = model.eager_optimizer_state_parameters()
+    assert len(eager_parameters) == 2
+    for parameter in eager_parameters:
+        assert set(method._student_optimizer.state[parameter]) == {
+            "step",
+            "exp_avg",
+            "exp_avg_sq",
+        }
+
+    loss_map, outputs, metrics = method.single_train_step(_raw_batch(), 0)
+
+    assert torch.isfinite(loss_map["total_loss"])
+    assert loss_map["flow_matching_loss"] is loss_map["total_loss"]
+    assert metrics == {}
+    method.backward(loss_map, outputs)
+    assert model.transformer.proj.weight.grad is not None
+    assert torch.count_nonzero(model.transformer.proj.weight.grad) > 0
+    assert model.transformer.empty_clip_feat.grad is not None
+    assert model.transformer.empty_sync_feat.grad is not None
+
+
+def test_flow_matching_method_compiles_inner_transformer_forward(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    config = _training_config()
+    config.model = ModelTrainingConfig(
+        compile_train_fn=True,
+        torch_compile_kwargs={"backend": "eager"},
+    )
+    monkeypatch.setattr(
+        "fastvideo.dataset.mmaudio_feature_dataset."
+        "build_mmaudio_feature_dataloader",
+        lambda *args, **kwargs: None,
+    )
+    compile_calls: list[tuple[object, dict[str, object]]] = []
+
+    def _compile(fn, **kwargs):
+        compile_calls.append((fn, kwargs))
+        return fn
+
+    monkeypatch.setattr(torch, "compile", _compile)
+    model = MMAudioModel(
+        init_from="unused",
+        training_config=config,
+        transformer=_TinyMMAudioTransformer(),
+    )
+    config.distributed = DistributedConfig(strategy="ddp")
+    cfg = RunConfig(
+        models={},
+        method={
+            "_target_": (
+                "fastvideo.train.methods.fine_tuning.flow_matching."
+                "FlowMatchingFineTuneMethod"
+            )
+        },
+        training=config,
+        callbacks={},
+        raw={},
+    )
+    method = FlowMatchingFineTuneMethod(
+        cfg=cfg,
+        role_models={"student": model},
+    )
+    method.cuda_generator = torch.Generator().manual_seed(19)
+
+    loss_map, _, _ = method.single_train_step(_raw_batch(), 0)
+
+    assert len(compile_calls) == 1
+    assert compile_calls[0][1] == {"backend": "eager", "fullgraph": True}
+    assert torch.isfinite(loss_map["total_loss"])
+
+
+def test_compile_train_fn_rejects_fsdp(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    config = _training_config()
+    config.model = ModelTrainingConfig(compile_train_fn=True)
+    monkeypatch.setattr(
+        "fastvideo.dataset.mmaudio_feature_dataset."
+        "build_mmaudio_feature_dataloader",
+        lambda *args, **kwargs: None,
+    )
+    model = MMAudioModel(
+        init_from="unused",
+        training_config=config,
+        transformer=_TinyMMAudioTransformer(),
+    )
+    cfg = RunConfig(
+        models={},
+        method={
+            "_target_": (
+                "fastvideo.train.methods.fine_tuning.flow_matching."
+                "FlowMatchingFineTuneMethod"
+            )
+        },
+        training=config,
+        callbacks={},
+        raw={},
+    )
+
+    with pytest.raises(ValueError, match="strategy=ddp"):
+        FlowMatchingFineTuneMethod(
+            cfg=cfg,
+            role_models={"student": model},
+        )
+
+
+def test_v2_training_is_rejected_by_default() -> None:
+    with pytest.raises(ValueError, match="does not support `_v2`"):
+        MMAudioModel(
+            init_from="unused",
+            training_config=_training_config(),
+            transformer=_TinyMMAudioTransformer(v2=True),
+        )
+
+
+def test_official_mmaudio_variant_architectures() -> None:
+    expected = {
+        "small_16k": (20, 448, 12, 8, 7, 250, False),
+        "small_44k": (40, 448, 12, 8, 7, 345, False),
+        "medium_44k": (40, 896, 12, 8, 14, 345, False),
+        "large_44k": (40, 896, 21, 14, 14, 345, False),
+        "large_44k_v2": (40, 896, 21, 14, 14, 345, True),
+    }
+    actual = {
+        name: (
+            values["latent_dim"],
+            values["hidden_dim"],
+            values["depth"],
+            values["fused_depth"],
+            values["num_heads"],
+            values["latent_seq_len"],
+            values["v2"],
+        )
+        for name, values in MMAUDIO_VARIANT_ARCHITECTURES.items()
+    }
+    assert actual == expected
+
+
+def test_from_scratch_builds_variant_without_loading_dit_checkpoint(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    captured: dict[str, object] = {}
+
+    monkeypatch.setattr(
+        "fastvideo.train.models.mmaudio.mmaudio._load_or_compute_latent_stats",
+        lambda *args, **kwargs: (
+            torch.zeros(1, 1, 40),
+            torch.ones(1, 1, 40),
+        ),
+    )
+    monkeypatch.setattr(
+        "fastvideo.train.models.mmaudio.mmaudio._load_empty_string_features",
+        lambda *args, **kwargs: torch.zeros(77, 1024),
+    )
+
+    def _build(**kwargs):
+        captured.update(kwargs)
+        return _TinyMMAudioTransformer()
+
+    monkeypatch.setattr(
+        "fastvideo.train.models.mmaudio.mmaudio.build_fsdp_model_from_scratch",
+        _build,
+    )
+    model = MMAudioModel(
+        training_config=_training_config(),
+        variant="small_44k",
+        from_scratch=True,
+        empty_string_features_path="unused.safetensors",
+    )
+
+    assert model._init_from == "scratch:small_44k"
+    assert isinstance(model.training_config.pipeline_config, MMAudioSmall44kV2AConfig)
+    assert captured["model_cls"].__name__ == "MMAudioTransformer"
+    assert captured["init_params"]["hf_config"]["v2"] is False
+
+
+def test_from_scratch_rejects_v2_variant_before_loading_assets() -> None:
+    with pytest.raises(ValueError, match="does not train 'large_44k_v2'"):
+        MMAudioModel(
+            training_config=_training_config(),
+            variant="large_44k_v2",
+            from_scratch=True,
+            empty_string_features_path="unused.safetensors",
+        )
+
+
+def test_44k_variants_share_cached_latent_statistics(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    payload = {
+        "variant": "small_44k",
+        "latent_mean": torch.zeros(1, 1, 40),
+        "latent_std": torch.ones(1, 1, 40),
+    }
+
+    class _ReceiverGroup:
+        is_first_rank = False
+
+        @staticmethod
+        def broadcast_object(value, src=0):
+            del value, src
+            return payload
+
+    monkeypatch.setattr(
+        "fastvideo.train.models.mmaudio.mmaudio.get_world_group",
+        lambda: _ReceiverGroup(),
+    )
+    mean, std = _load_or_compute_latent_stats(
+        _training_config(),
+        variant="medium_44k",
+        cache_path=None,
+        chunk_size=1,
+    )
+
+    torch.testing.assert_close(mean, payload["latent_mean"])
+    torch.testing.assert_close(std, payload["latent_std"])
+
+
+def test_16k_rejects_44k_cached_latent_statistics(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    payload = {
+        "variant": "small_44k",
+        "latent_mean": torch.zeros(1, 1, 40),
+        "latent_std": torch.ones(1, 1, 40),
+    }
+
+    class _ReceiverGroup:
+        is_first_rank = False
+
+        @staticmethod
+        def broadcast_object(value, src=0):
+            del value, src
+            return payload
+
+    monkeypatch.setattr(
+        "fastvideo.train.models.mmaudio.mmaudio.get_world_group",
+        lambda: _ReceiverGroup(),
+    )
+    with pytest.raises(ValueError, match="small_44k.*small_16k"):
+        _load_or_compute_latent_stats(
+            _training_config(),
+            variant="small_16k",
+            cache_path=None,
+            chunk_size=1,
+        )

fastvideo/tests/train/utils/test_config.py

@@ -54,6 +54,7 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None:
     cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
     t = cfg.training
 
+    assert t.distributed.strategy == "fsdp"
     assert t.distributed.num_gpus == 1
     assert t.distributed.tp_size == 1
     assert t.distributed.sp_size == 1
@@ -67,9 +68,13 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None:
 
     assert t.optimizer.learning_rate == 0.0
     assert t.optimizer.betas == (0.9, 0.999)
+    assert t.optimizer.eps == pytest.approx(1e-8)
     assert t.optimizer.weight_decay == 0.0
     assert t.optimizer.lr_scheduler == "constant"
     assert t.optimizer.min_lr_ratio == 0.5
+    assert t.optimizer.lr_milestones == ()
+    assert t.optimizer.lr_gamma == pytest.approx(0.1)
+    assert t.optimizer.fused is False
 
     assert t.loop.max_train_steps == 0
     assert t.loop.gradient_accumulation_steps == 1
@@ -87,6 +92,8 @@ def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None:
     assert t.model.weighting_scheme == "uniform"
     assert t.model.precondition_outputs is False
     assert t.model.moba_config == {}
+    assert t.model.compile_train_fn is False
+    assert t.model.torch_compile_kwargs == {}
 
     assert t.dit_precision == "fp32"
     assert t.vsa_sparsity == 0.0
@@ -106,6 +113,7 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
         },
         "data": {
             "data_path": "/some/path",
+            "preprocessed_data_type": "mmaudio_features",
             "train_batch_size": 2,
             "dataloader_num_workers": 4,
             "training_cfg_rate": 0.1,
@@ -118,10 +126,13 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
         "optimizer": {
             "learning_rate": 1e-4,
             "betas": [0.9, 0.95],
+            "eps": 1e-6,
             "weight_decay": 0.01,
             "lr_scheduler": "cosine",
             "lr_warmup_steps": 100,
             "min_lr_ratio": 0.1,
+            "lr_milestones": [240, 270],
+            "lr_gamma": 0.2,
         },
         "loop": {
             "max_train_steps": 1000,
@@ -149,6 +160,11 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
             "logit_mean": 0.5,
             "logit_std": 1.5,
             "precondition_outputs": True,
+            "compile_train_fn": True,
+            "torch_compile_kwargs": {
+                "backend": "eager",
+                "dynamic": False,
+            },
         },
         "dit_precision": "bf16",
     }
@@ -161,14 +177,18 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
 
     assert t.data.train_batch_size == 2
     assert t.data.data_path == "/some/path"
+    assert t.data.preprocessed_data_type == "mmaudio_features"
     assert t.data.num_frames == 33
     assert t.data.seed == 42
 
     assert t.optimizer.learning_rate == pytest.approx(1e-4)
     assert t.optimizer.betas == (0.9, 0.95)
+    assert t.optimizer.eps == pytest.approx(1e-6)
     assert t.optimizer.lr_scheduler == "cosine"
     assert t.optimizer.lr_warmup_steps == 100
     assert t.optimizer.min_lr_ratio == pytest.approx(0.1)
+    assert t.optimizer.lr_milestones == (240, 270)
+    assert t.optimizer.lr_gamma == pytest.approx(0.2)
 
     assert t.loop.max_train_steps == 1000
     assert t.loop.gradient_accumulation_steps == 4
@@ -185,6 +205,11 @@ def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
     assert t.vsa_sparsity == pytest.approx(0.5)
     assert t.model.weighting_scheme == "logit_normal"
     assert t.model.precondition_outputs is True
+    assert t.model.compile_train_fn is True
+    assert t.model.torch_compile_kwargs == {
+        "backend": "eager",
+        "dynamic": False,
+    }
     assert t.dit_precision == "bf16"
 
 
@@ -260,6 +285,20 @@ def test_betas_parses_list_and_string_forms(
     assert cfg.training.optimizer.betas == expected
 
 
+def test_distributed_strategy_parses_ddp(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    data["training"] = {"distributed": {"strategy": "DDP"}}
+    cfg = load_run_config(_write_yaml(tmp_path, data))
+    assert cfg.training.distributed.strategy == "ddp"
+
+
+def test_distributed_strategy_rejects_unknown_value(tmp_path: Path) -> None:
+    data = _minimal_yaml()
+    data["training"] = {"distributed": {"strategy": "zero3"}}
+    with pytest.raises(ValueError, match="strategy must be one of"):
+        load_run_config(_write_yaml(tmp_path, data))
+
+
 def test_pipeline_quant_config_resolves_for_load_and_export(tmp_path: Path) -> None:
     from fastvideo.layers.quantization.nvfp4_qat_train_config import (
         NVFP4QATTrainConfig, )

fastvideo/tests/train/utils/test_distributed_strategy.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU-only coverage for the modular trainer's DDP strategy."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.train.utils.distributed_strategy import (
+    build_replicated_model_from_scratch,
+    normalize_distributed_strategy,
+    unwrap_ddp_module,
+    wrap_module_ddp,
+)
+from fastvideo.training.checkpointing_utils import ModelWrapper
+
+
+class _TinyModel(torch.nn.Module):
+
+    def __init__(self, width: int = 2) -> None:
+        super().__init__()
+        self.linear = torch.nn.Linear(width, width)
+        self.public_value = 17
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        return self.linear(value)
+
+    def helper(self) -> int:
+        return self.public_value
+
+
+@pytest.fixture
+def local_process_group(tmp_path: Path):
+    owns_group = not dist.is_initialized()
+    if owns_group:
+        rendezvous = tmp_path / "ddp_rendezvous"
+        dist.init_process_group(
+            "gloo",
+            init_method=f"file://{rendezvous}",
+            rank=0,
+            world_size=1,
+        )
+    try:
+        yield
+    finally:
+        if owns_group:
+            dist.destroy_process_group()
+
+
+def test_strategy_normalization() -> None:
+    assert normalize_distributed_strategy(None) == "fsdp"
+    assert normalize_distributed_strategy(" DDP ") == "ddp"
+    with pytest.raises(ValueError, match="strategy must be one of"):
+        normalize_distributed_strategy("deepspeed")
+
+
+def test_replicated_initialization_is_seeded() -> None:
+    first = build_replicated_model_from_scratch(
+        _TinyModel,
+        {"width": 3},
+        device=torch.device("cpu"),
+        default_dtype=torch.float32,
+        seed=123,
+    )
+    second = build_replicated_model_from_scratch(
+        _TinyModel,
+        {"width": 3},
+        device=torch.device("cpu"),
+        default_dtype=torch.float32,
+        seed=123,
+    )
+    for left, right in zip(first.parameters(), second.parameters(), strict=True):
+        torch.testing.assert_close(left, right)
+
+
+def test_ddp_adapter_forward_delegation_and_checkpoint(
+    local_process_group,
+) -> None:
+    del local_process_group
+    wrapped = wrap_module_ddp(
+        _TinyModel(),
+        device=torch.device("cpu"),
+        broadcast_buffers=False,
+    )
+    assert wrapped.helper() == 17
+    assert unwrap_ddp_module(wrapped) is wrapped.module
+
+    output = wrapped(torch.ones(2, 2)).sum()
+    output.backward()
+    assert all(parameter.grad is not None for parameter in wrapped.parameters())
+
+    # torch.distributed.checkpoint canonicalizes DDP's ``module.`` prefix.
+    # ModelWrapper must apply the same canonicalization when filtering out
+    # frozen parameters, otherwise it silently emits an empty checkpoint.
+    state = ModelWrapper(wrapped).state_dict()
+    assert set(state) == {"linear.weight", "linear.bias"}

fastvideo/tests/train/utils/test_optimizer.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for FastVideo modular trainer optimizer construction."""
+
+import pytest
+import torch
+
+from fastvideo.train.utils.optimizer import (
+    build_optimizer_and_scheduler,
+    seed_adamw_parameter_state,
+)
+from fastvideo.train.utils.training_config import OptimizerConfig, TrainingLoopConfig
+from fastvideo.training.checkpointing_utils import SchedulerWrapper
+
+
+def test_multistep_with_warmup_matches_mmaudio_scheduler_composition() -> None:
+    parameter = torch.nn.Parameter(torch.zeros(()))
+    config = OptimizerConfig(
+        learning_rate=1e-4,
+        lr_scheduler="multistep_with_warmup",
+        lr_warmup_steps=2,
+        lr_milestones=(3, 5),
+        lr_gamma=0.1,
+    )
+    optimizer, scheduler = build_optimizer_and_scheduler(
+        params=[parameter],
+        optimizer_config=config,
+        loop_config=TrainingLoopConfig(max_train_steps=10),
+        learning_rate=config.learning_rate,
+        betas=config.betas,
+        scheduler_name=config.lr_scheduler,
+    )
+
+    assert isinstance(scheduler, torch.optim.lr_scheduler.SequentialLR)
+    initial_lr = optimizer.param_groups[0]["lr"]
+    assert initial_lr == pytest.approx(1e-4 / 3)
+
+    values = []
+    for _ in range(8):
+        optimizer.step()
+        scheduler.step()
+        values.append(optimizer.param_groups[0]["lr"])
+
+    assert values[1] == pytest.approx(1e-4)
+    assert values[4] == pytest.approx(1e-5)
+    assert values[6] == pytest.approx(1e-6)
+
+
+def test_multistep_with_warmup_requires_milestones() -> None:
+    parameter = torch.nn.Parameter(torch.zeros(()))
+    with pytest.raises(ValueError, match="lr_milestones"):
+        build_optimizer_and_scheduler(
+            params=[parameter],
+            optimizer_config=OptimizerConfig(
+                learning_rate=1e-4,
+                lr_scheduler="multistep_with_warmup",
+                lr_warmup_steps=2,
+            ),
+            loop_config=TrainingLoopConfig(max_train_steps=10),
+            learning_rate=1e-4,
+            betas=(0.9, 0.95),
+            scheduler_name="multistep_with_warmup",
+        )
+
+
+@pytest.mark.parametrize("fused", [False, True])
+def test_seed_adamw_state_uses_expected_step_device(fused: bool) -> None:
+    parameter = torch.nn.Parameter(torch.zeros(()))
+    optimizer = torch.optim.AdamW([parameter], fused=fused)
+    seed_adamw_parameter_state(optimizer, parameter)
+    expected = parameter.device if fused else torch.device("cpu")
+    assert optimizer.state[parameter]["step"].device == expected
+
+
+def test_seed_adamw_state_rejects_unowned_parameter() -> None:
+    owned = torch.nn.Parameter(torch.zeros(()))
+    unowned = torch.nn.Parameter(torch.zeros(()))
+    optimizer = torch.optim.AdamW([owned])
+    with pytest.raises(ValueError, match="unowned"):
+        seed_adamw_parameter_state(optimizer, unowned)
+
+
+def test_scheduler_wrapper_stabilizes_multistep_milestone_keys() -> None:
+    parameter = torch.nn.Parameter(torch.zeros(()))
+    config = OptimizerConfig(
+        learning_rate=1e-4,
+        lr_scheduler="multistep_with_warmup",
+        lr_warmup_steps=2,
+        lr_milestones=(3, 5),
+    )
+    _, scheduler = build_optimizer_and_scheduler(
+        params=[parameter],
+        optimizer_config=config,
+        loop_config=TrainingLoopConfig(max_train_steps=10),
+        learning_rate=config.learning_rate,
+        betas=config.betas,
+        scheduler_name=config.lr_scheduler,
+    )
+    wrapper = SchedulerWrapper(scheduler)
+    state = wrapper.state_dict()
+    milestones = state["scheduler"]["_schedulers"][1]["milestones"]
+    assert milestones == {"3": 1, "5": 1}
+
+    # Model what DCP returns after flattening mapping keys to strings.
+    wrapper.load_state_dict(state)
+    restored = scheduler.state_dict()["_schedulers"][1]["milestones"]
+    assert restored == {3: 1, 5: 1}

fastvideo/train/callbacks/ddp_posthoc_ema.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Official nitrous PostHocEMA lifecycle for native DDP training."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+from fastvideo.logger import init_logger
+from fastvideo.train.callbacks.callback import Callback
+from fastvideo.train.utils.distributed_strategy import unwrap_ddp_module
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+
+logger = init_logger(__name__)
+
+
+class DDPPostHocEMACallback(Callback):
+    """Run the exact ``nitrous_ema.PostHocEMA`` used by MMAudio.
+
+    Official MMAudio owns two complete CUDA EMA models on local rank zero and
+    checkpoints them every 5,000 optimizer steps. Other DDP ranks do not hold
+    EMA copies. This callback preserves that lifecycle while the ordinary
+    FastVideo FSDP callback continues to maintain rank-local CPU shards.
+    """
+
+    def __init__(
+        self,
+        *,
+        sigma_rels: list[float] | tuple[float, ...] = (0.05, 0.1),
+        update_every: int = 1,
+        checkpoint_every: int = 5000,
+        checkpoint_folder: str | None = None,
+        start_iter: int = 0,
+        default_output_sigma: float = 0.05,
+        step_size_correction: bool = True,
+    ) -> None:
+        self.sigma_rels = tuple(float(value) for value in sigma_rels)
+        self.update_every = max(1, int(update_every))
+        self.checkpoint_every = max(1, int(checkpoint_every))
+        self.checkpoint_folder = checkpoint_folder
+        self.start_iter = int(start_iter)
+        self.default_output_sigma = float(default_output_sigma)
+        self.step_size_correction = bool(step_size_correction)
+
+        self._ema: Any | None = None
+        self._calls = 0
+        self._rank = 0
+        self._checkpoint_root: Path | None = None
+
+    def on_train_start(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        del iteration
+        transformer = method.student.transformer
+        if not isinstance(transformer, DistributedDataParallel):
+            raise TypeError("DDPPostHocEMACallback requires a native DDP transformer")
+
+        self._rank = dist.get_rank() if dist.is_initialized() else 0
+        root = self.checkpoint_folder
+        if not root:
+            root = str(Path(self.training_config.checkpoint.output_dir) / "posthoc_ema" / "official_ddp")
+        self._checkpoint_root = Path(root).expanduser().resolve()
+
+        if self._rank == 0:
+            try:
+                from nitrous_ema import PostHocEMA
+            except ImportError as exc:
+                raise ImportError("DDP MMAudio training requires nitrous-ema. Install "
+                                  "FastVideo with the mmaudio-train extra.") from exc
+
+            module = unwrap_ddp_module(transformer)
+            self._ema = PostHocEMA(
+                module,
+                sigma_rels=self.sigma_rels,
+                update_every=self.update_every,
+                checkpoint_every_num_steps=self.checkpoint_every,
+                checkpoint_folder=str(self._checkpoint_root),
+                step_size_correction=self.step_size_correction,
+            ).to(method.student.device)
+            logger.info(
+                "Official DDP PostHocEMA enabled on rank 0: "
+                "sigma_rels=%s checkpoint_every=%d folder=%s",
+                self.sigma_rels,
+                self.checkpoint_every,
+                self._checkpoint_root,
+            )
+
+    def on_training_step_end(
+        self,
+        method: TrainingMethod,
+        loss_dict: dict[str, Any],
+        iteration: int = 0,
+    ) -> None:
+        del method, loss_dict
+        if iteration < self.start_iter:
+            return
+        self._calls += 1
+        if self._rank == 0:
+            if self._ema is None:
+                raise RuntimeError("DDP PostHocEMA was not initialized")
+            self._ema.update()
+            if int(self._ema.step) != self._calls:
+                raise RuntimeError("DDP PostHocEMA step diverged from the trainer step")
+
+    def _restore_checkpoint(self) -> None:
+        if self._rank != 0 or self._ema is None or self._calls <= 0:
+            return
+        if self._checkpoint_root is None:
+            raise RuntimeError("DDP PostHocEMA checkpoint folder is unset")
+        for index, ema_model in enumerate(self._ema.ema_models):
+            path = self._checkpoint_root / f"{index}.{self._calls}.pt"
+            if not path.is_file():
+                raise FileNotFoundError("Cannot resume DDP PostHocEMA: expected matching snapshot "
+                                        f"{path}. Keep training checkpoints aligned with "
+                                        "checkpoint_every.")
+            state = torch.load(
+                path,
+                map_location=method_device(ema_model),
+                weights_only=True,
+            )
+            ema_model.load_state_dict(state)
+        logger.info("Restored official DDP PostHocEMA at step %d", self._calls)
+
+    def state_dict(self) -> dict[str, Any]:
+        # Full nitrous states already live in rank-0 snapshots. Keeping only
+        # the common step in DCP gives every rank the same callback schema.
+        return {"calls": self._calls}
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+    ) -> None:
+        self._calls = int(state_dict.get("calls", 0))
+        self._restore_checkpoint()
+
+    def on_train_end(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        del method, iteration
+        if self._rank != 0 or self._ema is None or self._calls <= 0:
+            return
+        if self._calls % self.checkpoint_every != 0:
+            self._ema.checkpoint()
+        synthesized = self._ema.synthesize_ema_model(
+            sigma_rel=self.default_output_sigma,
+            step=self._calls,
+            device="cpu",
+        )
+        if self._checkpoint_root is None:
+            raise RuntimeError("DDP PostHocEMA checkpoint folder is unset")
+        sigma_name = str(self.default_output_sigma).replace(".", "p")
+        path = self._checkpoint_root / (f"mmaudio_ema_final_sigma_{sigma_name}_step_"
+                                        f"{self._calls:09d}.pth")
+        torch.save(synthesized.ema_model.state_dict(), path)
+        logger.info("Saved official synthesized DDP PostHocEMA to %s", path)
+
+
+def method_device(module: torch.nn.Module) -> torch.device:
+    """Return the device of a nitrous EMA module."""
+    return next(module.parameters()).device
+
+
+__all__ = ["DDPPostHocEMACallback"]

fastvideo/train/callbacks/mmaudio_validation.py

@@ -0,0 +1,509 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Validation loss and audio sampling for MMAudio feature training."""
+
+from __future__ import annotations
+
+import contextlib
+import math
+import re
+from pathlib import Path
+from typing import Any, TYPE_CHECKING
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from scipy.io import wavfile
+
+from fastvideo.distributed import get_local_torch_device, get_world_group
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.train.callbacks.callback import Callback
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+
+logger = init_logger(__name__)
+
+
+def _global_inference_indices(
+    total_samples: int,
+    world_size: int,
+    rank: int,
+) -> list[int | None]:
+    """Assign a fixed global sample budget with equal calls on every rank."""
+    if total_samples <= 0 or world_size <= 0 or not 0 <= rank < world_size:
+        raise ValueError("Invalid distributed MMAudio inference assignment")
+    calls_per_rank = math.ceil(total_samples / world_size)
+    return [
+        global_index if global_index < total_samples else None for local_index in range(calls_per_rank)
+        for global_index in (local_index * world_size + rank, )
+    ]
+
+
+def _inference_sample_index(
+    global_index: int | None,
+    local_index: int,
+    available: int,
+) -> int:
+    """Wrap an assigned global sample onto a row of this rank's feature batch.
+
+    ``available`` is the per-rank validation batch size, which is usually
+    smaller than the global sample budget.
+    """
+    return (global_index if global_index is not None else local_index) % available
+
+
+def _safe_sample_name(value: str) -> str:
+    value = re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._")
+    return value or "sample"
+
+
+class _IgnoredComponent:
+    """Placeholder for a required pipeline module that is never dereferenced.
+
+    ``MMAudioPipeline`` requires the CLIP text/vision and Synchformer encoders,
+    but its direct-feature conditioning path returns before touching them.
+    """
+
+    def __repr__(self) -> str:
+        return "<ignored component>"
+
+
+class MMAudioValidationCallback(Callback):
+    """Evaluate cached val features and periodically run native V2A inference.
+
+    Validation loss follows the official MMAudio ``val_fn``: sample a VAE
+    posterior latent, a logit-normal flow time, prior noise, and independent
+    video/text CFG masks. The RNG is reset for every pass, making values at
+    different training steps directly comparable.
+
+    Optional inference reuses the live FSDP transformer in
+    :class:`MMAudioPipeline`, while frozen VAE/vocoder weights are loaded from
+    ``inference_model_path``. Precomputed CLIP, Synchformer, and text features
+    go directly into the pipeline, so validation does not decode source video
+    or run feature encoders again.
+    """
+
+    def __init__(
+        self,
+        *,
+        data_path: str,
+        every_steps: int = 5000,
+        max_batches: int = 0,
+        batch_size: int = 8,
+        num_data_workers: int = 2,
+        run_at_start: bool = False,
+        use_ema: bool = False,
+        inference_every_steps: int = 20000,
+        inference_model_path: str = "",
+        inference_num_samples: int = 16,
+        inference_num_steps: int = 25,
+        inference_guidance_scale: float = 4.5,
+        inference_seed: int = 14159265,
+        inference_save_video: bool = True,
+        inference_log_to_tracker: bool = True,
+        output_dir: str | None = None,
+    ) -> None:
+        self.data_path = str(data_path)
+        self.every_steps = int(every_steps)
+        self.max_batches = max(0, int(max_batches))
+        self.batch_size = max(1, int(batch_size))
+        self.num_data_workers = max(0, int(num_data_workers))
+        self.run_at_start = bool(run_at_start)
+        self.use_ema = bool(use_ema)
+        self.inference_every_steps = max(0, int(inference_every_steps))
+        self.inference_model_path = str(inference_model_path)
+        self.inference_num_samples = max(1, int(inference_num_samples))
+        self.inference_num_steps = max(1, int(inference_num_steps))
+        self.inference_guidance_scale = float(inference_guidance_scale)
+        self.inference_seed = int(inference_seed)
+        self.inference_save_video = bool(inference_save_video)
+        self.inference_log_to_tracker = bool(inference_log_to_tracker)
+        self.output_dir = output_dir
+
+        self._dataloader: Any = None
+        self._pipeline: Any = None
+        self._rank = 0
+
+    def on_train_start(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        del iteration
+        if not self.data_path:
+            raise ValueError("MMAudio validation requires callbacks.validation.data_path")
+        self._rank = get_world_group().rank
+        transformer = method.student.transformer
+        from fastvideo.dataset.mmaudio_feature_dataset import (
+            build_mmaudio_feature_dataloader, )
+
+        feature_shapes = {
+            "latent_seq_len": int(transformer.latent_seq_len),
+            "latent_dim": int(transformer.latent_dim),
+            "clip_seq_len": int(transformer.clip_seq_len),
+            "clip_dim": int(transformer.config.arch_config.clip_dim),
+            "sync_seq_len": int(transformer.sync_seq_len),
+            "sync_dim": int(transformer.config.arch_config.sync_dim),
+            "text_seq_len": int(transformer.config.arch_config.text_seq_len),
+            "text_dim": int(transformer.config.arch_config.text_dim),
+        }
+        self._dataloader = build_mmaudio_feature_dataloader(
+            self.data_path,
+            batch_size=self.batch_size,
+            num_data_workers=self.num_data_workers,
+            seed=int(self.training_config.data.seed),
+            pin_memory=self.training_config.distributed.pin_cpu_memory,
+            feature_shapes=feature_shapes,
+            include_metadata=True,
+        )
+        logger.info("Initialized MMAudio validation features from %s", self.data_path)
+
+    def on_validation_begin(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        if iteration == 0 and not self.run_at_start:
+            return
+        run_loss = self.every_steps > 0 and iteration % self.every_steps == 0
+        run_inference = (self.inference_every_steps > 0 and bool(self.inference_model_path)
+                         and iteration % self.inference_every_steps == 0)
+        if not run_loss and not run_inference:
+            return
+
+        transformer = method.student.transformer
+        was_training = transformer.training
+        transformer.eval()
+        try:
+            ema_context = self._ema_context(transformer) if self.use_ema else contextlib.nullcontext(transformer)
+            with ema_context:
+                validation_batch = None
+                if run_loss:
+                    validation_batch = self._run_loss_validation(method, iteration)
+                if run_inference:
+                    if validation_batch is None:
+                        validation_batch = next(iter(self._dataloader))
+                    self._run_inference(method, validation_batch, iteration)
+        finally:
+            if was_training:
+                transformer.train()
+
+    def _ema_context(self, transformer: torch.nn.Module) -> Any:
+        callback_dict = getattr(self, "_callback_dict", None)
+        if callback_dict is not None:
+            for callback in callback_dict._callbacks.values():
+                if callback is self:
+                    continue
+                context = getattr(callback, "ema_context", None)
+                if callable(context):
+                    return context(transformer)
+        logger.warning("MMAudio validation requested EMA weights but no EMA callback was found")
+        return contextlib.nullcontext(transformer)
+
+    @torch.inference_mode()
+    def _run_loss_validation(
+        self,
+        method: TrainingMethod,
+        iteration: int,
+    ) -> dict[str, Any] | None:
+        if self._dataloader is None:
+            raise RuntimeError("MMAudio validation dataloader is not initialized")
+        device = get_local_torch_device()
+        generator = torch.Generator(device=device).manual_seed(int(self.training_config.data.seed) + self._rank, )
+        loss_sum = torch.zeros((), device=device, dtype=torch.float64)
+        sample_count = torch.zeros((), device=device, dtype=torch.float64)
+        first_batch: dict[str, Any] | None = None
+
+        for batch_index, raw_batch in enumerate(self._dataloader):
+            if self.max_batches and batch_index >= self.max_batches:
+                break
+            if first_batch is None:
+                first_batch = raw_batch
+            training_batch = method.student.prepare_batch(
+                raw_batch,
+                generator=generator,
+                latents_source="data",
+            )
+            noisy_latents = training_batch.noisy_model_input
+            timesteps = training_batch.timesteps
+            target = training_batch.training_target
+            if not all(isinstance(value, torch.Tensor) for value in (noisy_latents, timesteps, target)):
+                raise RuntimeError("MMAudio validation batch is incomplete")
+            assert isinstance(noisy_latents, torch.Tensor)
+            assert isinstance(timesteps, torch.Tensor)
+            assert isinstance(target, torch.Tensor)
+            prediction = method.student.predict_noise(
+                noisy_latents,
+                timesteps,
+                training_batch,
+                conditional=True,
+                attn_kind="dense",
+            )
+            batch_size = int(noisy_latents.shape[0])
+            loss_sum += F.mse_loss(
+                prediction.float(),
+                target.float(),
+                reduction="mean",
+            ).double() * batch_size
+            sample_count += batch_size
+
+        if torch.distributed.is_initialized():
+            torch.distributed.all_reduce(loss_sum)
+            torch.distributed.all_reduce(sample_count)
+        if sample_count.item() == 0:
+            raise RuntimeError("MMAudio validation did not produce any samples")
+        loss = float((loss_sum / sample_count).item())
+        if self._rank == 0:
+            tracker = getattr(method, "tracker", None)
+            if tracker is not None:
+                tracker.log({"validation/flow_matching_loss": loss}, iteration)
+            logger.info("MMAudio validation step=%d loss=%.8f", iteration, loss)
+        return first_batch
+
+    def _build_pipeline(self, transformer: torch.nn.Module) -> Any:
+        if self._pipeline is not None:
+            return self._pipeline
+        from fastvideo.pipelines.basic.mmaudio import MMAudioPipeline
+
+        ignored_component = _IgnoredComponent()
+        self._pipeline = MMAudioPipeline.from_pretrained(
+            self.inference_model_path,
+            inference_mode=True,
+            loaded_modules={
+                "transformer": transformer,
+                "text_encoder": ignored_component,
+                "tokenizer": ignored_component,
+                "image_encoder": ignored_component,
+                "image_encoder_2": ignored_component,
+            },
+            workload_type="v2a",
+            num_gpus=int(self.training_config.distributed.num_gpus),
+            tp_size=int(self.training_config.distributed.tp_size),
+            sp_size=int(self.training_config.distributed.sp_size),
+            hsdp_shard_dim=int(self.training_config.distributed.hsdp_shard_dim),
+            pin_cpu_memory=bool(self.training_config.distributed.pin_cpu_memory),
+            dit_cpu_offload=False,
+            dit_layerwise_offload=False,
+            vae_cpu_offload=False,
+        )
+        self._pipeline.fastvideo_args.pipeline_config = self.training_config.pipeline_config
+        for module_name in ("text_encoder", "tokenizer", "image_encoder", "image_encoder_2"):
+            module = self._pipeline.get_module(module_name)
+            assert isinstance(module,
+                              _IgnoredComponent), (f"MMAudio validation placeholder for {module_name} was not used")
+        return self._pipeline
+
+    @torch.inference_mode()
+    def _run_inference(
+        self,
+        method: TrainingMethod,
+        raw_batch: dict[str, Any],
+        iteration: int,
+    ) -> None:
+        pipeline = self._build_pipeline(method.student.transformer)
+        pc = self.training_config.pipeline_config
+        if pc is None:
+            raise RuntimeError("MMAudio inference validation requires pipeline_config")
+
+        text_features = raw_batch.get("text_features")
+        if not isinstance(text_features, torch.Tensor):
+            raise RuntimeError("MMAudio validation inference requires text features")
+        available = int(text_features.shape[0])
+        if available <= 0:
+            raise RuntimeError("MMAudio validation inference batch is empty")
+        world_size = (torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1)
+        assignments = _global_inference_indices(
+            self.inference_num_samples,
+            world_size,
+            self._rank,
+        )
+        output_root = Path(self.output_dir or
+                           self.training_config.checkpoint.output_dir, ) / "validation_audio" / f"step_{iteration:09d}"
+        output_root.mkdir(parents=True, exist_ok=True)
+
+        for local_index, global_index in enumerate(assignments):
+            # Every FSDP rank must execute the same number of forwards. Ranks
+            # outside a non-divisible global budget run one padded call but do
+            # not save it.
+            sample_index = _inference_sample_index(global_index, local_index, available)
+            output_index = (global_index if global_index is not None else self.inference_num_samples + self._rank)
+            batch = ForwardBatch(
+                data_type="video",
+                prompt="",
+                negative_prompt="",
+                audio_start_in_s=0.0,
+                audio_end_in_s=float(pc.duration_s),
+                num_inference_steps=self.inference_num_steps,
+                guidance_scale=self.inference_guidance_scale,
+                seed=self.inference_seed + output_index,
+                num_videos_per_prompt=1,
+                height=8,
+                width=8,
+                num_frames=1,
+                save_video=False,
+                return_frames=False,
+                extra={
+                    "mmaudio_clip_features": raw_batch["clip_features"][sample_index:sample_index + 1],
+                    "mmaudio_sync_features": raw_batch["sync_features"][sample_index:sample_index + 1],
+                    "mmaudio_text_features": raw_batch["text_features"][sample_index:sample_index + 1],
+                },
+            )
+            result = pipeline.forward(batch, pipeline.fastvideo_args)
+            audio = result.extra.get("audio")
+            sample_rate = result.extra.get("audio_sample_rate")
+            if not isinstance(audio, np.ndarray) or not isinstance(sample_rate, int):
+                raise RuntimeError("MMAudio validation inference did not produce decoded audio")
+            if global_index is None:
+                continue
+            sample_id = self._metadata_value(raw_batch, "sample_id", sample_index)
+            source_path = self._metadata_value(raw_batch, "source_path", sample_index)
+            stem = f"sample_{global_index:03d}_{_safe_sample_name(sample_id)}"
+            audio_path = output_root / f"{stem}.wav"
+            wavfile.write(audio_path, sample_rate, audio.astype(np.float32))
+            logger.info("Saved MMAudio validation audio to %s", audio_path)
+            if self.inference_save_video and source_path:
+                self._mux_source_video(
+                    source_path=Path(source_path),
+                    output_path=output_root / f"{stem}.mp4",
+                    audio=audio,
+                    sample_rate=sample_rate,
+                    duration_s=float(pc.duration_s),
+                )
+
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+        if self._rank == 0 and self.inference_log_to_tracker:
+            self._log_inference_media(method, output_root, iteration)
+
+    @staticmethod
+    def _metadata_value(
+        raw_batch: dict[str, Any],
+        key: str,
+        index: int,
+    ) -> str:
+        values = raw_batch.get(key)
+        if isinstance(values, list | tuple) and index < len(values):
+            return str(values[index])
+        return ""
+
+    @staticmethod
+    def _mux_source_video(
+        *,
+        source_path: Path,
+        output_path: Path,
+        audio: np.ndarray,
+        sample_rate: int,
+        duration_s: float,
+    ) -> Path | None:
+        if not source_path.is_file():
+            logger.warning("Cannot compose MMAudio validation video; missing %s", source_path)
+            return None
+        try:
+            import av
+            from av import AudioFrame
+
+            # Match official mmaudio.data.av_utils.read_frames: keep every
+            # decoded source frame through duration_s at the stream's guessed
+            # frame rate, including a frame exactly on the endpoint.
+            all_frames: list[np.ndarray] = []
+            with av.open(str(source_path)) as source:
+                input_stream = source.streams.video[0]
+                frame_rate = input_stream.guessed_rate
+                input_stream.thread_type = "AUTO"
+                for packet in source.demux(input_stream):
+                    reached_end = False
+                    for frame in packet.decode():
+                        if frame.time is not None and frame.time > duration_s:
+                            reached_end = True
+                            break
+                        all_frames.append(frame.to_ndarray(format="rgb24"))
+                    if reached_end:
+                        break
+            if not all_frames:
+                raise ValueError("source video produced no decoded frames")
+
+            # Match official reencode_with_audio: H.264 at the original
+            # guessed FPS, 10 Mbps, yuv420p, followed by mono float audio
+            # encoded as AAC at the model sample rate.
+            output = av.open(str(output_path), "w")
+            try:
+                video_stream = output.add_stream("h264", frame_rate)
+                video_stream.codec_context.bit_rate = int(10 * 1e6)
+                video_stream.width = int(all_frames[0].shape[1])
+                video_stream.height = int(all_frames[0].shape[0])
+                video_stream.pix_fmt = "yuv420p"
+                audio_stream = output.add_stream("aac", sample_rate)
+
+                for frame_array in all_frames:
+                    video_frame = av.VideoFrame.from_ndarray(frame_array)
+                    for packet in video_stream.encode(video_frame):
+                        output.mux(packet)
+                for packet in video_stream.encode():
+                    output.mux(packet)
+
+                audio_np = np.asarray(audio, dtype=np.float32)
+                if audio_np.ndim == 1:
+                    audio_np = audio_np[None, :]
+                elif audio_np.ndim == 2 and audio_np.shape[1] == 1:
+                    audio_np = audio_np.T
+                if audio_np.ndim != 2 or audio_np.shape[0] != 1:
+                    raise ValueError(f"Expected mono audio, got {audio_np.shape}")
+                audio_frame = AudioFrame.from_ndarray(
+                    np.ascontiguousarray(audio_np),
+                    format="flt",
+                    layout="mono",
+                )
+                audio_frame.sample_rate = sample_rate
+                for packet in audio_stream.encode(audio_frame):
+                    output.mux(packet)
+                for packet in audio_stream.encode():
+                    output.mux(packet)
+            finally:
+                output.close()
+        except Exception as exc:
+            logger.warning(
+                "Failed to compose MMAudio validation video %s: %s",
+                output_path,
+                exc,
+            )
+            return None
+        logger.info("Saved MMAudio validation video to %s", output_path)
+        return output_path
+
+    def _log_inference_media(
+        self,
+        method: TrainingMethod,
+        output_root: Path,
+        iteration: int,
+    ) -> None:
+        tracker = getattr(method, "tracker", None)
+        if tracker is None:
+            return
+        artifacts = []
+        for path in sorted(output_root.glob("sample_*.mp4"))[:self.inference_num_samples]:
+            artifact = tracker.video(
+                str(path),
+                caption=path.stem,
+                format="mp4",
+            )
+            if artifact is not None:
+                artifacts.append(artifact)
+        if artifacts:
+            tracker.log_artifacts(
+                {"validation/mmaudio_samples": artifacts},
+                iteration,
+            )
+
+    def on_train_end(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        del method, iteration
+        if self._pipeline is not None:
+            self._pipeline.close()
+            self._pipeline = None
+
+
+__all__ = ["MMAudioValidationCallback"]

fastvideo/train/callbacks/posthoc_ema.py

@@ -0,0 +1,296 @@
+# SPDX-License-Identifier: Apache-2.0
+"""FSDP2-compatible post-hoc EMA used by the official MMAudio recipe."""
+
+from __future__ import annotations
+
+import contextlib
+import os
+from collections.abc import Generator
+from pathlib import Path
+from typing import Any, TYPE_CHECKING
+
+import numpy as np
+import torch
+
+from fastvideo.logger import init_logger
+from fastvideo.train.callbacks.callback import Callback
+from fastvideo.training.training_utils import EMA_FSDP
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+
+logger = init_logger(__name__)
+
+
+def sigma_rel_to_gamma(sigma_rel: float) -> float:
+    """Algorithm 2 from Karras et al., matching ``nitrous-ema``."""
+    if sigma_rel <= 0:
+        raise ValueError("Post-hoc EMA sigma_rel must be positive")
+    t = sigma_rel**-2
+    return float(np.roots([1, 7, 16 - t, 12 - t]).real.max())
+
+
+def _p_dot_p(
+    t_a: torch.Tensor,
+    gamma_a: torch.Tensor,
+    t_b: torch.Tensor,
+    gamma_b: torch.Tensor,
+) -> torch.Tensor:
+    t_ratio = t_a / t_b
+    t_exp = torch.where(t_a < t_b, gamma_b, -gamma_a)
+    t_max = torch.maximum(t_a, t_b)
+    numerator = (gamma_a + 1) * (gamma_b + 1) * t_ratio**t_exp
+    denominator = (gamma_a + gamma_b + 1) * t_max
+    return numerator / denominator
+
+
+def solve_posthoc_weights(
+    timesteps: torch.Tensor,
+    gammas: torch.Tensor,
+    target_timestep: int,
+    target_gamma: float,
+) -> torch.Tensor:
+    """Algorithm 3 from Karras et al., matching ``nitrous-ema``."""
+    t_i = timesteps.double().reshape(-1, 1)
+    gamma_i = gammas.double().reshape(-1, 1)
+    t_j = timesteps.double().reshape(1, -1)
+    gamma_j = gammas.double().reshape(1, -1)
+    matrix = _p_dot_p(t_i, gamma_i, t_j, gamma_j)
+    target_t = torch.tensor([[target_timestep]], dtype=torch.float64)
+    target_g = torch.tensor([[target_gamma]], dtype=torch.float64)
+    rhs = _p_dot_p(t_i, gamma_i, target_t, target_g)
+    return torch.linalg.solve(matrix, rhs).squeeze(-1)
+
+
+class PostHocEMACallback(Callback):
+    """Maintain and checkpoint multiple Karras EMA profiles on local shards.
+
+    The upstream ``nitrous-ema`` implementation deep-copies a complete model
+    on rank zero. That is appropriate for MMAudio's DDP trainer but not for a
+    FastVideo FSDP2/HSDP transformer. This callback applies the same update and
+    synthesis equations independently to every local FSDP shard. Together the
+    rank-local snapshots represent the same full EMA model without gathering
+    it on every optimizer step.
+    """
+
+    def __init__(
+        self,
+        *,
+        sigma_rels: list[float] | tuple[float, ...] = (0.05, 0.1),
+        update_every: int = 1,
+        checkpoint_every: int = 5000,
+        checkpoint_folder: str | None = None,
+        start_iter: int = 0,
+        default_output_sigma: float = 0.05,
+        step_size_correction: bool = True,
+    ) -> None:
+        if len(sigma_rels) < 2:
+            raise ValueError("Post-hoc EMA requires at least two sigma profiles")
+        self.sigma_rels = tuple(float(value) for value in sigma_rels)
+        self.gammas = tuple(sigma_rel_to_gamma(value) for value in self.sigma_rels)
+        if len(set(self.gammas)) != len(self.gammas):
+            raise ValueError("Post-hoc EMA sigma profiles must be unique")
+        self.update_every = max(1, int(update_every))
+        self.checkpoint_every = max(1, int(checkpoint_every))
+        self.checkpoint_folder = checkpoint_folder
+        self.start_iter = int(start_iter)
+        self.default_output_sigma = float(default_output_sigma)
+        self.step_size_correction = bool(step_size_correction)
+
+        self._ema_models: list[EMA_FSDP] = []
+        self._calls = 0
+        self._initted = False
+        self._rank = 0
+        self._snapshot_root: Path | None = None
+
+    def on_train_start(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        del iteration
+        transformer = method.student.transformer
+        self._ema_models = [EMA_FSDP(transformer, decay=0.0, mode="local_shard") for _ in self.gammas]
+        if torch.distributed.is_initialized():
+            self._rank = torch.distributed.get_rank()
+        root = self.checkpoint_folder
+        if not root:
+            root = str(Path(self.training_config.checkpoint.output_dir) / "posthoc_ema")
+        self._snapshot_root = Path(root).expanduser().resolve()
+        (self._snapshot_root / f"rank_{self._rank:05d}").mkdir(
+            parents=True,
+            exist_ok=True,
+        )
+        logger.info(
+            "PostHocEMA enabled: sigma_rels=%s, checkpoint_every=%d, folder=%s",
+            self.sigma_rels,
+            self.checkpoint_every,
+            self._snapshot_root,
+        )
+
+    def _decay(self, gamma: float) -> float:
+        step = self._calls
+        if not self.step_size_correction:
+            return (1.0 - 1.0 / (step + 1))**(1.0 + gamma)
+        first = (1.0 - 1.0 / (step + 1))**(1.0 + gamma)
+        second = (1.0 - 1.0 / (step + 1 + self.update_every))**(1.0 + gamma)
+        return (first * second)**(self.update_every / 2)
+
+    def on_training_step_end(
+        self,
+        method: TrainingMethod,
+        loss_dict: dict[str, Any],
+        iteration: int = 0,
+    ) -> None:
+        del loss_dict
+        if iteration < self.start_iter or not self._ema_models:
+            return
+
+        previous_calls = self._calls
+        self._calls += 1
+        if previous_calls % self.update_every != 0:
+            return
+
+        transformer = method.student.transformer
+        if not self._initted:
+            for ema in self._ema_models:
+                ema._init_shadow(transformer)
+            self._initted = True
+        for gamma, ema in zip(self.gammas, self._ema_models, strict=True):
+            ema.decay = self._decay(gamma)
+            ema.update(transformer)
+
+        if self._calls % self.checkpoint_every == 0:
+            self._save_snapshots()
+
+    def _save_snapshots(self) -> None:
+        if self._snapshot_root is None:
+            raise RuntimeError("PostHocEMA callback has not been initialized")
+        rank_dir = self._snapshot_root / f"rank_{self._rank:05d}"
+        for index, ema in enumerate(self._ema_models):
+            path = rank_dir / f"{index}.{self._calls}.pt"
+            temporary = path.with_suffix(".pt.tmp")
+            torch.save(ema.state_dict(), temporary)
+            os.replace(temporary, path)
+        logger.info("Saved PostHocEMA shard snapshots at step %d", self._calls)
+
+    def _snapshot_records(
+        self,
+        *,
+        max_step: int | None,
+    ) -> list[tuple[Path, int, float]]:
+        if self._snapshot_root is None:
+            return []
+        rank_dir = self._snapshot_root / f"rank_{self._rank:05d}"
+        records: list[tuple[Path, int, float]] = []
+        for path in rank_dir.glob("*.pt"):
+            try:
+                profile_index, timestep = map(int, path.stem.split("."))
+                gamma = self.gammas[profile_index]
+            except (ValueError, IndexError):
+                continue
+            if max_step is None or timestep <= max_step:
+                records.append((path, timestep, gamma))
+        return sorted(records, key=lambda item: (item[1], item[0].name))
+
+    def synthesize_local_shard(
+        self,
+        *,
+        sigma_rel: float | None = None,
+        step: int | None = None,
+    ) -> dict[str, torch.Tensor] | None:
+        target_sigma = self.default_output_sigma if sigma_rel is None else float(sigma_rel)
+        records = self._snapshot_records(max_step=step)
+        if not records:
+            return None
+        target_step = max(record[1] for record in records) if step is None else int(step)
+        if target_step > max(record[1] for record in records):
+            raise ValueError("Cannot synthesize PostHocEMA beyond the newest snapshot")
+
+        weights = solve_posthoc_weights(
+            torch.tensor([record[1] for record in records]),
+            torch.tensor([record[2] for record in records]),
+            target_step,
+            sigma_rel_to_gamma(target_sigma),
+        )
+        synthesized: dict[str, torch.Tensor] = {}
+        for (path, _, _), weight in zip(records, weights.tolist(), strict=True):
+            state = torch.load(path, map_location="cpu", weights_only=True)
+            if not isinstance(state, dict):
+                raise ValueError(f"Invalid PostHocEMA snapshot: {path}")
+            for name, tensor in state.items():
+                if not isinstance(tensor, torch.Tensor):
+                    continue
+                if name not in synthesized:
+                    synthesized[name] = tensor.float().mul(float(weight))
+                else:
+                    synthesized[name].add_(tensor.float(), alpha=float(weight))
+        return synthesized
+
+    @contextlib.contextmanager
+    def ema_context(
+        self,
+        transformer: torch.nn.Module,
+    ) -> Generator[torch.nn.Module, None, None]:
+        synthesized = self.synthesize_local_shard()
+        if synthesized is None:
+            yield transformer
+            return
+        temporary_ema = EMA_FSDP(transformer, decay=0.0, mode="local_shard")
+        temporary_ema.load_state_dict(synthesized)
+        with temporary_ema.apply_to_model(transformer):
+            yield transformer
+
+    def on_train_end(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        del method, iteration
+        if not self._initted or self._snapshot_root is None:
+            return
+
+        # Upstream MMAudio synthesizes the sigma=0.05 model after training.
+        # Save the latest profiles even when a short/early-stopped run did not
+        # end exactly on checkpoint_every, then synthesize each FSDP shard.
+        self._save_snapshots()
+        synthesized = self.synthesize_local_shard(
+            sigma_rel=self.default_output_sigma,
+            step=self._calls,
+        )
+        if synthesized is None:
+            return
+        rank_dir = self._snapshot_root / f"rank_{self._rank:05d}"
+        sigma_name = str(self.default_output_sigma).replace(".", "p")
+        path = rank_dir / (f"synthesized_sigma_{sigma_name}_step_{self._calls:09d}.pt")
+        temporary = path.with_suffix(".pt.tmp")
+        torch.save(synthesized, temporary)
+        os.replace(temporary, path)
+        logger.info("Saved synthesized PostHocEMA shard to %s", path)
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "calls": self._calls,
+            "initted": self._initted,
+            "profiles": {
+                str(index): ema.state_dict()
+                for index, ema in enumerate(self._ema_models)
+            },
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self._calls = int(state_dict.get("calls", 0))
+        self._initted = bool(state_dict.get("initted", False))
+        profiles = state_dict.get("profiles", {})
+        if isinstance(profiles, dict):
+            for index, ema in enumerate(self._ema_models):
+                profile = profiles.get(str(index))
+                if isinstance(profile, dict):
+                    ema.load_state_dict(profile)
+
+
+__all__ = [
+    "PostHocEMACallback",
+    "sigma_rel_to_gamma",
+    "solve_posthoc_weights",
+]

fastvideo/train/methods/__init__.py

@@ -6,6 +6,7 @@ __all__ = [
     "TrainingMethod",
     "DMD2Method",
     "FineTuneMethod",
+    "FlowMatchingFineTuneMethod",
     "KDMethod",
     "SelfForcingMethod",
     "DiffusionForcingSFTMethod",
@@ -21,6 +22,10 @@ def __getattr__(name: str) -> object:
     if name == "FineTuneMethod":
         from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
         return FineTuneMethod
+    if name == "FlowMatchingFineTuneMethod":
+        from fastvideo.train.methods.fine_tuning.flow_matching import (
+            FlowMatchingFineTuneMethod, )
+        return FlowMatchingFineTuneMethod
     if name == "KDMethod":
         from fastvideo.train.methods.knowledge_distillation.kd import KDMethod
         return KDMethod

fastvideo/train/methods/base.py

@@ -12,6 +12,7 @@ from fastvideo import envs
 from fastvideo.logger import init_logger
 from fastvideo.train.models.base import ModelBase
 from fastvideo.train.utils.checkpoint import _RoleModuleContainer
+from fastvideo.train.utils.optimizer import seed_adamw_parameter_state
 from fastvideo.training.checkpointing_utils import (
     ModelWrapper,
     OptimizerWrapper,
@@ -189,11 +190,7 @@ class TrainingMethod(torch.nn.Module, ABC):
                         continue
                     if len(opt.state.get(p, {})) > 0:
                         continue
-                    opt.state[p] = {
-                        "step": torch.tensor(0.0),
-                        "exp_avg": torch.zeros_like(p),
-                        "exp_avg_sq": torch.zeros_like(p),
-                    }
+                    seed_adamw_parameter_state(opt, p)
 
     # -- Shared hooks (override in subclasses as needed) --
 

fastvideo/train/methods/fine_tuning/__init__.py

@@ -7,11 +7,13 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod
     from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
+    from fastvideo.train.methods.fine_tuning.flow_matching import FlowMatchingFineTuneMethod
     from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod
 
 __all__ = [
     "DiffusionForcingSFTMethod",
     "FineTuneMethod",
+    "FlowMatchingFineTuneMethod",
     "TeacherForcingSFTMethod",
 ]
 
@@ -27,6 +29,11 @@ def __getattr__(name: str) -> object:
         from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
 
         return FineTuneMethod
+    if name == "FlowMatchingFineTuneMethod":
+        from fastvideo.train.methods.fine_tuning.flow_matching import (
+            FlowMatchingFineTuneMethod, )
+
+        return FlowMatchingFineTuneMethod
     if name == "TeacherForcingSFTMethod":
         from fastvideo.train.methods.fine_tuning.tfsft import (
             TeacherForcingSFTMethod, )

fastvideo/train/methods/fine_tuning/finetune.py

@@ -11,7 +11,9 @@ import torch.nn.functional as F
 from fastvideo.train.methods.base import TrainingMethod, LogScalar
 from fastvideo.train.models.base import ModelBase, NoisePrediction
 from fastvideo.train.utils.optimizer import (
-    build_optimizer_and_scheduler, )
+    build_optimizer_and_scheduler,
+    seed_adamw_parameter_state,
+)
 
 
 def _compute_finetune_loss_map(
@@ -254,3 +256,16 @@ class FineTuneMethod(TrainingMethod):
             betas=student_betas,
             scheduler_name=student_sched,
         )
+        optimizer_parameters = {
+            parameter
+            for group in self._student_optimizer.param_groups
+            for parameter in group["params"]
+        }
+        for parameter in self.student.eager_optimizer_state_parameters():
+            if parameter not in optimizer_parameters:
+                raise ValueError("An eager optimizer-state parameter is not owned by the "
+                                 "student optimizer")
+            seed_adamw_parameter_state(
+                self._student_optimizer,
+                parameter,
+            )

fastvideo/train/methods/fine_tuning/flow_matching.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shape-agnostic conditional flow-matching fine-tuning."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from fastvideo.logger import init_logger
+from fastvideo.train.methods.base import LogScalar
+from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.distributed_strategy import is_ddp_strategy
+
+logger = init_logger(__name__)
+
+TrainStepResult = tuple[
+    dict[str, torch.Tensor],
+    dict[str, Any],
+    dict[str, LogScalar],
+]
+
+TensorTrainFn = Callable[..., torch.Tensor]
+
+
+class FlowMatchingFineTuneMethod(FineTuneMethod):
+    """Fine-tune a model against a velocity target supplied by its adapter.
+
+    Unlike :class:`FineTuneMethod`, this method does not assume a five-
+    dimensional video latent layout. The model plugin owns sampling the flow
+    path and stores its target in ``TrainingBatch.training_target``.
+    """
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, ModelBase],
+    ) -> None:
+        super().__init__(cfg=cfg, role_models=role_models)
+        flow_matching_forward = getattr(self.student, "flow_matching_forward", None)
+        if not callable(flow_matching_forward):
+            raise TypeError("FlowMatchingFineTuneMethod requires the student model to "
+                            "implement flow_matching_forward()")
+        self._flow_matching_forward: TensorTrainFn = flow_matching_forward
+        self._train_fn: TensorTrainFn = self._eager_train_fn
+
+        if bool(self.training_config.model.compile_train_fn):
+            if not is_ddp_strategy(self.training_config):
+                raise ValueError("training.model.compile_train_fn is currently supported "
+                                 "only with training.distributed.strategy=ddp")
+            compile_kwargs = dict(self.training_config.model.torch_compile_kwargs)
+            compile_kwargs.setdefault("fullgraph", True)
+            compile_training_forward = getattr(self.student, "compile_training_forward", None)
+            if not callable(compile_training_forward):
+                raise TypeError("compile_train_fn requires the student model to implement "
+                                "compile_training_forward()")
+            logger.info(
+                "Enabling fullgraph transformer compilation with kwargs=%s",
+                compile_kwargs,
+            )
+            compile_training_forward(compile_kwargs)
+
+    def _eager_train_fn(
+        self,
+        noisy_latents: torch.Tensor,
+        timesteps: torch.Tensor,
+        target: torch.Tensor,
+        clip_features: torch.Tensor,
+        sync_features: torch.Tensor,
+        text_features: torch.Tensor,
+    ) -> torch.Tensor:
+        prediction = self._flow_matching_forward(
+            noisy_latents,
+            timesteps,
+            clip_features,
+            sync_features,
+            text_features,
+        )
+        return F.mse_loss(prediction.float(), target.float())
+
+    def single_train_step(
+        self,
+        batch: dict[str, Any],
+        iteration: int,
+    ) -> TrainStepResult:
+        del iteration
+        if self.cuda_generator is None:
+            raise RuntimeError("on_train_start() must run before training")
+
+        # RNG, existence masks, and CFG masks stay eager so the explicit
+        # per-rank CUDA Generator advances exactly as in official MMAudio.
+        training_batch = self.student.prepare_batch(
+            batch,
+            generator=self.cuda_generator,
+            latents_source="data",
+        )
+        noisy_latents = training_batch.noisy_model_input
+        timesteps = training_batch.timesteps
+        target = training_batch.training_target
+        if not all(isinstance(value, torch.Tensor) for value in (noisy_latents, timesteps, target)):
+            raise RuntimeError("prepare_batch() returned incomplete flow inputs")
+        assert isinstance(noisy_latents, torch.Tensor)
+        assert isinstance(timesteps, torch.Tensor)
+        assert isinstance(target, torch.Tensor)
+        if target.shape != noisy_latents.shape:
+            raise RuntimeError("Flow-matching target and model input shapes must match, got "
+                               f"{tuple(target.shape)} and {tuple(noisy_latents.shape)}")
+
+        conditions = training_batch.conditional_dict
+        if not isinstance(conditions, dict):
+            raise RuntimeError("prepare_batch() must set conditional_dict")
+        condition_values = tuple(conditions.get(key) for key in ("clip_features", "sync_features", "text_features"))
+        if not all(isinstance(value, torch.Tensor) for value in condition_values):
+            raise RuntimeError("MMAudio flow conditions must all be tensors")
+        clip_features, sync_features, text_features = condition_values
+        assert isinstance(clip_features, torch.Tensor)
+        assert isinstance(sync_features, torch.Tensor)
+        assert isinstance(text_features, torch.Tensor)
+
+        loss = self._train_fn(
+            noisy_latents,
+            timesteps,
+            target,
+            clip_features,
+            sync_features,
+            text_features,
+        )
+        attn_metadata = (training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata)
+        loss_map = {
+            "total_loss": loss,
+            "flow_matching_loss": loss,
+        }
+        outputs: dict[str, Any] = {
+            "_fv_backward": (timesteps, attn_metadata),
+        }
+        metrics: dict[str, LogScalar] = {}
+        return loss_map, outputs, metrics
+
+
+__all__ = ["FlowMatchingFineTuneMethod"]

fastvideo/train/models/base.py

@@ -117,6 +117,15 @@ class ModelBase(ABC):
     def on_train_start(self) -> None:  # noqa: B027
         """Called once before the training loop begins."""
 
+    def eager_optimizer_state_parameters(self, ) -> tuple[torch.nn.Parameter, ...]:
+        """Parameters whose optimizer state must exist before their first grad.
+
+        Most models return an empty tuple. Models with intermittently used
+        trainable parameters can override this so strict distributed
+        checkpoints retain a stable optimizer-state schema across steps.
+        """
+        return ()
+
     def decode_latents(
         self,
         latents_b_t_c_h_w: torch.Tensor,

fastvideo/train/models/mmaudio/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.train.models.mmaudio.mmaudio import MMAudioModel
+
+__all__ = ["MMAudioModel"]

fastvideo/train/models/mmaudio/mmaudio.py

@@ -0,0 +1,609 @@
+# SPDX-License-Identifier: Apache-2.0
+"""MMAudio model plugin for the modular FastVideo trainer."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+
+from fastvideo.configs.models.dits import (
+    MMAUDIO_44K_TRAINING_VARIANTS,
+    MMAUDIO_TRAINING_VARIANTS,
+    MMAUDIO_VARIANT_ARCHITECTURES,
+    get_mmaudio_transformer_config,
+)
+from fastvideo.configs.pipelines.mmaudio import (
+    MMAUDIO_PIPELINE_CONFIGS,
+    MMAudioV2AConfig,
+)
+from fastvideo.distributed import get_local_torch_device, get_world_group
+from fastvideo.forward_context import set_forward_context
+from fastvideo.logger import init_logger
+from fastvideo.models.dits.mmaudio import MMAudioTransformer
+from fastvideo.models.loader.fsdp_load import build_fsdp_model_from_scratch
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.platforms import AttentionBackendEnum
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.distributed_strategy import (
+    build_replicated_model_from_scratch,
+    is_ddp_strategy,
+    unwrap_ddp_module,
+    wrap_module_ddp,
+)
+from fastvideo.train.utils.module_state import apply_trainable
+from fastvideo.train.utils.moduleloader import load_module_from_path
+from fastvideo.utils import PRECISION_TO_TYPE
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.lora import LoraConfig
+    from fastvideo.train.utils.training_config import TrainingConfig
+
+logger = init_logger(__name__)
+
+
+def _load_empty_string_features(path: str | Path) -> torch.Tensor:
+    """Load only the fixed CLIP empty-string embedding from an official asset."""
+    source = Path(path).expanduser().resolve()
+    if not source.exists():
+        raise FileNotFoundError(f"MMAudio empty-string feature source does not exist: {source}")
+
+    if source.is_dir():
+        preferred = (
+            source / "transformer" / "diffusion_pytorch_model.safetensors",
+            source / "diffusion_pytorch_model.safetensors",
+        )
+        candidates = [candidate for candidate in preferred if candidate.is_file()]
+        if not candidates:
+            candidates = sorted(source.rglob("*.safetensors"))
+        for candidate in candidates:
+            from safetensors import safe_open
+
+            with safe_open(candidate, framework="pt", device="cpu") as handle:
+                if "empty_string_feat" in handle:
+                    return handle.get_tensor("empty_string_feat").float()
+        raise ValueError(f"No 'empty_string_feat' tensor was found in safetensors under {source}")
+
+    if source.suffix == ".safetensors":
+        from safetensors import safe_open
+
+        with safe_open(source, framework="pt", device="cpu") as handle:
+            if "empty_string_feat" not in handle:
+                raise ValueError(f"MMAudio safetensors {source} has no 'empty_string_feat' tensor")
+            return handle.get_tensor("empty_string_feat").float()
+
+    state = torch.load(source, map_location="cpu", weights_only=True)
+    if isinstance(state, torch.Tensor):
+        tensor = state
+    elif isinstance(state, dict):
+        tensor = state.get("empty_string_feat")
+        if tensor is None and isinstance(state.get("state_dict"), dict):
+            tensor = state["state_dict"].get("empty_string_feat")
+        if tensor is None:
+            for key, value in state.items():
+                if str(key).endswith("empty_string_feat") and isinstance(value, torch.Tensor):
+                    tensor = value
+                    break
+    else:
+        tensor = None
+    if not isinstance(tensor, torch.Tensor):
+        raise ValueError(f"MMAudio asset {source} does not contain an empty-string tensor")
+    if tensor.ndim == 3 and tensor.shape[0] == 1:
+        tensor = tensor.squeeze(0)
+    return tensor.float()
+
+
+def _load_or_compute_latent_stats(
+    training_config: TrainingConfig,
+    *,
+    variant: str,
+    cache_path: str | None,
+    chunk_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute/cache official latent statistics once, then share across ranks."""
+    world_group = get_world_group()
+    payload: dict[str, Any] | None = None
+    if world_group.is_first_rank:
+        stats_path = Path(cache_path).expanduser().resolve() if cache_path else None
+        if stats_path is not None and stats_path.is_file():
+            loaded = torch.load(stats_path, map_location="cpu", weights_only=True)
+            if not isinstance(loaded, dict):
+                raise ValueError(f"Invalid MMAudio latent statistics file: {stats_path}")
+            payload = loaded
+            logger.info("Loaded MMAudio latent statistics from %s", stats_path)
+        else:
+            from fastvideo.dataset.mmaudio_feature_dataset import (
+                compute_mmaudio_latent_stats, )
+
+            architecture = MMAUDIO_VARIANT_ARCHITECTURES[variant]
+            logger.info("Computing MMAudio latent statistics from the first feature cache")
+            latent_mean, latent_std = compute_mmaudio_latent_stats(
+                training_config.data.data_path,
+                latent_seq_len=int(architecture["latent_seq_len"]),
+                latent_dim=int(architecture["latent_dim"]),
+                chunk_size=chunk_size,
+            )
+            payload = {
+                "variant": variant,
+                "latent_mean": latent_mean,
+                "latent_std": latent_std,
+            }
+            if stats_path is not None:
+                stats_path.parent.mkdir(parents=True, exist_ok=True)
+                torch.save(payload, stats_path)
+                logger.info("Saved MMAudio latent statistics to %s", stats_path)
+
+    payload = world_group.broadcast_object(payload, src=0)
+    if not isinstance(payload, dict):
+        raise RuntimeError("Failed to broadcast MMAudio latent statistics")
+    cached_variant = payload.get("variant")
+    same_44k_contract = (cached_variant in MMAUDIO_44K_TRAINING_VARIANTS and variant in MMAUDIO_44K_TRAINING_VARIANTS)
+    if cached_variant is not None and cached_variant != variant and not same_44k_contract:
+        raise ValueError(f"MMAudio latent statistics are for {cached_variant!r}, not {variant!r}")
+    latent_mean = payload.get("latent_mean")
+    latent_std = payload.get("latent_std")
+    if not isinstance(latent_mean, torch.Tensor) or not isinstance(latent_std, torch.Tensor):
+        raise ValueError("MMAudio latent statistics must contain tensor mean/std values")
+    expected_shape = (1, 1, int(MMAUDIO_VARIANT_ARCHITECTURES[variant]["latent_dim"]))
+    if tuple(latent_mean.shape) != expected_shape or tuple(latent_std.shape) != expected_shape:
+        raise ValueError(f"MMAudio latent statistics for {variant!r} must have shape "
+                         f"{expected_shape}, got {tuple(latent_mean.shape)} and "
+                         f"{tuple(latent_std.shape)}")
+    return latent_mean, latent_std
+
+
+class MMAudioModel(ModelBase):
+    """Official-compatible conditional flow-matching training adapter.
+
+    Only the MMAudio transformer participates in the training graph. Audio VAE,
+    DFN5B, and Synchformer outputs are read from an offline feature cache.
+    """
+
+    _transformer_cls_name = "MMAudioTransformer"
+
+    def __init__(
+        self,
+        *,
+        init_from: str | None = None,
+        training_config: TrainingConfig,
+        variant: str | None = None,
+        from_scratch: bool = False,
+        empty_string_features_path: str | None = None,
+        latent_statistics_path: str | None = None,
+        latent_statistics_chunk_size: int = 32,
+        trainable: bool = True,
+        allow_v2_training: bool = False,
+        lora: LoraConfig | dict[str, Any] | None = None,
+        attention_backend: AttentionBackendEnum | str | None = None,
+        transformer: torch.nn.Module | None = None,
+    ) -> None:
+        super().__init__(
+            trainable=trainable,
+            lora=lora,
+            attention_backend=attention_backend,
+        )
+        if int(training_config.distributed.sp_size or 1) != 1:
+            raise ValueError("MMAudio training does not yet support sequence parallelism; set sp_size=1")
+        if int(training_config.distributed.tp_size or 1) != 1:
+            raise ValueError("MMAudio training does not yet support tensor parallelism; set tp_size=1")
+        use_ddp = is_ddp_strategy(training_config)
+        if use_ddp and (int(training_config.distributed.hsdp_replicate_dim) != 1
+                        or int(training_config.distributed.hsdp_shard_dim) not in {-1, 1}):
+            raise ValueError("MMAudio DDP strategy does not use HSDP dimensions; "
+                             "set hsdp_replicate_dim=1 and hsdp_shard_dim=1")
+
+        self.training_config = training_config
+        if variant is not None:
+            try:
+                pipeline_config_cls = MMAUDIO_PIPELINE_CONFIGS[variant]
+            except KeyError as exc:
+                supported = ", ".join(MMAUDIO_PIPELINE_CONFIGS)
+                raise ValueError(f"Unknown MMAudio variant {variant!r}; expected one of: {supported}") from exc
+            self.training_config.pipeline_config = pipeline_config_cls()
+        elif self.training_config.pipeline_config is None:
+            self.training_config.pipeline_config = MMAudioV2AConfig()
+
+        if transformer is None:
+            if from_scratch:
+                if init_from:
+                    raise ValueError("MMAudio from_scratch=true cannot also set init_from; use "
+                                     "empty_string_features_path for the fixed CLIP embedding")
+                if variant is None:
+                    raise ValueError("MMAudio from-scratch training requires variant")
+                if variant not in MMAUDIO_TRAINING_VARIANTS:
+                    supported = ", ".join(MMAUDIO_TRAINING_VARIANTS)
+                    raise ValueError(f"The official MMAudio recipe does not train {variant!r}; "
+                                     f"supported from-scratch variants: {supported}")
+                if not empty_string_features_path:
+                    raise ValueError("MMAudio from-scratch training requires "
+                                     "empty_string_features_path")
+                latent_mean, latent_std = _load_or_compute_latent_stats(
+                    self.training_config,
+                    variant=variant,
+                    cache_path=latent_statistics_path,
+                    chunk_size=latent_statistics_chunk_size,
+                )
+                empty_string_feat = _load_empty_string_features(empty_string_features_path)
+                architecture = MMAUDIO_VARIANT_ARCHITECTURES[variant]
+                expected_text_shape = (77, 1024)
+                if tuple(empty_string_feat.shape) != expected_text_shape:
+                    raise ValueError("MMAudio empty-string features must have shape "
+                                     f"{expected_text_shape}, got {tuple(empty_string_feat.shape)}")
+                model_config = get_mmaudio_transformer_config(variant)
+                hf_config = {
+                    "_class_name": self._transformer_cls_name,
+                    **architecture,
+                    "clip_dim": 1024,
+                    "clip_seq_len": 64,
+                    "sync_dim": 768,
+                    "sync_seq_len": 192,
+                    "text_dim": 1024,
+                    "text_seq_len": 77,
+                    "mlp_ratio": 4.0,
+                }
+                default_dtype = PRECISION_TO_TYPE[self.training_config.dit_precision]
+                init_params = {
+                    "config": model_config,
+                    "hf_config": hf_config,
+                    "latent_mean": latent_mean,
+                    "latent_std": latent_std,
+                    "empty_string_feat": empty_string_feat,
+                }
+                if use_ddp:
+                    transformer = build_replicated_model_from_scratch(
+                        MMAudioTransformer,
+                        init_params,
+                        device=get_local_torch_device(),
+                        default_dtype=default_dtype,
+                        seed=self.training_config.data.seed,
+                    )
+                else:
+                    transformer = build_fsdp_model_from_scratch(
+                        model_cls=MMAudioTransformer,
+                        init_params=init_params,
+                        device=get_local_torch_device(),
+                        hsdp_replicate_dim=self.training_config.distributed.hsdp_replicate_dim,
+                        hsdp_shard_dim=self.training_config.distributed.hsdp_shard_dim,
+                        default_dtype=default_dtype,
+                        param_dtype=torch.bfloat16,
+                        reduce_dtype=torch.float32,
+                        seed=self.training_config.data.seed,
+                        pin_cpu_memory=self.training_config.distributed.pin_cpu_memory,
+                    )
+                self._init_from = f"scratch:{variant}"
+            else:
+                if not init_from:
+                    raise ValueError("MMAudio pretrained training requires init_from, or set "
+                                     "from_scratch=true with a variant")
+                if use_ddp:
+                    raise NotImplementedError("MMAudio DDP currently supports from_scratch=true only; "
+                                              "pretrained component loading remains on the FSDP path")
+                self._init_from = str(init_from)
+                transformer = load_module_from_path(
+                    model_path=self._init_from,
+                    module_type="transformer",
+                    training_config=self.training_config,
+                    override_transformer_cls_name=self._transformer_cls_name,
+                    attention_backend=self.attention_backend,
+                )
+        else:
+            self._init_from = str(init_from or f"provided:{variant or 'unknown'}")
+        self.transformer = transformer
+        if bool(getattr(self.transformer, "v2", False)) and not allow_v2_training:
+            raise ValueError("The official MMAudio training recipe does not support `_v2` "
+                             "checkpoints. Use small_44k, medium_44k, or large_44k.")
+
+        if not self._enable_lora_if_configured(self.transformer):
+            self.transformer = apply_trainable(self.transformer, trainable=self._trainable)
+        if self._trainable:
+            # These are checkpoint statistics/fixed text features in the
+            # official model, not optimization variables. ``apply_trainable``
+            # intentionally enables a whole module, so restore their contract.
+            for name in ("latent_mean", "latent_std", "empty_string_feat"):
+                parameter = getattr(self.transformer, name, None)
+                if isinstance(parameter, torch.nn.Parameter):
+                    parameter.requires_grad_(False)
+            # The two learned null-video tokens are trained by MMAudio.
+            for name in ("empty_clip_feat", "empty_sync_feat"):
+                parameter = getattr(self.transformer, name, None)
+                if isinstance(parameter, torch.nn.Parameter):
+                    parameter.requires_grad_(True)
+        if use_ddp:
+            # Official MMAudio wraps the fully initialized FP32 model after
+            # selecting trainable parameters. It disables buffer broadcasts;
+            # the fixed latent statistics and positional buffers are identical
+            # because every rank uses the same initialization seed.
+            self.transformer = wrap_module_ddp(
+                self.transformer,
+                device=get_local_torch_device(),
+                broadcast_buffers=False,
+            )
+
+        self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
+            shift=1.0,
+            invert_sigmas=True,
+            sigma_min=0.0,
+            use_reference_discrete_timesteps=True,
+        )
+        self.dataloader: Any = None
+        self.start_step = 0
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.transformer.parameters()).device
+
+    def _unshard_transformer(self) -> None:
+        unshard = getattr(self.transformer, "unshard", None)
+        if callable(unshard):
+            unshard()
+
+    def eager_optimizer_state_parameters(self, ) -> tuple[torch.nn.Parameter, ...]:
+        parameters = (
+            getattr(self.transformer, "empty_clip_feat", None),
+            getattr(self.transformer, "empty_sync_feat", None),
+        )
+        return tuple(parameter for parameter in parameters
+                     if isinstance(parameter, torch.nn.Parameter) and parameter.requires_grad)
+
+    def init_preprocessors(self, training_config: TrainingConfig) -> None:
+        if training_config.data.preprocessed_data_type != "mmaudio_features":
+            raise ValueError("MMAudioModel requires "
+                             "training.data.preprocessed_data_type=mmaudio_features")
+        from fastvideo.dataset.mmaudio_feature_dataset import (
+            build_mmaudio_feature_dataloader, )
+
+        feature_shapes = {
+            "latent_seq_len": int(self.transformer.latent_seq_len),
+            "latent_dim": int(self.transformer.latent_dim),
+            "clip_seq_len": int(self.transformer.clip_seq_len),
+            "clip_dim": int(self.transformer.config.arch_config.clip_dim),
+            "sync_seq_len": int(self.transformer.sync_seq_len),
+            "sync_dim": int(self.transformer.config.arch_config.sync_dim),
+            "text_seq_len": int(self.transformer.config.arch_config.text_seq_len),
+            "text_dim": int(self.transformer.config.arch_config.text_dim),
+        }
+        self.dataloader = build_mmaudio_feature_dataloader(
+            training_config.data.data_path,
+            batch_size=training_config.data.train_batch_size,
+            num_data_workers=training_config.data.dataloader_num_workers,
+            seed=training_config.data.seed,
+            pin_memory=training_config.distributed.pin_cpu_memory,
+            feature_shapes=feature_shapes,
+        )
+
+    @staticmethod
+    def _batch_tensor(raw_batch: dict[str, Any], primary: str, alias: str | None = None) -> torch.Tensor:
+        value = raw_batch.get(primary)
+        if value is None and alias is not None:
+            value = raw_batch.get(alias)
+        if not isinstance(value, torch.Tensor):
+            names = primary if alias is None else f"{primary!r} (or {alias!r})"
+            raise ValueError(f"MMAudio training batch is missing tensor {names}")
+        return value
+
+    @staticmethod
+    def _existence_mask(
+        raw_batch: dict[str, Any],
+        primary: str,
+        alias: str,
+        *,
+        batch_size: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        value = raw_batch.get(primary, raw_batch.get(alias))
+        if value is None:
+            return torch.ones(batch_size, device=device, dtype=torch.bool)
+        if not isinstance(value, torch.Tensor):
+            value = torch.as_tensor(value)
+        value = value.to(device=device, dtype=torch.bool).reshape(-1)
+        if value.numel() != batch_size:
+            raise ValueError(f"{primary} must have {batch_size} values, got {value.numel()}")
+        return value
+
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        # MMAudio uses transformer-owned latent statistics and learned null
+        # tokens while constructing the batch, before ``transformer.forward``
+        # would normally trigger FSDP2's pre-forward all-gather. Explicitly
+        # unshard here so these tensors can be broadcast/expanded safely. The
+        # following transformer forward performs the normal reshard.
+        self._unshard_transformer()
+
+        device = self.device
+        latent_mean = self._batch_tensor(raw_batch, "audio_latent_mean", "a_mean").to(device=device)
+        latent_std = self._batch_tensor(raw_batch, "audio_latent_std", "a_std").to(device=device)
+        clip_features = self._batch_tensor(raw_batch, "clip_features").to(device=device).clone()
+        sync_features = self._batch_tensor(raw_batch, "sync_features").to(device=device).clone()
+        text_features = self._batch_tensor(raw_batch, "text_features").to(device=device).clone()
+
+        if latent_mean.shape != latent_std.shape:
+            raise ValueError("MMAudio audio latent mean/std shapes must match, got "
+                             f"{tuple(latent_mean.shape)} and {tuple(latent_std.shape)}")
+        batch_size = int(latent_mean.shape[0])
+        expected_latent_shape = (
+            batch_size,
+            int(self.transformer.latent_seq_len),
+            int(self.transformer.latent_dim),
+        )
+        if tuple(latent_mean.shape) != expected_latent_shape:
+            raise ValueError(f"MMAudio audio latents must have shape {expected_latent_shape}, "
+                             f"got {tuple(latent_mean.shape)}")
+
+        video_exists = self._existence_mask(
+            raw_batch,
+            "video_exists",
+            "video_exist",
+            batch_size=batch_size,
+            device=device,
+        )
+        text_exists = self._existence_mask(
+            raw_batch,
+            "text_exists",
+            "text_exist",
+            batch_size=batch_size,
+            device=device,
+        )
+        empty_clip = self.transformer.get_empty_clip_sequence(1).to(dtype=clip_features.dtype)
+        empty_sync = self.transformer.get_empty_sync_sequence(1).to(dtype=sync_features.dtype)
+        empty_text = self.transformer.get_empty_string_sequence(1).to(dtype=text_features.dtype)
+        video_exists_expanded = video_exists[:, None, None]
+        text_exists_expanded = text_exists[:, None, None]
+        clip_features = torch.where(video_exists_expanded, clip_features, empty_clip)
+        sync_features = torch.where(video_exists_expanded, sync_features, empty_sync)
+        text_features = torch.where(text_exists_expanded, text_features, empty_text)
+
+        if latents_source == "data":
+            posterior_noise = torch.empty_like(latent_mean).normal_(generator=generator)
+            clean_latents = latent_mean + latent_std * posterior_noise
+            clean_latents = self.transformer.normalize(clean_latents)
+        elif latents_source == "zeros":
+            clean_latents = torch.zeros_like(latent_mean)
+        else:
+            raise ValueError(f"Unknown latents_source: {latents_source!r}")
+
+        logit_mean = float(self.training_config.model.logit_mean)
+        logit_std = float(self.training_config.model.logit_std)
+        timestep = torch.sigmoid(torch.randn(batch_size, device=device, generator=generator) * logit_std + logit_mean)
+        prior_noise = torch.empty_like(clean_latents).normal_(generator=generator)
+        timestep_expanded = timestep[:, None, None]
+        noisy_latents = (1.0 - timestep_expanded) * prior_noise + timestep_expanded * clean_latents
+
+        null_probability = float(self.training_config.data.training_cfg_rate)
+        if not 0.0 <= null_probability <= 1.0:
+            raise ValueError("training.data.training_cfg_rate must be in [0, 1]")
+        null_video = torch.rand(batch_size, device=device, generator=generator) < null_probability
+        null_video_expanded = null_video[:, None, None]
+        clip_features = torch.where(null_video_expanded, empty_clip, clip_features)
+        sync_features = torch.where(null_video_expanded, empty_sync, sync_features)
+        null_text = torch.rand(batch_size, device=device, generator=generator) < null_probability
+        null_text_expanded = null_text[:, None, None]
+        text_features = torch.where(null_text_expanded, empty_text, text_features)
+
+        # Keep learned null-video tokens in every step's autograd graph even
+        # when this batch samples no null conditions. AdamW then creates a
+        # stable optimizer-state schema, which is required for strict DCP
+        # checkpoint resume. Multiplication by zero leaves forward values
+        # unchanged; selected null conditions still receive their real grads.
+        clip_features = clip_features + empty_clip.sum() * 0.0
+        sync_features = sync_features + empty_sync.sum() * 0.0
+
+        batch = TrainingBatch()
+        batch.latents = clean_latents
+        batch.noisy_model_input = noisy_latents
+        batch.noise = prior_noise
+        batch.timesteps = timestep
+        batch.sigmas = timestep_expanded
+        batch.training_target = clean_latents - prior_noise
+        batch.raw_latent_shape = tuple(clean_latents.shape)
+        batch.conditional_dict = {
+            "clip_features": clip_features,
+            "sync_features": sync_features,
+            "text_features": text_features,
+        }
+        return batch
+
+    def add_noise(
+        self,
+        clean_latents: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        while timestep.ndim < clean_latents.ndim:
+            timestep = timestep.unsqueeze(-1)
+        return (1.0 - timestep) * noise + timestep * clean_latents
+
+    def compile_training_forward(
+        self,
+        compile_kwargs: dict[str, Any],
+    ) -> None:
+        """Compile transformer math while leaving DDP control flow eager."""
+        module = unwrap_ddp_module(self.transformer)
+        compiled_forward = torch.compile(module.forward, **compile_kwargs)
+        module.forward = compiled_forward
+        logger.info(
+            "Compiled inner MMAudio transformer forward with kwargs=%s",
+            compile_kwargs,
+        )
+
+    def flow_matching_forward(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        clip_features: torch.Tensor,
+        sync_features: torch.Tensor,
+        text_features: torch.Tensor,
+    ) -> torch.Tensor:
+        """Tensor-only adapter around the DDP-wrapped training forward."""
+        device_type = noisy_latents.device.type
+        with torch.autocast(
+                device_type=device_type,
+                dtype=torch.bfloat16,
+                enabled=device_type == "cuda",
+        ):
+            return self.transformer(
+                hidden_states=noisy_latents,
+                encoder_hidden_states={
+                    "clip_features": clip_features,
+                    "sync_features": sync_features,
+                    "text_features": text_features,
+                },
+                timestep=timestep,
+            )
+
+    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",
+    ) -> torch.Tensor:
+        if attn_kind != "dense":
+            raise ValueError("MMAudio training currently supports dense Torch SDPA only")
+        if conditional:
+            conditions = batch.conditional_dict
+            if conditions is None:
+                raise RuntimeError("MMAudio TrainingBatch is missing conditional_dict")
+        elif cfg_uncond is not None:
+            conditions = cfg_uncond
+        else:
+            self._unshard_transformer()
+            batch_size = noisy_latents.shape[0]
+            conditions = {
+                "clip_features": self.transformer.get_empty_clip_sequence(batch_size),
+                "sync_features": self.transformer.get_empty_sync_sequence(batch_size),
+                "text_features": self.transformer.get_empty_string_sequence(batch_size),
+            }
+
+        with set_forward_context(current_timestep=timestep, attn_metadata=None):
+            return self.flow_matching_forward(
+                noisy_latents,
+                timestep,
+                conditions["clip_features"],
+                conditions["sync_features"],
+                conditions["text_features"],
+            )
+
+    def backward(
+        self,
+        loss: torch.Tensor,
+        ctx: Any,
+        *,
+        grad_accum_rounds: int,
+    ) -> None:
+        timestep, attn_metadata = ctx
+        with set_forward_context(current_timestep=timestep, attn_metadata=attn_metadata):
+            (loss / max(1, int(grad_accum_rounds))).backward()
+
+
+__all__ = ["MMAudioModel"]

fastvideo/train/trainer.py

@@ -161,6 +161,7 @@ class Trainer:
         progress = tqdm(
             range(start_step + 1, max_steps + 1),
             initial=start_step,
+            total=max_steps,
             desc="Steps",
             disable=self.local_rank > 0,
         )

fastvideo/train/utils/config.py

@@ -344,7 +344,8 @@ def _build_training_config(
 
     ``preprocessed_data_type`` selects the dataloader's tensor contract:
     ``t2v`` carries video and text, ``t2va`` carries synchronized video,
-    audio, and text, and ``text_only`` carries prompt conditioning.
+    audio, and text, ``text_only`` carries prompt conditioning, and
+    ``mmaudio_features`` carries cached MMAudio V2A features.
     """
     d = dict(t.get("distributed", {}) or {})
     da = dict(t.get("data", {}) or {})
@@ -379,9 +380,9 @@ def _build_training_config(
         data_path = str(raw_data_path)
 
     preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower()
-    if preprocessed_data_type not in {"t2v", "t2va", "text_only"}:
+    if preprocessed_data_type not in {"t2v", "t2va", "text_only", "mmaudio_features"}:
         raise ValueError("training.data.preprocessed_data_type must be one of "
-                         "{'t2v', 't2va', 'text_only'}, got "
+                         "{'t2v', 't2va', 'text_only', 'mmaudio_features'}, got "
                          f"{preprocessed_data_type!r}")
 
     peak_tflops = get_optional_float(
@@ -393,8 +394,12 @@ def _build_training_config(
         raise ValueError("training.performance.peak_tflops_per_gpu must be "
                          f"> 0, got {peak_tflops}")
 
+    from fastvideo.train.utils.distributed_strategy import (
+        normalize_distributed_strategy, )
+
     return TrainingConfig(
         distributed=DistributedConfig(
+            strategy=normalize_distributed_strategy(str(d.get("strategy", "fsdp") or "fsdp")),
             num_gpus=num_gpus,
             tp_size=int(d.get("tp_size", 1) or 1),
             sp_size=int(d.get("sp_size", num_gpus) or num_gpus),
@@ -417,12 +422,16 @@ def _build_training_config(
         optimizer=OptimizerConfig(
             learning_rate=float(o.get("learning_rate", 0.0) or 0.0),
             betas=betas,
+            eps=float(o.get("eps", 1e-8) or 1e-8),
             weight_decay=float(o.get("weight_decay", 0.0) or 0.0),
             lr_scheduler=str(o.get("lr_scheduler", "constant") or "constant"),
             lr_warmup_steps=int(o.get("lr_warmup_steps", 0) or 0),
             lr_num_cycles=int(o.get("lr_num_cycles", 0) or 0),
             lr_power=float(o.get("lr_power", 0.0) or 0.0),
             min_lr_ratio=float(o.get("min_lr_ratio", 0.5) or 0.5),
+            lr_milestones=tuple(int(step) for step in (o.get("lr_milestones", []) or [])),
+            lr_gamma=float(o.get("lr_gamma", 0.1) or 0.1),
+            fused=bool(o.get("fused", False)),
         ),
         loop=TrainingLoopConfig(
             max_train_steps=int(lo.get("max_train_steps", 0) or 0),
@@ -454,6 +463,8 @@ def _build_training_config(
             precondition_outputs=bool(m.get("precondition_outputs", False)),
             moba_config=dict(m.get("moba_config", {}) or {}),
             enable_gradient_checkpointing_type=(m.get("enable_gradient_checkpointing_type")),
+            compile_train_fn=bool(m.get("compile_train_fn", False)),
+            torch_compile_kwargs=dict(m.get("torch_compile_kwargs", {}) or {}),
         ),
         pipeline_config=pipeline_config,
         model_path=model_path,

fastvideo/train/utils/distributed_strategy.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Distributed model strategies for the modular trainer."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Any, TypeVar
+from collections.abc import Generator
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+_ModelT = TypeVar("_ModelT", bound=torch.nn.Module)
+SUPPORTED_DISTRIBUTED_STRATEGIES = frozenset({"fsdp", "ddp"})
+
+
+def normalize_distributed_strategy(value: str | None) -> str:
+    """Normalize and validate a modular-training strategy name."""
+    strategy = str(value or "fsdp").strip().lower()
+    if strategy not in SUPPORTED_DISTRIBUTED_STRATEGIES:
+        choices = ", ".join(sorted(SUPPORTED_DISTRIBUTED_STRATEGIES))
+        raise ValueError(f"training.distributed.strategy must be one of {{{choices}}}, "
+                         f"got {strategy!r}")
+    return strategy
+
+
+def is_ddp_strategy(training_config: Any) -> bool:
+    distributed = getattr(training_config, "distributed", None)
+    return normalize_distributed_strategy(getattr(distributed, "strategy", "fsdp")) == "ddp"
+
+
+@contextmanager
+def _fork_model_rng(
+    device: torch.device,
+    seed: int,
+) -> Generator[None, None, None]:
+    fork_devices: list[int] = []
+    if device.type == "cuda":
+        fork_devices = [device.index if device.index is not None else torch.cuda.current_device()]
+    with torch.random.fork_rng(devices=fork_devices):
+        torch.manual_seed(int(seed))
+        yield
+
+
+def build_replicated_model_from_scratch(
+    model_cls: type[_ModelT],
+    init_params: dict[str, Any],
+    *,
+    device: torch.device,
+    default_dtype: torch.dtype,
+    seed: int,
+) -> _ModelT:
+    """Initialize one deterministic, full model replica on the local device.
+
+    Model plugins call :func:`wrap_module_ddp` only after applying their
+    trainable/frozen parameter policy. This is required because DDP builds its
+    gradient reducer from the parameters that require gradients at wrap time.
+    """
+    old_dtype = torch.get_default_dtype()
+    try:
+        torch.set_default_dtype(default_dtype)
+        # Official MMAudio constructs on CPU and then calls ``.cuda()`` before
+        # DDP. Preserve that RNG/initialization order for scratch parity.
+        with _fork_model_rng(device, seed), torch.device("cpu"):
+            model = model_cls(**init_params)
+    finally:
+        torch.set_default_dtype(old_dtype)
+    return model.to(device=device)
+
+
+class DelegatingDistributedDataParallel(DistributedDataParallel):
+    """DDP wrapper that preserves a model plugin's public attributes.
+
+    FastVideo model plugins access architecture attributes and helper methods
+    such as ``latent_seq_len`` and ``normalize`` through ``transformer``.
+    Native PyTorch DDP exposes these on ``module`` only; this adapter delegates
+    unknown attributes while retaining the standard DDP forward/reducer.
+    """
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            module = super().__getattr__("module")
+            return getattr(module, name)
+
+
+def wrap_module_ddp(
+    module: _ModelT,
+    *,
+    device: torch.device,
+    broadcast_buffers: bool = False,
+) -> DelegatingDistributedDataParallel:
+    """Wrap a full model replica like official MMAudio's DDP runner."""
+    if not dist.is_available() or not dist.is_initialized():
+        raise RuntimeError("DDP strategy requires an initialized process group")
+    if not any(parameter.requires_grad for parameter in module.parameters()):
+        raise ValueError("DDP strategy requires at least one trainable parameter")
+
+    kwargs: dict[str, Any] = {
+        "broadcast_buffers": bool(broadcast_buffers),
+    }
+    if device.type == "cuda":
+        if device.index is None:
+            raise ValueError("CUDA DDP requires a concrete local device index")
+        kwargs.update({
+            "device_ids": [device.index],
+            "output_device": device.index,
+        })
+    wrapped = DelegatingDistributedDataParallel(module, **kwargs)
+    logger.info(
+        "Wrapped %s with native DDP on %s (broadcast_buffers=%s)",
+        type(module).__name__,
+        device,
+        broadcast_buffers,
+    )
+    return wrapped
+
+
+def unwrap_ddp_module(module: _ModelT) -> torch.nn.Module:
+    """Return the underlying module for DDP, otherwise return the input."""
+    if isinstance(module, DistributedDataParallel):
+        return module.module
+    return module
+
+
+__all__ = [
+    "DelegatingDistributedDataParallel",
+    "SUPPORTED_DISTRIBUTED_STRATEGIES",
+    "build_replicated_model_from_scratch",
+    "is_ddp_strategy",
+    "normalize_distributed_strategy",
+    "unwrap_ddp_module",
+    "wrap_module_ddp",
+]

fastvideo/train/utils/optimizer.py

@@ -36,28 +36,80 @@ def build_optimizer_and_scheduler(
         raise ValueError("No trainable parameters passed to "
                          "build_optimizer_and_scheduler")
 
+    optimizer_kwargs = {}
+    if bool(getattr(optimizer_config, "fused", False)):
+        optimizer_kwargs["fused"] = True
     optimizer = torch.optim.AdamW(
         params,
         lr=float(learning_rate),
         betas=betas,
         weight_decay=float(optimizer_config.weight_decay),
-        eps=1e-8,
+        eps=float(optimizer_config.eps),
+        **optimizer_kwargs,
     )
 
-    scheduler = get_scheduler(
-        str(scheduler_name),
-        optimizer=optimizer,
-        num_warmup_steps=int(optimizer_config.lr_warmup_steps),
-        num_training_steps=int(loop_config.max_train_steps),
-        num_cycles=int(optimizer_config.lr_num_cycles),
-        power=float(optimizer_config.lr_power),
-        min_lr_ratio=float(optimizer_config.min_lr_ratio),
-        last_epoch=-1,
-    )
+    if str(scheduler_name) == "multistep_with_warmup":
+        milestones = list(optimizer_config.lr_milestones)
+        if not milestones:
+            raise ValueError("multistep_with_warmup requires training.optimizer.lr_milestones")
+        warmup_steps = int(optimizer_config.lr_warmup_steps)
+        if warmup_steps <= 0:
+            raise ValueError("multistep_with_warmup requires lr_warmup_steps > 0")
+
+        def warmup(current_step: int) -> float:
+            return (current_step + 1) / (warmup_steps + 1)
+
+        warmup_scheduler = torch.optim.lr_scheduler.LambdaLR(
+            optimizer,
+            lr_lambda=warmup,
+        )
+        step_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+            optimizer,
+            milestones=milestones,
+            gamma=float(optimizer_config.lr_gamma),
+        )
+        scheduler = torch.optim.lr_scheduler.SequentialLR(
+            optimizer,
+            schedulers=[warmup_scheduler, step_scheduler],
+            milestones=[warmup_steps],
+        )
+    else:
+        scheduler = get_scheduler(
+            str(scheduler_name),
+            optimizer=optimizer,
+            num_warmup_steps=int(optimizer_config.lr_warmup_steps),
+            num_training_steps=int(loop_config.max_train_steps),
+            num_cycles=int(optimizer_config.lr_num_cycles),
+            power=float(optimizer_config.lr_power),
+            min_lr_ratio=float(optimizer_config.min_lr_ratio),
+            last_epoch=-1,
+        )
 
     return optimizer, scheduler
 
 
+def seed_adamw_parameter_state(
+    optimizer: torch.optim.Optimizer,
+    parameter: torch.nn.Parameter,
+) -> None:
+    """Create AdamW state using PyTorch's fused/capturable device contract."""
+    if optimizer.state.get(parameter):
+        return
+    owner_group = next(
+        (group for group in optimizer.param_groups if any(candidate is parameter for candidate in group["params"])),
+        None,
+    )
+    if owner_group is None:
+        raise ValueError("Cannot seed optimizer state for an unowned parameter")
+    step_on_parameter = bool(owner_group.get("capturable", False) or owner_group.get("fused", False))
+    step_device = parameter.device if step_on_parameter else torch.device("cpu")
+    optimizer.state[parameter] = {
+        "step": torch.zeros((), dtype=torch.float32, device=step_device),
+        "exp_avg": torch.zeros_like(parameter),
+        "exp_avg_sq": torch.zeros_like(parameter),
+    }
+
+
 def clip_grad_norm_if_needed(
     module: torch.nn.Module,
     max_grad_norm: float,

fastvideo/train/utils/training_config.py

@@ -18,6 +18,8 @@ class DistributedConfig:
     hsdp_replicate_dim: int = 1
     hsdp_shard_dim: int = -1
     pin_cpu_memory: bool = False
+    # Keep this last so existing positional construction remains compatible.
+    strategy: str = "fsdp"
 
 
 @dataclass(slots=True)
@@ -38,12 +40,18 @@ class DataConfig:
 class OptimizerConfig:
     learning_rate: float = 0.0
     betas: tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
     weight_decay: float = 0.0
     lr_scheduler: str = "constant"
     lr_warmup_steps: int = 0
     lr_num_cycles: int = 0
     lr_power: float = 0.0
     min_lr_ratio: float = 0.5
+    lr_milestones: tuple[int, ...] = ()
+    lr_gamma: float = 0.1
+    # ``False`` preserves PyTorch's existing optimizer selection. When true,
+    # FastVideo explicitly requests CUDA fused AdamW.
+    fused: bool = False
 
 
 @dataclass(slots=True)
@@ -87,6 +95,10 @@ class ModelTrainingConfig:
     precondition_outputs: bool = False
     moba_config: dict = field(default_factory=dict)
     enable_gradient_checkpointing_type: str | None = None
+    # Compile the model adapter's tensor-only training forward. This is
+    # separate from loader-level model compilation and defaults off.
+    compile_train_fn: bool = False
+    torch_compile_kwargs: dict = field(default_factory=dict)
 
 
 @dataclass(slots=True)

fastvideo/training/checkpointing_utils.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import random
+from collections import Counter
+from collections.abc import Mapping
 from typing import Any
 
 import numpy as np
@@ -21,6 +23,8 @@ class ModelWrapper(torch.distributed.checkpoint.stateful.Stateful):
             k.replace("._checkpoint_wrapped_module.", ".")
             for k, v in self.model.named_parameters() if v.requires_grad
         }
+        if isinstance(self.model, torch.nn.parallel.DistributedDataParallel):
+            param_requires_grad = {key.removeprefix("module.") for key in param_requires_grad}
 
         filtered_state_dict = {k: v for k, v in state_dict.items() if k in param_requires_grad}
 
@@ -62,10 +66,35 @@ class SchedulerWrapper(torch.distributed.checkpoint.stateful.Stateful):
         self.scheduler = scheduler
 
     def state_dict(self) -> dict[str, Any]:
-        return {"scheduler": self.scheduler.state_dict()}
+        return {"scheduler": _scheduler_state_for_dcp(self.scheduler.state_dict())}
 
     def load_state_dict(self, state_dict: dict[str, Any]) -> None:
-        self.scheduler.load_state_dict(state_dict["scheduler"])
+        self.scheduler.load_state_dict(_scheduler_state_from_dcp(state_dict["scheduler"]))
+
+
+def _scheduler_state_for_dcp(value: Any, *, name: str = "") -> Any:
+    """Use stable string keys for scheduler mappings flattened by DCP."""
+    if isinstance(value, Mapping):
+        if name == "milestones":
+            return {str(key): item for key, item in value.items()}
+        return {key: _scheduler_state_for_dcp(item, name=str(key)) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_scheduler_state_for_dcp(item) for item in value]
+    return value
+
+
+def _scheduler_state_from_dcp(value: Any, *, name: str = "") -> Any:
+    """Restore MultiStepLR milestones after DCP stringifies mapping keys."""
+    if isinstance(value, Mapping):
+        if name == "milestones":
+            milestones: dict[int, Any] = {}
+            for key, item in value.items():
+                milestones[int(key)] = item
+            return Counter(milestones)
+        return {key: _scheduler_state_from_dcp(item, name=str(key)) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_scheduler_state_from_dcp(item) for item in value]
+    return value
 
 
 class RandomStateWrapper(torch.distributed.checkpoint.stateful.Stateful):

fastvideo/utils.py

@@ -657,8 +657,11 @@ def verify_model_config_and_directory(
     else:
         # Full snapshots keep the historical invariant that transformer/ is
         # present and every active manifest component exists locally.
+        # Explicitly marked offline preprocessor trees contain only frozen
+        # encoders and avoid duplicating a DiT that will never be loaded.
         transformer_dir = os.path.join(model_path, "transformer")
-        if not os.path.exists(transformer_dir):
+        preprocessor_only = config.get("_fastvideo_preprocessor_only") is True
+        if not preprocessor_only and not os.path.exists(transformer_dir):
             raise ValueError(f"Model directory {model_path} does not contain a transformer/ directory.")
 
     # Diffusers convention: component entries start with [library, class].

fastvideo/workflow/preprocess/preprocess_workflow.py

@@ -117,6 +117,9 @@ class PreprocessWorkflow(WorkflowBase):
         elif fastvideo_args.workload_type == WorkloadType.I2V:
             from fastvideo.workflow.preprocess.preprocess_workflow_i2v import (PreprocessWorkflowI2V)
             return cast(PreprocessWorkflow, PreprocessWorkflowI2V)
+        elif fastvideo_args.workload_type == WorkloadType.V2A:
+            from fastvideo.workflow.preprocess.preprocess_workflow_v2a import (PreprocessWorkflowV2A)
+            return cast(PreprocessWorkflow, PreprocessWorkflowV2A)
         else:
             raise ValueError(
                 f"Workload type: {fastvideo_args.workload_type} is not supported in preprocessing workflow.")

fastvideo/workflow/preprocess/preprocess_workflow_v2a.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared raw-media workflow for video-to-audio feature preprocessing."""
+
+from __future__ import annotations
+
+import json
+import math
+from collections import deque
+from concurrent.futures import Future, ThreadPoolExecutor
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+from collections.abc import Iterator
+
+import torch
+from torch.utils.data import DataLoader, Subset
+from tqdm import tqdm
+
+from fastvideo.configs.configs import DatasetType, VideoLoaderType
+from fastvideo.dataset.v2a_feature_cache import V2AFeatureShardWriter
+from fastvideo.dataset.vggsound import VGGSoundDataset
+from fastvideo.distributed.parallel_state import get_world_rank, get_world_size
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.pipeline_batch_info import PreprocessBatch
+from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import (MMAudioTorioRowPreprocessor,
+                                                                       PREPROCESSED_MEDIA_KEY, PREPROCESS_ERROR_KEY)
+from fastvideo.workflow.preprocess.preprocess_workflow import PreprocessWorkflow
+
+if TYPE_CHECKING:
+    from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase
+
+logger = init_logger(__name__)
+
+
+def _identity_collate(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    return rows
+
+
+def build_v2a_dataset(
+    dataset_type: DatasetType,
+    dataset_path: str,
+    *,
+    dataset_metadata_path: str = "",
+    dataset_split: str = "train",
+):
+    """Build a raw-media dataset for a V2A preprocessing workflow."""
+    if isinstance(dataset_type, str):
+        dataset_type = DatasetType.from_string(dataset_type)
+    if dataset_type is DatasetType.VGGSOUND:
+        metadata_path = dataset_metadata_path or None
+        return VGGSoundDataset(
+            dataset_path,
+            split=dataset_split,
+            metadata_path=metadata_path,
+        )
+    raise ValueError(f"V2A preprocessing currently supports dataset_type=vggsound; got {dataset_type.value!r}")
+
+
+class V2AForwardBatchBuilder:
+    """Translate standardized raw-media rows into a FastVideo batch."""
+
+    def __init__(self, seed: int) -> None:
+        self.seed = seed
+
+    def __call__(self, rows: list[dict[str, Any]]) -> PreprocessBatch:
+        extra: dict[str, Any] = {}
+        preprocessed_media = [row.get(PREPROCESSED_MEDIA_KEY) for row in rows]
+        preprocess_errors = [row.get(PREPROCESS_ERROR_KEY) for row in rows]
+        if any(media is not None for media in preprocessed_media):
+            extra[PREPROCESSED_MEDIA_KEY] = preprocessed_media
+        if any(error is not None for error in preprocess_errors):
+            extra[PREPROCESS_ERROR_KEY] = preprocess_errors
+        return PreprocessBatch(
+            data_type="video",
+            video_loader=[str(row["video_path"]) for row in rows],
+            video_file_name=[str(row["id"]) for row in rows],
+            prompt=[str(row["caption"]) for row in rows],
+            batch_size=len(rows),
+            generator=torch.Generator("cpu").manual_seed(self.seed),
+            extra=extra,
+        )
+
+
+class PreprocessWorkflowV2A(PreprocessWorkflow):
+    """Run model-specific V2A encoders with shared data/cache orchestration."""
+
+    training_dataloader: DataLoader
+    preprocess_pipeline: ComposedPipelineBase
+    forward_batch_builder: V2AForwardBatchBuilder
+    feature_writer: V2AFeatureShardWriter
+
+    def register_components(self) -> None:
+        self.torio_row_preprocessor: MMAudioTorioRowPreprocessor | None = None
+        self.torio_decode_workers = 0
+        self.initial_skipped = 0
+        assert self.fastvideo_args.preprocess_config is not None
+        config = self.fastvideo_args.preprocess_config
+        dataset = build_v2a_dataset(
+            config.dataset_type,
+            config.dataset_path,
+            dataset_metadata_path=config.dataset_metadata_path,
+            dataset_split=config.dataset_split,
+        )
+
+        rank = get_world_rank()
+        world_size = get_world_size()
+        feature_writer = V2AFeatureShardWriter(
+            config.dataset_output_dir,
+            rank=rank,
+            samples_per_shard=config.samples_per_file,
+        )
+        rank_indices = list(range(rank, len(dataset), world_size))
+        pending_indices = [index for index in rank_indices if not feature_writer.contains(str(dataset[index]["id"]))]
+        self.initial_skipped = len(rank_indices) - len(pending_indices)
+        rank_indices = pending_indices
+        rank_dataset = Subset(dataset, rank_indices)
+
+        dataloader_workers = config.dataloader_num_workers
+        if config.video_loader_type == VideoLoaderType.TORIO:
+            pc = self.fastvideo_args.pipeline_config
+            duration_s = float(pc.duration_s)
+            sample_rate = int(pc.sampling_rate)
+            alignment = int(pc.spectrogram_frame_rate * pc.latent_downsample_rate)
+            target_samples = math.ceil(duration_s * sample_rate / alignment) * alignment
+            self.torio_row_preprocessor = MMAudioTorioRowPreprocessor(
+                duration_s=duration_s,
+                target_sample_rate=sample_rate,
+                target_samples=target_samples,
+                normalize_audio=config.dataset_split.lower() == "train",
+                clip_fps=pc.clip_frame_rate,
+                sync_fps=pc.sync_frame_rate,
+                clip_size=pc.clip_image_size,
+                sync_size=pc.sync_image_size,
+            )
+            self.torio_decode_workers = max(0, dataloader_workers)
+            # Large video tensors cannot pass safely through a small /dev/shm.
+            # Threads also avoid tensor serialization while overlapping the
+            # next decode batch with GPU feature extraction.
+            dataloader_workers = 0
+            logger.info(
+                "Torio preprocessing uses %d background decode threads on rank %d",
+                self.torio_decode_workers,
+                rank,
+            )
+
+        dataloader = DataLoader(
+            rank_dataset,
+            batch_size=config.preprocess_video_batch_size,
+            num_workers=dataloader_workers,
+            collate_fn=_identity_collate,
+            shuffle=False,
+            persistent_workers=dataloader_workers > 0,
+        )
+        self.add_component("training_dataloader", dataloader)
+        self.add_component("forward_batch_builder", V2AForwardBatchBuilder(config.seed))
+        self.add_component("feature_writer", feature_writer)
+
+    def _iter_preprocessed_rows(self) -> Iterator[list[dict[str, Any]]]:
+        preprocessor = self.torio_row_preprocessor
+        if preprocessor is None:
+            yield from self.training_dataloader
+            return
+
+        if self.torio_decode_workers <= 0:
+            for rows in self.training_dataloader:
+                yield [preprocessor(row) for row in rows]
+            return
+
+        raw_batches = iter(self.training_dataloader)
+        queued_batches: deque[list[Future[dict[str, Any]]]] = deque()
+        with ThreadPoolExecutor(
+                max_workers=self.torio_decode_workers,
+                thread_name_prefix="mmaudio-torio",
+        ) as executor:
+
+            def submit_next() -> bool:
+                try:
+                    rows = next(raw_batches)
+                except StopIteration:
+                    return False
+                queued_batches.append([executor.submit(preprocessor, row) for row in rows])
+                return True
+
+            submit_next()
+            submit_next()
+            while queued_batches:
+                futures = queued_batches.popleft()
+                rows = [future.result() for future in futures]
+                submit_next()
+                yield rows
+
+    def prepare_system_environment(self) -> None:
+        assert self.fastvideo_args.preprocess_config is not None
+        output_root = Path(self.fastvideo_args.preprocess_config.dataset_output_dir).expanduser().resolve()
+        output_root.mkdir(parents=True, exist_ok=True)
+        self.output_root = output_root
+        self.failure_log = output_root / f"failures_rank_{get_world_rank():05d}.jsonl"
+        logger.info("V2A feature cache output: %s", output_root)
+
+    def _write_failures(self, failures: list[dict[str, str]]) -> None:
+        if not failures:
+            return
+        with self.failure_log.open("a", encoding="utf-8") as handle:
+            for failure in failures:
+                handle.write(json.dumps(failure, ensure_ascii=False) + "\n")
+
+    def run(self) -> None:
+        processed = 0
+        skipped = self.initial_skipped
+        failed = 0
+        try:
+            for rows in tqdm(
+                    self._iter_preprocessed_rows(),
+                    total=len(self.training_dataloader),
+                    desc="Preprocessing V2A training data",
+                    unit="batch",
+            ):
+                pending_rows = [row for row in rows if not self.feature_writer.contains(str(row["id"]))]
+                skipped += len(rows) - len(pending_rows)
+                if not pending_rows:
+                    continue
+                batch = self.forward_batch_builder(pending_rows)
+                batch = self.preprocess_pipeline.forward(batch, self.fastvideo_args)
+                failures = batch.extra.get("failed_samples", [])
+                failed += len(failures)
+                self._write_failures(failures)
+                features = batch.extra.get("precomputed_features", {})
+                metadata = batch.extra.get("precomputed_metadata", [])
+                self.feature_writer.append(features, metadata)
+                processed += len(metadata)
+        finally:
+            self.feature_writer.close()
+        logger.info(
+            "Finished V2A preprocessing on rank %s: processed=%s skipped=%s failed=%s output=%s",
+            get_world_rank(),
+            processed,
+            skipped,
+            failed,
+            self.output_root,
+        )
+
+
+__all__ = [
+    "PreprocessWorkflowV2A",
+    "V2AForwardBatchBuilder",
+    "build_v2a_dataset",
+]

pyproject.toml

@@ -147,6 +147,12 @@ swanlab = [
     "swanlab>=0.6.7",
 ]
 
+mmaudio-train = [
+    "librosa>=0.10.0",
+    "nitrous-ema==0.0.1",
+    "tensordict>=0.6.1",
+]
+
 lint = [
     "pre-commit==4.0.1",
     "types-aiofiles",

scripts/checkpoint_conversion/convert_mmaudio_to_diffusers.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Convert MMAudio ``large_44k_v2`` assets into a FastVideo component tree.
+"""Convert every official MMAudio variant into a FastVideo component tree.
 
 The converter is deliberately offline: every large source asset must already
 exist locally. It splits the shared DFN5B OpenCLIP checkpoint into native
@@ -16,6 +16,10 @@ Example::
       --dfn5b-dir official_weights/mmaudio/DFN5B-CLIP-ViT-H-14-384 \
       --bigvgan-dir official_weights/mmaudio/bigvgan_v2_44khz_128band_512x \
       --output converted_weights/mmaudio/large_44k_v2
+
+Use ``--variant small_16k`` with ``v1-16.pth`` and ``best_netG.pt`` for
+the 16 kHz model. The four 44.1 kHz variants share ``v1-44.pth`` and the
+NVIDIA BigVGAN-v2 directory.
 """
 
 from __future__ import annotations
@@ -30,24 +34,80 @@ import torch
 from safetensors.torch import save_file
 
 
-TRANSFORMER_CONFIG = {
+BASE_TRANSFORMER_CONFIG = {
     "_class_name": "MMAudioTransformer",
     "latent_dim": 40,
     "clip_dim": 1024,
     "sync_dim": 768,
     "text_dim": 1024,
-    "hidden_dim": 896,
-    "depth": 21,
-    "fused_depth": 14,
-    "num_heads": 14,
     "mlp_ratio": 4.0,
     "latent_seq_len": 345,
     "clip_seq_len": 64,
     "sync_seq_len": 192,
     "text_seq_len": 77,
-    "v2": True,
 }
 
+TRANSFORMER_VARIANTS = {
+    "small_16k": {
+        "latent_dim": 20,
+        "hidden_dim": 448,
+        "depth": 12,
+        "fused_depth": 8,
+        "num_heads": 7,
+        "latent_seq_len": 250,
+        "v2": False,
+    },
+    "small_44k": {
+        "hidden_dim": 448,
+        "depth": 12,
+        "fused_depth": 8,
+        "num_heads": 7,
+        "v2": False,
+    },
+    "medium_44k": {
+        "hidden_dim": 896,
+        "depth": 12,
+        "fused_depth": 8,
+        "num_heads": 14,
+        "v2": False,
+    },
+    "large_44k": {
+        "hidden_dim": 896,
+        "depth": 21,
+        "fused_depth": 14,
+        "num_heads": 14,
+        "v2": False,
+    },
+    "large_44k_v2": {
+        "hidden_dim": 896,
+        "depth": 21,
+        "fused_depth": 14,
+        "num_heads": 14,
+        "v2": True,
+    },
+}
+
+BIGVGAN_16K_CONFIG = {
+    "_class_name": "BigVGANV2",
+    "resblock": "1",
+    "num_mels": 80,
+    "upsample_rates": [4, 4, 2, 2, 2, 2],
+    "upsample_kernel_sizes": [8, 8, 4, 4, 4, 4],
+    "upsample_initial_channel": 1536,
+    "resblock_kernel_sizes": [3, 7, 11],
+    "resblock_dilation_sizes": [
+        [1, 3, 5],
+        [1, 3, 5],
+        [1, 3, 5],
+    ],
+    "activation": "snakebeta",
+    "snake_logscale": True,
+    "use_bias_at_final": True,
+    "use_tanh_at_final": True,
+    "weight_norm_removed": False,
+}
+
+
 TEXT_ENCODER_CONFIG = {
     "architectures": ["MMAudioDFNCLIPTextEncoder"],
     "vocab_size": 49408,
@@ -263,28 +323,112 @@ def _load_dfn5b_state(directory: Path) -> dict[str, torch.Tensor]:
 def convert(args: argparse.Namespace) -> None:
     output = args.output.resolve()
     output.mkdir(parents=True, exist_ok=True)
+    mode = "16k" if args.variant.endswith("16k") else "44k"
+    model_index = dict(MODEL_INDEX)
+    model_index["_fastvideo_mmaudio_variant"] = args.variant
+    model_index["_fastvideo_audio_sample_rate"] = (
+        16000 if mode == "16k" else 44100
+    )
 
-    transformer_state = _load_torch_state(args.transformer_checkpoint)
-    # Official ``MMAudio.load_weights`` discards this derived buffer. Keeping
-    # it would make a standard strict FastVideo component load fail.
-    transformer_state.pop("t_embed.freqs", None)
-    transformer_state.pop("latent_rot", None)
-    transformer_state.pop("clip_rot", None)
-    _write_component(output, "transformer", transformer_state, TRANSFORMER_CONFIG)
+    transformer_config_only = bool(getattr(args, "transformer_config_only", False))
+    if args.transformer_only and args.preprocessor_only:
+        raise ValueError(
+            "--transformer-only and --preprocessor-only are mutually exclusive")
+    if transformer_config_only and (args.transformer_only or args.preprocessor_only):
+        raise ValueError(
+            "--transformer-config-only cannot be combined with other partial conversion flags"
+        )
+
+    transformer_config = {
+        **BASE_TRANSFORMER_CONFIG,
+        **TRANSFORMER_VARIANTS[args.variant],
+    }
+    if transformer_config_only:
+        transformer_dir = output / "transformer"
+        transformer_dir.mkdir(parents=True, exist_ok=True)
+        _write_json(transformer_dir / "config.json", transformer_config)
+        transformer_model_index = {
+            key: value
+            for key, value in model_index.items()
+            if key.startswith("_") or key == "transformer"
+        }
+        _write_json(output / "model_index.json", transformer_model_index)
+        print(f"Wrote MMAudio {args.variant} transformer export skeleton to {output}")
+        return
+
+    if not args.preprocessor_only:
+        if args.transformer_checkpoint is None:
+            raise ValueError(
+                "--transformer-checkpoint is required unless --preprocessor-only is used")
+        transformer_state = _load_torch_state(args.transformer_checkpoint)
+        # Official ``MMAudio.load_weights`` discards these derived buffers.
+        transformer_state.pop("t_embed.freqs", None)
+        transformer_state.pop("latent_rot", None)
+        transformer_state.pop("clip_rot", None)
+        _write_component(output, "transformer", transformer_state,
+                         transformer_config)
+
+    if args.transformer_only:
+        transformer_model_index = {
+            key: value
+            for key, value in model_index.items()
+            if key.startswith("_") or key == "transformer"
+        }
+        _write_json(output / "model_index.json", transformer_model_index)
+        print(f"Converted MMAudio {args.variant} transformer to {output}")
+        return
+
+    required_assets = {
+        "--audio-vae-checkpoint": args.audio_vae_checkpoint,
+        "--synchformer-checkpoint": args.synchformer_checkpoint,
+        "--dfn5b-dir": args.dfn5b_dir,
+    }
+    if not args.preprocessor_only:
+        if mode == "16k":
+            required_assets["--bigvgan-checkpoint"] = args.bigvgan_checkpoint
+        else:
+            required_assets["--bigvgan-dir"] = args.bigvgan_dir
+    missing_assets = [name for name, path in required_assets.items() if path is None]
+    if missing_assets:
+        raise ValueError(
+            "Full pipeline conversion requires " + ", ".join(missing_assets)
+        )
+
+    assert args.audio_vae_checkpoint is not None
+    assert args.synchformer_checkpoint is not None
+    assert args.dfn5b_dir is not None
+    if not args.preprocessor_only:
+        if mode == "16k":
+            assert args.bigvgan_checkpoint is not None
+        else:
+            assert args.bigvgan_dir is not None
 
     vae_state = _load_torch_state(args.audio_vae_checkpoint)
-    decoder_state = {
-        key: tensor
-        for key, tensor in vae_state.items()
-        if key.startswith("decoder.") or key in {"data_mean", "data_std"}
-    }
-    if not decoder_state:
-        raise ValueError("Audio VAE checkpoint did not contain decoder weights")
+    if args.preprocessor_only:
+        audio_vae_state = vae_state
+        need_encoder = True
+        if not any(key.startswith("encoder.") for key in audio_vae_state):
+            raise ValueError(
+                "Audio VAE checkpoint did not contain encoder weights")
+    else:
+        audio_vae_state = {
+            key: tensor
+            for key, tensor in vae_state.items()
+            if key.startswith("decoder.") or key in {"data_mean", "data_std"}
+        }
+        need_encoder = False
+        if not audio_vae_state:
+            raise ValueError(
+                "Audio VAE checkpoint did not contain decoder weights")
     _write_component(
         output,
         "audio_vae",
-        decoder_state,
-        {"_class_name": "MMAudioVAE", "mode": "44k", "need_encoder": False},
+        audio_vae_state,
+        {
+            "_class_name": "MMAudioVAE",
+            "mode": mode,
+            "need_encoder": need_encoder
+        },
     )
 
     synchformer_state = _load_torch_state(args.synchformer_checkpoint)
@@ -304,14 +448,38 @@ def convert(args: argparse.Namespace) -> None:
     _write_component(output, "image_encoder", map_open_clip_vision_state(dfn_state), IMAGE_ENCODER_CONFIG)
     write_open_clip_tokenizer(output)
 
-    bigvgan_config_path = args.bigvgan_dir / "config.json"
-    if not bigvgan_config_path.is_file():
-        raise FileNotFoundError(bigvgan_config_path)
-    with bigvgan_config_path.open(encoding="utf-8") as handle:
-        bigvgan_config = json.load(handle)
-    bigvgan_config["_class_name"] = "BigVGANV2"
-    bigvgan_config["weight_norm_removed"] = False
-    bigvgan_state = _load_torch_state(args.bigvgan_dir / "bigvgan_generator.pt")
+    if args.preprocessor_only:
+        preprocess_components = {
+            key: value
+            for key, value in model_index.items()
+            if key.startswith("_") or key in {
+                "audio_vae",
+                "text_encoder",
+                "tokenizer",
+                "image_encoder",
+                "image_encoder_2",
+            }
+        }
+        preprocess_components["_fastvideo_preprocessor_only"] = True
+        _write_json(output / "model_index.json", preprocess_components)
+        print(f"Converted MMAudio {mode} preprocessing components to {output}")
+        return
+
+    if mode == "16k":
+        assert args.bigvgan_checkpoint is not None
+        bigvgan_config = dict(BIGVGAN_16K_CONFIG)
+        bigvgan_state = _load_torch_state(args.bigvgan_checkpoint)
+    else:
+        assert args.bigvgan_dir is not None
+        bigvgan_config_path = args.bigvgan_dir / "config.json"
+        if not bigvgan_config_path.is_file():
+            raise FileNotFoundError(bigvgan_config_path)
+        with bigvgan_config_path.open(encoding="utf-8") as handle:
+            bigvgan_config = json.load(handle)
+        bigvgan_config["_class_name"] = "BigVGANV2"
+        bigvgan_config["weight_norm_removed"] = False
+        bigvgan_state = _load_torch_state(
+            args.bigvgan_dir / "bigvgan_generator.pt")
     _write_component(output, "vocoder", bigvgan_state, bigvgan_config)
 
     _write_json(
@@ -325,17 +493,53 @@ def convert(args: argparse.Namespace) -> None:
             "use_reference_discrete_timesteps": True,
         },
     )
-    _write_json(output / "model_index.json", MODEL_INDEX)
-    print(f"Converted MMAudio large_44k_v2 components to {output}")
+    _write_json(output / "model_index.json", model_index)
+    print(f"Converted MMAudio {args.variant} components to {output}")
 
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
-    parser.add_argument("--transformer-checkpoint", type=Path, required=True)
-    parser.add_argument("--audio-vae-checkpoint", type=Path, required=True)
-    parser.add_argument("--synchformer-checkpoint", type=Path, required=True)
-    parser.add_argument("--dfn5b-dir", type=Path, required=True)
-    parser.add_argument("--bigvgan-dir", type=Path, required=True)
+    parser.add_argument(
+        "--variant",
+        choices=tuple(TRANSFORMER_VARIANTS),
+        default="large_44k_v2",
+        help="Transformer architecture. v1 variants are required for training.",
+    )
+    parser.add_argument("--transformer-checkpoint", type=Path)
+    parser.add_argument("--audio-vae-checkpoint", type=Path)
+    parser.add_argument("--synchformer-checkpoint", type=Path)
+    parser.add_argument("--dfn5b-dir", type=Path)
+    parser.add_argument(
+        "--bigvgan-dir",
+        type=Path,
+        help="NVIDIA BigVGAN-v2 directory used by 44.1 kHz variants.",
+    )
+    parser.add_argument(
+        "--bigvgan-checkpoint",
+        type=Path,
+        help="Official best_netG.pt used by the small_16k variant.",
+    )
+    parser.add_argument(
+        "--transformer-only",
+        action="store_true",
+        help="Write only the transformer component needed by training.",
+    )
+    parser.add_argument(
+        "--transformer-config-only",
+        action="store_true",
+        help=(
+            "Write a weight-free transformer component tree used as the "
+            "DCP export template for from-scratch training."
+        ),
+    )
+    parser.add_argument(
+        "--preprocessor-only",
+        action="store_true",
+        help=(
+            "Write only the VAE encoder, DFN5B text/vision encoders, "
+            "Synchformer, and tokenizer needed for offline feature extraction."
+        ),
+    )
     parser.add_argument("--output", type=Path, required=True)
     return parser.parse_args()
 

tests/local_tests/mmaudio/PORT_STATUS.md

@@ -6,8 +6,8 @@
 - workload_types: `V2A`, `T2A`
 - official_ref: `../MMAudio` at `974010a026c731054592d8f777218bd9d85a6c24`
 - first_variant: `large_44k_v2`
-- phase: `native_pipeline_complete`
-- status: `real_weight_parity_pass`
+- phase: `all_official_variants_and_scratch_training_integrated`
+- status: `large_v2_inference_parity_and_4gpu_scratch_training_pass`
 - last_updated: `2026-08-02`
 
 ## Native Components
@@ -24,8 +24,9 @@
 ## Pipeline Integration
 
 - Pipeline: `fastvideo/pipelines/basic/mmaudio/MMAudioPipeline`
-- Config: `fastvideo/configs/pipelines/mmaudio.py::MMAudioV2AConfig`
-- Preset: `mmaudio_large_44k_v2`
+- Config: five variant-specific classes in `fastvideo/configs/pipelines/mmaudio.py`
+- Presets: `mmaudio_small_16k`, `mmaudio_small_44k`,
+  `mmaudio_medium_44k`, `mmaudio_large_44k`, `mmaudio_large_44k_v2`
 - Registry: resolves both `WorkloadType.V2A` and `WorkloadType.T2A`
 - Required production components: `transformer`, `scheduler`, `text_encoder`,
   `tokenizer`, `image_encoder`, `image_encoder_2`, `audio_vae`, `vocoder`
@@ -35,7 +36,7 @@
   published training/default duration; longer and shorter inference is accepted
 - Existing T2V/I2V/T2I pipelines are not routed through these stages.
 
-The V2A preprocessing contract is identical to official MMAudio:
+The inference V2A preprocessing contract is identical to official MMAudio:
 
 1. timestamp sampling at 8 FPS for DFN5B and 25 FPS for Synchformer;
 2. DFN path: bicubic resize to `384x384`, float `[0,1]`, CLIP normalization;
@@ -50,16 +51,85 @@ The V2A preprocessing contract is identical to official MMAudio:
 - Total local artifact size: about 9.1 GB
 - Converted weights and official assets remain ignored/untracked.
 
+The v1 training smoke additionally uses a transformer-only component tree at
+`converted_weights/mmaudio/small_44k` (about 601 MB). The converter supports
+`small_16k`, `small_44k`, `medium_44k`, `large_44k`, and the inference-default
+`large_44k_v2`; `--transformer-only` avoids duplicating inference-only assets.
+
+## Training Integration
+
+- Model adapter: `fastvideo/train/models/mmaudio/MMAudioModel`
+- Shared method: shape-agnostic
+  `fastvideo/train/methods/fine_tuning/flow_matching.py`
+- Data: official-compatible TensorDict memmaps through
+  `fastvideo/dataset/mmaudio_feature_dataset.py`
+- Config: four `*_from_scratch.yaml` v1 configs plus the pretrained
+  `small_44k.yaml` fine-tuning config
+- Numerical recipe: posterior sampling, latent normalization, logit-normal
+  time, linear flow interpolation, independent video/text CFG dropout, and
+  velocity MSE match upstream `Runner.train_fn`
+- Distributed contract: real single-GPU pretrained FSDP step and real 4-GPU
+  scratch-initialized FSDP step passed; DCP save/restore passed on the pretrained path
+- Scratch contract: rank-0 latent statistics reduction/cache, fixed CLIP empty
+  string loading, random DiT initialization, flow loss/backward, AdamW,
+  official warmup/MultiStepLR composition, gradient clipping, and EMA are wired
+- Production code imports no `mmaudio.*` modules.
+
+## Native VGGSound Preprocessing
+
+- Shared orchestration: `fastvideo/workflow/preprocess/preprocess_workflow_v2a.py`
+- Dataset adapter: `fastvideo/dataset/vggsound.py`
+- MMAudio feature stage: `fastvideo/pipelines/preprocess/mmaudio/stages.py`
+- Sharded resumable cache: `fastvideo/dataset/v2a_feature_cache.py`
+- Input manifests use the official `id<TAB>label` contract; `label` is passed as
+  the full caption to the DFN5B text encoder.
+- `dataset_metadata_path` explicitly selects one manifest and `dataset_split`
+  preserves the official train-only peak normalization behavior.
+- Train, val, and test use one pipeline class and separate output roots.
+- Production preprocessing imports no modules from the reference checkout.
+- The reference training loader uses the removed
+  `torio.io.StreamingMediaDecoder`. The default FastVideo backend uses
+  `torchaudio.load` and a PyAV-hosted FFmpeg `fps` graph. An explicit `torio`
+  backend is available in a separate PyTorch 2.7 environment for exact
+  reference-media audits. Both feed the same native FastVideo feature stage;
+  inference keeps its separate timestamp sampler and remains unchanged.
+- The FFmpeg training sampler fixes fractional-FPS boundary clips that decoded
+  as 64 DFN frames but only 199 Synchformer frames with timestamp comparison.
+  Three previously failed VGGSound clips now produce exactly 64/200 frames, and
+  one passed the full FastVideo feature extraction and shard-writing path.
+- The isolated torio backend passed exact audio, CLIP-frame, and sync-frame
+  parity against the official VGGSound adapter plus a full FastVideo shard smoke.
+- Torio preprocessing now decodes with a bounded background thread pool per
+  rank and keeps two batches queued, overlapping CPU media decode with GPU
+  feature extraction without DataLoader shared-memory tensor transfer. Cached
+  IDs are filtered before decode. Four threads per rank were fastest in a
+  16-video decode check; a 4-GPU/16-video end-to-end smoke wrote 4/4 samples on
+  every rank and reduced the feature loop from 19.25 seconds on one GPU to a
+  6.65-second slowest rank on the four-GPU run.
+- Timestamp, FFmpeg, and torio caches must use separate output roots.
+
 ## Parity Evidence
 
 | Scope | Result |
 |---|---|
-| Official/FastVideo video preprocessing | exact (`clip max_abs=0`, `sync max_abs=0`) |
+| Official/FastVideo inference video preprocessing | exact (`clip max_abs=0`, `sync max_abs=0`) |
+| Training FFmpeg fps-filter boundary regression | pass (64 DFN frames, 200 Synchformer frames) |
+| Previously failed real VGGSound clips | 3/3 recovered; end-to-end feature extraction 1/1 |
+| Official/FastVideo torio training media preprocessing | exact audio/CLIP/sync tensors (`atol=0`, `rtol=0`) |
+| Torio native FastVideo feature/shard path | pass; all five feature tensors written with expected shapes |
 | Condition features, random latent, projected conditions | exact |
 | First flow prediction and 25-step final latent | exact |
 | Final 2-second V2A waveform (89,088 samples) | exact (`atol=0`, `rtol=0`) |
 | Real 10-second variable-duration V2A | pass (441,344 samples, 10.0078 s) |
 | Default FastVideo offload path | real one-step smoke pass |
+| Real `small_44k` v1 transformer forward | exact BF16 parity |
+| Real `small_44k` v1 FastVideo train step | pass (forward/backward/AdamW/grad clip) |
+| Real `small_44k` 4-GPU from-scratch train step | pass (5.87 s, one feature shard) |
+| Real DCP checkpoint and resume to next step | pass |
+| Local MMAudio parity suite | `21 passed, 1 skipped` (the skipped test is the opt-in full gate) |
+| Combined parity/dataset/training regression | `33 passed, 1 skipped` |
+| Real filtered-caption manifests | exact ID/order/caption/count parity: train 180,062; val 2,047; test 15,221 |
+| Filtered-caption tokenizer lengths | 0/197,330 exceed the 77-token contract; maximum 74 |
 | Local suite | full rerun pending; VAE FP32/BF16 random-weight and real-weight cases pass |
 | VAE component parity | state structure and FP32/BF16 random-weight decoder parity pass on GB200 (`atol=1e-6`, `rtol=1e-6`); real-weight encode/decode parity passes (`atol=1e-5`, `rtol=1e-5`) |
 | Exact-head GB200 V2A examples | two 5-second, 25-step, seed-42 generations pass; each output is mono PCM16 at 44.1 kHz with 221,184 samples |
@@ -92,7 +162,9 @@ official `large_44k_v2`, DFN5B, Synchformer, VAE, and BigVGAN assets.
 
 - Publishing the converted checkpoint and immutable source revisions.
 - Optional source-video mux/re-encode helper; the current V2A result is WAV/audio.
-- 16 kHz and small/medium variants.
+- Real-weight waveform parity for `small_16k`, `small_44k`, `medium_44k`, and
+  `large_44k`; their configs/converter/presets are implemented.
 - Sequence/tensor-parallel optimization.
-- Training integration. The official repository does not support training the
-  `_v2` variant; any training port should start from a v1 44.1 kHz checkpoint.
+- Full-dataset loss/quality convergence validation.
+- The training config intentionally keeps `data_path` empty until the train
+  feature cache has been produced.

tests/local_tests/mmaudio/README.md

@@ -37,8 +37,10 @@ only in the shell when a future gated asset requires one.
 - Duration: 8 seconds is the published training/default duration, but inference
   uses dynamic sequence lengths and accepts shorter or longer clips. As in the
   official demo, quality can fall when moving far away from 8 seconds.
-- Deferred until the base parity gate passes: 16 kHz, small/medium variants,
-  sequence/tensor parallel optimization, and quality-reference publication.
+- All five official architecture variants now have FastVideo configs, presets,
+  registry entries, and conversion support. Real-weight parity is complete for
+  `large_44k_v2`; the other four remain local follow-up gates. Sequence/tensor
+  parallel optimization and quality-reference publication are deferred.
 
 ## Shared Environment Setup
 
@@ -147,3 +149,93 @@ pytest tests/local_tests/mmaudio/test_mmaudio_pipeline_parity.py -v -s
   presented as runnable.
 - MMAudio checkpoints are documented upstream as CC-BY-NC 4.0; converted model
   publishing must preserve the applicable license and attribution.
+
+## Native VGGSound Feature Extraction
+
+FastVideo consumes the same `id<TAB>label` split manifests as the reference
+`mmaudio.data.extraction.vgg_sound.VGGSound` adapter. `label` is a full caption,
+not the original VGGSound category. The reference checkout is used only for
+local parity review and is not imported by production preprocessing.
+
+Training video frames are resampled with one PyAV-hosted FFmpeg filter graph:
+8 FPS for DFN5B and 25 FPS for Synchformer. This matches the reference
+`StreamingMediaDecoder.add_basic_video_stream(frame_rate=...)` behavior while
+remaining compatible with FastVideo's current PyTorch environment. The
+inference timestamp sampler is a separate default path and is not changed.
+
+```bash
+DATASET_PATH=/path/to/VGGSound SPLIT=train GPU_NUM=4 \
+  bash examples/training/finetune/mmaudio/preprocess_vggsound.sh
+```
+
+The launcher supports `train`, `val`, and `test` through the same FastVideo
+pipeline. The corresponding default manifests are
+`sets/filtered_caption/vgg-{split}-filtered-caption.tsv`, and caches are written
+under `mmaudio_features/{split}`. The cache is resumable: rerunning the command
+skips IDs already written and retries pending/failed IDs. Failure logs are
+append-only, so deduplicate by ID when auditing multiple runs.
+
+Caches produced with the earlier timestamp-based training sampler are valid
+TensorDict data but are not preprocessing-compatible with the FFmpeg sampler.
+Do not resume into such a cache. Rebuild into a new output root instead:
+
+```bash
+OUTPUT_ROOT=/path/to/VGGSound/mmaudio_features_ffmpeg \
+  ./examples/training/finetune/mmaudio/run_preprocess_vggsound.sh
+```
+
+### Optional official-compatible torio reader
+
+For audits that require the original training media reader, FastVideo also
+provides an opt-in backend in
+`fastvideo/pipelines/preprocess/mmaudio/torio_media_reader.py`. It reproduces
+the official `StreamingMediaDecoder` streams, audio normalization/resampling,
+and both torchvision transform chains. The rest of the operation remains the
+native FastVideo preprocessing pipeline: VAE, DFN5B, Synchformer, text encoder,
+distributed workflow, progress bar, resumable writer, and failure logging.
+Production code does not import `mmaudio.*` or execute the reference checkout.
+
+Torio was removed from current PyTorch releases, so keep its PyTorch 2.7 CUDA
+12.8 stack in a sibling environment rather than changing FastVideo's `.venv`:
+
+```text
+/path/to/workspace/FastVideo
+/path/to/workspace/envs/fastvideo-mmaudio-torio
+```
+
+The verified environment uses Python 3.12, PyTorch 2.7.0+cu128,
+torchvision/torchaudio 0.22.0/2.7.0, NumPy 1.26.4, and an editable
+`fastvideo --no-deps` install. Select it only through the dedicated launcher:
+
+```bash
+DATASET_PATH=/path/to/VGGSound \
+TORIO_VENV_PATH=/path/to/workspace/envs/fastvideo-mmaudio-torio \
+CUDA_VISIBLE_DEVICES=0,1,2,3 GPU_NUM=4 BATCH_SIZE=16 \
+DATALOADER_NUM_WORKERS=4 \
+  ./examples/training/finetune/mmaudio/run_preprocess_vggsound_torio.sh
+```
+
+By default this processes `train val test` with their existing split manifests
+and writes a separate cache under `mmaudio_features_torio/{split}`. Override
+`SPLITS`, `OUTPUT_ROOT`, or the usual batching variables exactly as with the
+standard launcher. Never resume a PyAV/FFmpeg cache into the torio output root.
+
+For torio, `DATALOADER_NUM_WORKERS` is the number of background decode threads
+per GPU rank, not DataLoader subprocesses. The default is 4 based on local
+VGGSound measurements. FastVideo keeps two batches queued, so CPU decoding of
+the next batch overlaps GPU feature extraction without sending large tensors
+through `/dev/shm`. IDs already present in the feature cache are filtered before
+decode and are not read again on resume.
+
+Exact real-media parity against the official VGGSound dataset adapter is
+covered by an opt-in local test:
+
+```bash
+PYTHONNOUSERSITE=1 \
+MMAUDIO_TORIO_TEST_VIDEO=/path/to/an-8s-or-longer-video.mp4 \
+/path/to/workspace/envs/fastvideo-mmaudio-torio/bin/pytest -q \
+  tests/local_tests/mmaudio/test_mmaudio_torio_preprocessing.py -s
+```
+
+This test imports the official dataset adapter only on the test side; it is not
+a runtime dependency of FastVideo.

tests/local_tests/mmaudio/test_mmaudio_audio_vae_parity.py

@@ -46,6 +46,18 @@ def test_mmaudio_44k_audio_vae_state_structure() -> None:
     }
 
 
+def test_mmaudio_44k_mel_converter_parity() -> None:
+    from mmaudio.ext.mel_converter import get_mel_converter
+
+    from fastvideo.models.audio.mmaudio_processing import build_mmaudio_mel_converter
+
+    waveform = torch.randn((2, 353_280), generator=torch.Generator().manual_seed(1234))
+    expected = get_mel_converter("44k")(waveform)
+    actual = build_mmaudio_mel_converter("44k")(waveform)
+
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"])
 def test_mmaudio_44k_audio_vae_decoder_implementation_parity(dtype: torch.dtype) -> None:
     if not torch.cuda.is_available():

tests/local_tests/mmaudio/test_mmaudio_torio_preprocessing.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Opt-in parity for FastVideo's isolated torio preprocessing backend."""
+
+from __future__ import annotations
+
+import os
+import sys
+import types
+from pathlib import Path
+
+import pytest
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+OFFICIAL_REPO = Path(
+    os.environ.get("MMAUDIO_OFFICIAL_REPO", REPO_ROOT.parent / "MMAudio")
+)
+
+
+@pytest.mark.skipif(
+    not os.environ.get("MMAUDIO_TORIO_TEST_VIDEO"),
+    reason="Set MMAUDIO_TORIO_TEST_VIDEO to run real media parity.",
+)
+def test_fastvideo_torio_reader_matches_official_vggsound_dataset(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    pytest.importorskip("torio.io")
+    if not OFFICIAL_REPO.is_dir():
+        pytest.skip(f"Official MMAudio checkout is missing: {OFFICIAL_REPO}")
+
+    video_path = Path(os.environ["MMAUDIO_TORIO_TEST_VIDEO"])
+    if not video_path.is_file():
+        pytest.skip(f"Test video is missing: {video_path}")
+
+    monkeypatch.syspath_prepend(str(OFFICIAL_REPO))
+    dist_utils = types.ModuleType("mmaudio.utils.dist_utils")
+    dist_utils.local_rank = 0
+    monkeypatch.setitem(sys.modules, "mmaudio.utils.dist_utils", dist_utils)
+
+    from mmaudio.data.extraction.vgg_sound import VGGSound
+
+    video_id = video_path.stem
+    videos = tmp_path / "videos"
+    videos.mkdir()
+    (videos / f"{video_id}.mp4").symlink_to(video_path.resolve())
+    manifest = tmp_path / "captions.tsv"
+    manifest.write_text(
+        f"id\tlabel\n{video_id}\tA reference preprocessing sample.\n",
+        encoding="utf-8",
+    )
+
+    official = VGGSound(
+        videos,
+        tsv_path=manifest,
+        sample_rate=44_100,
+        duration_sec=8.0,
+        audio_samples=353_280,
+        normalize_audio=True,
+    ).sample(0)
+
+    from fastvideo.pipelines.preprocess.mmaudio.torio_media_reader import (
+        preprocess_mmaudio_media_with_torio,
+    )
+
+    audio, clip, sync, duration = preprocess_mmaudio_media_with_torio(
+        video_path,
+        duration_s=8.0,
+        target_sample_rate=44_100,
+        target_samples=353_280,
+        normalize_audio=True,
+    )
+
+    torch.testing.assert_close(audio, official["audio"], atol=0, rtol=0)
+    torch.testing.assert_close(clip, official["clip_video"], atol=0, rtol=0)
+    torch.testing.assert_close(sync, official["sync_video"], atol=0, rtol=0)
+    assert duration == 8.0

tests/local_tests/mmaudio/test_mmaudio_transformer_parity.py

@@ -21,6 +21,12 @@ OFFICIAL_WEIGHTS = Path(
         REPO_ROOT.parent / "MMAudio/weights/mmaudio_large_44k_v2.pth",
     )
 )
+SMALL_44K_OFFICIAL_WEIGHTS = Path(
+    os.environ.get(
+        "MMAUDIO_SMALL_44K_TRANSFORMER_WEIGHTS",
+        REPO_ROOT.parent / "MMAudio/weights/mmaudio_small_44k.pth",
+    )
+)
 
 
 def _require_cuda_and_weights() -> None:
@@ -33,7 +39,8 @@ def _require_cuda_and_weights() -> None:
         )
 
 
-def test_mmaudio_transformer_implementation_parity() -> None:
+@pytest.mark.parametrize("v2", [False, True])
+def test_mmaudio_transformer_implementation_parity(v2: bool) -> None:
     from mmaudio.model.networks import MMAudio
 
     from fastvideo.configs.models.dits.mmaudio import (
@@ -56,7 +63,7 @@ def test_mmaudio_transformer_implementation_parity() -> None:
         "clip_seq_len": 3,
         "sync_seq_len": 8,
         "text_seq_len": 4,
-        "v2": True,
+        "v2": v2,
     }
     official = MMAudio(**kwargs)
     fastvideo = MMAudioTransformer(MMAudioTransformerConfig(arch_config=MMAudioArchConfig(**kwargs)), hf_config={})
@@ -78,6 +85,75 @@ def test_mmaudio_transformer_implementation_parity() -> None:
     torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6)
 
 
+def test_mmaudio_small_44k_v1_transformer_parity() -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("MMAudio full-transformer parity requires CUDA")
+    if not SMALL_44K_OFFICIAL_WEIGHTS.is_file():
+        pytest.skip(
+            "Official MMAudio small_44k weights are absent. Set "
+            "MMAUDIO_SMALL_44K_TRANSFORMER_WEIGHTS."
+        )
+
+    from mmaudio.model.networks import small_44k
+
+    from fastvideo.configs.models.dits.mmaudio import (
+        MMAudioArchConfig,
+        MMAudioTransformerConfig,
+    )
+    from fastvideo.models.dits.mmaudio import MMAudioTransformer
+    from fastvideo.models.loader.utils import set_default_torch_dtype
+
+    device = torch.device("cuda:0")
+    dtype = torch.bfloat16
+    state = torch.load(
+        SMALL_44K_OFFICIAL_WEIGHTS,
+        map_location="cpu",
+        weights_only=True,
+    )
+    arch = MMAudioArchConfig(
+        hidden_dim=448,
+        depth=12,
+        fused_depth=8,
+        num_heads=7,
+        v2=False,
+    )
+    with torch.device(device), set_default_torch_dtype(dtype):
+        official = small_44k()
+        fastvideo = MMAudioTransformer(
+            MMAudioTransformerConfig(arch_config=arch),
+            hf_config={},
+        )
+    official.load_weights(dict(state))
+    state.pop("latent_rot", None)
+    state.pop("clip_rot", None)
+    missing, unexpected = fastvideo.load_state_dict(state, strict=True)
+    assert missing == []
+    assert unexpected == []
+    official.eval()
+    fastvideo.eval()
+
+    generator = torch.Generator(device=device).manual_seed(5678)
+    latent = torch.randn(
+        (1, 345, 40), generator=generator, device=device, dtype=dtype
+    )
+    clip = torch.randn(
+        (1, 64, 1024), generator=generator, device=device, dtype=dtype
+    )
+    sync = torch.randn(
+        (1, 192, 768), generator=generator, device=device, dtype=dtype
+    )
+    text = torch.randn(
+        (1, 77, 1024), generator=generator, device=device, dtype=dtype
+    )
+    timestep = torch.tensor([0.375], device=device, dtype=dtype)
+
+    with torch.inference_mode(), torch.autocast("cuda", dtype=dtype):
+        expected = official(latent, clip, sync, text, timestep)
+        actual = fastvideo(latent, (clip, sync, text), timestep)
+
+    torch.testing.assert_close(actual, expected, atol=1e-3, rtol=1e-3)
+
+
 def test_mmaudio_large_44k_v2_transformer_parity() -> None:
     _require_cuda_and_weights()
     from mmaudio.model.networks import large_44k_v2