fix: land unmerged stacked NFT PRs

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)

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(

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))

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
+      ]
+    ]
+  }
+}

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",
+    ],
+)