diff --git a/miles/ray/rollout.py b/miles/ray/rollout.py index 2b49513a..b6a665f6 100644 --- a/miles/ray/rollout.py +++ b/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) diff --git a/scripts/run-diffusion-nft-sd3-pickscore.sh b/scripts/run-diffusion-nft-sd3-pickscore.sh index 947a071f..e40f8fa7 100755 --- a/scripts/run-diffusion-nft-sd3-pickscore.sh +++ b/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 \ diff --git a/tests/ci/fixtures/e2e_standards/test_sd3_pickscore_nft_2xGPU.json b/tests/ci/fixtures/e2e_standards/test_sd3_pickscore_nft_2xGPU.json new file mode 100644 index 00000000..ef1dd1ec --- /dev/null +++ b/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 + ] + ] + } +} diff --git a/tests/e2e/short/test_sd3_pickscore_nft_2xGPU.py b/tests/e2e/short/test_sd3_pickscore_nft_2xGPU.py new file mode 100644 index 00000000..08710458 --- /dev/null +++ b/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", + ], +) 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 ae8da9df..785b90fe 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py +++ b/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):