refactor(diffusion): require rollout scheduler sigmas, drop timesteps-derived fallbacks

miles/backends/fsdp_utils/actor.py

@@ -364,7 +364,6 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
             raise ValueError("rollout_data['train_data'] is empty")
 
         num_pairs = len(train_pairs)
-        num_train_timesteps = self.scheduler.config.num_train_timesteps
 
         ref_mode = self.args.ref_mode
         if ref_mode == "lora_base" and not all(hasattr(m, "disable_adapter") for m in self.models.values()):
@@ -376,7 +375,6 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
         scheduler_timesteps, scheduler_sigmas = scheduler_meta_from_rollout(
             rollout_data,
             device=device,
-            num_train_timesteps=num_train_timesteps,
         )
         self.scheduler.timesteps = scheduler_timesteps
         self.scheduler.sigmas = scheduler_sigmas

miles/ray/data_conversion_hub/flow_grpo.py

@@ -4,6 +4,7 @@
 
 import torch
 
+from miles.utils.train_data_utils import scheduler_meta_from_samples
 from miles.utils.types import RolloutDebugTensors, Sample
 
 
@@ -25,28 +26,10 @@ def _expand_samples_to_train_pairs(
     """Flat train pairs in sample-major order (all pairs for sample 0, then sample 1, ...)."""
     device = torch.device("cpu")
     train_data: list[dict[str, Any]] = []
-    first_traj = samples[0].dit_trajectory
-    scheduler_meta: dict[str, torch.Tensor] = {"scheduler_timesteps": first_traj.timesteps.detach().cpu().float()}
-
-    if first_traj.sigmas is not None:
-        scheduler_meta["scheduler_sigmas"] = first_traj.sigmas.detach().cpu().float()
+    scheduler_meta = scheduler_meta_from_samples(samples)
 
     for sample, rew, raw_r in zip(samples, rewards, raw_rewards, strict=True):
         traj, denoising_env, rollout_log_probs = _sample_required_inputs(sample)
-        if not torch.equal(traj.timesteps.detach().cpu().float(), scheduler_meta["scheduler_timesteps"]):
-            raise ValueError(
-                f"sample {sample.index} has different scheduler_timesteps than sample 0; "
-                "the converter assumes one shared schedule across the batch"
-            )
-        expected_sigmas = scheduler_meta.get("scheduler_sigmas")
-        traj_sigmas = None if traj.sigmas is None else traj.sigmas.detach().cpu().float()
-        if (expected_sigmas is None) != (traj_sigmas is None) or (
-            expected_sigmas is not None and not torch.equal(traj_sigmas, expected_sigmas)
-        ):
-            raise ValueError(
-                f"sample {sample.index} has different scheduler_sigmas than sample 0; "
-                "the converter assumes one shared schedule across the batch"
-            )
         per_sample_features = _build_per_sample_features(
             sample,
             reward=rew,

miles/ray/data_conversion_hub/nft.py

@@ -5,6 +5,7 @@
 
 import torch
 
+from miles.utils.train_data_utils import scheduler_meta_from_samples
 from miles.utils.types import Sample
 
 
@@ -51,20 +52,7 @@ def expand_samples_to_train_pairs(
             f"NFT convert length mismatch: samples={len(samples)} "
             f"rewards={len(rewards)} raw_rewards={len(raw_rewards)}"
         )
-    first_traj = samples[0].dit_trajectory
-    if first_traj is None:
-        raise ValueError("sample 0 missing dit_trajectory")
-    if first_traj.timesteps is None:
-        raise ValueError("NFT needs dit_trajectory.timesteps from rollout")
-    if first_traj.sigmas is not None:
-        scheduler_sigmas = first_traj.sigmas.detach().cpu().float()
-    else:
-        ts = first_traj.timesteps.detach().cpu().float()
-        scheduler_sigmas = torch.cat([ts / 1000.0, ts.new_zeros(1)])
-    scheduler_meta = {
-        "scheduler_timesteps": first_traj.timesteps.detach().cpu().float(),
-        "scheduler_sigmas": scheduler_sigmas,
-    }
+    scheduler_meta = scheduler_meta_from_samples(samples)
     sigmas = resolve_nft_sigmas(
         scheduler_meta["scheduler_sigmas"],
         training_timestep_fraction=args.diffusion_nft_timestep_fraction,

miles/utils/train_data_utils.py

@@ -22,20 +22,50 @@ def stack_train_pair_rollout_debug(
     return torch.stack([item["rollout_debug_tensors"][key] for item in batch], dim=0)
 
 
+def scheduler_meta_from_samples(samples: list) -> dict[str, torch.Tensor]:
+    """Build the batch scheduler meta from sample 0's trajectory, enforcing one shared schedule."""
+    first = samples[0].dit_trajectory
+    if first is None:
+        raise ValueError("sample 0 missing dit_trajectory")
+    if first.timesteps is None:
+        raise ValueError("sample 0 missing dit_trajectory.timesteps")
+    if first.sigmas is None:
+        raise ValueError("sample 0 missing dit_trajectory.sigmas; rollout engine must return the sigmas snapshot")
+    meta = {
+        "scheduler_timesteps": first.timesteps.detach().cpu().float(),
+        "scheduler_sigmas": first.sigmas.detach().cpu().float(),
+    }
+    for sample in samples[1:]:
+        traj = sample.dit_trajectory
+        if (
+            traj is None
+            or traj.timesteps is None
+            or not torch.equal(traj.timesteps.detach().cpu().float(), meta["scheduler_timesteps"])
+        ):
+            raise ValueError(
+                f"sample {sample.index} has different scheduler_timesteps than sample 0; "
+                "the converter assumes one shared schedule across the batch"
+            )
+        if traj.sigmas is None or not torch.equal(traj.sigmas.detach().cpu().float(), meta["scheduler_sigmas"]):
+            raise ValueError(
+                f"sample {sample.index} has different scheduler_sigmas than sample 0; "
+                "the converter assumes one shared schedule across the batch"
+            )
+    return meta
+
+
 def scheduler_meta_from_rollout(
     rollout_data: dict,
     *,
     device: torch.device,
-    num_train_timesteps: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Use rollout-side scheduler metadata for train/rollout alignment."""
     if "scheduler_timesteps" not in rollout_data:
         raise ValueError("rollout_data missing scheduler_timesteps")
+    if "scheduler_sigmas" not in rollout_data:
+        raise ValueError("rollout_data missing scheduler_sigmas; rollout engine must return the sigmas snapshot")
     timesteps = rollout_data["scheduler_timesteps"].to(device=device, dtype=torch.float32)
-    if "scheduler_sigmas" in rollout_data:
-        sigmas = rollout_data["scheduler_sigmas"].to(device=device, dtype=torch.float32)
-    else:
-        sigmas = torch.cat([timesteps / float(num_train_timesteps), timesteps.new_zeros(1)])
+    sigmas = rollout_data["scheduler_sigmas"].to(device=device, dtype=torch.float32)
     return timesteps, sigmas
 
 

miles/utils/types.py

@@ -47,11 +47,9 @@ class DenoisingEnv:
 class DiTTrajectory:
     latents: torch.Tensor | None = None
     timesteps: torch.Tensor | None = None
-    # Rollout's scheduler.sigmas snapshot [T+1] (post-shift, includes
-    # terminal 0). Use this on the training side instead of recomputing
-    # sigmas from `timesteps / num_train_timesteps` — that round-trips
-    # σ * 1000 / 1000 in fp32 and drifts 1-2 ULPs, amplifying to ~3e-5
-    # log_prob diff.
+    # Rollout's scheduler.sigmas snapshot [T+1] (post-shift, includes terminal 0).
+    # Required for training — converters raise if missing; recomputing from
+    # `timesteps / num_train_timesteps` drifts 1-2 ULPs (~3e-5 log_prob diff).
     sigmas: torch.Tensor | None = None
 
 

tests/fast/backends/fsdp_utils/test_loss_hub_nft.py

@@ -75,6 +75,50 @@ class _Env:
         assert out["train_data"][0]["advantage"] == rewards[0]
         assert out["train_data"][2]["advantage"] == rewards[1]
 
+    def test_convert_rejects_mismatched_scheduler_meta(self):
+        # One shared schedule per batch, same contract as the flow_grpo converter.
+        class _Traj:
+            def __init__(self):
+                self.timesteps = torch.tensor([999.0, 500.0, 0.0])
+                self.sigmas = torch.tensor([1.0, 0.5, 0.0])
+                self.latents = torch.zeros(3, 2, 2)
+
+        class _Env:
+            pos_cond_kwargs = {}
+            neg_cond_kwargs = None
+
+        samples = [
+            Sample(index=0, prompt="a", reward=1.0, dit_trajectory=_Traj(), denoising_env=_Env()),
+            Sample(index=1, prompt="b", reward=3.0, dit_trajectory=_Traj(), denoising_env=_Env()),
+        ]
+        samples[1].dit_trajectory.sigmas = samples[1].dit_trajectory.sigmas + 1.0  # tamper
+        try:
+            expand_samples_to_train_pairs(_args(), samples, [1.0, 2.0], [1.0, 2.0])
+        except ValueError as e:
+            assert "scheduler_sigmas" in str(e)
+        else:
+            raise AssertionError("expected ValueError for mismatched scheduler_sigmas")
+
+    def test_convert_requires_rollout_sigmas(self):
+        # Sigmas come from the rollout scheduler snapshot; no timesteps/1000 fallback.
+        class _Traj:
+            def __init__(self):
+                self.timesteps = torch.tensor([999.0, 500.0, 0.0])
+                self.sigmas = None
+                self.latents = torch.zeros(3, 2, 2)
+
+        class _Env:
+            pos_cond_kwargs = {}
+            neg_cond_kwargs = None
+
+        samples = [Sample(index=0, prompt="a", reward=1.0, dit_trajectory=_Traj(), denoising_env=_Env())]
+        try:
+            expand_samples_to_train_pairs(_args(), samples, [1.0], [1.0])
+        except ValueError as e:
+            assert "sigmas" in str(e)
+        else:
+            raise AssertionError("expected ValueError for missing dit_trajectory.sigmas")
+
 
 class TestEmaShadow:
     def _model(self):

tests/fast/utils/test_grouping_parity.py

@@ -32,7 +32,7 @@
 import torch
 
 from miles.ray.data_conversion_hub.flow_grpo import expand_samples_to_train_pairs
-from miles.utils.train_data_utils import TrainDataDPSplitter
+from miles.utils.train_data_utils import TrainDataDPSplitter, scheduler_meta_from_rollout
 
 
 # --------------------------------------------------------------------------------------
@@ -154,12 +154,16 @@ def test_l2_converter_pairs_match_direct_indexing():
             assert torch.equal(d["rollout_step_noise_std_dev"], s.rollout_debug_tensors.rollout_noise_std_devs[idx])
 
 
-def test_l2_converter_sigmas_optional():
+def test_l2_converter_requires_sigmas():
+    """A trajectory without the rollout sigmas snapshot must raise (no derived fallback)."""
     T, sde = 4, [0, 2]
     samples = [_mk_sample(i, T, sde, with_sigmas=False, with_debug=False) for i in range(2)]
-    out = expand_samples_to_train_pairs(None, samples, [1.0, 2.0], [1.0, 2.0])
-    assert "scheduler_sigmas" not in out
-    assert len(out["train_data"]) == 2 * len(sde)
+    try:
+        expand_samples_to_train_pairs(None, samples, [1.0, 2.0], [1.0, 2.0])
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("expected ValueError for missing dit_trajectory.sigmas")
 
 
 def test_l2_converter_rejects_mismatched_scheduler_timesteps():
@@ -186,6 +190,22 @@ def test_l2_converter_rejects_mismatched_scheduler_sigmas():
         raise AssertionError("expected ValueError for mismatched scheduler_sigmas")
 
 
+def test_scheduler_meta_from_rollout_requires_sigmas():
+    """The train actor consumes the rollout sigmas snapshot verbatim; no timesteps/N fallback."""
+    ts = torch.tensor([999.0, 500.0, 1.0])
+    sig = torch.tensor([1.0, 0.5, 0.001, 0.0])
+    out_ts, out_sig = scheduler_meta_from_rollout(
+        {"scheduler_timesteps": ts, "scheduler_sigmas": sig}, device=torch.device("cpu")
+    )
+    assert torch.equal(out_ts, ts) and torch.equal(out_sig, sig)
+    try:
+        scheduler_meta_from_rollout({"scheduler_timesteps": ts}, device=torch.device("cpu"))
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("expected ValueError for missing scheduler_sigmas")
+
+
 # --------------------------------------------------------------------------------------
 # L3 — Cond window padding: per-microbatch collate(pad_to_len) == legacy window-slice
 # --------------------------------------------------------------------------------------