diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 26fc494d..401ba28a 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/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, diff --git a/miles/backends/fsdp_utils/loss_hub/sft.py b/miles/backends/fsdp_utils/loss_hub/sft.py index 581605da..fd850941 100644 --- a/miles/backends/fsdp_utils/loss_hub/sft.py +++ b/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 diff --git a/miles/backends/fsdp_utils/loss_hub/types.py b/miles/backends/fsdp_utils/loss_hub/types.py index 33d4f767..74d252b4 100644 --- a/miles/backends/fsdp_utils/loss_hub/types.py +++ b/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 diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py b/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py index 30e36b4b..3a2b1fed 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py +++ b/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,21 +92,46 @@ 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(): @@ -106,6 +139,46 @@ def test_dual_expert_micro_batch_is_phase_pure(self): 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):