fix(nft): fix deterministic mode and add the SD3.5 DiffusionNFT e2e

miles/ray/rollout.py

@@ -51,6 +51,7 @@ def __init__(self, args, pg):
         logger.info("RolloutManager init start")
         self.args = args
         self.pg = pg
+        torch.manual_seed(args.seed)
         from miles.dashboard import hooks
 
         hooks.register_rollout_manager(args)

scripts/run-diffusion-nft-sd3-pickscore.sh

@@ -100,6 +100,7 @@ python -u "${ROOT_DIR}/train_diffusion.py" \
   --prompt-data "${PROMPT_DATA}" \
   --input-key input \
   "${BATCH_ARGS[@]}" \
+  --deterministic-mode \
   --gradient-checkpointing \
   --colocate \
   --use-miles-router \

tests/ci/fixtures/e2e_standards/test_sd3_pickscore_nft_2xGPU.json

@@ -0,0 +1,128 @@
+{
+  "meta": {
+    "commit": "773c686e00709d9bcb66c1cc320d7d83fcbb0cb6",
+    "source": "test_sd3_pickscore_nft_2xGPU.py"
+  },
+  "metrics": {
+    "rollout/reward/raw_mean": [
+      [
+        0,
+        0.74586021900177
+      ],
+      [
+        1,
+        0.769538402557373
+      ]
+    ],
+    "rollout/reward/raw_median": [
+      [
+        0,
+        0.7397080659866333
+      ],
+      [
+        1,
+        0.7668601274490356
+      ]
+    ],
+    "rollout/reward/raw_num_samples": [
+      [
+        0,
+        64.0
+      ],
+      [
+        1,
+        64.0
+      ]
+    ],
+    "rollout/reward/raw_std": [
+      [
+        0,
+        0.04408378154039383
+      ],
+      [
+        1,
+        0.06347957998514175
+      ]
+    ],
+    "train/adv_abs_mean": [
+      [
+        1.0,
+        0.5435407111524708
+      ],
+      [
+        2.0,
+        0.36624840438081363
+      ]
+    ],
+    "train/nft_adv_mean": [
+      [
+        1.0,
+        8.50866652197308e-08
+      ],
+      [
+        2.0,
+        5.76837919652462e-08
+      ]
+    ],
+    "train/nft_loss": [
+      [
+        1.0,
+        14.975260416666666
+      ],
+      [
+        2.0,
+        16.517578125
+      ]
+    ],
+    "train/nft_loss_per_pair": [
+      [
+        1.0,
+        1.6639472113715277
+      ],
+      [
+        2.0,
+        1.8352796766493056
+      ]
+    ],
+    "train/nft_neg_loss": [
+      [
+        1.0,
+        0.3327924940321181
+      ],
+      [
+        2.0,
+        0.3670535617404514
+      ]
+    ],
+    "train/nft_pos_loss": [
+      [
+        1.0,
+        0.3327924940321181
+      ],
+      [
+        2.0,
+        0.3670535617404514
+      ]
+    ],
+    "train/nft_r_mean": [
+      [
+        1.0,
+        0.5000000074505806
+      ],
+      [
+        2.0,
+        0.5000000053809749
+      ]
+    ],
+    "train/nft_t_mean": [
+      [
+        1.0,
+        0.7304325178265572
+      ],
+      [
+        2.0,
+        0.7304325215518475
+      ]
+    ]
+  }
+}

tests/e2e/short/test_sd3_pickscore_nft_2xGPU.py

@@ -0,0 +1,43 @@
+"""E2E: SD3.5-medium PickScore DiffusionNFT, 2 train GPUs colocated with 2 sglang
+rollout engines + 1 dedicated PickScore GPU, 2 rollouts — runs the example script
+itself and checks its metric series against the registered standard
+(tests/ci/fixtures/e2e_standards/). Runs with --deterministic-mode (torch
+deterministic algorithms + NCCL/CUBLAS determinism), so every metric is compared
+strictly, bit for bit.
+
+DiffusionNFT is a dual-policy x0-MSE objective, not a ratio/log-prob one, so the
+train-side metrics here are the NFT branch losses rather than the Flow-GRPO
+`log_prob_*` series (which nft_loss_formula never emits). nft_pos_loss vs
+nft_neg_loss is the NFT analogue of the log-prob drift check: it tracks how far
+the trained policy has moved from the EMA reference.
+
+train/grad_norm is deliberately absent: fp16 forward enables ShardedGradScaler,
+whose initial scale overflows on NFT's much larger loss, so step 1 logs nan and
+nan == nan is False for any standard, with or without a tolerance.
+"""
+
+from tests.ci.e2e_metrics_registry import register_e2e_ci
+
+register_e2e_ci(
+    est_time=600,
+    suite="stage-c-3-gpu-h200",
+    script="scripts/run-diffusion-nft-sd3-pickscore.sh",
+    env={
+        "NUM_ROLLOUT": "2",
+        "CUDA_VISIBLE_DEVICES": "0,1,2",
+    },
+    metrics=[
+        "rollout/reward/raw_num_samples",
+        "rollout/reward/raw_mean",
+        "rollout/reward/raw_median",
+        "rollout/reward/raw_std",
+        "train/nft_loss",
+        "train/nft_loss_per_pair",
+        "train/nft_pos_loss",
+        "train/nft_neg_loss",
+        "train/nft_r_mean",
+        "train/nft_adv_mean",
+        "train/adv_abs_mean",
+        "train/nft_t_mean",
+    ],
+)

tests/fast/backends/fsdp_utils/test_loss_hub_nft.py

@@ -75,6 +75,35 @@ class _Env:
         assert out["train_data"][0]["advantage"] == rewards[0]
         assert out["train_data"][2]["advantage"] == rewards[1]
 
+    def test_shuffled_timesteps_reproduce_under_a_seed(self):
+        # The shuffle draws from the global RNG, so RolloutManager seeds it; without that seed
+        # the pair order differs per run and no bitwise e2e standard is possible.
+        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)
+
+        class _Env:
+            pos_cond_kwargs = {}
+            neg_cond_kwargs = None
+
+        def order():
+            samples = [
+                Sample(index=i, prompt="p", reward=1.0, dit_trajectory=_Traj(), denoising_env=_Env()) for i in range(4)
+            ]
+            out = expand_samples_to_train_pairs(
+                _args(diffusion_nft_shuffle_timesteps=True), samples, [0.0] * 4, [1.0] * 4
+            )
+            return [p["timestep"] for p in out["train_data"]]
+
+        torch.manual_seed(1234)
+        first = order()
+        torch.manual_seed(1234)
+        assert order() == first
+        # Guard the test itself: an unshuffled schedule would pass trivially.
+        assert first != sorted(first, reverse=True)
+
 
 class TestEmaShadow:
     def _model(self):