From 98f761ec453d1f3400b8a657ced562cebdeff7b4 Mon Sep 17 00:00:00 2001
From: William Lin <SolitaryThinker@users.noreply.github.com>
Date: Sun, 23 Aug 2026 01:09:26 -0500
Subject: [PATCH] [bugfix] validation: inherit the trained denoising ladder
 (#1738)

---
 .../test_validation_denoising_ladder.py       | 63 +++++++++++++++++++
 fastvideo/train/callbacks/validation.py       | 38 +++++++++++
 2 files changed, 101 insertions(+)
 create mode 100644 fastvideo/tests/train/callbacks/test_validation_denoising_ladder.py

diff --git a/fastvideo/tests/train/callbacks/test_validation_denoising_ladder.py b/fastvideo/tests/train/callbacks/test_validation_denoising_ladder.py
new file mode 100644
index 000000000..7c8baf458
--- /dev/null
+++ b/fastvideo/tests/train/callbacks/test_validation_denoising_ladder.py
@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Validation samples on the ladder the method trains against.
+
+``sampling_steps`` only sets ``num_inference_steps``; the scheduler expands N
+of those into an N-point sigma grid, i.e. N-1 forwards on its own spacing. A
+few-step recipe therefore validates off its operating point unless
+``sampling_timesteps`` repeats the trained ladder exactly.
+"""
+
+from __future__ import annotations
+
+import types
+
+import pytest
+
+from fastvideo.train.callbacks.validation import ValidationCallback
+
+LADDER = [1000, 750, 500, 250]
+
+
+def _callback(**kwargs):
+    return ValidationCallback(
+        pipeline_target="fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline",
+        dataset_file="unused.json",
+        **kwargs,
+    )
+
+
+def _method(ladder=LADDER):
+    cfg = {} if ladder is None else {"dmd_denoising_steps": list(ladder)}
+    return types.SimpleNamespace(method_config=cfg)
+
+
+def test_ladder_is_inherited_when_unset() -> None:
+    cb = _callback(sampling_steps=[4])
+    assert cb.sampling_timesteps is None
+    cb._adopt_training_denoising_ladder(_method())
+    assert cb.sampling_timesteps == LADDER
+
+
+def test_matching_ladder_is_accepted() -> None:
+    cb = _callback(sampling_steps=[4], sampling_timesteps=LADDER)
+    cb._adopt_training_denoising_ladder(_method())
+    assert cb.sampling_timesteps == LADDER
+
+
+def test_diverging_ladder_raises() -> None:
+    cb = _callback(sampling_steps=[4], sampling_timesteps=[999, 749, 500])
+    with pytest.raises(ValueError, match="disagrees with the trained ladder"):
+        cb._adopt_training_denoising_ladder(_method())
+
+
+@pytest.mark.parametrize("ladder", [None, []])
+def test_methods_without_a_ladder_are_left_alone(ladder) -> None:
+    cb = _callback(sampling_steps=[40])
+    cb._adopt_training_denoising_ladder(_method(ladder))
+    assert cb.sampling_timesteps is None
+
+
+def test_method_without_config_is_tolerated() -> None:
+    cb = _callback(sampling_steps=[40])
+    cb._adopt_training_denoising_ladder(types.SimpleNamespace())
+    assert cb.sampling_timesteps is None
diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py
index 4b635e827..e83775ed0 100644
--- a/fastvideo/train/callbacks/validation.py
+++ b/fastvideo/train/callbacks/validation.py
@@ -240,6 +240,43 @@ class ValidationCallback(Callback):
     # Callback hooks
     # ----------------------------------------------------------
 
+    def _adopt_training_denoising_ladder(
+        self,
+        method: TrainingMethod,
+    ) -> None:
+        """Keep validation on the timestep ladder the method trains against.
+
+        ``sampling_steps`` only sets ``num_inference_steps``, which the
+        scheduler turns into an N-point sigma grid -- N-1 forwards on its own
+        spacing, not the trained ladder. Validation reaches the trained
+        operating point only when ``sampling_timesteps`` repeats
+        ``method.dmd_denoising_steps`` exactly, so inherit it by default and
+        refuse to run when the two disagree.
+        """
+        method_config = getattr(method, "method_config", None)
+        if not isinstance(method_config, dict):
+            return
+        raw = method_config.get("dmd_denoising_steps")
+        if not isinstance(raw, list) or not raw:
+            return
+        trained = [int(s) for s in raw]
+
+        if self.sampling_timesteps is None:
+            self.sampling_timesteps = trained
+            logger.info(
+                "validation: inheriting the trained denoising ladder %s "
+                "(%d forwards)",
+                trained,
+                len(trained),
+            )
+            return
+        if self.sampling_timesteps != trained:
+            raise ValueError("callbacks.validation.sampling_timesteps "
+                             f"{self.sampling_timesteps} disagrees with the trained ladder "
+                             f"{trained} (method.dmd_denoising_steps). Validation would "
+                             "sample off the operating point the student is trained for. "
+                             "Drop the override to inherit the ladder, or align the two.")
+
     def on_train_start(
         self,
         method: TrainingMethod,
@@ -247,6 +284,7 @@ class ValidationCallback(Callback):
     ) -> None:
         self.method = method
         tc = self.training_config
+        self._adopt_training_denoising_ladder(method)
 
         self.world_group = get_world_group()
         self.sp_group = get_sp_group()