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..37234935 100644 --- a/miles/backends/fsdp_utils/configs/qwen_image.py +++ b/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", diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 790db46c..e55df3fe 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/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, *, diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index d067a2c5..4775804e 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/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: diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index 7ac28fe7..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_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, 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..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 @@ -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) 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..9132b800 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,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)