[ci]: collect fastvideo/tests/stages in the unit-test lane (was never run by any CI job)

.agents/skills/reseed-performance-baseline/SKILL.md

@@ -69,7 +69,7 @@ approval, then upload reviewed accepted baseline records.
 | `model_id` | Yes | Benchmark id, e.g. `wan-t2v-1.3b-2gpu`. This maps to the HF subdirectory after `sanitize(model_id)`. |
 | `gpu_type` | Yes | Exact GPU device string from the performance record, e.g. the L40S device name emitted by CI. Baselines are GPU-specific. |
 | `source_results` | Yes | One or more local paths or Buildkite artifact URLs for accepted shifted performance JSONs. Prefer normalized `normalized_perf_*.json` artifacts emitted by `compare_baseline.py`. Accept `source_result` as an alias only for a single JSON. |
-| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `0.05` (5%). |
+| `max_intra_batch_regression` | No | Maximum allowed regression of any source JSON against the source batch median. Default: `PERF_MAX_REGRESSION` if set, otherwise `0.05` (5%). |
 | `intent_rationale` | Yes | One-line explanation for why the baseline shift is legitimate. This is written into provenance and should be reused in the PR. |
 
 Hardcoded defaults:
@@ -148,7 +148,8 @@ For each metric with at least two non-null source values:
 4. Stop if any source record regresses against the batch median by more than
    `max_intra_batch_regression`.
 
-Default `max_intra_batch_regression` to `0.05`. Print a table with per-source values, batch median, and
+Default `max_intra_batch_regression` to `PERF_MAX_REGRESSION` when set,
+otherwise `0.05`. Print a table with per-source values, batch median, and
 worst intra-batch regression.
 
 This check prevents uploading a mixed batch where one JSON is materially
@@ -182,7 +183,7 @@ present, that run is not a valid source for baseline reseeding.
 
 ### 2. Sync and back up existing HF records under /tmp
 
-Use `fastvideo/performance/hf_store.py` helpers directly. Do **not** use
+Use `fastvideo/tests/performance/hf_store.py` helpers directly. Do **not** use
 `compare_baseline.py` as a sync shortcut; on full main runs it can persist
 records, while this step must only fetch and back up existing history.
 
@@ -191,7 +192,7 @@ The sync command pattern is:
 ```bash
 export PERFORMANCE_TRACKING_ROOT="${PERFORMANCE_TRACKING_ROOT:-/tmp/perf-tracking}"
 export HF_REPO_ID="${HF_REPO_ID:-FastVideo/performance-tracking}"
-python -c 'from fastvideo.performance.hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
+PYTHONPATH=fastvideo/tests/performance python -c 'from hf_store import sync_from_hf; import os; sync_from_hf(os.environ["PERFORMANCE_TRACKING_ROOT"], strict=True)'
 ```
 
 Then back up only the sanitized model directory under `/tmp`:
@@ -199,8 +200,8 @@ Then back up only the sanitized model directory under `/tmp`:
 ```bash
 SHORT_COMMIT=$(git rev-parse --short=12 HEAD)
 TIMESTAMP=$(date -u +%Y%m%d_%H%M%S)
-MODEL_SAFE=$(python - <<'PY'
-from fastvideo.performance.hf_store import sanitize
+MODEL_SAFE=$(PYTHONPATH=fastvideo/tests/performance python - <<'PY'
+from hf_store import sanitize
 print(sanitize("<model_id>"))
 PY
 )
@@ -234,7 +235,7 @@ first baseline seed. Continue, but report that baseline history was empty.
 Load the last 5 successful records for the target:
 
 ```python
-from fastvideo.performance.hf_store import load_records_for_model
+from hf_store import load_records_for_model
 
 records = load_records_for_model(
     "/tmp/perf-tracking",
@@ -371,7 +372,7 @@ prepared records plus backup on disk.
 Use the shared storage helper so the path and repo type match CI:
 
 ```python
-from fastvideo.performance.hf_store import upload_record
+from hf_store import upload_record
 
 upload_record("<local_record_path>", record, strict=True)
 ```
@@ -459,7 +460,7 @@ directories created for this reseed. Never remove unrelated `/tmp` contents.
   intentional baseline replacement.
 - `fastvideo/tests/performance/compare_baseline.py` — normalization, rolling
   median comparison, and persistence rules.
-- `fastvideo/performance/hf_store.py` — HF sync, record loading,
+- `fastvideo/tests/performance/hf_store.py` — HF sync, record loading,
   `sanitize()`, and `upload_record()`.
 - `fastvideo/tests/performance/test_inference_performance.py` — source result
   JSON schema.

.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"

apps/performance_dashboard/README.md

@@ -12,17 +12,13 @@ Defaults:
 
 - `HF_REPO_ID=FastVideo/performance-tracking`
 - `PERFORMANCE_TRACKING_ROOT=/tmp/fastvideo-perf-dashboard`
+- `PERF_MAX_REGRESSION=0.05`
 
-Records can include source metadata and rolling-baseline policy context:
+Records can include source metadata:
 
 - `run_source`: `pr`, `local`, `scheduled_main`, or `unknown`
 - `baseline_eligible`: only successful scheduled-main records should be true
 - Buildkite metadata such as branch, PR number, build URL, build ID, and job ID
-- `regression_thresholds`: per-metric rolling-baseline percent and absolute
-  floors used for recomputed status context
-
-Dashboard/API metric payloads expose `threshold_exceeded` for raw threshold
-crossings; `regressed` remains the gated CI-failure signal.
 
 Set one of `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` if the
 configured dataset repo requires authenticated access:
@@ -93,8 +89,7 @@ Trend charts show metric-specific axes and exact point details on hover/focus:
 - PR number, branch, and Buildkite URL when present
 
 The latest status table uses the stored JSON `success` value. Recomputed
-baseline context applies each metric's percent and absolute regression floors
-and does not override stored status.
+baseline context is shown separately and does not override stored status.
 
 ## API
 

apps/performance_dashboard/frontend/src/App.tsx

@@ -440,8 +440,6 @@ export default function App() {
                   <th>Throughput</th>
                   <th>Memory</th>
                   <th>Worst</th>
-                  <th>Exceeded</th>
-                  <th>Failing</th>
                 </tr>
               </thead>
               <tbody>
@@ -467,12 +465,6 @@ export default function App() {
                     <td>{formatNumber(row.metrics.throughput?.current, 3)}</td>
                     <td>{formatNumber(row.metrics.memory?.current, 1)}</td>
                     <td>{formatNumber(row.worst_regression_pct, 1)}%</td>
-                    <td>
-                      {row.threshold_exceeded_metrics.length
-                        ? row.threshold_exceeded_metrics.join(", ")
-                        : "none"}
-                    </td>
-                    <td>{row.failing_metrics.length ? row.failing_metrics.join(", ") : "none"}</td>
                   </tr>
                 ))}
               </tbody>

apps/performance_dashboard/frontend/src/api.ts

@@ -2,12 +2,6 @@ export type MetricValue = {
   current: number | null;
   baseline: number | null;
   regression_pct: number | null;
-  absolute_delta: number | null;
-  threshold_percent: number;
-  threshold_absolute: number;
-  gated: boolean;
-  threshold_exceeded: boolean;
-  regressed: boolean;
   label: string;
   lower_is_better: boolean;
   precision: number;
@@ -21,8 +15,7 @@ export type SummaryRow = {
   success: boolean;
   baseline_n: number;
   worst_regression_pct: number | null;
-  threshold_exceeded_metrics: string[];
-  failing_metrics: string[];
+  regression_threshold_pct: number;
   computed_regression_status: "pass" | "fail";
   status: "pass" | "fail";
   run_source: RunSource;

docs/contributing/performance_benchmarks.md

@@ -72,10 +72,7 @@ fastvideo/tests/performance/
     │           writes Markdown summary + (optionally) uploads new records
     ├── dashboard.py
     │       └── builds time-series Plotly HTML from HF history
-
-fastvideo/performance/
-    ├── hf_store.py               # shared HF I/O + DataFrame helpers
-    └── metric_policy.py          # shared rolling-baseline threshold policy
+    └── hf_store.py               # shared HF I/O + DataFrame helpers
 ```
 
 The HF dataset (`FastVideo/performance-tracking` by default) holds one
@@ -95,28 +92,25 @@ and recipe changes instead of treating all records for a model as equivalent.
 
 ## Metrics
 
-Each benchmark records six metrics. The rolling-baseline comparator also has a
-per-metric policy with direction, percent threshold, absolute threshold, and a
-`gated` flag.
+Each benchmark records six metrics:
 
-| Metric | Raw key | Normalized key | Direction | Default rolling policy |
-|---|---|---|---|---|
-| End-to-end generation latency | `avg_generation_time_s` | `latency` | Lower is better | 8% and 0.5 s |
-| Video throughput | `throughput_fps` | `throughput` | Higher is better | 8% and 0.05 FPS |
-| Peak GPU memory | `max_peak_memory_mb` | `memory` | Lower is better | 5% and 256 MB |
-| Text encoder time | `text_encoder_time_s` | `text_encoder_time_s` | Lower is better | 5% and 0.25 s |
-| DiT denoising time | `dit_time_s` | `dit_time_s` | Lower is better | 5% and 0.25 s |
-| VAE decode time | `vae_decode_time_s` | `vae_decode_time_s` | Lower is better | 5% and 0.25 s |
+| Metric | Raw key | Normalized key | Direction |
+|---|---|---|---|
+| End-to-end generation latency | `avg_generation_time_s` | `latency` | Lower is better |
+| Video throughput | `throughput_fps` | `throughput` | Higher is better |
+| Peak GPU memory | `max_peak_memory_mb` | `memory` | Lower is better |
+| Text encoder time | `text_encoder_time_s` | `text_encoder_time_s` | Lower is better |
+| DiT denoising time | `dit_time_s` | `dit_time_s` | Lower is better |
+| VAE decode time | `vae_decode_time_s` | `vae_decode_time_s` | Lower is better |
 
 `test_inference_performance.py` temporarily sets `FASTVIDEO_STAGE_LOGGING=1`
 while it runs so pipeline stage execution times are available in
 `generate_video(...).logging_info`. Stage logs use pipeline-unique keys such as
 `prompt_encoding_stage` so duplicate stage classes do not collide. For
-`PipelineStage` entries, shared component stage bases emit a stable
-`component_metric`: text encoding stages map to `text_encoder_time_s`,
-denoising stages and subclasses map to `dit_time_s`, and decoding stages map to
-`vae_decode_time_s`. The extractor falls back to known `stage_class` names for
-older logs that do not include `component_metric` or that used the class name as
+`PipelineStage` entries, the extractor maps the `stage_class` field:
+`TextEncodingStage` maps to `text_encoder_time_s`, `DenoisingStage` and
+`DmdDenoisingStage` map to `dit_time_s`, and `DecodingStage` maps to
+`vae_decode_time_s`, with a fallback for older logs that used the class name as
 the stage key. Generator-side timings such as `PostDecodeFrameProcessStage`,
 `VideoSaveStage`, and `AudioMuxStage` are intentionally ignored. If a pipeline
 does not report one of the mapped stages, that component metric is stored as
@@ -162,22 +156,9 @@ headroom and almost never need touching.
 
 `compare_baseline.py` loads the last 5 successful, baseline-eligible records
 for the same `(model_id, gpu_type)` from the HF dataset, computes the median
-for each available metric, and evaluates the current run with the metric's
-rolling regression policy. For latency, memory, and component times, higher
-values are regressions. For throughput, lower values are regressions.
-
-A metric exceeds its rolling threshold when both of these are true:
-
-```text
-percent_delta > threshold_percent
-absolute_delta > threshold_absolute
-```
-
-Gated metrics fail CI when that threshold crossing happens. Set `gated: false`
-for metrics that should remain visible in reports and the dashboard without
-failing CI. Dashboard/API payloads expose `threshold_exceeded` separately from
-`regressed`, where `regressed` means a gated CI failure. Missing or `null`
-metrics are skipped.
+for each available metric, and fails if the current run regresses by more than
+`PERF_MAX_REGRESSION` (default 5%). For latency, memory, and component times,
+higher values are regressions. For throughput, lower values are regressions.
 
 This is the **drift detector** — it catches sub-threshold regressions that
 slowly add up. Only scheduled-main successful records are baseline eligible.
@@ -258,13 +239,6 @@ Written by `test_inference_performance.py`. One file per benchmark run.
     "max_dit_time_s": 10.0,
     "max_vae_decode_time_s": 10.0
   },
-  "regression_thresholds": {
-    "latency": {
-      "threshold_percent": 0.10,
-      "threshold_absolute": 1.0,
-      "gated": true
-    }
-  },
   "commit": "<full sha>",
   "pr_number": "1234",
   "timestamp": "2026-05-08T22:00:00+00:00",
@@ -291,13 +265,6 @@ result, used as the rolling-baseline source of truth.
   "text_encoder_time_s": 2.141,
   "dit_time_s": 8.437,
   "vae_decode_time_s": 3.208,
-  "regression_thresholds": {
-    "latency": {
-      "threshold_percent": 0.08,
-      "threshold_absolute": 0.5,
-      "gated": true
-    }
-  },
   "success": true
 }
 ```
@@ -314,17 +281,18 @@ successful main/full-suite uploads and remain eligible for rolling baselines.
 
 | Variable | Default | Used by | Purpose |
 |---|---|---|---|
+| `PERF_MAX_REGRESSION` | `0.05` | `compare_baseline.py` | Per-metric regression fraction that fails the build. |
 | `PERFORMANCE_TRACKING_ROOT` | `/tmp/perf-tracking` | `compare_baseline.py`, `dashboard.py` | Local directory the HF dataset is synced to. |
 | `PERF_REPORTS_DIR` | `/root/data/perf_reports` | `compare_baseline.py`, `dashboard.py` | Where the Markdown summary and Plotly HTML get written for Buildkite to pick up. |
-| `HF_REPO_ID` | `FastVideo/performance-tracking` | `fastvideo/performance/hf_store.py` | HF dataset repo holding rolling-baseline records. |
-| `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `fastvideo/performance/hf_store.py` | Required for upload or private dataset reads. |
+| `HF_REPO_ID` | `FastVideo/performance-tracking` | `hf_store.py` | HF dataset repo holding rolling-baseline records. |
+| `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, `HF_TOKEN` | unset | `hf_store.py` | Required for upload or private dataset reads. |
 | `PERF_RUN_SOURCE` | inferred | `compare_baseline.py` | Source metadata for uploaded records: `pr`, `local`, `scheduled_main`, or `unknown`. |
 | `PERF_UPLOAD_POLICY` | `never` | `compare_baseline.py` | Upload policy: `never`, `pass`, or `always`. |
 | `PERF_PYTEST_RC` | unset | `compare_baseline.py` | Static-threshold pytest exit code, used so scheduled-main failures can be uploaded with `success=false`. |
 | `TEST_SCOPE` | unset | `compare_baseline.py` | CI context used to infer scheduled-main runs together with `BUILDKITE_BRANCH=main`. |
 | `BUILDKITE_BRANCH`, `BUILDKITE_COMMIT`, `BUILDKITE_PULL_REQUEST` | unset | `compare_baseline.py`, `test_inference_performance.py` | CI metadata stamped into records. |
 | `DASHBOARD_DAYS` | `30` | `dashboard.py` | Lookback window for the Plotly trend pages. |
-| `PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS` | `3600` | `fastvideo/performance/hf_store.py` | Freshness window for reusing an existing HF sync when requested by dashboard consumers. |
+| `PERFORMANCE_TRACKING_SYNC_REUSE_TTL_SECONDS` | `3600` | `hf_store.py` | Freshness window for reusing an existing HF sync when requested by dashboard consumers. |
 | `FASTVIDEO_STAGE_LOGGING` | set by the pytest test | `test_inference_performance.py` | Enables pipeline stage timing capture for component metrics during benchmark runs. |
 
 ## CI integration
@@ -336,11 +304,9 @@ Buildkite artifact upload is in
 `.buildkite/scripts/pr_test.sh:upload_performance_artifacts`.
 
 Each performance build runs pytest first. If that fixed-threshold phase fails,
-PR/direct runs skip `compare_baseline.py` because they only upload passing
-records. Scheduled-main runs still execute `compare_baseline.py` with
-`PERF_PYTEST_RC` set so the failed canonical attempt is visible in normalized
-JSON and dashboard history. The dashboard runs best-effort for observability.
-When the rolling-baseline phase runs, it emits:
+`compare_baseline.py` is skipped, so Markdown summaries and normalized JSON
+artifacts are not emitted. The dashboard still runs best-effort for
+observability. When pytest passes, the rolling-baseline phase emits:
 
 * **Markdown summary** — appended to `$GITHUB_STEP_SUMMARY` when that variable
   is set, and written as `perf_<sha>_<ts>.md` for Buildkite upload. Contains a
@@ -381,13 +347,9 @@ When the rolling-baseline phase runs, it emits:
          "max_vae_decode_time_s": 10.0
        },
        "default": { "max_generation_time_s": 120.0, "max_peak_memory_mb": 30000.0 }
-     },
-     "regression_thresholds": {
-       "latency": { "threshold_percent": 0.10, "threshold_absolute": 1.0, "gated": true }
      }
-  }
-
-  ```
+   }
+   ```
 
    Legacy v1 configs without `config_schema_version` still load, but should not
    gain v2 identity or metadata fields until they are migrated to
@@ -410,13 +372,6 @@ When the rolling-baseline phase runs, it emits:
    a useful fixed gate. The rolling baseline will still track component times
    when static component thresholds are omitted.
 
-6. Omit `regression_thresholds` to use the default rolling-baseline policy, or
-   include only benchmark-specific deviations. Tune these independently from
-   the fixed thresholds when a metric is noisy or should be informational. The
-   fixed `thresholds` block is an absolute pytest ceiling. The
-   `regression_thresholds` block controls rolling-baseline comparisons against
-   recent scheduled-main records.
-
 ## Troubleshooting
 
 **"No baseline for ... Initializing"** — first run for this `(model_id,
@@ -433,5 +388,5 @@ pipelines that did not report a mapped component stage.
 
 **Component timing is `null`** — the generated result did not include a mapped
 stage in `logging_info.stages`. Check that the pipeline emits stage logging
-and that the stage emits `component_metric` or is covered by the legacy
-`STAGE_METRIC_MAP` fallback in `test_inference_performance.py`.
+and that the stage name is listed in `STAGE_METRIC_MAP` in
+`test_inference_performance.py`.

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/performance/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Performance benchmark and dashboard utilities."""

fastvideo/performance/metric_policy.py

@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Metric policy for rolling performance baseline comparisons."""
-
-from __future__ import annotations
-
-from collections.abc import Mapping
-from dataclasses import dataclass
-from typing import Any
-
-
-@dataclass(frozen=True)
-class MetricPolicy:
-    key: str
-    label: str
-    precision: int
-    lower_is_better: bool
-    threshold_percent: float
-    threshold_absolute: float
-    gated: bool = True
-
-
-@dataclass(frozen=True)
-class MetricDelta:
-    absolute: float
-    percent: float
-    threshold_exceeded: bool
-    regressed: bool
-
-
-DEFAULT_METRIC_POLICIES: tuple[MetricPolicy, ...] = (
-    MetricPolicy("latency", "Latency", 3, True, 0.08, 0.5),
-    MetricPolicy("throughput", "Throughput", 3, False, 0.08, 0.05),
-    MetricPolicy("memory", "Memory", 1, True, 0.05, 256.0),
-    MetricPolicy("text_encoder_time_s", "Text Enc", 3, True, 0.05, 0.25),
-    MetricPolicy("dit_time_s", "DiT", 3, True, 0.05, 0.25),
-    MetricPolicy("vae_decode_time_s", "VAE Decode", 3, True, 0.05, 0.25),
-)
-
-
-def _optional_float(value: Any) -> float | None:
-    if value is None or isinstance(value, bool):
-        return None
-    try:
-        return float(value)
-    except (TypeError, ValueError):
-        return None
-
-
-def _optional_bool(value: Any) -> bool | None:
-    if isinstance(value, bool):
-        return value
-    if isinstance(value, str):
-        normalized = value.strip().lower()
-        if normalized in {"1", "true", "yes", "on"}:
-            return True
-        if normalized in {"0", "false", "no", "off"}:
-            return False
-    return None
-
-
-def resolve_metric_policies(threshold_overrides: Mapping[str, Any] | None, ) -> tuple[MetricPolicy, ...]:
-    """Return default metric policies with optional per-metric overrides."""
-
-    if not isinstance(threshold_overrides, Mapping):
-        threshold_overrides = {}
-    policies: list[MetricPolicy] = []
-    for base_policy in DEFAULT_METRIC_POLICIES:
-        raw_override = threshold_overrides.get(base_policy.key, {})
-        if not isinstance(raw_override, Mapping):
-            raw_override = {}
-
-        threshold_percent = _optional_float(raw_override.get("threshold_percent"))
-        threshold_absolute = _optional_float(raw_override.get("threshold_absolute"))
-        gated = _optional_bool(raw_override.get("gated"))
-
-        policies.append(
-            MetricPolicy(
-                key=base_policy.key,
-                label=base_policy.label,
-                precision=base_policy.precision,
-                lower_is_better=base_policy.lower_is_better,
-                threshold_percent=(base_policy.threshold_percent if threshold_percent is None else threshold_percent),
-                threshold_absolute=(base_policy.threshold_absolute
-                                    if threshold_absolute is None else threshold_absolute),
-                gated=base_policy.gated if gated is None else gated,
-            ))
-    return tuple(policies)
-
-
-def serialize_metric_thresholds(policies: tuple[MetricPolicy, ...], ) -> dict[str, dict[str, float | bool]]:
-    return {
-        policy.key: {
-            "threshold_percent": policy.threshold_percent,
-            "threshold_absolute": policy.threshold_absolute,
-            "gated": policy.gated,
-        }
-        for policy in policies
-    }
-
-
-def regression_delta(
-    policy: MetricPolicy,
-    current: float,
-    baseline: float,
-) -> MetricDelta | None:
-    if baseline <= 0:
-        return None
-    absolute_delta = current - baseline if policy.lower_is_better else baseline - current
-    percent_delta = absolute_delta / baseline
-    threshold_exceeded = (percent_delta > policy.threshold_percent and absolute_delta > policy.threshold_absolute)
-    return MetricDelta(
-        absolute=absolute_delta,
-        percent=percent_delta,
-        threshold_exceeded=threshold_exceeded,
-        regressed=policy.gated and threshold_exceeded,
-    )

fastvideo/performance_dashboard/api.py

@@ -13,7 +13,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 
-from fastvideo.performance import hf_store
+from fastvideo.tests.performance import hf_store
 
 from .service import build_latest_summary, build_trends, filter_records
 
@@ -150,6 +150,7 @@ def create_app(store: PerformanceDataStore | None = None) -> FastAPI:
         filtered = filter_records(loaded, model_id=model_id, gpu_type=gpu_type)
         rows = build_latest_summary(
             filtered,
+            max_regression=float(os.environ.get("PERF_MAX_REGRESSION", "0.05")),
             run_source=run_source,
         )
         return {

fastvideo/performance_dashboard/metrics.py

@@ -1,8 +1,26 @@
 # SPDX-License-Identifier: Apache-2.0
 """Metric definitions shared by the performance dashboard backend."""
 
-from fastvideo.performance.metric_policy import DEFAULT_METRIC_POLICIES
+from __future__ import annotations
 
-METRICS = DEFAULT_METRIC_POLICIES
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class MetricDefinition:
+    key: str
+    label: str
+    precision: int
+    lower_is_better: bool
+
+
+METRICS: tuple[MetricDefinition, ...] = (
+    MetricDefinition("latency", "Latency", 3, True),
+    MetricDefinition("throughput", "Throughput", 3, False),
+    MetricDefinition("memory", "Memory", 1, True),
+    MetricDefinition("text_encoder_time_s", "Text Encoder", 3, True),
+    MetricDefinition("dit_time_s", "DiT", 3, True),
+    MetricDefinition("vae_decode_time_s", "VAE Decode", 3, True),
+)
 
 METRIC_BY_KEY = {metric.key: metric for metric in METRICS}

fastvideo/performance_dashboard/service.py

@@ -13,8 +13,9 @@ from collections import defaultdict
 from datetime import datetime, timezone
 from typing import Any
 
-from fastvideo.performance.hf_store import is_baseline_eligible_record, safe_float
-from fastvideo.performance.metric_policy import regression_delta, resolve_metric_policies
+from fastvideo.tests.performance.hf_store import is_baseline_eligible_record, safe_float
+
+from .metrics import METRICS
 
 Record = dict[str, Any]
 
@@ -94,9 +95,19 @@ def baseline_value(records: list[Record], metric_key: str) -> float | None:
     return float(statistics.median(values))
 
 
+def regression_percent(metric_key: str, current: float | None, baseline: float | None) -> float | None:
+    if current is None or baseline is None or baseline <= 0:
+        return None
+    metric = next(metric for metric in METRICS if metric.key == metric_key)
+    if metric.lower_is_better:
+        return (current - baseline) / baseline * 100.0
+    return (baseline - current) / baseline * 100.0
+
+
 def build_latest_summary(records: list[Record],
                          *,
                          baseline_window: int = 5,
+                         max_regression: float = 0.05,
                          run_source: str | None = None) -> list[Record]:
     rows: list[Record] = []
     for (model_id, gpu_type), group in group_by_model_gpu(records).items():
@@ -112,58 +123,52 @@ def build_latest_summary(records: list[Record],
             if record is not latest and record.get("success", True) and is_baseline_eligible_record(record)
         ]
         baseline_records = baseline_pool[-baseline_window:]
-        metric_policies = resolve_metric_policies(latest.get("regression_thresholds"))
 
         metrics: dict[str, Record] = {}
         regressions: list[float] = []
-        failing_metrics: list[str] = []
-        threshold_exceeded_metrics: list[str] = []
-        for policy in metric_policies:
-            current = safe_float(latest.get(policy.key))
-            baseline = baseline_value(baseline_records, policy.key)
-            delta = None
-            if current is not None and baseline is not None:
-                delta = regression_delta(policy, current, baseline)
-            regression = None if delta is None else delta.percent * 100.0
-            metrics[policy.key] = {
+        for metric in METRICS:
+            current = safe_float(latest.get(metric.key))
+            baseline = baseline_value(baseline_records, metric.key)
+            regression = regression_percent(metric.key, current, baseline)
+            metrics[metric.key] = {
                 "current": current,
                 "baseline": baseline,
                 "regression_pct": regression,
-                "absolute_delta": None if delta is None else delta.absolute,
-                "threshold_percent": policy.threshold_percent * 100.0,
-                "threshold_absolute": policy.threshold_absolute,
-                "gated": policy.gated,
-                "threshold_exceeded": False if delta is None else delta.threshold_exceeded,
-                "regressed": False if delta is None else delta.regressed,
-                "label": policy.label,
-                "lower_is_better": policy.lower_is_better,
-                "precision": policy.precision,
+                "label": metric.label,
+                "lower_is_better": metric.lower_is_better,
+                "precision": metric.precision,
             }
             if regression is not None:
                 regressions.append(regression)
-            if delta is not None and delta.threshold_exceeded:
-                threshold_exceeded_metrics.append(policy.key)
-            if delta is not None and delta.regressed:
-                failing_metrics.append(policy.key)
 
         worst_regression = max(regressions) if regressions else None
         success = bool(latest.get("success", True))
         status = "pass" if success else "fail"
 
         rows.append({
-            "model_id": model_id,
-            "gpu_type": gpu_type,
-            "timestamp": latest.get("timestamp"),
-            "commit_sha": latest.get("commit_sha"),
+            "model_id":
+            model_id,
+            "gpu_type":
+            gpu_type,
+            "timestamp":
+            latest.get("timestamp"),
+            "commit_sha":
+            latest.get("commit_sha"),
             **record_metadata(latest),
-            "success": success,
-            "baseline_n": len(baseline_records),
-            "worst_regression_pct": worst_regression,
-            "threshold_exceeded_metrics": threshold_exceeded_metrics,
-            "failing_metrics": failing_metrics,
-            "computed_regression_status": "fail" if failing_metrics else "pass",
-            "status": status,
-            "metrics": metrics,
+            "success":
+            success,
+            "baseline_n":
+            len(baseline_records),
+            "worst_regression_pct":
+            worst_regression,
+            "regression_threshold_pct":
+            max_regression * 100.0,
+            "computed_regression_status":
+            "fail" if worst_regression is not None and worst_regression > max_regression * 100.0 else "pass",
+            "status":
+            status,
+            "metrics":
+            metrics,
         })
 
     return sorted(rows, key=lambda row: (row["status"] != "fail", row["model_id"], row["gpu_type"]))
@@ -174,15 +179,14 @@ def build_trends(records: list[Record]) -> list[Record]:
     for (model_id, gpu_type), group in group_by_model_gpu(records).items():
         points = []
         for record in group:
-            metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
             point = {
                 "timestamp": record.get("timestamp"),
                 "commit_sha": record.get("commit_sha"),
                 **record_metadata(record),
                 "success": bool(record.get("success", True)),
                 "metrics": {
-                    policy.key: safe_float(record.get(policy.key))
-                    for policy in metric_policies
+                    metric.key: safe_float(record.get(metric.key))
+                    for metric in METRICS
                 },
             }
             points.append(point)

fastvideo/pipelines/stages/base.py

@@ -34,7 +34,6 @@ class PipelineStage(ABC):
     composed with other stages to create a complete pipeline. Each stage is responsible
     for a specific part of the process, such as prompt encoding, latent preparation, etc.
     """
-    performance_component_metric: str | None = None
 
     def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
         """
@@ -156,9 +155,6 @@ class PipelineStage(ABC):
                 logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000)
                 batch.logging_info.add_stage_execution_time(stage_key, execution_time)
                 batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name)
-                component_metric = self.performance_component_metric
-                if component_metric is not None:
-                    batch.logging_info.add_stage_metric(stage_key, "component_metric", component_metric)
             except Exception as e:
                 torch.cuda.synchronize()
                 execution_time = time.perf_counter() - start_time

fastvideo/pipelines/stages/decoding.py

@@ -28,7 +28,6 @@ class DecodingStage(PipelineStage):
     This stage handles the decoding of latent representations into the final
     output format (e.g., pixel values).
     """
-    performance_component_metric = "vae_decode_time_s"
 
     def __init__(self, vae, pipeline=None) -> None:
         self.vae: ParallelTiledVAE = vae

fastvideo/pipelines/stages/denoising.py

@@ -51,7 +51,6 @@ class DenoisingStage(PipelineStage):
     This stage handles the iterative denoising process that transforms
     the initial noise into the final output.
     """
-    performance_component_metric = "dit_time_s"
 
     def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None:
         super().__init__()
@@ -1188,7 +1187,6 @@ class Cosmos25V2WDenoisingStage(Cosmos25DenoisingStage):
 
 class Cosmos25AutoDenoisingStage(PipelineStage):
     """Route Cosmos 2.5 denoising to T2W vs V2W/I2W."""
-    performance_component_metric = "dit_time_s"
 
     def __init__(self, transformer, scheduler) -> None:
         super().__init__()

fastvideo/pipelines/stages/text_encoding.py

@@ -24,7 +24,6 @@ class TextEncodingStage(PipelineStage):
     This stage handles the encoding of text prompts into the embedding space
     expected by the diffusion model.
     """
-    performance_component_metric = "text_encoder_time_s"
 
     def __init__(self, text_encoders, tokenizers) -> None:
         """
@@ -351,7 +350,6 @@ class Cosmos25TextEncodingStage(PipelineStage):
     Cosmos 2.5 uses Reason1 (Qwen2.5-VL) and relies on the encoder's
     `compute_text_embeddings_online()`.
     """
-    performance_component_metric = "text_encoder_time_s"
 
     def __init__(self, text_encoder) -> None:
         super().__init__()

fastvideo/tests/contract/test_ci_test_collection.py

@@ -1,85 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Guard: every test directory must be collected by some CI lane or be on
-the explicit allowlist below.
-
-Three separate incidents on 2026-07-05 found test files that no CI lane
-ever collects (fastvideo/tests/stages/, tests/local_tests/ additions in
-PR #1509, and this sweep found seven dark directories in total): the tests
-pass review, merge, and then silently never run. This test makes going
-dark an explicit, reviewed decision instead of an accident: adding a new
-test directory fails CI until it is either wired into a lane or
-allowlisted here with a reason.
-
-Pure text analysis — no fastvideo imports, no GPU, no torch.
-"""
-from pathlib import Path
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-TESTS_ROOT = REPO_ROOT / "fastvideo" / "tests"
-
-# Files whose text constitutes "a CI lane references this directory".
-CI_SOURCES = [
-    TESTS_ROOT / "modal" / "pr_test.py",
-    TESTS_ROOT / "modal" / "ssim_test.py",
-    *sorted((REPO_ROOT / ".buildkite").rglob("*.yml")),
-    *sorted((REPO_ROOT / ".buildkite").rglob("*.sh")),
-]
-
-# Directories that intentionally have no CI lane today. Every entry needs a
-# reason; remove the entry when the directory gets wired into a lane.
-# State as found on 2026-07-05 — these SHOULD shrink over time, not grow.
-ALLOWLIST = {
-    "attention": "no lane yet — GPU attention-backend tests, run manually",
-    "audio": "no lane yet — audio encoder tests, run manually",
-    "distributed": "no lane yet — multi-GPU torchrun tests, run manually",
-    "hooks": "no lane yet — run manually",
-    "layers": "no lane yet — torchrun FSDP dispatch tests, run manually",
-    "nightly": "by design: nightly cadence, not per-PR",
-    "ops": "no lane yet — GPU op tests, run manually",
-    "modal": "CI infrastructure itself, not a test suite",
-}
-
-
-def _dirs_with_tests() -> list[str]:
-    dirs = []
-    for child in sorted(TESTS_ROOT.iterdir()):
-        if child.is_dir() and any(child.rglob("test_*.py")):
-            dirs.append(child.name)
-    return dirs
-
-
-def _ci_text() -> str:
-    return "\n".join(
-        src.read_text(errors="replace") for src in CI_SOURCES if src.exists())
-
-
-def test_every_test_directory_is_collected_or_allowlisted():
-    ci_text = _ci_text()
-    dark = [
-        name for name in _dirs_with_tests()
-        if f"tests/{name}" not in ci_text and name not in ALLOWLIST
-    ]
-    assert not dark, (
-        f"Test directories not referenced by any CI lane and not "
-        f"allowlisted: {dark}. Wire them into a lane in "
-        f"fastvideo/tests/modal/pr_test.py (or a Buildkite step), or add an "
-        f"allowlist entry with a reason in {__file__}.")
-
-
-def test_local_tests_stays_out_of_ci():
-    # tests/local_tests/ (repo root) is developer-local by design (author
-    # decision, 2026-07-05): parity scaffolds and machine-specific checks
-    # that must never gate CI. Fail if any CI source starts collecting it.
-    assert "tests/local_tests" not in _ci_text(), (
-        "tests/local_tests/ is local-only by design; remove the CI "
-        "reference or move the tests into a fastvideo/tests/ lane.")
-
-
-def test_allowlist_entries_are_still_real_directories():
-    # A stale allowlist hides regressions; entries must track reality.
-    missing = [
-        name for name in ALLOWLIST
-        if name != "modal" and not (TESTS_ROOT / name).is_dir()
-    ]
-    assert not missing, (
-        f"Allowlisted directories no longer exist — remove them: {missing}")

fastvideo/tests/performance/compare_baseline.py

@@ -8,8 +8,8 @@ This script:
    baseline-eligible successful records (filtered by gpu_type),
 4) writes normalized records back to the HF dataset repo according to
    PERF_UPLOAD_POLICY,
-5) exits non-zero if any gated metric exceeds both its percent and absolute
-   regression floors.
+5) exits non-zero if any metric regresses by more than PERF_MAX_REGRESSION
+   (default 5%).
 """
 
 import glob
@@ -21,36 +21,21 @@ from datetime import datetime, timezone
 from typing import Any
 
 try:
-    from fastvideo.performance.hf_store import (
+    from .hf_store import (
         load_records_for_model,
         safe_float,
         sanitize,
         sync_from_hf,
         upload_record,
     )
-    from fastvideo.performance.metric_policy import (
-        MetricPolicy,
-        regression_delta,
-        resolve_metric_policies,
-        serialize_metric_thresholds,
-    )
 except ImportError:
-    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
-    if repo_root not in sys.path:
-        sys.path.insert(0, repo_root)
-    from fastvideo.performance.hf_store import (
+    from hf_store import (
         load_records_for_model,
         safe_float,
         sanitize,
         sync_from_hf,
         upload_record,
     )
-    from fastvideo.performance.metric_policy import (
-        MetricPolicy,
-        regression_delta,
-        resolve_metric_policies,
-        serialize_metric_thresholds,
-    )
 
 RESULTS_DIR = os.path.join(
     os.path.dirname(os.path.abspath(__file__)),
@@ -61,9 +46,25 @@ TRACKING_ROOT = os.environ.get(
     "/tmp/perf-tracking",
 )
 PERF_REPORTS_DIR = os.environ.get("PERF_REPORTS_DIR", "/root/data/perf_reports")
+MAX_REGRESSION = float(os.environ.get("PERF_MAX_REGRESSION", "0.05"))
 UPLOAD_POLICY = os.environ.get("PERF_UPLOAD_POLICY", "never").strip().lower()
 VALID_UPLOAD_POLICIES = {"never", "pass", "always"}
 VALID_RUN_SOURCES = {"pr", "local", "scheduled_main", "unknown"}
+METRICS = (
+    ("latency", "Latency", 3),
+    ("throughput", "Throughput", 3),
+    ("memory", "Memory", 1),
+    ("text_encoder_time_s", "Text Enc", 3),
+    ("dit_time_s", "DiT", 3),
+    ("vae_decode_time_s", "VAE Decode", 3),
+)
+LOWER_IS_BETTER_METRICS = {
+    "latency",
+    "memory",
+    "text_encoder_time_s",
+    "dit_time_s",
+    "vae_decode_time_s",
+}
 
 
 def _should_persist_tracking() -> bool:
@@ -167,7 +168,6 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
     text_encoder_time = safe_float(result.get("text_encoder_time_s"))
     dit_time = safe_float(result.get("dit_time_s"))
     vae_decode_time = safe_float(result.get("vae_decode_time_s"))
-    metric_policies = resolve_metric_policies(result.get("regression_thresholds"))
 
     return {
         "model_id": model_id,
@@ -180,7 +180,6 @@ def normalize_performance_result(result: dict[str, Any]) -> dict[str, Any]:
         "text_encoder_time_s": text_encoder_time,
         "dit_time_s": dit_time,
         "vae_decode_time_s": vae_decode_time,
-        "regression_thresholds": serialize_metric_thresholds(metric_policies),
         "success": True,
         **_record_metadata(_detect_run_source(), result),
     }
@@ -230,41 +229,55 @@ def _baseline_metric(records: list[dict[str, Any]], key: str) -> float | None:
     return statistics.median(values)
 
 
-def _metric_policy_summary(policy: MetricPolicy) -> str:
-    gated = "gated" if policy.gated else "info"
-    return (
-        f"{gated}, >{policy.threshold_percent * 100:.1f}% "
-        f"and >{policy.threshold_absolute:.{policy.precision}f}"
-    )
-
-
 def _check_regressions(
     current: dict[str, Any],
     baseline_records: list[dict[str, Any]],
-    metric_policies: tuple[MetricPolicy, ...],
+    max_regression: float,
 ) -> list[str]:
     failures: list[str] = []
 
-    for policy in metric_policies:
-        baseline = _baseline_metric(baseline_records, policy.key)
-        curr = safe_float(current.get(policy.key))
-        if baseline is None or curr is None:
+    for metric, _label, _precision in METRICS:
+        if metric not in LOWER_IS_BETTER_METRICS:
             continue
-        delta = regression_delta(policy, curr, baseline)
-        if delta is None or not delta.regressed:
+        baseline = _baseline_metric(baseline_records, metric)
+        curr = safe_float(current.get(metric))
+        if baseline is None or curr is None or baseline <= 0:
             continue
-        failures.append(
-            f"{current['model_id']} {policy.key} regressed by "
-            f"{delta.percent * 100:.1f}% and "
-            f"{delta.absolute:.{policy.precision}f} "
-            f"(current={curr:.{policy.precision}f}, "
-            f"baseline_median={baseline:.{policy.precision}f}, "
-            f"threshold={_metric_policy_summary(policy)})"
-        )
+        regression = (curr - baseline) / baseline
+        if regression > max_regression:
+            failures.append(f"{current['model_id']} {metric} regressed by "
+                            f"{regression * 100:.1f}% "
+                            f"(current={curr:.3f}, baseline_median={baseline:.3f})")
+
+    baseline_tp = _baseline_metric(baseline_records, "throughput")
+    curr_tp = safe_float(current.get("throughput"))
+    if baseline_tp is not None and curr_tp is not None and baseline_tp > 0:
+        regression = (baseline_tp - curr_tp) / baseline_tp
+        if regression > max_regression:
+            failures.append(f"{current['model_id']} throughput regressed by "
+                            f"{regression * 100:.1f}% "
+                            f"(current={curr_tp:.3f}, baseline_median={baseline_tp:.3f})")
 
     return failures
 
 
+def _metric_delta_percent(
+    metric: str,
+    current: dict[str, Any],
+    baseline_records: list[dict[str, Any]],
+) -> float | None:
+    curr = safe_float(current.get(metric))
+    baseline = _baseline_metric(baseline_records, metric)
+    if curr is None or baseline is None or baseline <= 0:
+        return None
+
+    if metric in LOWER_IS_BETTER_METRICS:
+        return (curr - baseline) / baseline * 100.0
+    if metric == "throughput":
+        return (baseline - curr) / baseline * 100.0
+    return None
+
+
 def _compact_value(value: float | None, precision: int = 3) -> str:
     if value is None:
         return "n/a"
@@ -274,42 +287,23 @@ def _compact_value(value: float | None, precision: int = 3) -> str:
 def _build_summary_row(
     record: dict[str, Any],
     baseline_records: list[dict[str, Any]],
-    metric_policies: tuple[MetricPolicy, ...],
     has_failed: bool,
 ) -> dict[str, Any]:
     """Format a single benchmark result as a row for the Markdown table."""
 
-    metric_values: dict[str, dict[str, Any]] = {}
+    metric_values: dict[str, dict[str, float | None]] = {}
     regressions: list[float] = []
-    failing_metrics: list[str] = []
-    threshold_exceeded_metrics: list[str] = []
-    for policy in metric_policies:
-        curr = safe_float(record.get(policy.key))
-        baseline = _baseline_metric(baseline_records, policy.key)
-        delta = (
-            regression_delta(policy, curr, baseline)
-            if curr is not None and baseline is not None
-            else None
-        )
-        regression = None if delta is None else delta.percent * 100.0
-        absolute_delta = None if delta is None else delta.absolute
-        metric_values[policy.key] = {
+    for metric, _label, _precision in METRICS:
+        curr = safe_float(record.get(metric))
+        baseline = _baseline_metric(baseline_records, metric)
+        regression = _metric_delta_percent(metric, record, baseline_records)
+        metric_values[metric] = {
             "curr": curr,
             "base": baseline,
             "regression_pct": regression,
-            "absolute_delta": absolute_delta,
-            "threshold_percent": policy.threshold_percent * 100.0,
-            "threshold_absolute": policy.threshold_absolute,
-            "gated": policy.gated,
-            "threshold_exceeded": False if delta is None else delta.threshold_exceeded,
-            "regressed": False if delta is None else delta.regressed,
         }
         if regression is not None:
             regressions.append(regression)
-        if delta is not None and delta.threshold_exceeded:
-            threshold_exceeded_metrics.append(policy.key)
-        if delta is not None and delta.regressed:
-            failing_metrics.append(policy.key)
 
     worst_regression_pct = max(regressions) if regressions else None
 
@@ -319,50 +313,40 @@ def _build_summary_row(
         "baseline_n": len(baseline_records),
         "metrics": metric_values,
         "worst_regression_pct": worst_regression_pct,
-        "threshold_exceeded_metrics": threshold_exceeded_metrics,
-        "failing_metrics": failing_metrics,
         "failed": has_failed,
     }
 
 
 def _build_markdown_summary(
     summary_rows: list[dict[str, Any]],
-    metric_policies: tuple[MetricPolicy, ...],
+    max_regression: float,
 ) -> str:
     lines = [
         "## Performance Baseline Comparison",
         "",
-        "Threshold: gated metrics fail only when both percent and absolute "
-        "regression floors are exceeded.",
+        f"Threshold: regressions greater than {max_regression * 100:.1f}% fail",
         "",
         ("| Model | GPU | Baseline N | Latency (curr/base) | "
          "Throughput (curr/base) | Memory (curr/base) | "
          "Text Enc (curr/base) | DiT (curr/base) | "
-         "VAE Decode (curr/base) | Worst Regression | Exceeded Metrics | "
-         "Failing Metrics | Status |"),
-        "|---|---|---:|---|---|---|---|---|---|---:|---|---|---|",
+         "VAE Decode (curr/base) | Worst Regression | Status |"),
+        "|---|---|---:|---|---|---|---|---|---|---:|---|",
     ]
 
     for row in summary_rows:
         metric_cells = []
-        for policy in metric_policies:
-            values = row["metrics"][policy.key]
-            metric_cells.append(f"{_compact_value(values['curr'], policy.precision)} / "
-                                f"{_compact_value(values['base'], policy.precision)}")
+        for metric, _label, precision in METRICS:
+            values = row["metrics"][metric]
+            metric_cells.append(f"{_compact_value(values['curr'], precision)} / "
+                                f"{_compact_value(values['base'], precision)}")
 
         worst_reg = ("n/a" if row["worst_regression_pct"] is None else f"{row['worst_regression_pct']:.1f}%")
-        exceeded_metrics = (
-            ", ".join(row["threshold_exceeded_metrics"])
-            if row["threshold_exceeded_metrics"]
-            else "none"
-        )
-        failing_metrics = ", ".join(row["failing_metrics"]) if row["failing_metrics"] else "none"
         status = "FAIL" if row["failed"] else "PASS"
 
         lines.append(f"| {row['model_id']} | {row['gpu_type']} | "
                      f"{row['baseline_n']} | "
                      f"{' | '.join(metric_cells)} | "
-                     f"{worst_reg} | {exceeded_metrics} | {failing_metrics} | {status} |")
+                     f"{worst_reg} | {status} |")
 
     return "\n".join(lines) + "\n"
 
@@ -416,7 +400,6 @@ def main() -> int:
 
     for raw in current_results:
         record = _normalize_record(raw)
-        metric_policies = resolve_metric_policies(record.get("regression_thresholds"))
 
         baseline_records = load_records_for_model(
             TRACKING_ROOT,
@@ -433,7 +416,7 @@ def main() -> int:
             failures: list[str] = []
             record["success"] = True
         else:
-            failures = _check_regressions(record, baseline_records, metric_policies)
+            failures = _check_regressions(record, baseline_records, MAX_REGRESSION)
         if static_threshold_failed:
             failures.append(f"{record['model_id']} fixed-threshold phase failed "
                             f"(PERF_PYTEST_RC={os.environ.get('PERF_PYTEST_RC')})")
@@ -451,10 +434,10 @@ def main() -> int:
             print("Tracking upload skipped for "
                   f"{record['model_id']} ({record['run_source']}, success={record['success']})")
 
-        summary_rows.append(_build_summary_row(record, baseline_records, metric_policies, bool(failures)))
+        summary_rows.append(_build_summary_row(record, baseline_records, bool(failures)))
 
     commit_sha = os.environ.get("BUILDKITE_COMMIT", "unknown")[:7]
-    markdown = _build_markdown_summary(summary_rows, resolve_metric_policies(None))
+    markdown = _build_markdown_summary(summary_rows, MAX_REGRESSION)
     _emit_markdown_summary(markdown, commit_sha)
 
     if all_failures:

fastvideo/tests/performance/dashboard.py

@@ -1,19 +1,12 @@
 # SPDX-License-Identifier: Apache-2.0
 import os
-import sys
 from html import escape
 from datetime import datetime
 
 import plotly.express as px
 import pandas as pd
 
-try:
-    from fastvideo.performance.hf_store import load_as_dataframe, sync_from_hf
-except ImportError:
-    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
-    if repo_root not in sys.path:
-        sys.path.insert(0, repo_root)
-    from fastvideo.performance.hf_store import load_as_dataframe, sync_from_hf
+from hf_store import sync_from_hf, load_as_dataframe
 
 TRACKING_ROOT = os.environ.get(
     "PERFORMANCE_TRACKING_ROOT",

fastvideo/tests/performance/hf_store.py

@@ -57,7 +57,9 @@ def is_baseline_eligible_record(record: dict[str, Any]) -> bool:
     """
     if record.get("baseline_eligible") is True:
         return True
-    return "baseline_eligible" not in record and "run_source" not in record
+    if "baseline_eligible" not in record and "run_source" not in record:
+        return True
+    return False
 
 
 def resolve_hf_token() -> str | None:
@@ -310,12 +312,8 @@ def load_records_for_model(
 # ---------------------------------------------------------------------------
 
 _NUMERIC_COLS = (
-    "latency",
-    "throughput",
-    "memory",
-    "text_encoder_time_s",
-    "dit_time_s",
-    "vae_decode_time_s",
+    "latency", "throughput", "memory",
+    "text_encoder_time_s", "dit_time_s", "vae_decode_time_s",
 )
 
 

fastvideo/tests/performance/test_compare_baseline_policy.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from fastvideo.tests.performance import compare_baseline
-from fastvideo.performance.metric_policy import resolve_metric_policies
 
 
 def _raw_result():
@@ -73,157 +72,9 @@ def test_normalized_record_includes_source_metadata(monkeypatch):
     assert record["job_id"] == "job-1"
 
 
-def test_normalized_record_includes_effective_regression_thresholds(monkeypatch):
-    monkeypatch.setenv("PERF_RUN_SOURCE", "pr")
-    raw = _raw_result()
-    raw["regression_thresholds"] = {
-        "latency": {
-            "threshold_percent": 0.09,
-            "threshold_absolute": 0.75,
-            "gated": True,
-        },
-        "throughput": {
-            "gated": False,
-        },
-    }
-
-    record = compare_baseline.normalize_performance_result(raw)
-
-    assert record["regression_thresholds"]["latency"] == {
-        "threshold_percent": 0.09,
-        "threshold_absolute": 0.75,
-        "gated": True,
-    }
-    assert record["regression_thresholds"]["throughput"]["gated"] is False
-
-
-def test_invalid_regression_threshold_container_uses_defaults():
-    policies = resolve_metric_policies(["not", "a", "mapping"])
-
-    latency = next(policy for policy in policies if policy.key == "latency")
-    assert latency.threshold_percent == 0.08
-    assert latency.threshold_absolute == 0.5
-    assert latency.gated is True
-
-
-def test_boolean_regression_threshold_values_are_ignored():
-    policies = resolve_metric_policies({
-        "latency": {
-            "threshold_percent": True,
-            "threshold_absolute": False,
-            "gated": "false",
-        }
-    })
-
-    latency = next(policy for policy in policies if policy.key == "latency")
-    assert latency.threshold_percent == 0.08
-    assert latency.threshold_absolute == 0.5
-    assert latency.gated is False
-
-
 def test_baseline_eligibility_only_for_successful_scheduled_main():
     assert compare_baseline._is_baseline_eligible("scheduled_main", True) is True
     assert compare_baseline._is_baseline_eligible("scheduled_main", False) is False
     assert compare_baseline._is_baseline_eligible("pr", True) is False
     assert compare_baseline._is_baseline_eligible("local", True) is False
 
-
-def test_latency_regression_requires_percent_and_absolute_floors():
-    baseline = [{"latency": 10.0}]
-    current = {"model_id": "wan", "latency": 10.6}
-
-    percent_only = resolve_metric_policies({
-        "latency": {
-            "threshold_percent": 0.05,
-            "threshold_absolute": 0.75,
-        }
-    })
-    absolute_only = resolve_metric_policies({
-        "latency": {
-            "threshold_percent": 0.10,
-            "threshold_absolute": 0.5,
-        }
-    })
-    both = resolve_metric_policies({
-        "latency": {
-            "threshold_percent": 0.05,
-            "threshold_absolute": 0.5,
-        }
-    })
-
-    assert compare_baseline._check_regressions(current, baseline, percent_only) == []
-    assert compare_baseline._check_regressions(current, baseline, absolute_only) == []
-
-    failures = compare_baseline._check_regressions(current, baseline, both)
-    assert len(failures) == 1
-    assert "latency regressed by 6.0% and 0.600" in failures[0]
-
-
-def test_throughput_regression_uses_higher_is_better_direction():
-    baseline = [{"throughput": 10.0}]
-    current = {"model_id": "wan", "throughput": 9.0}
-    policies = resolve_metric_policies({
-        "throughput": {
-            "threshold_percent": 0.05,
-            "threshold_absolute": 0.5,
-        }
-    })
-
-    failures = compare_baseline._check_regressions(current, baseline, policies)
-
-    assert len(failures) == 1
-    assert "throughput regressed by 10.0% and 1.000" in failures[0]
-
-
-def test_memory_regression_uses_metric_specific_absolute_floor():
-    baseline = [{"memory": 10000.0}]
-    current = {"model_id": "wan", "memory": 10600.0}
-    policies = resolve_metric_policies({
-        "memory": {
-            "threshold_percent": 0.05,
-            "threshold_absolute": 256.0,
-        }
-    })
-
-    failures = compare_baseline._check_regressions(current, baseline, policies)
-
-    assert len(failures) == 1
-    assert "memory regressed by 6.0% and 600.0" in failures[0]
-
-
-def test_component_metric_can_gate_independently():
-    baseline = [{"dit_time_s": 8.0}]
-    current = {"model_id": "wan", "dit_time_s": 8.6}
-    policies = resolve_metric_policies({
-        "dit_time_s": {
-            "threshold_percent": 0.05,
-            "threshold_absolute": 0.25,
-        }
-    })
-
-    failures = compare_baseline._check_regressions(current, baseline, policies)
-
-    assert len(failures) == 1
-    assert "dit_time_s regressed by 7.5% and 0.600" in failures[0]
-
-
-def test_informational_metric_remains_visible_without_failing():
-    baseline = [{"throughput": 10.0}]
-    current = {"model_id": "wan", "gpu_type": "NVIDIA L40S", "throughput": 8.0}
-    policies = resolve_metric_policies({
-        "throughput": {
-            "threshold_percent": 0.01,
-            "threshold_absolute": 0.01,
-            "gated": False,
-        }
-    })
-
-    row = compare_baseline._build_summary_row(current, baseline, policies, False)
-
-    assert compare_baseline._check_regressions(current, baseline, policies) == []
-    assert row["metrics"]["throughput"]["regression_pct"] == 20.0
-    assert row["metrics"]["throughput"]["gated"] is False
-    assert row["metrics"]["throughput"]["threshold_exceeded"] is True
-    assert row["metrics"]["throughput"]["regressed"] is False
-    assert row["threshold_exceeded_metrics"] == ["throughput"]
-    assert row["failing_metrics"] == []

fastvideo/tests/performance/test_dashboard_api.py

@@ -68,8 +68,6 @@ def test_summary_endpoint_returns_latest_group_status():
     assert body["count"] == 1
     assert body["status_counts"] == {"pass": 1, "fail": 0}
     assert body["rows"][0]["metrics"]["latency"]["baseline"] == 10.0
-    assert body["rows"][0]["metrics"]["latency"]["threshold_exceeded"] is True
-    assert body["rows"][0]["threshold_exceeded_metrics"] == ["latency", "throughput"]
     assert body["rows"][0]["computed_regression_status"] == "fail"
 
 

fastvideo/tests/performance/test_dashboard_service.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
-from fastvideo.performance import hf_store
 from fastvideo.performance_dashboard.service import build_latest_summary, build_trends, filter_records
+from fastvideo.tests.performance import hf_store
 
 
 def _record(ts, commit, latency, throughput, success=True, **metadata):
@@ -28,23 +28,16 @@ def test_build_latest_summary_uses_previous_successful_records_for_baseline():
         _record("2026-01-03T00:00:00+00:00", "c" * 40, 11.0, 9.0),
     ]
 
-    rows = build_latest_summary(records)
+    rows = build_latest_summary(records, max_regression=0.05)
 
     assert len(rows) == 1
     row = rows[0]
     assert row["baseline_n"] == 1
     assert row["metrics"]["latency"]["baseline"] == 10.0
     assert row["metrics"]["latency"]["regression_pct"] == 10.0
-    assert row["metrics"]["latency"]["absolute_delta"] == 1.0
-    assert row["metrics"]["latency"]["threshold_percent"] == 8.0
-    assert row["metrics"]["latency"]["threshold_absolute"] == 0.5
-    assert row["metrics"]["latency"]["threshold_exceeded"] is True
-    assert row["metrics"]["latency"]["regressed"] is True
     assert row["metrics"]["throughput"]["regression_pct"] == 10.0
     assert row["status"] == "pass"
     assert row["computed_regression_status"] == "fail"
-    assert row["threshold_exceeded_metrics"] == ["latency", "throughput"]
-    assert row["failing_metrics"] == ["latency", "throughput"]
 
 
 def test_build_latest_summary_status_uses_latest_record_success_field():
@@ -80,7 +73,7 @@ def test_build_latest_summary_run_source_filter_keeps_canonical_baseline():
         ),
     ]
 
-    rows = build_latest_summary(records, run_source="pr")
+    rows = build_latest_summary(records, max_regression=0.05, run_source="pr")
 
     assert len(rows) == 1
     assert rows[0]["run_source"] == "pr"
@@ -90,60 +83,6 @@ def test_build_latest_summary_run_source_filter_keeps_canonical_baseline():
     assert rows[0]["computed_regression_status"] == "fail"
 
 
-def test_build_latest_summary_requires_absolute_floor_for_computed_regression():
-    records = [
-        _record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
-        _record(
-            "2026-01-02T00:00:00+00:00",
-            "b" * 40,
-            10.6,
-            10.0,
-            regression_thresholds={
-                "latency": {
-                    "threshold_percent": 0.05,
-                    "threshold_absolute": 0.75,
-                    "gated": True,
-                }
-            },
-        ),
-    ]
-
-    rows = build_latest_summary(records)
-
-    assert round(rows[0]["metrics"]["latency"]["regression_pct"], 1) == 6.0
-    assert round(rows[0]["metrics"]["latency"]["absolute_delta"], 3) == 0.6
-    assert rows[0]["metrics"]["latency"]["threshold_exceeded"] is False
-    assert rows[0]["metrics"]["latency"]["regressed"] is False
-    assert rows[0]["computed_regression_status"] == "pass"
-
-
-def test_build_latest_summary_separates_informational_threshold_crossing():
-    records = [
-        _record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),
-        _record(
-            "2026-01-02T00:00:00+00:00",
-            "b" * 40,
-            10.6,
-            10.0,
-            regression_thresholds={
-                "latency": {
-                    "threshold_percent": 0.05,
-                    "threshold_absolute": 0.5,
-                    "gated": False,
-                }
-            },
-        ),
-    ]
-
-    rows = build_latest_summary(records)
-
-    assert rows[0]["metrics"]["latency"]["threshold_exceeded"] is True
-    assert rows[0]["metrics"]["latency"]["regressed"] is False
-    assert rows[0]["threshold_exceeded_metrics"] == ["latency"]
-    assert rows[0]["failing_metrics"] == []
-    assert rows[0]["computed_regression_status"] == "pass"
-
-
 def test_filter_records_and_trends_preserve_metric_points():
     records = [
         _record("2026-01-01T00:00:00+00:00", "a" * 40, 10.0, 10.0),

fastvideo/tests/performance/test_inference_performance.py

@@ -179,11 +179,7 @@ def _extract_component_times(result: dict) -> dict[str, float | None]:
             logger.debug("Skipping malformed stage '%s' data: %r", stage_name, stage_data)
             continue
         stage_class = stage_data.get("stage_class", stage_name)
-        component_metric = stage_data.get("component_metric")
-        if isinstance(component_metric, str) and component_metric in component_times:
-            metric_key = component_metric
-        else:
-            metric_key = STAGE_METRIC_MAP.get(stage_class)
+        metric_key = STAGE_METRIC_MAP.get(stage_class)
         if metric_key is None:
             logger.debug("Unmapped stage '%s' class '%s' (%.3fs)",
                          stage_name,
@@ -313,7 +309,6 @@ def _run_benchmark(cfg):
         "max_peak_memory_mb": round(max_peak_memory, 1),
         "individual_peak_memories_mb": [round(m, 1) for m in peak_memories],
         "thresholds": thresholds,
-        "regression_thresholds": cfg.get("regression_thresholds", {}),
         "commit": os.environ.get("BUILDKITE_COMMIT", ""),
         "pr_number": os.environ.get("BUILDKITE_PULL_REQUEST", ""),
         "timestamp": datetime.now(timezone.utc).isoformat(),

fastvideo/tests/performance/test_inference_performance_component_times.py

@@ -1,20 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from fastvideo.pipelines.pipeline_batch_info import PipelineLoggingInfo
-from fastvideo.pipelines.stages.denoising import Cosmos25AutoDenoisingStage, DenoisingStage
-from fastvideo.pipelines.stages.text_encoding import Cosmos25TextEncodingStage
 from fastvideo.tests.performance.test_inference_performance import _extract_component_times
 
 
-class SubclassStyleDenoisingStage(DenoisingStage):
-    pass
-
-
 def test_extract_component_times_handles_pipeline_logging_info_object():
     logging_info = PipelineLoggingInfo()
     logging_info.add_stage_execution_time("prompt_encoding_stage", 1.25)
     logging_info.add_stage_metric("prompt_encoding_stage", "stage_class", "TextEncodingStage")
-    logging_info.add_stage_metric("prompt_encoding_stage", "component_metric", "text_encoder_time_s")
 
     assert _extract_component_times({"logging_info": logging_info}) == {
         "text_encoder_time_s": 1.25,
@@ -52,35 +45,6 @@ def test_extract_component_times_uses_stage_class_for_pipeline_stage_keys():
     }
 
 
-def test_denoising_stage_subclasses_inherit_component_metric():
-    assert SubclassStyleDenoisingStage.performance_component_metric == "dit_time_s"
-
-
-def test_cosmos25_direct_pipeline_stages_define_component_metrics():
-    assert Cosmos25TextEncodingStage.performance_component_metric == "text_encoder_time_s"
-    assert Cosmos25AutoDenoisingStage.performance_component_metric == "dit_time_s"
-
-
-def test_extract_component_times_uses_component_metric_for_stage_subclasses():
-    result = {
-        "logging_info": {
-            "stages": {
-                "denoising_stage": {
-                    "execution_time": 4.2,
-                    "stage_class": "CosmosDenoisingStage",
-                    "component_metric": "dit_time_s",
-                },
-            },
-        },
-    }
-
-    assert _extract_component_times(result) == {
-        "text_encoder_time_s": None,
-        "dit_time_s": 4.2,
-        "vae_decode_time_s": None,
-    }
-
-
 def test_extract_component_times_keeps_legacy_class_name_keys():
     # Backward-compatibility check for logs produced before pipeline-unique
     # stage keys carried a separate stage_class field.

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."""