fix(sft): scope deterministic training randomness

miles/backends/fsdp_utils/actor.py

@@ -357,7 +357,7 @@ def train(self, rollout_id: int, rollout_data_ref) -> None:  # type: ignore[over
 
     def _train_core(self, rollout_id: int, rollout_data) -> None:
         """Run the shared diffusion training loop."""
-        device = torch.cuda.current_device()
+        device = torch.device("cuda", torch.cuda.current_device())
 
         train_pairs: list = rollout_data["train_data"]
         if not train_pairs:
@@ -414,6 +414,8 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
             args=self.args,
             forward_dtype=self._forward_dtype,
             device=device,
+            rollout_id=rollout_id,
+            dp_rank=self.parallel_state.dp_rank,
         )
 
         # ------------- Recompute old log-probs (impl-consistent PPO ratio) -------------
@@ -422,9 +424,13 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
                 # write_old_log_prob returns before recording; this is never reduced.
                 unused_metrics = new_metric_buffer(self.parallel_state.dp_group, device, self.models)
                 # Skip window 0: its training forward runs on the same pre-update weights and doubles as the recompute.
+                # Start the id after window 0's micro-batches to stay aligned with the training loop.
+                microbatch_id = len(microbatch_schedule[0])
                 for microbatch_ranges in microbatch_schedule[1:]:
                     legacy_pad_to_len = self._maybe_legacy_window_pad_len(train_pairs, microbatch_ranges)
                     for pair_lo, pair_hi in microbatch_ranges:
+                        loss_ctx.microbatch_id = microbatch_id
+                        microbatch_id += 1
                         self._forward_train_pair_batch(
                             loss_ctx,
                             train_pairs[pair_lo:pair_hi],
@@ -435,6 +441,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
 
         # ------------- Forward / Backward -------------
         with timer("actor_train"):
+            microbatch_id = 0
             for optim_step_idx, microbatch_ranges in enumerate(microbatch_schedule):
                 self.optimizer.zero_grad(set_to_none=True)
 
@@ -449,6 +456,8 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
 
                 for pair_lo, pair_hi in microbatch_ranges:
                     chunk = train_pairs[pair_lo:pair_hi]
+                    loss_ctx.microbatch_id = microbatch_id
+                    microbatch_id += 1
                     loss_sum = self._forward_train_pair_batch(
                         loss_ctx,
                         chunk,

miles/backends/fsdp_utils/loss_hub/sft.py

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import hashlib
+
 import torch
 import torch.nn as nn
 
@@ -10,21 +12,52 @@
 from miles.utils.metric_buffer import MetricBuffer
 
 
-def sample_grid_indices(ctx: DiffusionLossContext, bsz: int) -> tuple[str, nn.Module, torch.Tensor]:
-    """Pick one DiT component per micro-batch (phase-pure), then grid indices within its range."""
+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,
+    *,
+    generator: torch.Generator,
+) -> tuple[str, nn.Module, torch.Tensor]:
+    """Pick one rank-aligned DiT component, then per-sample grid indices."""
     num_grid = len(ctx.scheduler.timesteps)
+    config = ctx.train_pipeline_config
     if len(ctx.models) == 1:
         component_name, model = next(iter(ctx.models.items()))
-        return component_name, model, torch.randint(num_grid, (bsz,))
-
-    # A uniform anchor index picks each expert with probability equal to its share of the
-    # grid, keeping the marginal over indices uniform while the micro-batch stays single-expert.
-    num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps)
-    config = ctx.train_pipeline_config
-    components = [config.component_for_timestep(float(t), num_train_timesteps) for t in ctx.scheduler.timesteps]
-    component_name = components[int(torch.randint(num_grid, (1,)))]
-    pool = torch.tensor([i for i, name in enumerate(components) if name == component_name])
-    return component_name, ctx.models[component_name], pool[torch.randint(len(pool), (bsz,))]
+        component_for_timestep = getattr(config, "component_for_timestep", None)
+        if component_for_timestep is None:
+            pool = torch.arange(num_grid, device=generator.device)
+        else:
+            num_train_timesteps = int(ctx.scheduler.config.num_train_timesteps)
+            pool = torch.tensor(
+                [
+                    i
+                    for i, timestep in enumerate(ctx.scheduler.timesteps)
+                    if component_for_timestep(float(timestep), num_train_timesteps) == component_name
+                ],
+                device=generator.device,
+            )
+    else:
+        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)
+        )
+        component_name = components[int(torch.randint(num_grid, (1,), generator=expert_generator))]
+        model = ctx.models[component_name]
+        pool = torch.tensor(
+            [i for i, name in enumerate(components) if name == component_name],
+            device=generator.device,
+        )
+
+    draw = torch.randint(len(pool), (bsz,), device=generator.device, generator=generator)
+    return component_name, model, pool[draw]
 
 
 def prepare_sft_batch(
@@ -39,12 +72,14 @@ def prepare_sft_batch(
     bsz = len(batch)
 
     x0 = torch.stack([pair["latent"] for pair in batch]).to(device=device, dtype=torch.float32)
-    component_name, model, idx = sample_grid_indices(ctx, bsz)
-    idx = idx.to(device)
+    sample_generator = torch.Generator(device=device).manual_seed(
+        _seed("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)
     sigmas = ctx.scheduler.sigmas[idx].to(dtype=torch.float32)
 
-    noise = torch.randn(x0.shape, device=device, dtype=torch.float32)
+    noise = torch.randn(x0.shape, device=device, dtype=torch.float32, generator=sample_generator)
     sigma_exp = sigmas.view(bsz, *([1] * (x0.ndim - 1)))
     latents = (1.0 - sigma_exp) * x0 + sigma_exp * noise
 

miles/backends/fsdp_utils/loss_hub/types.py

@@ -21,6 +21,9 @@ class DiffusionLossContext:
     args: Namespace
     forward_dtype: torch.dtype
     device: torch.device
+    rollout_id: int = 0
+    microbatch_id: int = 0
+    dp_rank: int = 0
 
 
 @dataclass

tests/fast/backends/fsdp_utils/test_loss_hub_sft.py

@@ -26,6 +26,11 @@ def component_for_timestep(self, timestep, num_train_timesteps):
         return "transformer" if timestep >= 0.875 * num_train_timesteps else "transformer_2"
 
 
+class _SingleConfig(_Config):
+    def component_for_timestep(self, timestep, num_train_timesteps):
+        return "transformer"
+
+
 def _scheduler():
     sigmas = torch.linspace(1.0, 1.0 / NUM_GRID, NUM_GRID)
     return Namespace(
@@ -35,15 +40,18 @@ def _scheduler():
     )
 
 
-def _ctx(models):
+def _ctx(models, rollout_id=3, microbatch_id=0, dp_rank=0, config=None):
     return DiffusionLossContext(
         models=models,
-        train_pipeline_config=_Config(),
+        train_pipeline_config=config if config is not None else _Config(),
         sde_backend=None,
         scheduler=_scheduler(),
-        args=Namespace(),
+        args=Namespace(seed=42),
         forward_dtype=torch.float32,
         device=torch.device("cpu"),
+        rollout_id=rollout_id,
+        microbatch_id=microbatch_id,
+        dp_rank=dp_rank,
     )
 
 
@@ -84,28 +92,93 @@ def test_timestep_scaling(self):
         assert torch.allclose(prepared.timesteps_for_model, prepared.timesteps / NUM_TRAIN_TIMESTEPS)
 
     def test_single_model_indices_cover_grid_uniformly(self):
-        torch.manual_seed(0)
-        ctx = _ctx({"transformer": nn.Identity()})
-        _, _, idx = sample_grid_indices(ctx, bsz=20000)
+        ctx = _ctx({"transformer": nn.Identity()}, config=_SingleConfig())
+        _, _, idx = sample_grid_indices(
+            ctx,
+            bsz=20000,
+            generator=torch.Generator().manual_seed(1),
+        )
         counts = torch.bincount(idx, minlength=NUM_GRID).float()
         assert counts.min() > 0
         assert ((counts / 20000) - 1 / NUM_GRID).abs().max() < 0.02
 
+    def test_single_wan_expert_only_samples_its_timesteps(self):
+        config = _Config()
+        timesteps = _scheduler().timesteps
+        for component_name in ("transformer", "transformer_2"):
+            ctx = _ctx({component_name: nn.Identity()}, config=config)
+            name, _, idx = sample_grid_indices(
+                ctx,
+                bsz=1000,
+                generator=torch.Generator().manual_seed(1),
+            )
+            expected = {
+                i
+                for i, timestep in enumerate(timesteps)
+                if config.component_for_timestep(float(timestep), NUM_TRAIN_TIMESTEPS) == component_name
+            }
+            assert name == component_name
+            assert set(idx.tolist()) == expected
+
     def test_dual_expert_micro_batch_is_phase_pure(self):
-        torch.manual_seed(0)
         models = {"transformer": nn.Identity(), "transformer_2": nn.Identity()}
         ctx = _ctx(models)
         config = ctx.train_pipeline_config
         picked = set()
-        for _ in range(20):
-            name, model, idx = sample_grid_indices(ctx, bsz=4)
+        for call in range(20):
+            ctx.microbatch_id = call
+            name, model, idx = sample_grid_indices(
+                ctx,
+                bsz=4,
+                generator=torch.Generator().manual_seed(call + 100),
+            )
             picked.add(name)
             assert model is models[name]
             for i in idx.tolist():
                 t = float(ctx.scheduler.timesteps[i])
                 assert config.component_for_timestep(t, NUM_TRAIN_TIMESTEPS) == name
         assert picked == {"transformer", "transformer_2"}
 
+    def test_prepare_is_independent_of_global_rng_state(self):
+        batch = _batch()
+        torch.manual_seed(0)
+        global_state = torch.get_rng_state()
+
+        first = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+        assert torch.equal(torch.get_rng_state(), global_state)
+
+        torch.rand(1000)
+        second = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+        assert torch.equal(first.timesteps, second.timesteps)
+        assert torch.equal(first.latents, second.latents)
+        assert torch.equal(first.extras["target"], second.extras["target"])
+
+    def test_samples_within_microbatch_use_different_noise(self):
+        batch = _batch(bsz=4)
+        prepared = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+        x0 = torch.stack([pair["latent"] for pair in batch]).float()
+        noise = prepared.extras["target"] + x0
+        assert all(not torch.equal(noise[0], noise[i]) for i in range(1, len(batch)))
+
+    def test_identity_changes_draws(self):
+        batch = _batch()
+        base = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+        other_rollout = prepare_sft_batch(_ctx({"transformer": nn.Identity()}, rollout_id=4), batch)
+        other_slot = prepare_sft_batch(_ctx({"transformer": nn.Identity()}, microbatch_id=1), batch)
+        other_dp_rank = prepare_sft_batch(_ctx({"transformer": nn.Identity()}, dp_rank=1), batch)
+        assert not torch.equal(base.extras["target"], other_rollout.extras["target"])
+        assert not torch.equal(base.extras["target"], other_slot.extras["target"])
+        assert not torch.equal(base.extras["target"], other_dp_rank.extras["target"])
+
+    def test_dual_expert_choice_is_rank_aligned(self):
+        models = {"transformer": nn.Identity(), "transformer_2": nn.Identity()}
+        for slot in range(10):
+            # Different DP ranks choose the same expert but use independent sample RNG.
+            a = prepare_sft_batch(_ctx(models, microbatch_id=slot, dp_rank=0), _batch())
+            b = prepare_sft_batch(_ctx(models, microbatch_id=slot, dp_rank=1), _batch())
+            assert a.component_name == b.component_name
+            assert not torch.equal(a.extras["target"], b.extras["target"])
+
 
 class TestSftLossFormula:
     def test_zero_loss_on_exact_velocity(self):