[feat]: B200 SSIM full-quality support and resource-aware local runner

Add B200 to the device-reference mappings across the SSIM suite and migrate
test_gen3c_similarity.py onto the shared reference_utils helpers (device
resolution, tiered output/reference layout, full-quality param selection,
executor shutdown in try/finally). Add run_parallel.py, a local scheduler
that mirrors the Modal static discovery and packs one-, two-, and four-GPU
SSIM jobs onto disjoint CUDA_VISIBLE_DEVICES sets, with per-task logs, JUnit
XML, and a summary.json under generated_videos/<tier>/_runs/. Verified: all
files compile and --dry-run discovers all 19 tasks with correct GPU sizing
against current main. Full-quality B200 references generated on GB200 are
uploaded to the HF reference repo separately.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

.github/workflows/ci-precommit.yml

@@ -1,11 +1,7 @@
 name: pre-commit
 
 on:
-  # pull_request_target instead of pull_request: the workflow definition and
-  # the hook config are always taken from the BASE branch, so fork /
-  # first-time-contributor PRs run immediately without a maintainer clicking
-  # "Approve and run". The PR head is checked out as data only.
-  pull_request_target:
+  pull_request:
     branches: [main]
   workflow_call:
     inputs:
@@ -19,25 +15,12 @@ permissions:
 
 jobs:
   pre-commit:
-    if: github.event.pull_request.draft != true
+    if: github.event_name == 'workflow_call' || github.event.pull_request.draft != true
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v4
       with:
         ref: ${{ inputs.ref || '' }}
-    # For PR events, lint the PR head — but keep the hook definitions from
-    # the base branch so an untrusted PR cannot alter what gets executed.
-    - name: Save trusted hook config
-      if: github.event_name == 'pull_request_target'
-      run: cp .pre-commit-config.yaml "$RUNNER_TEMP/trusted-pre-commit-config.yaml"
-    - uses: actions/checkout@v4
-      if: github.event_name == 'pull_request_target'
-      with:
-        ref: ${{ github.event.pull_request.head.sha }}
-        persist-credentials: false
-    - name: Restore trusted hook config
-      if: github.event_name == 'pull_request_target'
-      run: cp "$RUNNER_TEMP/trusted-pre-commit-config.yaml" .pre-commit-config.yaml
     - uses: actions/setup-python@v5
       with:
         python-version: "3.12"

examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml

@@ -1,78 +0,0 @@
-# Causal Consistency Distillation: Wan 2.1 T2V 1.3B Causal
-#
-# ODE-data-free distillation. A frozen teacher takes a single CFG Euler step;
-# the student matches an EMA copy of itself at the next timestep, all under
-# clean-history teacher forcing.
-#
-# All three roles initialize from the SAME checkpoint (the teacher-forcing
-# AR-diffusion model). Point init_from at that checkpoint for a real run.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-  teacher:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-  ema:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-
-method:
-  _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod
-  discrete_cd_N: 48
-  guidance_scale: 3.0
-  ema_decay: 0.99
-  ema_start_step: 200
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/Wan-Syn_77x448x832_600k
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 3000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_causal_cd
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_cd_shift5
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-
-pipeline:
-  flow_shift: 5

examples/train/configs/fine_tuning/wan/dfsft_causal_t2v_framewise.yaml

@@ -1,73 +0,0 @@
-# DFSFT (Diffusion-Forcing SFT), frame-wise: Wan 2.1 T2V 1.3B Causal
-#
-# - Student: trainable causal Wan model with a block size of 1 frame
-# - Training: each frame gets its own independent noise level (frame-wise
-#   diffusion forcing), versus the chunk-wise variant that shares one noise
-#   level across num_frames_per_block frames.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    num_frames_per_block: 1
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
-  chunk_size: 1
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/Wan-Syn_77x448x832_600k
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 4000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_causal_dfsft_framewise
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_dfsft_framewise_shift5_gauss_weight
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-  validation:
-    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
-    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
-    every_steps: 50
-    sampling_steps: [40]
-    guidance_scale: 6.0
-    num_frames: 69
-
-pipeline:
-  flow_shift: 5

examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml

@@ -1,72 +0,0 @@
-# TFSFT (Teacher-Forcing SFT): Wan 2.1 T2V 1.3B Causal
-#
-# - Student: trainable causal Wan model
-# - Training: inhomogeneous timesteps per chunk, but the causal transformer
-#   denoises the current block while attending to *clean* history (clean_x),
-#   not its own noisy rollout.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
-  chunk_size: 3
-
-training:
-  distributed:
-    num_gpus: 8
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 8
-
-  data:
-    data_path: data/Wan-Syn_77x448x832_600k
-    dataloader_num_workers: 4
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    seed: 1000
-    num_latent_t: 18
-    num_height: 448
-    num_width: 832
-    num_frames: 69
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 4000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/wan2.1_causal_tfsft
-    training_state_checkpointing_steps: 1000
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: distillation_wan_r
-    run_name: wan2.1_causal_tfsft_shift5_gauss_weight
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 1.0
-  validation:
-    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
-    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
-    every_steps: 50
-    sampling_steps: [40]
-    guidance_scale: 6.0
-    num_frames: 69
-
-pipeline:
-  flow_shift: 5

fastvideo/models/dits/causal_wanvideo.py

@@ -437,7 +437,6 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         # Causal-specific
         self.block_mask = None
-        self.teacher_forcing_block_mask = None
         self.num_frame_per_block = config.arch_config.num_frames_per_block
         assert self.num_frame_per_block <= 3
         self.independent_first_frame = False
@@ -501,70 +500,6 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         return block_mask
 
-    @staticmethod
-    def _prepare_teacher_forcing_mask(
-        device: torch.device | str, num_frames: int = 21,
-        frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1
-    ) -> BlockMask:
-        """Attention mask for the teacher-forcing ``[clean | noisy]`` sequence.
-
-        A noisy token attends to its own block plus the clean context of all
-        strictly previous blocks; clean tokens are block-wise causal.
-        """
-        if local_attn_size != -1:
-            raise NotImplementedError(
-                f"Teacher forcing ignores local_attn_size={local_attn_size}: "
-                "unlike the block-wise causal mask, this mask always attends "
-                "to the full clean context. Windowed teacher forcing is not "
-                "implemented; use local_attn_size=-1 for teacher-forcing "
-                "training.")
-        total_length = num_frames * frame_seqlen * 2
-        padded_length = math.ceil(total_length / 128) * 128 - total_length
-
-        clean_ends = num_frames * frame_seqlen
-        context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-        noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long)
-
-        attention_block_size = frame_seqlen * num_frame_per_block
-        frame_indices = torch.arange(
-            start=0, end=num_frames * frame_seqlen,
-            step=attention_block_size, device=device, dtype=torch.long
-        )
-        for start in frame_indices:
-            context_ends[start:start + attention_block_size] = start + attention_block_size
-
-        noisy_image_start_list = torch.arange(
-            num_frames * frame_seqlen, total_length,
-            step=attention_block_size, device=device, dtype=torch.long
-        )
-        noisy_image_end_list = noisy_image_start_list + attention_block_size
-        for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)):
-            noise_noise_starts[start:end] = start
-            noise_noise_ends[start:end] = end
-            noise_context_ends[start:end] = block_index * attention_block_size
-
-        def attention_mask(b, h, q_idx, kv_idx):
-            clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx])
-            c1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx])
-            c2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx])
-            noise_mask = (q_idx >= clean_ends) & (c1 | c2)
-            eye_mask = q_idx == kv_idx
-            return eye_mask | clean_mask | noise_mask
-
-        block_mask = create_block_mask(
-            attention_mask, B=None, H=None,
-            Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length,
-            _compile=False, device=device)
-
-        if not dist.is_initialized() or dist.get_rank() == 0:
-            print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames")
-            print(block_mask)
-
-        return block_mask
-
     def _forward_inference(
                 self,
                 hidden_states: torch.Tensor,
@@ -693,12 +628,9 @@ class CausalWanTransformer3DModel(BaseDiT):
                 encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
                 | None = None,
                 start_frame: int = 0,
-                clean_x: torch.Tensor | None = None,
-                aug_t: torch.Tensor | None = None,
                 **kwargs) -> torch.Tensor:
 
         orig_dtype = hidden_states.dtype
-        teacher_forcing = clean_x is not None
         if not isinstance(encoder_hidden_states, torch.Tensor):
             encoder_hidden_states = encoder_hidden_states[0]
         if isinstance(encoder_hidden_states_image,
@@ -731,26 +663,15 @@ class CausalWanTransformer3DModel(BaseDiT):
         freqs_cis = (freqs_cos,
                      freqs_sin) if freqs_cos is not None else None
 
-        if teacher_forcing:
-            if self.teacher_forcing_block_mask is None:
-                self.teacher_forcing_block_mask = self._prepare_teacher_forcing_mask(
-                    device=hidden_states.device,
-                    num_frames=num_frames,
-                    frame_seqlen=post_patch_height * post_patch_width,
-                    num_frame_per_block=self.num_frame_per_block,
-                    local_attn_size=self.local_attn_size,
-                )
-            block_mask = self.teacher_forcing_block_mask
-        else:
-            if self.block_mask is None:
-                self.block_mask = self._prepare_blockwise_causal_attn_mask(
-                    device=hidden_states.device,
-                    num_frames=num_frames,
-                    frame_seqlen=post_patch_height * post_patch_width,
-                    num_frame_per_block=self.num_frame_per_block,
-                    local_attn_size=self.local_attn_size
-                )
-            block_mask = self.block_mask
+        # Construct blockwise causal attn mask
+        if self.block_mask is None:
+            self.block_mask = self._prepare_blockwise_causal_attn_mask(
+                device=hidden_states.device,
+                num_frames=num_frames,
+                frame_seqlen=post_patch_height * post_patch_width,
+                num_frame_per_block=self.num_frame_per_block,
+                local_attn_size=self.local_attn_size
+            )
 
         hidden_states = self.patch_embedding(hidden_states)
         grid_sizes = torch.stack(
@@ -758,7 +679,6 @@ class CausalWanTransformer3DModel(BaseDiT):
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
         encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1)
-        encoder_hidden_states_text = encoder_hidden_states
 
         temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
                         timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image)
@@ -774,35 +694,18 @@ class CausalWanTransformer3DModel(BaseDiT):
 
         assert encoder_hidden_states.dtype == orig_dtype
 
-        if teacher_forcing:
-            # Tile RoPE/modulation so clean frame i and noisy frame i share a position.
-            clean_tokens = self.patch_embedding(clean_x).flatten(2).transpose(1, 2)
-            hidden_states = torch.cat([clean_tokens, hidden_states], dim=1)
-            if aug_t is None:
-                aug_t = torch.zeros_like(timestep)
-            _, timestep_proj_clean, _, _ = self.condition_embedder(
-                aug_t.flatten(), encoder_hidden_states_text, None)
-            timestep_proj_clean = timestep_proj_clean.unflatten(
-                1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape)
-            timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1)
-            freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0),
-                         torch.cat([freqs_sin, freqs_sin], dim=0))
-
         # 4. Transformer blocks
         if torch.is_grad_enabled() and self.gradient_checkpointing:
             for block in self.blocks:
                 hidden_states = self._gradient_checkpointing_func(
                     block, hidden_states, encoder_hidden_states,
                     timestep_proj, freqs_cis,
-                    block_mask=block_mask)
+                    block_mask=self.block_mask)
         else:
             for block in self.blocks:
                 hidden_states = block(hidden_states, encoder_hidden_states,
                                         timestep_proj, freqs_cis,
-                                        block_mask=block_mask)
-
-        if teacher_forcing:
-            hidden_states = hidden_states[:, hidden_states.shape[1] // 2:]
+                                        block_mask=self.block_mask)
 
         # 5. Output norm, projection & unpatchify
         temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2)

fastvideo/tests/ssim/README.md

@@ -34,6 +34,28 @@ run `pytest fastvideo/tests/ssim/ -vs --ssim-full-quality` to use the
 `*_FULL_QUALITY_PARAMS` configs (default run keeps the original shortened test
 configs).
 
+For a local machine with mixed one-, two-, and four-GPU SSIM jobs, use the
+resource-aware runner instead of plain pytest:
+
+```bash
+.venv/bin/python fastvideo/tests/ssim/run_parallel.py \
+  --ssim-full-quality \
+  --hf-home /mnt/.cache \
+  --gpu-ids 0,1,2,3
+```
+
+The runner assigns disjoint `CUDA_VISIBLE_DEVICES` sets, keeps scheduling tasks
+that fit the currently free GPUs, and continues after individual failures by
+default. Per-task logs, JUnit XML, and `summary.json` are written under
+`generated_videos/<quality-tier>/_runs/<timestamp>/`. Use `--dry-run` to inspect
+the task/GPU plan, `--fail-fast` to stop after the first failure, or repeat
+`--test-file` / `--model-id` to select a subset. Worker core dumps are disabled
+by default; pass `--allow-core-dumps` only when a native crash needs debugging.
+
+Full-quality references are device-specific. A missing reference can therefore
+make pytest fail after generation even though the generated video or latent was
+successfully preserved under `generated_videos/full_quality/`.
+
 to override the HF reference repo at test time:
 `pytest fastvideo/tests/ssim/ -vs --ssim-reference-repo <org/repo>`
 

fastvideo/tests/ssim/run_parallel.py

@@ -0,0 +1,802 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run SSIM tests locally with exclusive, mixed-size GPU allocations.
+
+The SSIM suite contains one-, two-, and four-GPU tests. Plain pytest runs
+them serially, while pytest-xdist cannot safely infer how many GPUs each test
+needs. This runner mirrors the Modal scheduler's static discovery and assigns
+disjoint ``CUDA_VISIBLE_DEVICES`` sets to concurrent pytest subprocesses.
+
+Example:
+
+.. code-block:: bash
+
+    .venv/bin/python fastvideo/tests/ssim/run_parallel.py \
+        --ssim-full-quality \
+        --hf-home /mnt/.cache \
+        --gpu-ids 0,1,2,3
+
+Non-zero pytest exits do not stop other tasks by default. This is important
+for visual-generation audits where a missing device-specific reference may
+fail after the generated artifact has already been saved.
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import json
+import os
+import re
+import resource
+import signal
+import shutil
+import subprocess
+import sys
+import time
+import xml.etree.ElementTree as ET
+from dataclasses import asdict, dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+SSIM_DIR = Path(__file__).resolve().parent
+DEFAULT_HF_HOME = Path("/mnt/.cache")
+DEFAULT_ALLOC_CONF = "expandable_segments:False"
+HF_TOKEN_ENV_KEYS = ("HF_API_KEY", "HUGGINGFACE_HUB_TOKEN", "HF_TOKEN")
+HANDLED_SIGNALS = (signal.SIGHUP, signal.SIGINT, signal.SIGTERM)
+
+
+def _disable_core_dumps() -> None:
+    """Prevent routine worker failures from writing multi-gigabyte core files."""
+    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
+
+
+def _raise_keyboard_interrupt(_signum: int, _frame: Any) -> None:
+    """Route termination signals through the normal worker cleanup path."""
+    raise KeyboardInterrupt
+
+
+def _unblock_task_signals() -> None:
+    """Reset the launcher's short-lived blocked mask before child exec."""
+    signal.pthread_sigmask(signal.SIG_UNBLOCK, HANDLED_SIGNALS)
+
+
+@dataclass(frozen=True)
+class SSIMTask:
+    task_id: int
+    test_file: Path
+    required_gpus: int
+    model_id: str | None = None
+    test_function: str | None = None
+
+    @property
+    def node_id(self) -> str:
+        relative_file = self.test_file.relative_to(REPO_ROOT).as_posix()
+        if self.test_function is None:
+            return relative_file
+        return f"{relative_file}::{self.test_function}"
+
+    @property
+    def display_name(self) -> str:
+        parts = [self.test_file.name]
+        if self.test_function is not None:
+            parts.append(self.test_function)
+        if self.model_id is not None:
+            parts.append(self.model_id)
+        return "::".join(parts)
+
+    @property
+    def sort_key(self) -> tuple[str, str, str]:
+        return (
+            self.test_file.name,
+            self.test_function or "",
+            self.model_id or "",
+        )
+
+
+@dataclass
+class RunningTask:
+    task: SSIMTask
+    process: subprocess.Popen[str]
+    gpu_ids: list[str]
+    log_path: Path
+    junit_path: Path
+    log_handle: Any
+    started_at: float
+    command: list[str]
+
+
+@dataclass
+class TaskResult:
+    task_id: int
+    test_name: str
+    node_id: str
+    model_id: str | None
+    required_gpus: int
+    gpu_ids: list[str]
+    status: str
+    returncode: int | None
+    duration_seconds: float
+    tests: int
+    passed: int
+    failures: int
+    errors: int
+    skipped: int
+    log_path: str
+    junit_path: str
+    command: list[str]
+
+
+def _assignment_names(node: ast.Assign | ast.AnnAssign) -> list[str]:
+    if isinstance(node, ast.AnnAssign):
+        if isinstance(node.target, ast.Name):
+            return [node.target.id]
+        return []
+    return [target.id for target in node.targets if isinstance(target, ast.Name)]
+
+
+def _extract_required_gpus(module_ast: ast.Module, filepath: Path) -> int:
+    for node in module_ast.body:
+        if not isinstance(node, (ast.Assign, ast.AnnAssign)):
+            continue
+        if "REQUIRED_GPUS" not in _assignment_names(node):
+            continue
+        if isinstance(node.value, ast.Constant) and isinstance(node.value.value, int):
+            required_gpus = node.value.value
+            if required_gpus < 1:
+                raise ValueError(f"{filepath}: REQUIRED_GPUS must be positive")
+            return required_gpus
+        raise ValueError(f"{filepath}: REQUIRED_GPUS must be an integer literal")
+    return 1
+
+
+def _extract_string_constants(module_ast: ast.Module) -> dict[str, str]:
+    constants: dict[str, str] = {}
+    for node in module_ast.body:
+        if not isinstance(node, (ast.Assign, ast.AnnAssign)):
+            continue
+        if not isinstance(node.value, ast.Constant) or not isinstance(node.value.value, str):
+            continue
+        for name in _assignment_names(node):
+            constants[name] = node.value.value
+    return constants
+
+
+def _extract_model_ids(module_ast: ast.Module) -> list[str]:
+    """Extract literal model IDs from all ``*_MODEL_TO_PARAMS`` maps.
+
+    Both default and full-quality maps are inspected and de-duplicated. Named
+    string constants are supported as dictionary keys (for example SD3.5's
+    ``MODEL_ID``), fixing a limitation in the current Modal discovery helper.
+    """
+    string_constants = _extract_string_constants(module_ast)
+    model_ids: list[str] = []
+    seen: set[str] = set()
+    for node in module_ast.body:
+        if not isinstance(node, (ast.Assign, ast.AnnAssign)):
+            continue
+        if not any(name.endswith("MODEL_TO_PARAMS") for name in _assignment_names(node)):
+            continue
+        if not isinstance(node.value, ast.Dict):
+            continue
+        for key in node.value.keys:
+            model_id: str | None = None
+            if isinstance(key, ast.Constant) and isinstance(key.value, str):
+                model_id = key.value
+            elif isinstance(key, ast.Name):
+                model_id = string_constants.get(key.id)
+            if model_id is not None and model_id not in seen:
+                seen.add(model_id)
+                model_ids.append(model_id)
+    return model_ids
+
+
+def _extract_test_functions(module_ast: ast.Module) -> list[str]:
+    return [
+        node.name for node in module_ast.body
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name.startswith("test_")
+    ]
+
+
+def _normalize_selected_files(values: list[str]) -> set[str]:
+    selected: set[str] = set()
+    for value in values:
+        for part in value.split(","):
+            cleaned = part.strip()
+            if cleaned:
+                selected.add(Path(cleaned).name)
+    return selected
+
+
+def _normalize_selected_models(values: list[str]) -> set[str]:
+    selected: set[str] = set()
+    for value in values:
+        selected.update(part.strip() for part in value.split(",") if part.strip())
+    return selected
+
+
+def discover_tasks(
+    *,
+    selected_test_files: set[str] | None = None,
+    selected_model_ids: set[str] | None = None,
+) -> list[SSIMTask]:
+    selected_test_files = selected_test_files or set()
+    selected_model_ids = selected_model_ids or set()
+    matched_files: set[str] = set()
+    matched_models: set[str] = set()
+    task_specs: list[tuple[Path, int, str | None, str | None]] = []
+
+    for filepath in sorted(SSIM_DIR.glob("test_*_similarity.py")):
+        if selected_test_files and filepath.name not in selected_test_files:
+            continue
+        matched_files.add(filepath.name)
+        module_ast = ast.parse(filepath.read_text(encoding="utf-8"), filename=str(filepath))
+        required_gpus = _extract_required_gpus(module_ast, filepath)
+        model_ids = _extract_model_ids(module_ast)
+
+        if model_ids:
+            for model_id in model_ids:
+                if selected_model_ids and model_id not in selected_model_ids:
+                    continue
+                matched_models.add(model_id)
+                task_specs.append((filepath, required_gpus, model_id, None))
+            continue
+
+        if selected_model_ids:
+            continue
+        test_functions = _extract_test_functions(module_ast)
+        if not test_functions:
+            task_specs.append((filepath, required_gpus, None, None))
+        else:
+            # Function-level splitting gives independent LongCat T2V/I2V/VC
+            # tasks, keeping three one-GPU workers busy instead of serializing
+            # the whole file on one GPU.
+            for test_function in test_functions:
+                task_specs.append((filepath, required_gpus, None, test_function))
+
+    missing_files = sorted(selected_test_files - matched_files)
+    if missing_files:
+        raise ValueError(f"Requested SSIM test file(s) not found: {', '.join(missing_files)}")
+    missing_models = sorted(selected_model_ids - matched_models)
+    if missing_models:
+        raise ValueError(f"Requested SSIM model ID(s) not found: {', '.join(missing_models)}")
+
+    task_specs.sort(key=lambda spec: (spec[0].name, spec[3] or "", spec[2] or ""))
+    return [
+        SSIMTask(
+            task_id=index,
+            test_file=filepath,
+            required_gpus=required_gpus,
+            model_id=model_id,
+            test_function=test_function,
+        )
+        for index, (filepath, required_gpus, model_id, test_function) in enumerate(task_specs)
+    ]
+
+
+def _safe_name(value: str) -> str:
+    return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("_")
+
+
+def _parse_gpu_ids(value: str) -> list[str]:
+    gpu_ids: list[str] = []
+    seen: set[str] = set()
+    for item in value.split(","):
+        gpu_id = item.strip()
+        if not gpu_id or gpu_id in seen:
+            continue
+        seen.add(gpu_id)
+        gpu_ids.append(gpu_id)
+    if not gpu_ids:
+        raise ValueError("--gpu-ids must contain at least one GPU ID")
+    return gpu_ids
+
+
+def _preflight_gpus(gpu_ids: list[str]) -> list[str]:
+    probe = (
+        "import json, torch\n"
+        "names=[]\n"
+        "for i in range(torch.cuda.device_count()):\n"
+        "    names.append(torch.cuda.get_device_name(i))\n"
+        "    torch.empty(1, device=f'cuda:{i}')\n"
+        "torch.cuda.synchronize() if names else None\n"
+        "print(json.dumps({'count': torch.cuda.device_count(), 'names': names}))\n"
+    )
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = ",".join(gpu_ids)
+    result = subprocess.run(
+        [sys.executable, "-c", probe],
+        cwd=REPO_ROOT,
+        env=env,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    if result.returncode != 0:
+        detail = (result.stderr or result.stdout).strip()
+        raise RuntimeError(f"CUDA preflight failed: {detail or f'exit code {result.returncode}'}")
+    try:
+        payload = json.loads(result.stdout.strip().splitlines()[-1])
+    except (IndexError, json.JSONDecodeError) as error:
+        raise RuntimeError(f"CUDA preflight returned invalid output: {result.stdout!r}") from error
+    if payload.get("count") != len(gpu_ids):
+        recovery_hint = ""
+        if payload.get("count", 0) == 0:
+            recovery_hint = (
+                " No pytest task was launched. Verify `nvidia-smi` and recreate "
+                "the container/workspace with NVIDIA GPU device grants if it "
+                "reports `Failed to initialize NVML`."
+            )
+        raise RuntimeError(
+            f"Requested GPU IDs {gpu_ids}, but PyTorch can initialize only "
+            f"{payload.get('count', 0)} device(s). stderr: {result.stderr.strip()}"
+            f"{recovery_hint}"
+        )
+    return [str(name) for name in payload.get("names", [])]
+
+
+def _read_junit_counts(junit_path: Path) -> dict[str, int]:
+    counts = {"tests": 0, "failures": 0, "errors": 0, "skipped": 0, "passed": 0}
+    if not junit_path.exists():
+        return counts
+    try:
+        root = ET.parse(junit_path).getroot()
+    except (ET.ParseError, OSError):
+        return counts
+
+    suites = [root] if root.tag == "testsuite" else list(root.findall("testsuite"))
+    for suite in suites:
+        for key in ("tests", "failures", "errors", "skipped"):
+            counts[key] += int(suite.attrib.get(key, 0))
+    counts["passed"] = max(0, counts["tests"] - counts["failures"] - counts["errors"] - counts["skipped"])
+    return counts
+
+
+def _classify_status(returncode: int, counts: dict[str, int]) -> str:
+    if returncode != 0:
+        return "failed"
+    if counts["tests"] > 0 and counts["skipped"] == counts["tests"]:
+        return "skipped"
+    return "passed"
+
+
+def _build_command(
+    task: SSIMTask,
+    *,
+    junit_path: Path,
+    full_quality: bool,
+    skip_reference_download: bool,
+    pytest_args: list[str],
+) -> list[str]:
+    command = [
+        sys.executable,
+        "-m",
+        "pytest",
+        task.node_id,
+        "-vs",
+        f"--junitxml={junit_path}",
+        "-o",
+        "junit_family=xunit2",
+    ]
+    if full_quality:
+        command.append("--ssim-full-quality")
+    if skip_reference_download:
+        command.append("--skip-ssim-reference-download")
+    command.extend(pytest_args)
+    return command
+
+
+def _launch_task(
+    task: SSIMTask,
+    *,
+    gpu_ids: list[str],
+    run_dir: Path,
+    hf_home: Path,
+    full_quality: bool,
+    skip_reference_download: bool,
+    pytest_args: list[str],
+) -> RunningTask:
+    safe_name = _safe_name(f"{task.task_id:03d}_{task.display_name}")
+    log_path = run_dir / "logs" / f"{safe_name}.log"
+    junit_path = run_dir / "junit" / f"{safe_name}.xml"
+    log_path.parent.mkdir(parents=True, exist_ok=True)
+    junit_path.parent.mkdir(parents=True, exist_ok=True)
+    command = _build_command(
+        task,
+        junit_path=junit_path,
+        full_quality=full_quality,
+        skip_reference_download=skip_reference_download,
+        pytest_args=pytest_args,
+    )
+    env = os.environ.copy()
+    env.update({
+        "CUDA_VISIBLE_DEVICES": ",".join(gpu_ids),
+        "HF_HOME": str(hf_home),
+        "PYTHONUNBUFFERED": "1",
+        "PYTORCH_CUDA_ALLOC_CONF": env.get("PYTORCH_CUDA_ALLOC_CONF", DEFAULT_ALLOC_CONF),
+        "TOKENIZERS_PARALLELISM": "false",
+    })
+    if full_quality:
+        env["FASTVIDEO_SSIM_FULL_QUALITY"] = "1"
+    else:
+        env.pop("FASTVIDEO_SSIM_FULL_QUALITY", None)
+    if task.model_id is None:
+        env.pop("FASTVIDEO_SSIM_MODEL_ID", None)
+    else:
+        env["FASTVIDEO_SSIM_MODEL_ID"] = task.model_id
+
+    log_handle = log_path.open("w", encoding="utf-8")
+    try:
+        process = subprocess.Popen(
+            command,
+            cwd=REPO_ROOT,
+            env=env,
+            stdout=log_handle,
+            stderr=subprocess.STDOUT,
+            text=True,
+            preexec_fn=_unblock_task_signals,
+            start_new_session=True,
+        )
+    except BaseException:
+        log_handle.close()
+        raise
+    return RunningTask(
+        task=task,
+        process=process,
+        gpu_ids=gpu_ids,
+        log_path=log_path,
+        junit_path=junit_path,
+        log_handle=log_handle,
+        started_at=time.monotonic(),
+        command=command,
+    )
+
+
+def _terminate_running(running: list[RunningTask], timeout_seconds: float = 30.0) -> None:
+    for item in running:
+        if item.process.poll() is None:
+            try:
+                os.killpg(item.process.pid, signal.SIGTERM)
+            except ProcessLookupError:
+                pass
+    deadline = time.monotonic() + timeout_seconds
+    while time.monotonic() < deadline and any(item.process.poll() is None for item in running):
+        time.sleep(0.25)
+    for item in running:
+        if item.process.poll() is None:
+            try:
+                os.killpg(item.process.pid, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+
+
+def _finalize_task(item: RunningTask, returncode: int) -> TaskResult:
+    item.log_handle.close()
+    counts = _read_junit_counts(item.junit_path)
+    return TaskResult(
+        task_id=item.task.task_id,
+        test_name=item.task.display_name,
+        node_id=item.task.node_id,
+        model_id=item.task.model_id,
+        required_gpus=item.task.required_gpus,
+        gpu_ids=item.gpu_ids,
+        status=_classify_status(returncode, counts),
+        returncode=returncode,
+        duration_seconds=round(time.monotonic() - item.started_at, 3),
+        tests=counts["tests"],
+        passed=counts["passed"],
+        failures=counts["failures"],
+        errors=counts["errors"],
+        skipped=counts["skipped"],
+        log_path=str(item.log_path),
+        junit_path=str(item.junit_path),
+        command=item.command,
+    )
+
+
+def run_tasks(
+    tasks: list[SSIMTask],
+    *,
+    gpu_ids: list[str],
+    run_dir: Path,
+    hf_home: Path,
+    full_quality: bool,
+    skip_reference_download: bool,
+    pytest_args: list[str],
+    fail_fast: bool,
+) -> tuple[list[TaskResult], bool]:
+    pending = list(tasks)
+    running: list[RunningTask] = []
+    results: list[TaskResult] = []
+    free_gpu_ids = list(gpu_ids)
+    gpu_order = {gpu_id: index for index, gpu_id in enumerate(gpu_ids)}
+    interrupted = False
+    previous_handlers = {
+        signum: signal.getsignal(signum) for signum in HANDLED_SIGNALS
+    }
+    for signum in HANDLED_SIGNALS:
+        signal.signal(signum, _raise_keyboard_interrupt)
+
+    try:
+        while pending or running:
+            while pending:
+                next_index = next(
+                    (index for index, task in enumerate(pending) if task.required_gpus <= len(free_gpu_ids)),
+                    None,
+                )
+                if next_index is None:
+                    break
+                task = pending.pop(next_index)
+                assigned = free_gpu_ids[:task.required_gpus]
+                del free_gpu_ids[:task.required_gpus]
+                previous_mask = signal.pthread_sigmask(
+                    signal.SIG_BLOCK,
+                    HANDLED_SIGNALS,
+                )
+                try:
+                    item = _launch_task(
+                        task,
+                        gpu_ids=assigned,
+                        run_dir=run_dir,
+                        hf_home=hf_home,
+                        full_quality=full_quality,
+                        skip_reference_download=skip_reference_download,
+                        pytest_args=pytest_args,
+                    )
+                    running.append(item)
+                finally:
+                    signal.pthread_sigmask(signal.SIG_SETMASK, previous_mask)
+                print(
+                    f"START [{task.task_id:03d}] {task.display_name} "
+                    f"on GPU(s) {','.join(assigned)} -> {item.log_path}",
+                    flush=True,
+                )
+
+            completed = [item for item in running if item.process.poll() is not None]
+            if not completed:
+                time.sleep(0.5)
+                continue
+
+            stop_after_failure = False
+            for item in completed:
+                returncode = item.process.wait()
+                result = _finalize_task(item, returncode)
+                results.append(result)
+                running.remove(item)
+                free_gpu_ids.extend(item.gpu_ids)
+                free_gpu_ids.sort(key=lambda gpu_id: gpu_order[gpu_id])
+                print(
+                    f"DONE  [{item.task.task_id:03d}] {item.task.display_name}: "
+                    f"{result.status} rc={returncode} in {result.duration_seconds:.1f}s",
+                    flush=True,
+                )
+                stop_after_failure = stop_after_failure or result.status == "failed"
+
+            if stop_after_failure and fail_fast:
+                print("Fail-fast enabled; terminating active tasks.", flush=True)
+                _terminate_running(running)
+                break
+    except KeyboardInterrupt:
+        interrupted = True
+        print("Interrupted; terminating active SSIM tasks.", flush=True)
+        for signum in HANDLED_SIGNALS:
+            signal.signal(signum, signal.SIG_IGN)
+        _terminate_running(running)
+    except BaseException:
+        print("Runner error; terminating active SSIM tasks.", flush=True)
+        for signum in HANDLED_SIGNALS:
+            signal.signal(signum, signal.SIG_IGN)
+        _terminate_running(running)
+        raise
+    finally:
+        for signum, previous_handler in previous_handlers.items():
+            signal.signal(signum, previous_handler)
+        for item in list(running):
+            returncode = item.process.wait()
+            results.append(_finalize_task(item, returncode))
+        for task in pending:
+            results.append(
+                TaskResult(
+                    task_id=task.task_id,
+                    test_name=task.display_name,
+                    node_id=task.node_id,
+                    model_id=task.model_id,
+                    required_gpus=task.required_gpus,
+                    gpu_ids=[],
+                    status="not_run",
+                    returncode=None,
+                    duration_seconds=0.0,
+                    tests=0,
+                    passed=0,
+                    failures=0,
+                    errors=0,
+                    skipped=0,
+                    log_path="",
+                    junit_path="",
+                    command=[],
+                )
+            )
+    results.sort(key=lambda result: result.task_id)
+    return results, interrupted
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--gpu-ids",
+        default=os.environ.get("CUDA_VISIBLE_DEVICES", "0,1,2,3"),
+        help="Comma-separated physical GPU IDs or UUIDs (default: 0,1,2,3).",
+    )
+    parser.add_argument(
+        "--hf-home",
+        type=Path,
+        default=DEFAULT_HF_HOME,
+        help=f"Shared Hugging Face cache root (default: {DEFAULT_HF_HOME}).",
+    )
+    parser.add_argument("--ssim-full-quality", action="store_true", help="Select full-quality SSIM parameters.")
+    parser.add_argument(
+        "--skip-reference-download",
+        action="store_true",
+        help="Pass --skip-ssim-reference-download to every pytest task.",
+    )
+    parser.add_argument("--fail-fast", action="store_true", help="Stop after the first failed task.")
+    parser.add_argument(
+        "--allow-core-dumps",
+        action="store_true",
+        help="Keep the shell's core-dump limit instead of disabling worker core dumps.",
+    )
+    parser.add_argument("--test-file", action="append", default=[], help="Run only this test file (repeatable/CSV).")
+    parser.add_argument("--model-id", action="append", default=[], help="Run only this model ID (repeatable/CSV).")
+    parser.add_argument("--pytest-arg", action="append", default=[], help="Additional argument passed to pytest.")
+    parser.add_argument("--run-dir", type=Path, help="Directory for task logs, JUnit XML, and summary.json.")
+    parser.add_argument(
+        "--keep-existing-generated",
+        action="store_true",
+        help="Do not archive existing generated artifacts before the run.",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print discovery/allocation requirements without probing GPUs or launching pytest.",
+    )
+    return parser
+
+
+def _default_run_dir(full_quality: bool) -> Path:
+    tier = "full_quality" if full_quality else "default"
+    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+    return SSIM_DIR / "generated_videos" / tier / "_runs" / timestamp
+
+
+def _archive_existing_generated(full_quality: bool, run_dir: Path) -> list[Path]:
+    """Move stale generated artifacts aside before starting a new run.
+
+    Video output helpers add numeric suffixes when a target already exists,
+    while several tests later inspect the unsuffixed path. Archiving avoids a
+    stale artifact being mistaken for the result of the current run.
+    """
+    tier = "full_quality" if full_quality else "default"
+    tier_root = SSIM_DIR / "generated_videos" / tier
+    if not tier_root.exists():
+        return []
+
+    archive_root = run_dir / "previous_generated"
+    archived: list[Path] = []
+    for path in sorted(tier_root.iterdir()):
+        if path.name == "_runs":
+            continue
+        destination = archive_root / path.name
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        shutil.move(str(path), str(destination))
+        archived.append(destination)
+    return archived
+
+
+def _write_summary(
+    path: Path,
+    *,
+    started_at: str,
+    finished_at: str,
+    hf_home: Path,
+    gpu_ids: list[str],
+    gpu_names: list[str],
+    full_quality: bool,
+    interrupted: bool,
+    results: list[TaskResult],
+) -> None:
+    payload = {
+        "started_at": started_at,
+        "finished_at": finished_at,
+        "hf_home": str(hf_home),
+        "gpu_ids": gpu_ids,
+        "gpu_names": gpu_names,
+        "full_quality": full_quality,
+        "interrupted": interrupted,
+        "results": [asdict(result) for result in results],
+    }
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _build_parser().parse_args(argv)
+    if not args.allow_core_dumps:
+        _disable_core_dumps()
+    gpu_ids = _parse_gpu_ids(args.gpu_ids)
+    selected_files = _normalize_selected_files(args.test_file)
+    selected_models = _normalize_selected_models(args.model_id)
+    tasks = discover_tasks(
+        selected_test_files=selected_files,
+        selected_model_ids=selected_models,
+    )
+    if not tasks:
+        raise RuntimeError("No SSIM tasks matched the requested filters")
+    max_required = max(task.required_gpus for task in tasks)
+    if max_required > len(gpu_ids):
+        raise RuntimeError(
+            f"A selected task requires {max_required} GPUs, but --gpu-ids provides only {len(gpu_ids)}"
+        )
+
+    print(f"Discovered {len(tasks)} SSIM task(s) across {len(gpu_ids)} GPU slot(s):", flush=True)
+    for task in tasks:
+        print(f"  [{task.task_id:03d}] {task.required_gpus} GPU(s)  {task.display_name}", flush=True)
+    if args.dry_run:
+        return 0
+
+    hf_home = args.hf_home.expanduser().resolve()
+    hf_home.mkdir(parents=True, exist_ok=True)
+    if not any(os.environ.get(key, "").strip() for key in HF_TOKEN_ENV_KEYS):
+        print(
+            "WARNING: no Hugging Face token is set; gated model tasks may fail.",
+            flush=True,
+        )
+    run_dir = (args.run_dir or _default_run_dir(args.ssim_full_quality)).expanduser().resolve()
+    gpu_names = _preflight_gpus(gpu_ids)
+    if not args.keep_existing_generated:
+        archived = _archive_existing_generated(args.ssim_full_quality, run_dir)
+        if archived:
+            print(f"Archived {len(archived)} existing generated path(s) under {run_dir}", flush=True)
+    run_dir.mkdir(parents=True, exist_ok=True)
+    print(f"CUDA preflight passed: {dict(zip(gpu_ids, gpu_names, strict=True))}", flush=True)
+    print(f"HF_HOME={hf_home}", flush=True)
+    print(f"Run artifacts: {run_dir}", flush=True)
+
+    started_at = datetime.now(timezone.utc).isoformat()
+    results, interrupted = run_tasks(
+        tasks,
+        gpu_ids=gpu_ids,
+        run_dir=run_dir,
+        hf_home=hf_home,
+        full_quality=args.ssim_full_quality,
+        skip_reference_download=args.skip_reference_download,
+        pytest_args=args.pytest_arg,
+        fail_fast=args.fail_fast,
+    )
+    finished_at = datetime.now(timezone.utc).isoformat()
+    summary_path = run_dir / "summary.json"
+    _write_summary(
+        summary_path,
+        started_at=started_at,
+        finished_at=finished_at,
+        hf_home=hf_home,
+        gpu_ids=gpu_ids,
+        gpu_names=gpu_names,
+        full_quality=args.ssim_full_quality,
+        interrupted=interrupted,
+        results=results,
+    )
+
+    status_counts: dict[str, int] = {}
+    for result in results:
+        status_counts[result.status] = status_counts.get(result.status, 0) + 1
+    print(f"Summary: {status_counts}", flush=True)
+    print(f"Detailed results: {summary_path}", flush=True)
+    if interrupted:
+        return 130
+    return 1 if any(result.status in {"failed", "not_run"} for result in results) else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

fastvideo/tests/ssim/test_gamecraft_similarity.py

@@ -47,6 +47,7 @@ device_reference_folder = resolve_device_reference_folder(
         ("A100", "A100"),
         ("H100", "H100"),
         ("H200", "H200"),
+        ("B200", "B200"),
     ),
     device_name=get_cuda_device_name(),
     unknown_device_prefix="Unknown",

fastvideo/tests/ssim/test_gen3c_similarity.py

@@ -20,10 +20,16 @@ import glob
 from pathlib import Path
 
 import pytest
-import torch
 
 from fastvideo import VideoGenerator
 from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.reference_utils import (
+    build_generated_output_dir,
+    build_reference_folder_path,
+    get_cuda_device_name,
+    resolve_device_reference_folder,
+    select_ssim_params,
+)
 from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results
 from fastvideo.worker.multiproc_executor import MultiprocExecutor
 
@@ -50,16 +56,16 @@ def _resolve_gen3c_test_image_path() -> str:
 # Device detection
 # ---------------------------------------------------------------------------
 
-device_name = torch.cuda.get_device_name() if torch.cuda.is_available() else "cpu"
-device_reference_folder_suffix = "_reference_videos"
-
-if "A40" in device_name:
-    device_reference_folder = "A40" + device_reference_folder_suffix
-elif "L40S" in device_name:
-    device_reference_folder = "L40S" + device_reference_folder_suffix
-else:
-    device_reference_folder = None
-    logger.warning(f"Unsupported device for GEN3C SSIM tests: {device_name}")
+device_name = get_cuda_device_name()
+device_reference_folder = resolve_device_reference_folder(
+    (
+        ("A40", "A40"),
+        ("L40S", "L40S"),
+        ("B200", "B200"),
+    ),
+    device_name=device_name,
+    logger=logger,
+)
 
 # ---------------------------------------------------------------------------
 # GEN3C generation parameters
@@ -86,9 +92,16 @@ GEN3C_T2V_PARAMS = {
 MODEL_TO_PARAMS = {
     "GEN3C-Cosmos-7B": GEN3C_T2V_PARAMS,
 }
+FULL_QUALITY_MODEL_TO_PARAMS = {
+    # GEN3C has one production recipe today. Keeping an explicit map makes the
+    # full-quality tier selection and reference/output layout consistent with
+    # the rest of the SSIM suite without claiming a different recipe.
+    "GEN3C-Cosmos-7B": dict(GEN3C_T2V_PARAMS),
+}
 
 TEST_PROMPTS = [
-    "A camera slowly orbits around a young woman sitting at a table with a coffee mug with coffee in it in front of her, "
+    "A camera slowly orbits around a young woman sitting at a table with a "
+    "coffee mug with coffee in it in front of her, "
     "looking away naturally. Soft indoor lighting, cinematic framing, shallow depth of field, smooth camera motion.",
 ]
 
@@ -115,12 +128,17 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND
 
     script_dir = os.path.dirname(os.path.abspath(__file__))
-    base_output_dir = os.path.join(script_dir, "generated_videos", model_id)
-    output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND)
+    output_dir = build_generated_output_dir(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
+    )
     output_video_name = CANDIDATE_VIDEO_NAME
     os.makedirs(output_dir, exist_ok=True)
 
-    BASE_PARAMS = MODEL_TO_PARAMS[model_id]
+    params_map = select_ssim_params(MODEL_TO_PARAMS, FULL_QUALITY_MODEL_TO_PARAMS)
+    BASE_PARAMS = params_map[model_id]
     num_inference_steps = BASE_PARAMS["num_inference_steps"]
     model_path = BASE_PARAMS["model_path"]
 
@@ -178,18 +196,23 @@ def test_gen3c_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
     for stale_video in glob.glob(stale_pattern):
         os.remove(stale_video)
 
-    generator = VideoGenerator.from_pretrained(
-        model_path=model_path, **init_kwargs
-    )
-    generator.generate_video(prompt, **generation_kwargs)
-
-    if isinstance(generator.executor, MultiprocExecutor):
-        generator.executor.shutdown()
+    generator = None
+    try:
+        generator = VideoGenerator.from_pretrained(
+            model_path=model_path, **init_kwargs
+        )
+        generator.generate_video(prompt, **generation_kwargs)
+    finally:
+        if generator is not None and isinstance(generator.executor, MultiprocExecutor):
+            generator.executor.shutdown()
 
     assert os.path.exists(output_dir), f"Output not generated at {output_dir}"
 
-    reference_folder = os.path.join(
-        script_dir, device_reference_folder, model_id, ATTENTION_BACKEND
+    reference_folder = build_reference_folder_path(
+        script_dir,
+        device_reference_folder,
+        model_id,
+        ATTENTION_BACKEND,
     )
     if not os.path.exists(reference_folder):
         raise FileNotFoundError(

fastvideo/tests/ssim/test_lingbot_similarity.py

@@ -55,6 +55,7 @@ device_reference_folder = resolve_device_reference_folder(
         ("L40S", "L40S"),
         ("H100", "H100"),
         ("H200", "H200"),
+        ("B200", "B200"),
     ),
     device_name=device_name,
     logger=logger,

fastvideo/tests/ssim/test_longcat_similarity.py

@@ -42,6 +42,7 @@ device_reference_folder = resolve_device_reference_folder(
         ("L40S", "L40S"),
         ("H100", "H100"),
         ("H200", "H200"),
+        ("B200", "B200"),
     ),
     device_name=get_cuda_device_name(),
     logger=logger,

fastvideo/tests/ssim/test_matrixgame2_similarity.py

@@ -30,6 +30,7 @@ device_reference_folder = resolve_device_reference_folder(
         ("A40", "A40"),
         ("L40S", "L40S"),
         ("H200", "H200"),
+        ("B200", "B200"),
     ),
     device_name=get_cuda_device_name(),
     logger=logger,

fastvideo/tests/ssim/test_matrixgame3_similarity.py

@@ -28,6 +28,7 @@ device_reference_folder = resolve_device_reference_folder(
         ("A40", "A40"),
         ("L40S", "L40S"),
         ("H200", "H200"),
+        ("B200", "B200"),
     ),
     device_name=get_cuda_device_name(),
     fallback_device_prefix="L40S",

fastvideo/tests/ssim/test_sd35_similarity.py

@@ -31,6 +31,7 @@ device_reference_folder = resolve_device_reference_folder(
         ("L40S", "L40S"),
         ("H100", "H100"),
         ("H200", "H200"),
+        ("B200", "B200"),
         ("RTX 4090", "RTX4090"),
         ("4090", "RTX4090"),
     ),

fastvideo/tests/train/fixtures/wan_causal_t2v_causal_cd_min.yaml

@@ -1,56 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Minimum config to run a single training step of
-# CausalConsistencyDistillationMethod on WanCausalModel for the
-# per-method smoke test. Uses the real Wan 2.1 1.3B checkpoint for both
-# the trainable student and the frozen teacher (AR Euler-step target),
-# with tiny synthetic latents.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-  teacher:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-  ema:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: false
-
-method:
-  _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod
-  discrete_cd_N: 12
-  guidance_scale: 3.0
-  ema_decay: 0.95
-
-training:
-  dit_precision: bf16
-
-  distributed:
-    num_gpus: 1
-    sp_size: 1
-    tp_size: 1
-
-  data:
-    seed: 42
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    num_latent_t: 6
-    num_height: 64
-    num_width: 64
-    num_frames: 21
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 1
-    gradient_accumulation_steps: 1
-
-pipeline: {}

fastvideo/tests/train/fixtures/wan_causal_t2v_dfsft_framewise_min.yaml

@@ -1,46 +0,0 @@
-# Minimum config to run a single training step of
-# DiffusionForcingSFTMethod on a frame-wise WanCausalModel
-# (num_frames_per_block=1, chunk_size=1) for the per-method smoke
-# test.  Uses the real Wan 2.1 1.3B checkpoint with tiny synthetic
-# latents.
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    num_frames_per_block: 1
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
-  chunk_size: 1
-
-training:
-  dit_precision: bf16
-
-  distributed:
-    num_gpus: 1
-    sp_size: 1
-    tp_size: 1
-
-  data:
-    seed: 42
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    num_latent_t: 6
-    num_height: 64
-    num_width: 64
-    num_frames: 21
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 1
-    gradient_accumulation_steps: 1
-
-pipeline: {}

fastvideo/tests/train/fixtures/wan_causal_t2v_tfsft_min.yaml

@@ -1,45 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# Minimum config to run a single training step of
-# TeacherForcingSFTMethod on WanCausalModel for the per-method smoke
-# test. Identical to wan_causal_t2v_dfsft_min.yaml except the method,
-# which feeds clean history to the causal transformer (teacher forcing).
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.WanCausalModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-
-method:
-  _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod
-  chunk_size: 3
-
-training:
-  dit_precision: bf16
-
-  distributed:
-    num_gpus: 1
-    sp_size: 1
-    tp_size: 1
-
-  data:
-    seed: 42
-    train_batch_size: 1
-    training_cfg_rate: 0.0
-    num_latent_t: 6
-    num_height: 64
-    num_width: 64
-    num_frames: 21
-
-  optimizer:
-    learning_rate: 2.0e-6
-    betas: [0.9, 0.999]
-    weight_decay: 0.01
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 1
-    gradient_accumulation_steps: 1
-
-pipeline: {}

fastvideo/tests/train/methods/test_wan_causal_cd.py

@@ -1,147 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Per-method GPU smoke test: ``WanCausalModel`` + ``CausalConsistencyDistillationMethod``.
-
-Mirrors ``test_wan_causal_dfsft.py``. Causal consistency distillation
-bootstraps a consistency MSE between the student's ``x0`` at ``t`` and an EMA
-copy of the student at ``t_next``, where ``t_next`` is produced online by a
-single CFG Euler step of a frozen teacher (all under clean-history teacher
-forcing). This test exercises the full step: finite loss, nonzero student
-gradients, frozen teacher, and a post-step EMA update.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29519")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.train.methods.consistency_model.causal_cd import (
-    CausalConsistencyDistillationMethod, )
-from fastvideo.train.models.wan import WanCausalModel
-from fastvideo.train.utils.config import load_run_config
-
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures"
-    / "wan_causal_t2v_causal_cd_min.yaml")
-
-
-def _build_synthetic_batch(
-    device: torch.device,
-    dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    batch_size = 1
-    return {
-        "text_embedding":
-        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
-        "text_attention_mask":
-        torch.ones(batch_size, 16, device=device),
-        "vae_latent":
-        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
-    }
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_wan_causal_cd_single_train_step(
-        monkeypatch: pytest.MonkeyPatch) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    cfg = load_run_config(_FIXTURE)
-
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-
-    monkeypatch.setattr(
-        "fastvideo.train.utils.dataloader."
-        "build_parquet_t2v_train_dataloader",
-        lambda *args, **kwargs: None,
-    )
-
-    student = WanCausalModel(
-        init_from=cfg.models["student"]["init_from"],
-        training_config=cfg.training,
-        trainable=True,
-    )
-    student.transformer = student.transformer.to(device=device, dtype=dtype)
-
-    teacher = WanCausalModel(
-        init_from=cfg.models["teacher"]["init_from"],
-        training_config=cfg.training,
-        trainable=False,
-    )
-    teacher.transformer = teacher.transformer.to(device=device, dtype=dtype)
-
-    ema = WanCausalModel(
-        init_from=cfg.models["ema"]["init_from"],
-        training_config=cfg.training,
-        trainable=False,
-    )
-    ema.transformer = ema.transformer.to(device=device, dtype=dtype)
-
-    method = CausalConsistencyDistillationMethod(
-        cfg=cfg,
-        role_models={"student": student, "teacher": teacher, "ema": ema},
-    )
-    method.on_train_start()
-
-    batch = _build_synthetic_batch(device, dtype)
-    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
-
-    loss = loss_map["total_loss"]
-    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
-    assert torch.isfinite(loss).item(), (
-        f"total_loss is not finite: {loss.item()}")
-
-    method.backward(loss_map, outputs, grad_accum_rounds=1)
-
-    blocks = student.transformer.blocks
-    assert blocks is not None and len(blocks) > 0
-    layer0 = blocks[0]
-    trainable = [p for p in layer0.parameters() if p.requires_grad]
-    assert len(trainable) > 0, "student layer 0 has no trainable parameters"
-    for i, p in enumerate(trainable):
-        assert p.grad is not None, f"student layer 0 param[{i}] has None grad"
-        assert torch.isfinite(p.grad).all().item(), (
-            f"student layer 0 param[{i}] grad contains NaN/Inf")
-    assert any(p.grad.detach().float().norm().item() > 0.0 for p in trainable), (
-        "all student layer-0 grads are exactly zero; consistency loss "
-        "did not reach the first transformer block")
-
-    # Teacher must stay frozen.
-    assert all(not p.requires_grad for p in teacher.transformer.parameters()), (
-        "teacher must be frozen for Causal-CD")
-
-    # The EMA model and student start from the same checkpoint, so the first
-    # parameter must match before any update. FSDP fully_shard params are
-    # DTensors; compare the local shards (torch.equal is unsupported on
-    # DTensor).
-    def _local(p: torch.Tensor) -> torch.Tensor:
-        return p.to_local() if hasattr(p, "to_local") else p
-
-    ema_param = next(ema.transformer.parameters())
-    student_param = next(student.transformer.parameters())
-    assert torch.equal(_local(ema_param), _local(student_param)), (
-        "EMA model should start identical to the student (same checkpoint)")
-
-    # The EMA update must move EMA toward the student. Apply a visibly large
-    # perturbation so the bf16 lerp is well above rounding noise (the real
-    # optimizer step at lr=2e-6 would be sub-ULP in bf16).
-    with torch.no_grad():
-        student_param.add_(1.0)
-    before = _local(ema_param).detach().float().clone()
-    method._update_ema()
-    after = _local(ema_param).detach().float()
-    assert not torch.equal(before, after), (
-        "EMA weights did not move after _update_ema")
-    # EMA = decay*ema + (1-decay)*student moves ~ (1-decay) of the gap.
-    expected = before + (1.0 - method._ema_decay) * (
-        _local(student_param).detach().float() - before)
-    assert torch.allclose(after, expected, atol=1e-2), (
-        "EMA update did not follow the expected lerp")

fastvideo/tests/train/methods/test_wan_causal_dfsft_framewise.py

@@ -1,109 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Per-method GPU smoke test: frame-wise ``WanCausalModel`` + ``DiffusionForcingSFTMethod``.
-
-Mirrors ``test_wan_causal_dfsft.py`` but with a block size of 1 frame
-(``num_frames_per_block=1`` on the model, ``chunk_size=1`` on the method),
-so each frame gets its own independent noise level. The test asserts the
-override took effect and runs one train step: forward, finite loss, and
-nonzero gradients reaching the first transformer block.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29520")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.train.methods.fine_tuning.dfsft import (
-    DiffusionForcingSFTMethod, )
-from fastvideo.train.models.wan import WanCausalModel
-from fastvideo.train.utils.config import load_run_config
-
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures"
-    / "wan_causal_t2v_dfsft_framewise_min.yaml")
-
-
-def _build_synthetic_batch(
-    device: torch.device,
-    dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    batch_size = 1
-    return {
-        "text_embedding":
-        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
-        "text_attention_mask":
-        torch.ones(batch_size, 16, device=device),
-        "vae_latent":
-        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
-    }
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_wan_causal_dfsft_framewise_single_train_step(
-        monkeypatch: pytest.MonkeyPatch) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    cfg = load_run_config(_FIXTURE)
-
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-
-    monkeypatch.setattr(
-        "fastvideo.train.utils.dataloader."
-        "build_parquet_t2v_train_dataloader",
-        lambda *args, **kwargs: None,
-    )
-
-    student_cfg = cfg.models["student"]
-    model = WanCausalModel(
-        init_from=student_cfg["init_from"],
-        training_config=cfg.training,
-        trainable=True,
-        num_frames_per_block=student_cfg.get("num_frames_per_block"),
-    )
-    assert model.transformer.num_frame_per_block == 1, (
-        "frame-wise override did not reach the transformer")
-    model.transformer = model.transformer.to(device=device, dtype=dtype)
-
-    method = DiffusionForcingSFTMethod(
-        cfg=cfg,
-        role_models={"student": model},
-    )
-    method.on_train_start()
-
-    batch = _build_synthetic_batch(device, dtype)
-    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
-
-    loss = loss_map["total_loss"]
-    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
-    assert torch.isfinite(loss).item(), (
-        f"total_loss is not finite: {loss.item()}")
-
-    method.backward(loss_map, outputs, grad_accum_rounds=1)
-
-    blocks = getattr(model.transformer, "blocks", None)
-    assert blocks is not None and len(blocks) > 0
-    layer0 = blocks[0]
-
-    trainable = [p for p in layer0.parameters() if p.requires_grad]
-    assert len(trainable) > 0, "layer 0 has no trainable parameters"
-
-    for i, p in enumerate(trainable):
-        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
-        assert torch.isfinite(p.grad).all().item(), (
-            f"layer 0 param[{i}] grad contains NaN/Inf")
-
-    any_nonzero = any(
-        p.grad.detach().float().norm().item() > 0.0 for p in trainable)
-    assert any_nonzero, (
-        "all layer-0 grads are exactly zero; backward did not "
-        "reach the first transformer block")

fastvideo/tests/train/methods/test_wan_causal_tfsft.py

@@ -1,113 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Per-method GPU smoke test: ``WanCausalModel`` + ``TeacherForcingSFTMethod``.
-
-Mirrors ``test_wan_causal_dfsft.py``. Teacher forcing concatenates a clean
-context copy of every frame inside the causal transformer (``clean_x``) and
-denoises the current block while attending to *clean* history. This test
-exercises the ``clean_x`` path end-to-end: forward, finite loss, and nonzero
-gradients reaching the first transformer block.
-"""
-
-from __future__ import annotations
-
-import os
-
-os.environ.setdefault("MASTER_ADDR", "localhost")
-os.environ.setdefault("MASTER_PORT", "29518")
-
-from pathlib import Path
-
-import pytest
-import torch
-
-from fastvideo.train.methods.fine_tuning.tfsft import (
-    TeacherForcingSFTMethod, )
-from fastvideo.train.models.wan import WanCausalModel
-from fastvideo.train.utils.config import load_run_config
-
-
-_FIXTURE = str(
-    Path(__file__).resolve().parent.parent / "fixtures"
-    / "wan_causal_t2v_tfsft_min.yaml")
-
-
-def _build_synthetic_batch(
-    device: torch.device,
-    dtype: torch.dtype,
-) -> dict[str, torch.Tensor]:
-    batch_size = 1
-    return {
-        "text_embedding":
-        torch.randn(batch_size, 16, 4096, device=device, dtype=dtype),
-        "text_attention_mask":
-        torch.ones(batch_size, 16, device=device),
-        "vae_latent":
-        torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype),
-    }
-
-
-@pytest.mark.usefixtures("distributed_setup")
-def test_wan_causal_tfsft_single_train_step(
-        monkeypatch: pytest.MonkeyPatch) -> None:
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-
-    cfg = load_run_config(_FIXTURE)
-
-    device = torch.device("cuda:0")
-    dtype = torch.bfloat16
-
-    monkeypatch.setattr(
-        "fastvideo.train.utils.dataloader."
-        "build_parquet_t2v_train_dataloader",
-        lambda *args, **kwargs: None,
-    )
-
-    model = WanCausalModel(
-        init_from=cfg.models["student"]["init_from"],
-        training_config=cfg.training,
-        trainable=True,
-    )
-    model.transformer = model.transformer.to(device=device, dtype=dtype)
-
-    method = TeacherForcingSFTMethod(
-        cfg=cfg,
-        role_models={"student": model},
-    )
-    method.on_train_start()
-
-    batch = _build_synthetic_batch(device, dtype)
-    loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0)
-
-    loss = loss_map["total_loss"]
-    assert torch.is_tensor(loss), "total_loss must be a torch.Tensor"
-    assert torch.isfinite(loss).item(), (
-        f"total_loss is not finite: {loss.item()}")
-
-    method.backward(loss_map, outputs, grad_accum_rounds=1)
-
-    blocks = getattr(model.transformer, "blocks", None)
-    assert blocks is not None and len(blocks) > 0, (
-        "CausalWanTransformer is expected to expose ``.blocks``")
-    layer0 = blocks[0]
-
-    trainable = [p for p in layer0.parameters() if p.requires_grad]
-    assert len(trainable) > 0, "layer 0 has no trainable parameters"
-
-    for i, p in enumerate(trainable):
-        assert p.grad is not None, f"layer 0 param[{i}] has None grad"
-        assert torch.isfinite(p.grad).all().item(), (
-            f"layer 0 param[{i}] grad contains NaN/Inf")
-
-    any_nonzero = any(
-        p.grad.detach().float().norm().item() > 0.0 for p in trainable)
-    assert any_nonzero, (
-        "all layer-0 grads are exactly zero; backward did not "
-        "reach the first transformer block")
-
-    # Teacher forcing must build its own (concatenated) attention mask and
-    # must not have constructed the diffusion-forcing mask.
-    assert model.transformer.teacher_forcing_block_mask is not None, (
-        "teacher-forcing mask was not constructed")
-    assert model.transformer.block_mask is None, (
-        "diffusion-forcing mask should not be built on the TF path")

fastvideo/train/callbacks/validation.py

@@ -456,16 +456,11 @@ class ValidationCallback(Callback):
             None,
         )
 
-        loaded_modules: dict[str, Any] = {"transformer": transformer}
-        # Distillation methods build the flow-match scheduler their few-step DMD
-        # sampler needs; inject it so the pipeline doesn't fall back to UniPC.
-        method_scheduler = getattr(self.method, "_sf_scheduler", None)
-        if method_scheduler is not None:
-            loaded_modules["scheduler"] = method_scheduler
-
         kwargs: dict[str, Any] = {
             "inference_mode": True,
-            "loaded_modules": loaded_modules,
+            "loaded_modules": {
+                "transformer": transformer,
+            },
             "tp_size": tc.distributed.tp_size,
             "sp_size": tc.distributed.sp_size,
             "num_gpus": tc.distributed.num_gpus,
@@ -482,6 +477,12 @@ class ValidationCallback(Callback):
             **kwargs,
         )
 
+        scheduler = self._pipeline.get_module("scheduler")
+        if (scheduler is not None and type(scheduler).__name__ == "SelfForcingFlowMatchScheduler"):
+            scheduler.sigma_min = 0.0
+            scheduler.extra_one_step = True
+            scheduler.set_timesteps(num_inference_steps=1000, training=True)
+
         self._pipeline_key = key
         return self._pipeline
 

fastvideo/train/methods/__init__.py

@@ -9,8 +9,6 @@ __all__ = [
     "KDMethod",
     "SelfForcingMethod",
     "DiffusionForcingSFTMethod",
-    "TeacherForcingSFTMethod",
-    "CausalConsistencyDistillationMethod",
 ]
 
 
@@ -30,11 +28,4 @@ def __getattr__(name: str) -> object:
     if name == "DiffusionForcingSFTMethod":
         from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod
         return DiffusionForcingSFTMethod
-    if name == "TeacherForcingSFTMethod":
-        from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod
-        return TeacherForcingSFTMethod
-    if name == "CausalConsistencyDistillationMethod":
-        from fastvideo.train.methods.consistency_model.causal_cd import (
-            CausalConsistencyDistillationMethod, )
-        return CausalConsistencyDistillationMethod
     raise AttributeError(name)

fastvideo/train/methods/consistency_model/__init__.py

@@ -1,22 +1,3 @@
 # SPDX-License-Identifier: Apache-2.0
 
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from fastvideo.train.methods.consistency_model.causal_cd import (
-        CausalConsistencyDistillationMethod, )
-
-__all__ = [
-    "CausalConsistencyDistillationMethod",
-]
-
-
-def __getattr__(name: str) -> object:
-    if name == "CausalConsistencyDistillationMethod":
-        from fastvideo.train.methods.consistency_model.causal_cd import (
-            CausalConsistencyDistillationMethod, )
-
-        return CausalConsistencyDistillationMethod
-    raise AttributeError(name)
+__all__: list[str] = []

fastvideo/train/methods/consistency_model/causal_cd.py

@@ -1,237 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Causal consistency distillation method (algorithm layer)."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-import torch.nn.functional as F
-
-from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (SelfForcingFlowMatchScheduler)
-from fastvideo.train.methods.base import LogScalar, TrainingMethod
-from fastvideo.train.models.base import ModelBase
-from fastvideo.train.utils.checkpoint import _FullModelState
-from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler
-
-
-class CausalConsistencyDistillationMethod(TrainingMethod):
-
-    def __init__(
-        self,
-        *,
-        cfg: Any,
-        role_models: dict[str, ModelBase],
-    ) -> None:
-        super().__init__(cfg=cfg, role_models=role_models)
-
-        for role in ("student", "teacher", "ema"):
-            if role not in role_models:
-                raise ValueError(f"Causal-CD requires role {role!r} "
-                                 "(student trainable; teacher + ema frozen, "
-                                 "both initialized from the student's "
-                                 "checkpoint)")
-        if not self.student._trainable:
-            raise ValueError("Causal-CD requires student to be trainable")
-        self.teacher = role_models["teacher"]
-        self.ema_model = role_models["ema"]
-
-        self._attn_kind = self._infer_attn_kind()
-        self._guidance_scale = float(self.method_config.get("guidance_scale", 3.0))
-        self._discrete_cd_n = int(self.method_config.get("discrete_cd_N", 48))
-        if self._discrete_cd_n < 2:
-            raise ValueError("method.discrete_cd_N must be >= 2")
-        self._ema_decay = float(self.method_config.get("ema_decay", 0.99))
-        self._ema_start_step = int(self.method_config.get("ema_start_step", 200))
-        shift = getattr(self.training_config.pipeline_config, "flow_shift", None)
-        self._flow_shift = float(shift) if shift else 5.0
-
-        self.student.init_preprocessors(self.training_config)
-        self._sf_scheduler = SelfForcingFlowMatchScheduler(
-            num_inference_steps=self._discrete_cd_n,
-            num_train_timesteps=int(self.student.num_train_timesteps),
-            shift=self._flow_shift,
-            sigma_min=0.0,
-            sigma_max=1.0,
-            extra_one_step=True,
-            training=False,
-        )
-        self._init_optimizers_and_schedulers()
-
-    # ------------------------------------------------------------------
-
-    @property
-    def _optimizer_dict(self) -> dict[str, Any]:
-        return {"student": self._student_optimizer}
-
-    @property
-    def _lr_scheduler_dict(self) -> dict[str, Any]:
-        return {"student": self._student_lr_scheduler}
-
-    def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]:
-        del iteration
-        return [self._student_optimizer]
-
-    def get_lr_schedulers(self, iteration: int) -> list[Any]:
-        del iteration
-        return [self._student_lr_scheduler]
-
-    def checkpoint_state(self) -> dict[str, Any]:
-        # The EMA role is frozen (so the base class skips it) but mutated by
-        # _update_ema every step; without persisting it a resume reloads the
-        # EMA from init_from and the consistency target snaps back to the
-        # base checkpoint. Mirrors DiffusionNFT's frozen "old" role.
-        states = super().checkpoint_state()
-        states["roles.ema.transformer"] = _FullModelState(self.ema_model.transformer)
-        return states
-
-    # ------------------------------------------------------------------
-
-    def single_train_step(
-        self,
-        batch: dict[str, Any],
-        iteration: int,
-    ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]:
-        del iteration
-        training_batch = self.student.prepare_batch(
-            batch,
-            generator=self.cuda_generator,
-            latents_source="data",
-        )
-        clean_latents = training_batch.latents
-        if not torch.is_tensor(clean_latents) or clean_latents.ndim != 5:
-            raise ValueError("Causal-CD expects [B, T, C, H, W] latents")
-
-        batch_size, num_latents = int(clean_latents.shape[0]), int(clean_latents.shape[1])
-        device = clean_latents.device
-
-        sigmas = self._sf_scheduler.sigmas.to(device)
-        timesteps = self._sf_scheduler.timesteps.to(device)
-        idx = torch.randint(0, self._discrete_cd_n - 1, (1, ), generator=self.cuda_generator, device=device).squeeze(0)
-        t, t_next = timesteps[idx], timesteps[idx + 1]
-        sigma_t, sigma_t_next = sigmas[idx], sigmas[idx + 1]
-        t_pf = t * torch.ones(batch_size, num_latents, device=device)
-        t_next_pf = t_next * torch.ones(batch_size, num_latents, device=device)
-
-        noise = torch.randn(
-            clean_latents.shape,
-            generator=self.cuda_generator,
-            device=device,
-            dtype=clean_latents.dtype,
-        )
-        latent_t = (1.0 - sigma_t) * clean_latents + sigma_t * noise
-
-        # Set before any forward: predict_noise feeds batch.timesteps into
-        # set_forward_context (VSA sparsity gating), so the teacher CFG
-        # passes below must not see the stale timesteps from prepare_batch.
-        training_batch.timesteps = t_pf
-
-        with torch.no_grad():
-            v_cond = self._predict_flow(self.teacher,
-                                        latent_t,
-                                        t_pf,
-                                        training_batch,
-                                        conditional=True,
-                                        clean_x=clean_latents)
-            v_uncond = self._predict_flow(self.teacher,
-                                          latent_t,
-                                          t_pf,
-                                          training_batch,
-                                          conditional=False,
-                                          clean_x=clean_latents)
-            v_pred = v_uncond + self._guidance_scale * (v_cond - v_uncond)
-            dt = ((t - t_next) / float(self.student.num_train_timesteps))
-            latent_t_next = latent_t - dt * v_pred
-
-        flow_student = self._predict_flow(self.student,
-                                          latent_t,
-                                          t_pf,
-                                          training_batch,
-                                          conditional=True,
-                                          clean_x=clean_latents)
-        x0_t = latent_t - sigma_t * flow_student
-
-        with torch.no_grad():
-            flow_ema = self._predict_flow(self.ema_model,
-                                          latent_t_next,
-                                          t_next_pf,
-                                          training_batch,
-                                          conditional=True,
-                                          clean_x=clean_latents)
-            x0_t_next = latent_t_next - sigma_t_next * flow_ema
-
-        loss = F.mse_loss(x0_t.float(), x0_t_next.float())
-
-        loss_map = {"total_loss": loss, "causal_cd_loss": loss}
-        attn_metadata = (training_batch.attn_metadata_vsa if self._attn_kind == "vsa" else training_batch.attn_metadata)
-        outputs: dict[str, Any] = {"_fv_backward": (t_pf, attn_metadata)}
-        metrics: dict[str, LogScalar] = {}
-        return loss_map, outputs, metrics
-
-    # ------------------------------------------------------------------
-
-    def backward(
-        self,
-        loss_map: dict[str, torch.Tensor],
-        outputs: dict[str, Any],
-        *,
-        grad_accum_rounds: int = 1,
-    ) -> None:
-        grad_accum_rounds = max(1, int(grad_accum_rounds))
-        ctx = outputs.get("_fv_backward")
-        if ctx is None:
-            super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
-            return
-        self.student.backward(loss_map["total_loss"], ctx, grad_accum_rounds=grad_accum_rounds)
-
-    def optimizers_schedulers_step(self, iteration: int) -> None:
-        super().optimizers_schedulers_step(iteration)
-        if iteration >= self._ema_start_step:
-            self._update_ema()
-
-    # ------------------------------------------------------------------
-    # Internal helpers
-    # ------------------------------------------------------------------
-
-    def _predict_flow(
-        self,
-        model: ModelBase,
-        latents: torch.Tensor,
-        timestep: torch.Tensor,
-        batch: Any,
-        *,
-        conditional: bool,
-        clean_x: torch.Tensor,
-    ) -> torch.Tensor:
-        return model.predict_noise(latents,
-                                   timestep,
-                                   batch,
-                                   conditional=conditional,
-                                   cfg_uncond=None,
-                                   attn_kind=self._attn_kind,
-                                   clean_x=clean_x)
-
-    @torch.no_grad()
-    def _update_ema(self) -> None:
-        decay = self._ema_decay
-        for ema_p, p in zip(self.ema_model.transformer.parameters(), self.student.transformer.parameters(),
-                            strict=True):
-            ema_p.mul_(decay).add_(p.detach().to(ema_p.dtype), alpha=1.0 - decay)
-
-    def _init_optimizers_and_schedulers(self) -> None:
-        tc = self.training_config
-        student_lr = float(tc.optimizer.learning_rate)
-        if student_lr <= 0.0:
-            raise ValueError("training.learning_rate must be > 0 for causal-cd")
-        student_params = [p for p in self.student.transformer.parameters() if p.requires_grad]
-        (
-            self._student_optimizer,
-            self._student_lr_scheduler,
-        ) = build_optimizer_and_scheduler(
-            params=student_params,
-            optimizer_config=tc.optimizer,
-            loop_config=tc.loop,
-            learning_rate=student_lr,
-            betas=tc.optimizer.betas,
-            scheduler_name=str(tc.optimizer.lr_scheduler),
-        )

fastvideo/train/methods/fine_tuning/__init__.py

@@ -7,12 +7,10 @@ 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.tfsft import TeacherForcingSFTMethod
 
 __all__ = [
     "DiffusionForcingSFTMethod",
     "FineTuneMethod",
-    "TeacherForcingSFTMethod",
 ]
 
 
@@ -27,9 +25,4 @@ def __getattr__(name: str) -> object:
         from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
 
         return FineTuneMethod
-    if name == "TeacherForcingSFTMethod":
-        from fastvideo.train.methods.fine_tuning.tfsft import (
-            TeacherForcingSFTMethod, )
-
-        return TeacherForcingSFTMethod
     raise AttributeError(name)

fastvideo/train/methods/fine_tuning/dfsft.py

@@ -135,11 +135,12 @@ class DiffusionForcingSFTMethod(TrainingMethod):
             t_inhom.flatten(),
         )
 
-        pred = self._predict_noise(
+        pred = self.student.predict_noise(
             noisy_latents,
             t_inhom,
             training_batch,
-            clean_latents,
+            conditional=True,
+            attn_kind=self._attn_kind,
         )
 
         if bool(self.training_config.model.precondition_outputs):
@@ -177,23 +178,6 @@ class DiffusionForcingSFTMethod(TrainingMethod):
         metrics: dict[str, LogScalar] = {}
         return loss_map, outputs, metrics
 
-    def _predict_noise(
-        self,
-        noisy_latents: torch.Tensor,
-        timestep: torch.Tensor,
-        training_batch: Any,
-        clean_latents: torch.Tensor,
-    ) -> torch.Tensor:
-        # Unused here; the teacher-forcing subclass overrides this to pass clean_x.
-        del clean_latents
-        return self.student.predict_noise(
-            noisy_latents,
-            timestep,
-            training_batch,
-            conditional=True,
-            attn_kind=self._attn_kind,
-        )
-
     # TrainingMethod override: backward
     def backward(
         self,

fastvideo/train/methods/fine_tuning/tfsft.py

@@ -1,30 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Teacher-forcing SFT method (TFSFT; algorithm layer)."""
-
-from __future__ import annotations
-
-from typing import Any
-
-import torch
-
-from fastvideo.train.methods.fine_tuning.dfsft import (
-    DiffusionForcingSFTMethod, )
-
-
-class TeacherForcingSFTMethod(DiffusionForcingSFTMethod):
-
-    def _predict_noise(
-        self,
-        noisy_latents: torch.Tensor,
-        timestep: torch.Tensor,
-        training_batch: Any,
-        clean_latents: torch.Tensor,
-    ) -> torch.Tensor:
-        return self.student.predict_noise(
-            noisy_latents,
-            timestep,
-            training_batch,
-            conditional=True,
-            attn_kind=self._attn_kind,
-            clean_x=clean_latents,
-        )

fastvideo/train/methods/rl/diffusion_nft.py

@@ -10,6 +10,12 @@ from typing import Any
 
 import torch
 import torch.distributed as dist
+from torch.distributed.checkpoint.state_dict import (
+    StateDictOptions,
+    get_model_state_dict,
+    set_model_state_dict,
+)
+from torch.distributed.checkpoint.stateful import Stateful
 from tqdm.auto import tqdm
 
 from fastvideo.dataset.parquet_dataset_map_style import (
@@ -31,7 +37,6 @@ from fastvideo.train.methods.rl.common import (
     validation_caption,
     validation_shard_indices,
 )
-from fastvideo.train.utils.checkpoint import _FullModelState
 from fastvideo.train.utils.config import (
     get_optional_float,
     get_optional_int,
@@ -73,6 +78,31 @@ class _DiffusionNFTEMAState:
         self._method._ema_update_count = int(update_count)
 
 
+class _FullModelState(Stateful):
+    """DCP wrapper that saves frozen model parameters too.
+
+    The shared ``ModelWrapper`` intentionally filters to ``requires_grad``
+    parameters. DiffusionNFT's old policy is frozen but must be restored on
+    resume, so it needs full model state.
+    """
+
+    def __init__(self, model: torch.nn.Module) -> None:
+        self.model = model
+
+    def state_dict(self) -> dict[str, Any]:
+        return get_model_state_dict(self.model)  # type: ignore[no-any-return]
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+    ) -> None:
+        set_model_state_dict(
+            self.model,
+            model_state_dict=state_dict,
+            options=StateDictOptions(strict=False),
+        )
+
+
 class DiffusionNFTMethod(TrainingMethod):
     """DiffusionNFT-style RL for diffusion models.
 

fastvideo/train/models/wan/wan.py

@@ -320,8 +320,6 @@ class WanModel(ModelBase):
         conditional: bool,
         cfg_uncond: dict[str, Any] | None = None,
         attn_kind: Literal["dense", "vsa"] = "dense",
-        clean_x: torch.Tensor | None = None,
-        aug_t: torch.Tensor | None = None,
     ) -> torch.Tensor:
         device_type = self.device.type
         dtype = self._get_training_dtype()
@@ -349,11 +347,7 @@ class WanModel(ModelBase):
                 current_timestep=batch.timesteps,
                 attn_metadata=attn_metadata,
         ):
-            input_kwargs = (self._build_distill_input_kwargs(noisy_latents,
-                                                             timestep,
-                                                             text_dict,
-                                                             clean_x=clean_x,
-                                                             aug_t=aug_t))
+            input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
             transformer = self._get_transformer(timestep)
             pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
         return pred_noise
@@ -536,24 +530,17 @@ class WanModel(ModelBase):
         noise_input: torch.Tensor,
         timestep: torch.Tensor,
         text_dict: dict[str, torch.Tensor] | None,
-        clean_x: torch.Tensor | None = None,
-        aug_t: torch.Tensor | None = None,
     ) -> dict[str, Any]:
         if text_dict is None:
             raise ValueError("text_dict cannot be None for "
                              "Wan distillation")
-        kwargs: dict[str, Any] = {
+        return {
             "hidden_states": noise_input.permute(0, 2, 1, 3, 4),
             "encoder_hidden_states": text_dict["encoder_hidden_states"],
             "encoder_attention_mask": text_dict["encoder_attention_mask"],
             "timestep": timestep,
             "return_dict": False,
         }
-        if clean_x is not None:
-            # Teacher forcing: clean context latents (+ optional aug timestep).
-            kwargs["clean_x"] = clean_x.permute(0, 2, 1, 3, 4)
-            kwargs["aug_t"] = aug_t
-        return kwargs
 
     def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module:
         return self.transformer

fastvideo/train/models/wan/wan_causal.py

@@ -49,7 +49,6 @@ class WanCausalModel(WanModel, CausalModelBase):
         transformer_override_safetensor: str
         | None = None,
         lora: LoraConfig | dict[str, Any] | None = None,
-        num_frames_per_block: int | None = None,
     ) -> None:
         super().__init__(
             init_from=init_from,
@@ -63,16 +62,6 @@ class WanCausalModel(WanModel, CausalModelBase):
         )
         self._streaming_caches: (dict[tuple[int, str], _StreamingCaches]) = {}
 
-        if num_frames_per_block is not None:
-            num_frames_per_block = int(num_frames_per_block)
-            if not 1 <= num_frames_per_block <= 3:
-                # Same bound as CausalWanTransformer3DModel's config path
-                # (assert num_frame_per_block <= 3); this override must not
-                # bypass it.
-                raise ValueError("num_frames_per_block must be between 1 and 3, "
-                                 f"got {num_frames_per_block}")
-            self.transformer.num_frame_per_block = num_frames_per_block
-
     # --- CausalModelBase override: clear_caches ---
     def clear_caches(
         self,

fastvideo/train/utils/checkpoint.py

@@ -15,13 +15,6 @@ import numpy as np
 import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
-from torch.distributed.checkpoint.state_dict import (
-    StateDictOptions,
-    get_model_state_dict,
-    set_model_state_dict,
-)
-from torch.distributed.checkpoint.stateful import Stateful
-
 from fastvideo.logger import init_logger
 
 logger = init_logger(__name__)
@@ -138,32 +131,6 @@ class _RoleModuleContainer(torch.nn.Module):
             self.add_module(name, module)
 
 
-class _FullModelState(Stateful):
-    """DCP wrapper that saves frozen model parameters too.
-
-    The shared ``ModelWrapper`` intentionally filters to ``requires_grad``
-    parameters. Frozen-but-mutated roles (e.g. DiffusionNFT's old policy,
-    causal-CD's EMA target) must still be restored on resume, so they need
-    full model state.
-    """
-
-    def __init__(self, model: torch.nn.Module) -> None:
-        self.model = model
-
-    def state_dict(self) -> dict[str, Any]:
-        return get_model_state_dict(self.model)  # type: ignore[no-any-return]
-
-    def load_state_dict(
-        self,
-        state_dict: dict[str, Any],
-    ) -> None:
-        set_model_state_dict(
-            self.model,
-            model_state_dict=state_dict,
-            options=StateDictOptions(strict=False),
-        )
-
-
 class _CallbackStateWrapper:
     """Wraps a CallbackDict for DCP save/load."""