fix(nft): give NFT its own timestep hook instead of sharing process_timestep_as_input

miles/backends/fsdp_utils/configs/ltx.py

@@ -15,7 +15,6 @@
 class LTXTrainPipelineConfig(TrainPipelineConfig):
     """LTX-2.3 video GRPO: unguided velocity forward over ltx_core."""
 
-    needs_timestep_scaling = False
     supports_cfg_training = False
     # Rollout stores σ×1000 in trajectory timesteps; ltx_core AdaLN uses σ∈[0,1].
     sde_timestep_divisor = 1000.0

miles/backends/fsdp_utils/configs/qwen_image.py

@@ -59,9 +59,15 @@ class QwenImageTrainPipelineConfig(TrainPipelineConfig):
     hf_ckpt_name_patterns = ("qwen-image",)
     cfg_batching = False
 
+    # 1000 is the model's normalizer, not the scheduler range: sglang-d's DiT divides by the same
+    # literal. Reading num_train_timesteps here would diverge from the rollout at any other range.
     def process_timestep_as_input(self, timesteps):
         return timesteps / 1000.0
 
+    def process_sigma_as_timesteps_input(self, sigmas, *, num_train_timesteps):
+        # Identity only while the scheduler range equals the 1000 above; else sigma * N / 1000.
+        return sigmas
+
     lora_target_modules = [
         "to_q",
         "to_k",

miles/backends/fsdp_utils/configs/train_pipeline_config.py

@@ -106,6 +106,12 @@ def process_timestep_as_input(self, timesteps: torch.Tensor) -> torch.Tensor:
         sglang-d DiT rescales it -- the arithmetic has to match, not just the value."""
         return timesteps
 
+    def process_sigma_as_timesteps_input(self, sigmas: torch.Tensor, *, num_train_timesteps: int) -> torch.Tensor:
+        """NFT's sigma as this family's DiT takes its timesteps input. Separate from
+        ``process_timestep_as_input`` rather than pre-multiplying into it: for a family that
+        divides there, the composition is a multiply and a divide that do not cancel in fp32."""
+        return sigmas * float(num_train_timesteps)
+
     def compute_noise_pred(
         self,
         *,

miles/backends/fsdp_utils/configs/wan2_2.py

@@ -15,7 +15,6 @@ class Wan2_2TrainPipelineConfig(TrainPipelineConfig):
     # High-noise expert ("transformer") handles t >= boundary, low-noise expert
     # ("transformer_2") the rest.
     boundary_ratio = 0.875
-    # Wan DiT expects raw scheduler timesteps (0..num_train_timesteps), no /1000 scaling.
 
     def component_for_timestep(self, timestep: float, num_train_timesteps: int) -> str:
         if timestep >= self.boundary_ratio * num_train_timesteps:

miles/backends/fsdp_utils/loss_hub/nft.py

@@ -45,7 +45,7 @@ def prepare_nft_batch(
     return PreparedBatch(
         latents=xt,
         timesteps=t,
-        timesteps_for_model=config.process_timestep_as_input(t * float(num_train_timesteps)),
+        timesteps_for_model=config.process_sigma_as_timesteps_input(t, num_train_timesteps=num_train_timesteps),
         model=model,
         component_name=component_name,
         guidance_scale=0.0,

tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py

@@ -107,3 +107,28 @@ def test_qwen_image_divides_by_the_model_normalizer(self):
         out = QwenImageTrainPipelineConfig.process_timestep_as_input(QwenImageTrainPipelineConfig, self.TIMESTEPS)
         # One division, like the rollout: any rewrite of the expression drifts ULPs.
         assert torch.equal(out, self.TIMESTEPS / 1000.0)
+
+
+class TestProcessSigmaAsTimestepsInput:
+    # The NFT counterpart: each family rescales the opposite way from above.
+    NUM_TRAIN_TIMESTEPS = 1000
+    # 0.8474... does not survive a multiply then divide by 1000 in fp32.
+    SIGMAS = torch.tensor([0.8474337458610535, 0.5])
+
+    @pytest.mark.parametrize("config_cls", [SD3TrainPipelineConfig, Wan2_2TrainPipelineConfig])
+    def test_scales_up_to_the_scheduler_range(self, config_cls):
+        out = config_cls.process_sigma_as_timesteps_input(
+            config_cls, self.SIGMAS, num_train_timesteps=self.NUM_TRAIN_TIMESTEPS
+        )
+        assert torch.equal(out, self.SIGMAS * float(self.NUM_TRAIN_TIMESTEPS))
+
+    def test_qwen_image_passes_the_sigma_through(self):
+        out = QwenImageTrainPipelineConfig.process_sigma_as_timesteps_input(
+            QwenImageTrainPipelineConfig, self.SIGMAS, num_train_timesteps=self.NUM_TRAIN_TIMESTEPS
+        )
+        assert torch.equal(out, self.SIGMAS)
+        # Asserted so the equal() above keeps its teeth.
+        round_tripped = QwenImageTrainPipelineConfig.process_timestep_as_input(
+            QwenImageTrainPipelineConfig, self.SIGMAS * float(self.NUM_TRAIN_TIMESTEPS)
+        )
+        assert not torch.equal(round_tripped, self.SIGMAS)

tests/fast/backends/fsdp_utils/test_loss_hub_nft.py

@@ -8,8 +8,11 @@
 
 import torch
 
+from miles.backends.fsdp_utils.configs.qwen_image import QwenImageTrainPipelineConfig
+from miles.backends.fsdp_utils.configs.train_pipeline_config import TrainPipelineConfig
 from miles.backends.fsdp_utils.ema import EmaShadow
-from miles.backends.fsdp_utils.loss_hub.nft import corrupt, nft_r_from_advantages
+from miles.backends.fsdp_utils.loss_hub.nft import corrupt, nft_r_from_advantages, prepare_nft_batch
+from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext
 from miles.ray.data_conversion_hub.nft import expand_samples_to_train_pairs, resolve_nft_sigmas
 from miles.utils.types import Sample
 
@@ -120,6 +123,85 @@ class _Env:
             raise AssertionError("expected ValueError for missing dit_trajectory.sigmas")
 
 
+class _StubConfig:
+    """Cond plumbing stubbed out. Binds both hooks so a prepare wired to the wrong one fails on
+    the numbers, not on a missing attribute."""
+
+    process_timestep_as_input = TrainPipelineConfig.process_timestep_as_input
+    process_sigma_as_timesteps_input = TrainPipelineConfig.process_sigma_as_timesteps_input
+
+    def prepare_cond_kwargs(self, cond, device):
+        return {}
+
+    def collate_cond_for_sample_batch(self, per_sample_cond_kwargs, device, pad_to_len=None):
+        return {}
+
+
+class _Sd3StyleConfig(_StubConfig):
+    pass
+
+
+class _QwenStyleConfig(_StubConfig):
+    process_timestep_as_input = QwenImageTrainPipelineConfig.process_timestep_as_input
+    process_sigma_as_timesteps_input = QwenImageTrainPipelineConfig.process_sigma_as_timesteps_input
+
+
+class TestPrepareNftBatch:
+    NUM_TRAIN_TIMESTEPS = 1000
+    # 0.8474... does not survive a multiply then divide by 1000 in fp32.
+    SIGMAS = [0.8474337458610535, 0.5]
+
+    class _Env:
+        pos_cond_kwargs = None
+        neg_cond_kwargs = None
+
+    def _ctx(self, config):
+        return DiffusionLossContext(
+            models={"transformer": torch.nn.Identity()},
+            train_pipeline_config=config,
+            sde_backend=None,
+            scheduler=Namespace(config=Namespace(num_train_timesteps=self.NUM_TRAIN_TIMESTEPS)),
+            args=Namespace(seed=42),
+            forward_dtype=torch.float32,
+            device=torch.device("cpu"),
+        )
+
+    def _batch(self):
+        return [
+            {
+                "x0": torch.zeros(2, 2),
+                "timestep": sigma,
+                "denoising_env": self._Env(),
+                "advantage": 1.0,
+                "nft_num_timesteps": len(self.SIGMAS),
+            }
+            for sigma in self.SIGMAS
+        ]
+
+    def test_raw_sigma_reaches_the_family_hook(self):
+        seen = {}
+
+        class _Recording(_StubConfig):
+            def process_sigma_as_timesteps_input(self, sigmas, *, num_train_timesteps):
+                seen["sigmas"] = sigmas.clone()
+                seen["num_train_timesteps"] = num_train_timesteps
+                return sigmas
+
+        prepare_nft_batch(self._ctx(_Recording()), self._batch())
+        assert torch.equal(seen["sigmas"], torch.tensor(self.SIGMAS))
+        assert seen["num_train_timesteps"] == self.NUM_TRAIN_TIMESTEPS
+
+    def test_sd3_style_family_gets_the_scheduler_range(self):
+        prepared = prepare_nft_batch(self._ctx(_Sd3StyleConfig()), self._batch())
+        assert torch.equal(prepared.timesteps, torch.tensor(self.SIGMAS))
+        assert torch.equal(prepared.timesteps_for_model, torch.tensor(self.SIGMAS) * float(self.NUM_TRAIN_TIMESTEPS))
+
+    def test_qwen_style_family_gets_the_sigma_bit_exactly(self):
+        prepared = prepare_nft_batch(self._ctx(_QwenStyleConfig()), self._batch())
+        # equal, not allclose: the wrong hook would still pass allclose.
+        assert torch.equal(prepared.timesteps_for_model, torch.tensor(self.SIGMAS))
+
+
 class TestEmaShadow:
     def _model(self):
         return torch.nn.Linear(4, 4, bias=False)