fix(nft): seed the sigma shuffle and the corruption noise so NFT runs reproduce

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,

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)

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(
                 {

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)

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)