diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py index a8a0818e..89b4332c 100644 --- a/miles/backends/fsdp_utils/loss_hub/nft.py +++ b/miles/backends/fsdp_utils/loss_hub/nft.py @@ -5,6 +5,7 @@ import torch from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch +from miles.utils.hash_utils import stable_hash from miles.utils.metric_buffer import MetricBuffer @@ -41,7 +42,10 @@ def prepare_nft_batch( num_train_timesteps = ctx.scheduler.config.num_train_timesteps - xt = corrupt(x0, t, sample_noise(x0)) + noise_generator = torch.Generator(device=device).manual_seed( + stable_hash("nft_corrupt", int(ctx.args.seed), ctx.rollout_id, ctx.microbatch_id, ctx.dp_rank) + ) + xt = corrupt(x0, t, sample_noise(x0, generator=noise_generator)) return PreparedBatch( latents=xt, timesteps=t, diff --git a/miles/backends/fsdp_utils/loss_hub/sft.py b/miles/backends/fsdp_utils/loss_hub/sft.py index dc5c80ae..00daaddd 100644 --- a/miles/backends/fsdp_utils/loss_hub/sft.py +++ b/miles/backends/fsdp_utils/loss_hub/sft.py @@ -2,22 +2,14 @@ from __future__ import annotations -import hashlib - import torch import torch.nn as nn from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch +from miles.utils.hash_utils import stable_hash from miles.utils.metric_buffer import MetricBuffer -def _seed(scope: str, *parts: int) -> int: - """Create a stable seed for one named random stream.""" - payload = ":".join([scope, *(str(part) for part in parts)]).encode() - digest = hashlib.blake2b(payload, digest_size=8).digest() - return int.from_bytes(digest, "little") & (2**63 - 1) - - def sample_grid_indices( ctx: DiffusionLossContext, bsz: int, @@ -46,7 +38,7 @@ def sample_grid_indices( num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps) components = [config.component_for_timestep(float(t), num_train_timesteps) for t in ctx.scheduler.timesteps] expert_generator = torch.Generator().manual_seed( - _seed("expert", int(ctx.args.seed), ctx.rollout_id, ctx.microbatch_id) + stable_hash("expert", int(ctx.args.seed), ctx.rollout_id, ctx.microbatch_id) ) component_name = components[int(torch.randint(num_grid, (1,), generator=expert_generator))] model = ctx.models[component_name] @@ -72,7 +64,7 @@ def prepare_sft_batch( x0 = torch.stack([pair["latent"] for pair in batch]).to(device=device, dtype=torch.float32) sample_generator = torch.Generator(device=device).manual_seed( - _seed("sample", int(ctx.args.seed), ctx.rollout_id, ctx.microbatch_id, ctx.dp_rank) + stable_hash("sample", int(ctx.args.seed), ctx.rollout_id, ctx.microbatch_id, ctx.dp_rank) ) component_name, model, idx = sample_grid_indices(ctx, bsz, generator=sample_generator) timesteps = ctx.scheduler.timesteps[idx].to(dtype=torch.float32) diff --git a/miles/ray/data_conversion_hub/nft.py b/miles/ray/data_conversion_hub/nft.py index f871d9fa..3364c490 100644 --- a/miles/ray/data_conversion_hub/nft.py +++ b/miles/ray/data_conversion_hub/nft.py @@ -5,6 +5,7 @@ import torch +from miles.utils.hash_utils import stable_hash from miles.utils.train_data_utils import scheduler_meta_from_samples from miles.utils.types import Sample @@ -60,11 +61,21 @@ def expand_samples_to_train_pairs( num_timesteps = int(sigmas.numel()) train_data: list[dict[str, Any]] = [] - for sample, adv, raw in zip(samples, rewards, raw_rewards, strict=True): + for position, (sample, adv, raw) in enumerate(zip(samples, rewards, raw_rewards, strict=True)): if sample.denoising_env is None: raise ValueError(f"sample {sample.index} missing denoising_env") x0 = _clean_x0_from_sample(sample) - sample_sigmas = sigmas[torch.randperm(num_timesteps)] if args.diffusion_nft_shuffle_timesteps else sigmas + # Keyed on the sample's global index, which the data source advances across rollouts, + # so each sample draws its own permutation and the run reproduces. + stream = sample.index if sample.index is not None else position + shuffle_generator = torch.Generator().manual_seed( + stable_hash("nft_sigma_shuffle", int(args.seed), int(stream)) + ) + sample_sigmas = ( + sigmas[torch.randperm(num_timesteps, generator=shuffle_generator)] + if args.diffusion_nft_shuffle_timesteps + else sigmas + ) for t in sample_sigmas.tolist(): train_data.append( { diff --git a/miles/utils/hash_utils.py b/miles/utils/hash_utils.py new file mode 100644 index 00000000..cbf5cc2b --- /dev/null +++ b/miles/utils/hash_utils.py @@ -0,0 +1,18 @@ +"""Hashes that hold still across processes and runs.""" + +from __future__ import annotations + +import hashlib + + +def stable_hash(*parts: object) -> int: + """A 63-bit hash of ``parts``, identical in every process and every run. + + ``hash()`` is not: PYTHONHASHSEED randomises it per process, so anything derived from + it -- an RNG seed, a shard assignment, a cache key compared across ranks -- silently + stops agreeing. Callers name the parts that make a value distinct and get the same + integer back forever. + """ + payload = ":".join(str(part) for part in parts).encode() + digest = hashlib.blake2b(payload, digest_size=8).digest() + return int.from_bytes(digest, "little") & (2**63 - 1) diff --git a/scripts/run_diffusion_nft_sd3_pickscore.py b/scripts/run_diffusion_nft_sd3_pickscore.py index 6d4f559a..311cbc02 100644 --- a/scripts/run_diffusion_nft_sd3_pickscore.py +++ b/scripts/run_diffusion_nft_sd3_pickscore.py @@ -129,6 +129,7 @@ def execute(args: ScriptArgs, data_dir: str) -> None: "--rollout-num-gpus-per-engine 1 " f"--num-gpus-per-node {2 if args.smoke else 3} " "--colocate " + "--deterministic-mode " ) U.execute_train( diff --git a/tests/ci/e2e_metrics_registry.py b/tests/ci/e2e_metrics_registry.py index 9de73126..c1a17ad9 100644 --- a/tests/ci/e2e_metrics_registry.py +++ b/tests/ci/e2e_metrics_registry.py @@ -106,6 +106,11 @@ def load_series(jsonl_path: str | Path, metrics: list[str]) -> dict[str, list[li def _values_match(got: float, want: float, tol: dict | None) -> bool: + # A recorded NaN is a value like any other -- fp16 runs emit grad_norm=nan on the step the + # grad scaler overflows its init scale, deterministically. Neither == nor isclose matches + # NaN against itself, so without this a standard containing one could never pass. + if math.isnan(got) or math.isnan(want): + return math.isnan(got) and math.isnan(want) if tol is None: return got == want return math.isclose(got, want, rel_tol=tol.get("rtol", 0.0), abs_tol=tol.get("atol", 0.0)) diff --git a/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json b/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json new file mode 100644 index 00000000..2e2038f1 --- /dev/null +++ b/tests/ci/fixtures/e2e_standards/test_sd3_nft_pickscore_3xGPU.json @@ -0,0 +1,188 @@ +{ + "meta": { + "commit": "00eb32246d1bc8182383f4c02376a1303cbb6f0e", + "source": "test_sd3_nft_pickscore_3xGPU.py" + }, + "metrics": { + "rollout/reward/raw_mean": [ + [ + 0, + 0.74586021900177 + ], + [ + 1, + 0.769538402557373 + ], + [ + 2, + 0.748466968536377 + ], + [ + 3, + 0.7546428442001343 + ] + ], + "rollout/reward/raw_median": [ + [ + 0, + 0.7397080659866333 + ], + [ + 1, + 0.7668601274490356 + ], + [ + 2, + 0.744907796382904 + ], + [ + 3, + 0.7552561163902283 + ] + ], + "rollout/reward/raw_num_samples": [ + [ + 0, + 64.0 + ], + [ + 1, + 64.0 + ], + [ + 2, + 64.0 + ], + [ + 3, + 64.0 + ] + ], + "rollout/reward/raw_std": [ + [ + 0, + 0.04408378154039383 + ], + [ + 1, + 0.06347957998514175 + ], + [ + 2, + 0.03868388384580612 + ], + [ + 3, + 0.04985775798559189 + ] + ], + "train/grad_norm": [ + [ + 1.0, + NaN + ], + [ + 2.0, + 0.024831360206007957 + ], + [ + 3.0, + 0.027672801166772842 + ], + [ + 4.0, + 0.06675068289041519 + ] + ], + "train/nft_loss": [ + [ + 1.0, + 15.153754340277779 + ], + [ + 2.0, + 16.450520833333332 + ], + [ + 3.0, + 17.195638020833332 + ], + [ + 4.0, + 19.235731336805557 + ] + ], + "train/nft_neg_loss": [ + [ + 1.0, + 0.3367631700303819 + ], + [ + 2.0, + 0.365570068359375 + ], + [ + 3.0, + 0.3821360270182292 + ], + [ + 4.0, + 0.4274359809027778 + ] + ], + "train/nft_pos_loss": [ + [ + 1.0, + 0.3367631700303819 + ], + [ + 2.0, + 0.365570068359375 + ], + [ + 3.0, + 0.38216400146484375 + ], + [ + 4.0, + 0.42747921413845485 + ] + ], + "train/nft_r_mean": [ + [ + 1.0, + 0.5000000074505806 + ], + [ + 2.0, + 0.5000000053809749 + ], + [ + 3.0, + 0.5000000260770321 + ], + [ + 4.0, + 0.4999999875823657 + ] + ], + "train/nft_t_mean": [ + [ + 1.0, + 0.7304325223796897 + ], + [ + 2.0, + 0.7304325221727291 + ], + [ + 3.0, + 0.7304325236214532 + ], + [ + 4.0, + 0.7304325207240052 + ] + ] + } +} diff --git a/tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py b/tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py new file mode 100644 index 00000000..17e1c6f0 --- /dev/null +++ b/tests/e2e/short/test_sd3_nft_pickscore_3xGPU.py @@ -0,0 +1,38 @@ +"""E2E: SD3.5-medium DiffusionNFT with PickScore, 3-GPU (2 colocated FSDP DP=2 + sglang +rollout engines, 1 dedicated reward GPU) — runs the example script's real configuration, +not a reduced one, and checks its metric series against the registered standard +(tests/ci/fixtures/e2e_standards/). Runs with --deterministic-mode, so every metric is +compared strictly, bit for bit. + +Only --num-rollout is cut down, 100 -> 4, which is the shortest run that still reaches the +behaviour worth guarding. Step 1 overflows the fp16 grad scaler's 65536 init scale and is +skipped, so the weights do not move; step 2 is the first that lands, so through it the EMA +reference is still identical to the policy and NFT's two loss branches are algebraically +one value. They separate at step 3. Two rollouts would exercise none of that. + +The NFT-side metrics are what the GRPO e2e cannot cover: NFT has no log-prob ratio, its +loss is the two-branch x0-MSE, and nft_t_mean tracks the sigma grid the prepare hook feeds +the DiT. +""" + +from tests.ci.e2e_metrics_registry import register_e2e_ci + +register_e2e_ci( + est_time=900, + suite="stage-c-3-gpu-h200", + script="scripts/run_diffusion_nft_sd3_pickscore.py", + args=["--num-rollout", "4"], + labels=["e2e"], + metrics=[ + "rollout/reward/raw_num_samples", + "rollout/reward/raw_mean", + "rollout/reward/raw_median", + "rollout/reward/raw_std", + "train/grad_norm", + "train/nft_loss", + "train/nft_pos_loss", + "train/nft_neg_loss", + "train/nft_r_mean", + "train/nft_t_mean", + ], +) 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 9132b800..87872b39 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py @@ -26,6 +26,7 @@ def _args(**overrides): reward_key=None, diffusion_nft_timestep_fraction=1.0, diffusion_nft_shuffle_timesteps=False, + seed=42, custom_prepare_train_batch_path=None, custom_loss_function_path=None, ) @@ -155,7 +156,7 @@ class _Env: pos_cond_kwargs = None neg_cond_kwargs = None - def _ctx(self, config): + def _ctx(self, config, microbatch_id=0): return DiffusionLossContext( models={"transformer": torch.nn.Identity()}, train_pipeline_config=config, @@ -164,6 +165,7 @@ def _ctx(self, config): args=Namespace(seed=42), forward_dtype=torch.float32, device=torch.device("cpu"), + microbatch_id=microbatch_id, ) def _batch(self): @@ -202,6 +204,61 @@ def test_qwen_style_family_gets_the_sigma_bit_exactly(self): assert torch.equal(prepared.timesteps_for_model, torch.tensor(self.SIGMAS)) +class TestNftDeterminism: + # NFT draws two random streams -- the sigma permutation in the converter and the + # corruption noise in prepare. Both used the global RNG, so two runs of the same + # configuration trained on different data and no metric reproduced. + def _traj(self): + class _Traj: + def __init__(self): + self.timesteps = torch.tensor([999.0, 750.0, 500.0, 250.0, 0.0]) + self.sigmas = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + self.latents = torch.zeros(5, 2, 2) + + return _Traj() + + def _samples(self): + class _Env: + pos_cond_kwargs = None + neg_cond_kwargs = None + + return [ + Sample(index=i, prompt=p, reward=r, dit_trajectory=self._traj(), denoising_env=_Env()) + for i, (p, r) in enumerate([("a", 1.0), ("b", 3.0)]) + ] + + def test_sigma_shuffle_reproduces(self): + args = _args(diffusion_nft_shuffle_timesteps=True) + first = expand_samples_to_train_pairs(args, self._samples(), [-1.0, 1.0], [1.0, 3.0]) + second = expand_samples_to_train_pairs(args, self._samples(), [-1.0, 1.0], [1.0, 3.0]) + got = [p["timestep"] for p in first["train_data"]] + assert got == [p["timestep"] for p in second["train_data"]] + # Shuffled, not just handed back in scheduler order. + assert got[: len(got) // 2] != sorted(got[: len(got) // 2], reverse=True) + + def test_each_sample_draws_its_own_permutation(self): + args = _args(diffusion_nft_shuffle_timesteps=True) + out = expand_samples_to_train_pairs(args, self._samples(), [-1.0, 1.0], [1.0, 3.0]) + per_sample = {} + for pair in out["train_data"]: + per_sample.setdefault(pair["sample_index"], []).append(pair["timestep"]) + assert len(per_sample) == 2 + assert list(per_sample.values())[0] != list(per_sample.values())[1] + + def test_corruption_noise_reproduces(self): + ctx = TestPrepareNftBatch()._ctx(_Sd3StyleConfig()) + batch = TestPrepareNftBatch()._batch() + first = prepare_nft_batch(ctx, batch) + second = prepare_nft_batch(ctx, batch) + assert torch.equal(first.latents, second.latents) + + def test_a_different_microbatch_draws_different_noise(self): + harness = TestPrepareNftBatch() + first = prepare_nft_batch(harness._ctx(_Sd3StyleConfig()), harness._batch()) + second = prepare_nft_batch(harness._ctx(_Sd3StyleConfig(), microbatch_id=1), harness._batch()) + assert not torch.equal(first.latents, second.latents) + + class TestEmaShadow: def _model(self): return torch.nn.Linear(4, 4, bias=False) diff --git a/tests/fast/ci/test_e2e_metrics_registry.py b/tests/fast/ci/test_e2e_metrics_registry.py index 7418e880..932c3615 100644 --- a/tests/fast/ci/test_e2e_metrics_registry.py +++ b/tests/fast/ci/test_e2e_metrics_registry.py @@ -58,6 +58,24 @@ def test_strict_mismatch_fails_but_tolerance_passes(sandbox, monkeypatch): reg.check_or_update("test_foo.py", drifted, ["m"], tolerances={"m": {"atol": 1e-3}}) +def test_recorded_nan_matches_itself_but_not_a_number(sandbox, monkeypatch): + # fp16 runs emit grad_norm=nan on the step the scaler overflows its init scale. + nan_run = sandbox / "nan.jsonl" + finite_run = sandbox / "finite.jsonl" + _write_jsonl(nan_run, [{"step": 1, "m": float("nan")}, {"step": 2, "m": 0.25}]) + _write_jsonl(finite_run, [{"step": 1, "m": 0.5}, {"step": 2, "m": 0.25}]) + monkeypatch.setenv("MILES_E2E_METRICS_UPDATE", "1") + reg.check_or_update("test_foo.py", nan_run, ["m"]) + monkeypatch.delenv("MILES_E2E_METRICS_UPDATE") + reg.check_or_update("test_foo.py", nan_run, ["m"]) + # A run that stops overflowing is a real change, not a match. + with pytest.raises(AssertionError, match="strict"): + reg.check_or_update("test_foo.py", finite_run, ["m"]) + # Nor does a tolerance let a number pass against a recorded NaN. + with pytest.raises(AssertionError, match="tol"): + reg.check_or_update("test_foo.py", finite_run, ["m"], tolerances={"m": {"atol": 1e9}}) + + def test_series_shape_mismatches_fail(sandbox, monkeypatch): std = sandbox / "std.jsonl" _write_jsonl(std, [{"step": 1, "m": 0.5}, {"step": 2, "m": 0.25}])