From b8b5ac855b248a974da04c90f5d03ebb56a912ad Mon Sep 17 00:00:00 2001 From: rockdu Date: Tue, 11 Aug 2026 14:42:50 -0700 Subject: [PATCH] fix(nft): seed the sigma shuffle and the corruption noise so NFT runs reproduce --- miles/backends/fsdp_utils/loss_hub/nft.py | 6 +- miles/backends/fsdp_utils/loss_hub/sft.py | 14 +---- miles/ray/data_conversion_hub/nft.py | 15 ++++- miles/utils/hash_utils.py | 18 ++++++ .../backends/fsdp_utils/test_loss_hub_nft.py | 59 ++++++++++++++++++- 5 files changed, 97 insertions(+), 15 deletions(-) create mode 100644 miles/utils/hash_utils.py 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/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)