From 56f31e2b721affa62c552575998189acc91a9d64 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 11 Aug 2026 11:07:04 -0700 Subject: [PATCH 1/4] fix(nft): give NFT its own timestep hook instead of sharing process_timestep_as_input PR #125 routed NFT's prepare through process_timestep_as_input, but NFT's input is not a trajectory timestep. Its pairs carry a sigma in [0, 1] read straight off scheduler.sigmas, so the rescaling runs the opposite direction from flow-GRPO and SFT: the families that pass the timestep through unchanged there have to scale up to the scheduler range here, and qwen_image, which divides there, must pass through here. The old needs_timestep_scaling branch in nft.py encoded exactly that inversion; folding it into the shared hook required pre-multiplying by num_train_timesteps, which qwen_image then divides back out. That round trip is not the identity in float32. Multiplying by 1000 and dividing by 1000 lands off by a ULP on ~2% of sigmas, against a hook whose whole purpose is reproducing the rollout DiT's arithmetic bit-for-bit. Add process_sigma_as_input(sigmas, *, num_train_timesteps): the base scales up to the scheduler range (sd3, wan2_2, ltx), qwen_image passes the sigma through. Every family recovers its pre-#125 tensor exactly. Also drop the now-dead needs_timestep_scaling from ltx.py and reattach the wan2_2 comment that #125 orphaned. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/configs/ltx.py | 1 - .../backends/fsdp_utils/configs/qwen_image.py | 4 + .../configs/train_pipeline_config.py | 9 ++ miles/backends/fsdp_utils/configs/wan2_2.py | 3 +- miles/backends/fsdp_utils/loss_hub/nft.py | 2 +- .../test_train_pipeline_config_registry.py | 30 +++++++ .../backends/fsdp_utils/test_loss_hub_nft.py | 88 ++++++++++++++++++- 7 files changed, 133 insertions(+), 4 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/ltx.py b/miles/backends/fsdp_utils/configs/ltx.py index 3f78d290..f50b4acc 100644 --- a/miles/backends/fsdp_utils/configs/ltx.py +++ b/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 diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 18268af5..ae0318e1 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -62,6 +62,10 @@ class QwenImageTrainPipelineConfig(TrainPipelineConfig): def process_timestep_as_input(self, timesteps): return timesteps / 1000.0 + def process_sigma_as_input(self, sigmas, *, num_train_timesteps): + # NFT's sigma is already the normalized quantity this DiT takes; hand it over untouched. + return sigmas + lora_target_modules = [ "to_q", "to_k", diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 790db46c..d9d00d75 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -106,6 +106,15 @@ 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_input(self, sigmas: torch.Tensor, *, num_train_timesteps: int) -> torch.Tensor: + """The NFT training sigma as this family's DiT takes it. NFT draws its own grid off + ``scheduler.sigmas``, so it starts from sigma in [0, 1] rather than a trajectory + timestep in [0, num_train_timesteps): the families that need no rescaling above need + one here, and vice versa. Separate hook, not a reuse of process_timestep_as_input -- + composing the two round-trips the value through a multiply and a divide that do not + cancel in float32.""" + return sigmas * float(num_train_timesteps) + def compute_noise_pred( self, *, diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index d067a2c5..fcc71c91 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -15,8 +15,9 @@ 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. + # Wan DiT expects raw scheduler timesteps (0..num_train_timesteps), no /1000 scaling: it + # inherits the identity process_timestep_as_input and the scaling process_sigma_as_input. def component_for_timestep(self, timestep: float, num_train_timesteps: int) -> str: if timestep >= self.boundary_ratio * num_train_timesteps: return "transformer" diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index 7ac28fe7..11b38830 100644 --- a/miles/backends/fsdp_utils/loss_hub/nft.py +++ b/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_input(t, num_train_timesteps=num_train_timesteps), model=model, component_name=component_name, guidance_scale=0.0, diff --git a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py index f8156a4d..66f12c43 100644 --- a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py +++ b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py @@ -107,3 +107,33 @@ 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 TestProcessSigmaAsInput: + # What a family hands its DiT for NFT, whose grid is sigma in [0, 1] off scheduler.sigmas + # rather than a trajectory timestep. The rescaling runs the opposite direction from + # process_timestep_as_input, which is why NFT needs its own hook: + # + # sd3, wan2_2 sigma * N the DiT wants the raw scheduler range back + # qwen_image sigma already the normalized quantity the DiT takes + NUM_TRAIN_TIMESTEPS = 1000 + # 0.8474337458610535 is one of the ~2% of float32 sigmas where a multiply by 1000 followed + # by a divide by 1000 does not land back on the input. + 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_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_input( + QwenImageTrainPipelineConfig, self.SIGMAS, num_train_timesteps=self.NUM_TRAIN_TIMESTEPS + ) + # Bit-exact, not merely close: routing this through process_timestep_as_input instead + # would multiply then divide by 1000 and drift a ULP on the first element. + assert torch.equal(out, self.SIGMAS) + round_tripped = QwenImageTrainPipelineConfig.process_timestep_as_input( + QwenImageTrainPipelineConfig, self.SIGMAS * float(self.NUM_TRAIN_TIMESTEPS) + ) + assert not torch.equal(round_tripped, self.SIGMAS) diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py index 5b6406f8..79d36f6f 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/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,89 @@ class _Env: raise AssertionError("expected ValueError for missing dit_trajectory.sigmas") +class _StubConfig: + """Cond plumbing stubbed out; both timestep hooks are bound so that a prepare hook wired to + the wrong one fails on the numbers rather than on a missing attribute.""" + + process_timestep_as_input = TrainPipelineConfig.process_timestep_as_input + process_sigma_as_input = TrainPipelineConfig.process_sigma_as_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): + # sd3, wan2_2 and ltx all inherit both base hooks. + pass + + +class _QwenStyleConfig(_StubConfig): + process_timestep_as_input = QwenImageTrainPipelineConfig.process_timestep_as_input + process_sigma_as_input = QwenImageTrainPipelineConfig.process_sigma_as_input + + +class TestPrepareNftBatch: + # NFT's pair timestep is a sigma in [0, 1] off scheduler.sigmas, so prepare has to route it + # through process_sigma_as_input. Sending it through process_timestep_as_input instead means + # scaling up by N only for the normalizing families to divide it straight back out. + NUM_TRAIN_TIMESTEPS = 1000 + 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_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()) + # The un-rescaled sigma, not sigma * N: the family decides which direction to go. + 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()) + # Bit-exact: the multiply-then-divide round trip drifts a ULP on the first sigma. + assert torch.equal(prepared.timesteps_for_model, torch.tensor(self.SIGMAS)) + + class TestEmaShadow: def _model(self): return torch.nn.Linear(4, 4, bias=False) From 39a429778fb1588893e3272da5415ff4c00dc9c6 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 11 Aug 2026 11:14:23 -0700 Subject: [PATCH 2/4] docs(configs): record that qwen_image's 1000 is the model normalizer, not the scheduler range sglang-d's DiT divides by the literal 1000 (runtime/models/dits/qwen_image.py) while its flow-match scheduler converts sigma with num_train_timesteps. The two constants are different facts that happen to share a value; note which is which so neither hook gets 'fixed' into reading the other. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/configs/qwen_image.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index ae0318e1..8c2678ee 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -59,11 +59,18 @@ class QwenImageTrainPipelineConfig(TrainPipelineConfig): hf_ckpt_name_patterns = ("qwen-image",) cfg_batching = False + # 1000 is the model's own normalizer, not the scheduler range: sglang-d's DiT divides by the + # literal 1000 (runtime/models/dits/qwen_image.py) before the shared Timesteps(scale=1000) + # undoes it, and diffusers leaves that division to the caller. Reading num_train_timesteps + # here instead would diverge from the rollout the moment a scheduler ran with any other range. def process_timestep_as_input(self, timesteps): return timesteps / 1000.0 def process_sigma_as_input(self, sigmas, *, num_train_timesteps): - # NFT's sigma is already the normalized quantity this DiT takes; hand it over untouched. + # NFT's sigma is already the normalized quantity this DiT takes, because the scheduler + # range and the model normalizer are both 1000: rollout sends sigma * num_train_timesteps + # and the DiT divides by 1000. Only exact while those two agree, which they do for every + # qwen_image schedule we run; a range other than 1000 would need sigma * N / 1000. return sigmas lora_target_modules = [ From 21e2370c8d6af9a0f025449fbda686fe3f92e9e0 Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 11 Aug 2026 11:22:09 -0700 Subject: [PATCH 3/4] style(nft): tighten the timestep-hook comments Move the full rationale into the process_sigma_as_input docstring and leave one-line constraints at the call sites; drop the narrating lines and the wan2_2 comment that outlived the attribute it annotated and had drifted onto component_for_timestep. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/configs/qwen_image.py | 13 +++++-------- .../fsdp_utils/configs/train_pipeline_config.py | 13 +++++++------ miles/backends/fsdp_utils/configs/wan2_2.py | 2 -- .../test_train_pipeline_config_registry.py | 15 +++++++-------- .../fast/backends/fsdp_utils/test_loss_hub_nft.py | 14 ++++++-------- 5 files changed, 25 insertions(+), 32 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 8c2678ee..6133f47c 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -59,18 +59,15 @@ class QwenImageTrainPipelineConfig(TrainPipelineConfig): hf_ckpt_name_patterns = ("qwen-image",) cfg_batching = False - # 1000 is the model's own normalizer, not the scheduler range: sglang-d's DiT divides by the - # literal 1000 (runtime/models/dits/qwen_image.py) before the shared Timesteps(scale=1000) - # undoes it, and diffusers leaves that division to the caller. Reading num_train_timesteps - # here instead would diverge from the rollout the moment a scheduler ran with any other range. + # 1000 is the model's normalizer, not the scheduler range: sglang-d's DiT divides by the + # literal 1000 and diffusers leaves that division to its caller. Reading num_train_timesteps + # here would diverge from the rollout under any other scheduler range. def process_timestep_as_input(self, timesteps): return timesteps / 1000.0 def process_sigma_as_input(self, sigmas, *, num_train_timesteps): - # NFT's sigma is already the normalized quantity this DiT takes, because the scheduler - # range and the model normalizer are both 1000: rollout sends sigma * num_train_timesteps - # and the DiT divides by 1000. Only exact while those two agree, which they do for every - # qwen_image schedule we run; a range other than 1000 would need sigma * N / 1000. + # Exact only because the two 1000s above coincide: rollout sends sigma * num_train_timesteps + # and the DiT divides by 1000. A range other than 1000 would need sigma * N / 1000. return sigmas lora_target_modules = [ diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index d9d00d75..ba0293b6 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -107,12 +107,13 @@ def process_timestep_as_input(self, timesteps: torch.Tensor) -> torch.Tensor: return timesteps def process_sigma_as_input(self, sigmas: torch.Tensor, *, num_train_timesteps: int) -> torch.Tensor: - """The NFT training sigma as this family's DiT takes it. NFT draws its own grid off - ``scheduler.sigmas``, so it starts from sigma in [0, 1] rather than a trajectory - timestep in [0, num_train_timesteps): the families that need no rescaling above need - one here, and vice versa. Separate hook, not a reuse of process_timestep_as_input -- - composing the two round-trips the value through a multiply and a divide that do not - cancel in float32.""" + """The NFT training sigma as this family's DiT takes it. + + NFT draws its grid straight off ``scheduler.sigmas``, so it starts from sigma in [0, 1] + rather than a trajectory timestep -- the rescaling runs the opposite direction from + ``process_timestep_as_input``, and a family that is the identity there scales here. + Kept separate rather than pre-multiplying into that hook: for a family that divides, + the composition is a multiply and a divide that do not cancel in float32.""" return sigmas * float(num_train_timesteps) def compute_noise_pred( diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index fcc71c91..4775804e 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -16,8 +16,6 @@ class Wan2_2TrainPipelineConfig(TrainPipelineConfig): # ("transformer_2") the rest. boundary_ratio = 0.875 - # Wan DiT expects raw scheduler timesteps (0..num_train_timesteps), no /1000 scaling: it - # inherits the identity process_timestep_as_input and the scaling process_sigma_as_input. def component_for_timestep(self, timestep: float, num_train_timesteps: int) -> str: if timestep >= self.boundary_ratio * num_train_timesteps: return "transformer" diff --git a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py index 66f12c43..e2bdf297 100644 --- a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py +++ b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py @@ -110,15 +110,15 @@ def test_qwen_image_divides_by_the_model_normalizer(self): class TestProcessSigmaAsInput: - # What a family hands its DiT for NFT, whose grid is sigma in [0, 1] off scheduler.sigmas - # rather than a trajectory timestep. The rescaling runs the opposite direction from - # process_timestep_as_input, which is why NFT needs its own hook: + # The NFT counterpart, whose input is sigma off scheduler.sigmas. Each family rescales the + # opposite way from above: identity there means sigma * N here, and a divide there means + # pass-through here. # - # sd3, wan2_2 sigma * N the DiT wants the raw scheduler range back + # sd3, wan2_2 sigma * N the DiT wants the scheduler range back # qwen_image sigma already the normalized quantity the DiT takes NUM_TRAIN_TIMESTEPS = 1000 - # 0.8474337458610535 is one of the ~2% of float32 sigmas where a multiply by 1000 followed - # by a divide by 1000 does not land back on the input. + # 0.8474337458610535 is one of the ~2% of float32 sigmas that a multiply by 1000 followed by + # a divide by 1000 does not return unchanged. SIGMAS = torch.tensor([0.8474337458610535, 0.5]) @pytest.mark.parametrize("config_cls", [SD3TrainPipelineConfig, Wan2_2TrainPipelineConfig]) @@ -130,9 +130,8 @@ def test_qwen_image_passes_the_sigma_through(self): out = QwenImageTrainPipelineConfig.process_sigma_as_input( QwenImageTrainPipelineConfig, self.SIGMAS, num_train_timesteps=self.NUM_TRAIN_TIMESTEPS ) - # Bit-exact, not merely close: routing this through process_timestep_as_input instead - # would multiply then divide by 1000 and drift a ULP on the first element. assert torch.equal(out, self.SIGMAS) + # The composition this hook exists to avoid; asserted so the equal() above keeps its teeth. round_tripped = QwenImageTrainPipelineConfig.process_timestep_as_input( QwenImageTrainPipelineConfig, self.SIGMAS * float(self.NUM_TRAIN_TIMESTEPS) ) diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py index 79d36f6f..ff595d71 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py @@ -124,8 +124,8 @@ class _Env: class _StubConfig: - """Cond plumbing stubbed out; both timestep hooks are bound so that a prepare hook wired to - the wrong one fails on the numbers rather than on a missing attribute.""" + """Cond plumbing stubbed out. Binds both timestep hooks so that a prepare wired to the wrong + one fails on the numbers rather than on a missing attribute.""" process_timestep_as_input = TrainPipelineConfig.process_timestep_as_input process_sigma_as_input = TrainPipelineConfig.process_sigma_as_input @@ -138,7 +138,7 @@ def collate_cond_for_sample_batch(self, per_sample_cond_kwargs, device, pad_to_l class _Sd3StyleConfig(_StubConfig): - # sd3, wan2_2 and ltx all inherit both base hooks. + # Stands in for sd3, wan2_2 and ltx, which all inherit both base hooks. pass @@ -148,10 +148,9 @@ class _QwenStyleConfig(_StubConfig): class TestPrepareNftBatch: - # NFT's pair timestep is a sigma in [0, 1] off scheduler.sigmas, so prepare has to route it - # through process_sigma_as_input. Sending it through process_timestep_as_input instead means - # scaling up by N only for the normalizing families to divide it straight back out. NUM_TRAIN_TIMESTEPS = 1000 + # 0.8474337458610535 is one of the ~2% of float32 sigmas that a multiply by 1000 followed by + # a divide by 1000 does not return unchanged. SIGMAS = [0.8474337458610535, 0.5] class _Env: @@ -191,7 +190,6 @@ def process_sigma_as_input(self, sigmas, *, num_train_timesteps): return sigmas prepare_nft_batch(self._ctx(_Recording()), self._batch()) - # The un-rescaled sigma, not sigma * N: the family decides which direction to go. assert torch.equal(seen["sigmas"], torch.tensor(self.SIGMAS)) assert seen["num_train_timesteps"] == self.NUM_TRAIN_TIMESTEPS @@ -202,7 +200,7 @@ def test_sd3_style_family_gets_the_scheduler_range(self): def test_qwen_style_family_gets_the_sigma_bit_exactly(self): prepared = prepare_nft_batch(self._ctx(_QwenStyleConfig()), self._batch()) - # Bit-exact: the multiply-then-divide round trip drifts a ULP on the first sigma. + # equal, not allclose: routing through process_timestep_as_input would still pass allclose. assert torch.equal(prepared.timesteps_for_model, torch.tensor(self.SIGMAS)) From bc29eb134bc0f75aa06b89eb63a06cb6bfd5ca5b Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 11 Aug 2026 11:36:02 -0700 Subject: [PATCH 4/4] refactor(nft): rename to process_sigma_as_timesteps_input and cut the comments back The name now carries what the comments were spelling out: sigma in, timesteps input out. Leaves one line per non-local constraint -- why the hook is separate, and why qwen_image's 1000 must not become num_train_timesteps. Co-Authored-By: Claude Fable 5 --- .../backends/fsdp_utils/configs/qwen_image.py | 10 ++++------ .../configs/train_pipeline_config.py | 12 ++++------- miles/backends/fsdp_utils/loss_hub/nft.py | 2 +- .../test_train_pipeline_config_registry.py | 20 ++++++++----------- .../backends/fsdp_utils/test_loss_hub_nft.py | 16 +++++++-------- 5 files changed, 24 insertions(+), 36 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/qwen_image.py b/miles/backends/fsdp_utils/configs/qwen_image.py index 6133f47c..37234935 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/miles/backends/fsdp_utils/configs/qwen_image.py @@ -59,15 +59,13 @@ 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 - # literal 1000 and diffusers leaves that division to its caller. Reading num_train_timesteps - # here would diverge from the rollout under any other scheduler range. + # 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_input(self, sigmas, *, num_train_timesteps): - # Exact only because the two 1000s above coincide: rollout sends sigma * num_train_timesteps - # and the DiT divides by 1000. A range other than 1000 would need sigma * N / 1000. + 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 = [ diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index ba0293b6..e55df3fe 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -106,14 +106,10 @@ 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_input(self, sigmas: torch.Tensor, *, num_train_timesteps: int) -> torch.Tensor: - """The NFT training sigma as this family's DiT takes it. - - NFT draws its grid straight off ``scheduler.sigmas``, so it starts from sigma in [0, 1] - rather than a trajectory timestep -- the rescaling runs the opposite direction from - ``process_timestep_as_input``, and a family that is the identity there scales here. - Kept separate rather than pre-multiplying into that hook: for a family that divides, - the composition is a multiply and a divide that do not cancel in float32.""" + 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( diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index 11b38830..a8a0818e 100644 --- a/miles/backends/fsdp_utils/loss_hub/nft.py +++ b/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_sigma_as_input(t, num_train_timesteps=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, diff --git a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py index e2bdf297..c517efad 100644 --- a/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py +++ b/tests/fast/backends/fsdp_utils/configs/test_train_pipeline_config_registry.py @@ -109,29 +109,25 @@ def test_qwen_image_divides_by_the_model_normalizer(self): assert torch.equal(out, self.TIMESTEPS / 1000.0) -class TestProcessSigmaAsInput: - # The NFT counterpart, whose input is sigma off scheduler.sigmas. Each family rescales the - # opposite way from above: identity there means sigma * N here, and a divide there means - # pass-through here. - # - # sd3, wan2_2 sigma * N the DiT wants the scheduler range back - # qwen_image sigma already the normalized quantity the DiT takes +class TestProcessSigmaAsTimestepsInput: + # The NFT counterpart: each family rescales the opposite way from above. NUM_TRAIN_TIMESTEPS = 1000 - # 0.8474337458610535 is one of the ~2% of float32 sigmas that a multiply by 1000 followed by - # a divide by 1000 does not return unchanged. + # 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_input(config_cls, self.SIGMAS, num_train_timesteps=self.NUM_TRAIN_TIMESTEPS) + 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_input( + out = QwenImageTrainPipelineConfig.process_sigma_as_timesteps_input( QwenImageTrainPipelineConfig, self.SIGMAS, num_train_timesteps=self.NUM_TRAIN_TIMESTEPS ) assert torch.equal(out, self.SIGMAS) - # The composition this hook exists to avoid; asserted so the equal() above keeps its teeth. + # Asserted so the equal() above keeps its teeth. round_tripped = QwenImageTrainPipelineConfig.process_timestep_as_input( QwenImageTrainPipelineConfig, self.SIGMAS * float(self.NUM_TRAIN_TIMESTEPS) ) diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py index ff595d71..9132b800 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py @@ -124,11 +124,11 @@ class _Env: class _StubConfig: - """Cond plumbing stubbed out. Binds both timestep hooks so that a prepare wired to the wrong - one fails on the numbers rather than on a missing attribute.""" + """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_input = TrainPipelineConfig.process_sigma_as_input + process_sigma_as_timesteps_input = TrainPipelineConfig.process_sigma_as_timesteps_input def prepare_cond_kwargs(self, cond, device): return {} @@ -138,19 +138,17 @@ def collate_cond_for_sample_batch(self, per_sample_cond_kwargs, device, pad_to_l class _Sd3StyleConfig(_StubConfig): - # Stands in for sd3, wan2_2 and ltx, which all inherit both base hooks. pass class _QwenStyleConfig(_StubConfig): process_timestep_as_input = QwenImageTrainPipelineConfig.process_timestep_as_input - process_sigma_as_input = QwenImageTrainPipelineConfig.process_sigma_as_input + process_sigma_as_timesteps_input = QwenImageTrainPipelineConfig.process_sigma_as_timesteps_input class TestPrepareNftBatch: NUM_TRAIN_TIMESTEPS = 1000 - # 0.8474337458610535 is one of the ~2% of float32 sigmas that a multiply by 1000 followed by - # a divide by 1000 does not return unchanged. + # 0.8474... does not survive a multiply then divide by 1000 in fp32. SIGMAS = [0.8474337458610535, 0.5] class _Env: @@ -184,7 +182,7 @@ def test_raw_sigma_reaches_the_family_hook(self): seen = {} class _Recording(_StubConfig): - def process_sigma_as_input(self, sigmas, *, num_train_timesteps): + def process_sigma_as_timesteps_input(self, sigmas, *, num_train_timesteps): seen["sigmas"] = sigmas.clone() seen["num_train_timesteps"] = num_train_timesteps return sigmas @@ -200,7 +198,7 @@ def test_sd3_style_family_gets_the_scheduler_range(self): def test_qwen_style_family_gets_the_sigma_bit_exactly(self): prepared = prepare_nft_batch(self._ctx(_QwenStyleConfig()), self._batch()) - # equal, not allclose: routing through process_timestep_as_input would still pass allclose. + # equal, not allclose: the wrong hook would still pass allclose. assert torch.equal(prepared.timesteps_for_model, torch.tensor(self.SIGMAS))