refactor(fsdp): DiffusionLossContext stops carrying loop state

docs/user-guide/customization.md

@@ -279,24 +279,30 @@ def convert_samples_to_train_data(args, samples) -> dict:
 ### `--custom-prepare-train-batch-path`
 
 ```python
-def prepare(ctx, batch, *, pad_to_len=None) -> PreparedBatch:
+def prepare(ctx, batch, *, microbatch_id, optim_step_idx, pad_to_len=None) -> PreparedBatch:
     ...
 ```
 
 Read only under `--loss-type custom_loss`; the other loss types take their prepare from the loss hub's table
-above. Builds DiT inputs from train pairs.
+above. Builds DiT inputs from train pairs. `microbatch_id` numbers the micro-batch within the rollout (training
+counts from 0; the old-log-prob recompute of windows 1.. passes their training forwards' numbers) and
+`optim_step_idx` is its optimizer step; the built-in hooks seed their noise from `microbatch_id`.
 `PreparedBatch.latents`, `timesteps` and
 `timesteps_for_model` are keyed by latent stream name: `"visual"`, plus e.g.
 `"audio"` for joint audio-video models.
 
 ### `--custom-loss-function-path`
 
 ```python
-def loss_formula(ctx, batch, prepared, *, new_pred, old_pred, ref_pred, metrics, **kwargs):
+def loss_formula(
+    ctx, batch, prepared, *, new_pred, old_pred, ref_pred, metrics, optim_step_idx, train_output_inputs, **kwargs
+):
     ...
 ```
 
-Read only under `--loss-type custom_loss`. `new_pred`, `old_pred` and `ref_pred` are keyed by the same stream
+Read only under `--loss-type custom_loss`. `optim_step_idx` is the micro-batch's optimizer step within the rollout;
+`train_output_inputs` is the list the formula may fill for the train-output hook, empty at the start of each train
+call. `new_pred`, `old_pred` and `ref_pred` are keyed by the same stream
 names as `prepared.latents`. `old_pred` (the EMA forward) is `None` for a custom formula, since the trunk runs it
 only under `--loss-type nft`; `ref_pred` is `None` with `--ref-mode none`.
 
@@ -317,16 +323,14 @@ and the train-output hook see the returned pairs.
 ### `--custom-train-output-function-path`
 
 ```python
-def train_output(ctx, train_pairs) -> Any:
+def train_output(ctx, train_pairs, train_output_inputs) -> Any:
     ...
 ```
 
 Read only under `--loss-type custom_loss`. Runs without grad after this trainer group's train steps on a
-rollout. Every trainer group that
+rollout, with the list the loss formula filled. Every trainer group that
 trains later on the same rollout reads the return value of its paired rank (same rank, same
-train pairs) as `ctx.external_data[<role>]`. Prepare and loss hooks see the current optimizer step
-as `ctx.optim_step_idx` and can leave inputs for this hook in `ctx.train_output_inputs`, which
-starts empty on each train call.
+train pairs) as `ctx.external_data[<role>]`.
 
 ### `--sde-step-backend-path`
 

miles/backends/fsdp_utils/actor.py

@@ -370,6 +370,7 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An
         )
         if self.loss_hooks.train_input is not None:
             train_pairs = self.loss_hooks.train_input(loss_ctx, train_pairs)
+        train_output_inputs: list = []  # one list per train call
 
         num_pairs = len(train_pairs)
 
@@ -409,27 +410,28 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An
                     sigma_buckets=self.args.log_loss_sigma_bucket,
                 )
                 # 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.
+                # Each later window's micro-batches get the ids and step index of their training forwards.
                 microbatch_id = len(microbatch_schedule[0])
-                for microbatch_ranges in microbatch_schedule[1:]:
+                for optim_step_idx, microbatch_ranges in enumerate(microbatch_schedule[1:], 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],
+                            microbatch_id=microbatch_id,
+                            optim_step_idx=optim_step_idx,
+                            train_output_inputs=train_output_inputs,
                             metrics=unused_metrics,
                             pad_to_len=legacy_pad_to_len,
                             write_old_log_prob=True,
                         )
+                        microbatch_id += 1
 
         # ------------- 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)
-                loss_ctx.optim_step_idx = optim_step_idx
 
                 old_log_prob_from_new = self.args.diffusion_recompute_old_log_prob and optim_step_idx == 0
 
@@ -448,15 +450,17 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An
 
                 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,
+                        microbatch_id=microbatch_id,
+                        optim_step_idx=optim_step_idx,
+                        train_output_inputs=train_output_inputs,
                         metrics=metrics,
                         pad_to_len=legacy_pad_to_len,
                         old_log_prob_from_new=old_log_prob_from_new,
                     )
+                    microbatch_id += 1
                     if not self.args.debug_skip_optimizer_step:
                         # ShardedGradScaler keeps fp16 policy grads from underflowing
                         # (required for SD3.5 fp16 forward); no-op for bf16/fp32.
@@ -488,7 +492,7 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An
 
         if self.loss_hooks.train_output is not None:
             with torch.no_grad():
-                return self.loss_hooks.train_output(loss_ctx, train_pairs)
+                return self.loss_hooks.train_output(loss_ctx, train_pairs, train_output_inputs)
 
     def _maybe_legacy_window_pad_len(self, train_pairs: list, microbatch_ranges: list) -> int | None:
         """LEGACY 2D parity: the whole-window max cond seq_len (like the legacy tile path), or
@@ -509,13 +513,19 @@ def _forward_train_pair_batch(
         ctx: DiffusionLossContext,
         batch: list,
         *,
+        microbatch_id: int,
+        optim_step_idx: int,
+        train_output_inputs: list,
         metrics: MetricBuffer,
         pad_to_len: int | None = None,
         write_old_log_prob: bool = False,
         old_log_prob_from_new: bool = False,
     ) -> torch.Tensor | None:
-        """Run one prepared diffusion micro-batch."""
-        prepared = self.loss_hooks.prepare(ctx, batch, pad_to_len=pad_to_len)
+        """Run one prepared diffusion micro-batch: number ``microbatch_id`` of the rollout, in optimizer step
+        ``optim_step_idx``."""
+        prepared = self.loss_hooks.prepare(
+            ctx, batch, microbatch_id=microbatch_id, optim_step_idx=optim_step_idx, pad_to_len=pad_to_len
+        )
         train_pipeline_config = self.train_pipeline_config
         forward_dtype = self._forward_dtype
 
@@ -574,6 +584,8 @@ def _compute_noise_pred(model) -> dict[str, torch.Tensor]:
             old_pred=old_pred,
             ref_pred=ref_pred,
             metrics=metrics,
+            optim_step_idx=optim_step_idx,
+            train_output_inputs=train_output_inputs,
             write_old_log_prob=write_old_log_prob,
             old_log_prob_from_new=old_log_prob_from_new,
         )

miles/backends/fsdp_utils/loss_hub/__init__.py

@@ -30,15 +30,18 @@
 class LossHooks:
     """The hooks one trainer group runs on a rollout.
 
-    ``prepare`` builds a micro-batch's DiT inputs, ``formula`` turns the predictions into its loss, ``train_input``
-    preprocesses the rank's train pairs before micro-batching, and ``train_output`` runs after the train steps and
-    hands its return value to the groups that train later.
+    ``prepare(ctx, batch, *, microbatch_id, optim_step_idx, pad_to_len=None)`` builds a micro-batch's DiT inputs,
+    ``formula(ctx, batch, prepared, *, new_pred, old_pred, ref_pred, metrics, optim_step_idx, train_output_inputs,
+    ...)`` turns the predictions into its loss, ``train_input`` preprocesses the rank's train pairs before
+    micro-batching, and ``train_output(ctx, train_pairs, train_output_inputs)`` runs after the train steps and hands
+    its return value to the groups that train later; ``train_output_inputs`` is the list the formulas filled during
+    this train call.
     """
 
     prepare: Callable[..., PreparedBatch]
     formula: Callable[..., torch.Tensor | None]
     train_input: Callable[[DiffusionLossContext, list], list] | None = None
-    train_output: Callable[[DiffusionLossContext, list], Any] | None = None
+    train_output: Callable[[DiffusionLossContext, list, list], Any] | None = None
 
 
 LOSS_HOOKS: dict[tuple[str, str], LossHooks] = {

miles/backends/fsdp_utils/loss_hub/dmd2.py

@@ -18,6 +18,8 @@ def prepare_critic_batch(
     ctx: DiffusionLossContext,
     batch: list[dict],
     *,
+    microbatch_id: int,
+    optim_step_idx: int,
     pad_to_len: int | None = None,
 ) -> PreparedBatch:
     """SFT flow matching on the student's x0 predictions, conditioned on their rollout prompts.
@@ -33,12 +35,19 @@ def prepare_critic_batch(
         for pair in batch
     ]
     grid_index_range = None
-    if _is_scoring_step(ctx):
+    if _is_scoring_step(ctx, optim_step_idx):
         min_sigma, max_sigma = ctx.args.dmd2_score_sigma_range
         num_train_timesteps = config.num_train_timesteps
         # Grid index i carries the unshifted sigma 1 - i / num_train_timesteps.
         grid_index_range = (round(num_train_timesteps * (1 - max_sigma)), round(num_train_timesteps * (1 - min_sigma)))
-    return prepare_sft_batch(ctx, sft_batch, pad_to_len=pad_to_len, grid_index_range=grid_index_range)
+    return prepare_sft_batch(
+        ctx,
+        sft_batch,
+        microbatch_id=microbatch_id,
+        optim_step_idx=optim_step_idx,
+        pad_to_len=pad_to_len,
+        grid_index_range=grid_index_range,
+    )
 
 
 def critic_loss_formula(
@@ -50,6 +59,8 @@ def critic_loss_formula(
     old_pred: dict[str, torch.Tensor] | None,
     ref_pred: dict[str, torch.Tensor] | None,
     metrics: MetricBuffer,
+    optim_step_idx: int,
+    train_output_inputs: list,
     write_old_log_prob: bool = False,
     old_log_prob_from_new: bool = False,
 ) -> torch.Tensor:
@@ -59,9 +70,9 @@ def critic_loss_formula(
     """
     latents = prepared.latents["visual"]
     sigmas = prepared.extras["sigmas"]["visual"].view(len(batch), *([1] * (latents.ndim - 1)))
-    if _is_scoring_step(ctx):
+    if _is_scoring_step(ctx, optim_step_idx):
         x0_fake = latents - sigmas * new_pred["visual"].detach().float()
-        ctx.train_output_inputs.append(
+        train_output_inputs.append(
             (batch, prepared.component_name, latents, prepared.timesteps["visual"], sigmas, x0_fake)
         )
     if ctx.args.dmd2_critic_loss_space == "x0":
@@ -77,15 +88,19 @@ def critic_loss_formula(
         old_pred=old_pred,
         ref_pred=ref_pred,
         metrics=metrics,
+        optim_step_idx=optim_step_idx,
+        train_output_inputs=train_output_inputs,
         write_old_log_prob=write_old_log_prob,
         old_log_prob_from_new=old_log_prob_from_new,
     )
 
 
-def score_last_step_samples(ctx: DiffusionLossContext, train_pairs: list[dict]) -> dict[int, dict]:
+def score_last_step_samples(
+    ctx: DiffusionLossContext, train_pairs: list[dict], train_output_inputs: list
+) -> dict[int, dict]:
     """The teacher's x0 at each scored sample's noisy input, next to the critic's x0 kept there, by sample."""
     predictions = {}
-    for batch, component_name, latents, timesteps, sigmas, x0_fake in ctx.train_output_inputs:
+    for batch, component_name, latents, timesteps, sigmas, x0_fake in train_output_inputs:
         teacher = ctx.teacher_models[component_name]
         real_pred = _predict_velocity(ctx, teacher, batch, latents, timesteps, ctx.args.dmd2_teacher_guidance_scale)
         for pair, x0_real, fake in zip(batch, latents - sigmas * real_pred, x0_fake, strict=True):
@@ -97,8 +112,8 @@ def score_last_step_samples(ctx: DiffusionLossContext, train_pairs: list[dict])
     return predictions
 
 
-def _is_scoring_step(ctx: DiffusionLossContext) -> bool:
-    return ctx.optim_step_idx == ctx.args.num_steps_per_rollout - 1
+def _is_scoring_step(ctx: DiffusionLossContext, optim_step_idx: int) -> bool:
+    return optim_step_idx == ctx.args.num_steps_per_rollout - 1
 
 
 def _predict_velocity(
@@ -171,9 +186,15 @@ def prepare_generator_batch(
     ctx: DiffusionLossContext,
     batch: list[dict],
     *,
+    microbatch_id: int,
+    optim_step_idx: int,
     pad_to_len: int | None = None,
 ) -> PreparedBatch:
-    """The student's step ``i`` on its own rollout input ``x_i``, unguided."""
+    """The student's step ``i`` on its own rollout input ``x_i``, unguided.
+
+    The pairs carry their own noise, so ``microbatch_id`` and ``optim_step_idx`` are part of every prepare hook's
+    signature and unused here.
+    """
     if len(ctx.models) != 1:
         raise ValueError("DMD2 currently supports a single DiT component")
     device = ctx.device
@@ -211,6 +232,8 @@ def dmd2_loss_formula(
     old_pred: dict[str, torch.Tensor] | None,
     ref_pred: dict[str, torch.Tensor] | None,
     metrics: MetricBuffer,
+    optim_step_idx: int,
+    train_output_inputs: list,
     write_old_log_prob: bool = False,
     old_log_prob_from_new: bool = False,
 ) -> torch.Tensor:

miles/backends/fsdp_utils/loss_hub/flow_grpo.py

@@ -19,9 +19,15 @@ def prepare_flow_grpo_batch(
     ctx: DiffusionLossContext,
     batch: list[dict],
     *,
+    microbatch_id: int,
+    optim_step_idx: int,
     pad_to_len: int | None = None,
 ) -> PreparedBatch:
-    """Stack SDE-pair fields and build CFG conditioning (guidance from args)."""
+    """Stack SDE-pair fields and build CFG conditioning (guidance from args).
+
+    The pairs carry their own noise, so ``microbatch_id`` and ``optim_step_idx`` are part of every prepare hook's
+    signature and unused here.
+    """
     args = ctx.args
     device = ctx.device
     config = ctx.train_pipeline_config
@@ -115,6 +121,8 @@ def flow_grpo_loss_formula(
     old_pred: dict[str, torch.Tensor] | None,
     ref_pred: dict[str, torch.Tensor] | None,
     metrics: MetricBuffer,
+    optim_step_idx: int,
+    train_output_inputs: list,
     write_old_log_prob: bool = False,
     old_log_prob_from_new: bool = False,
 ) -> torch.Tensor | None:

miles/backends/fsdp_utils/loss_hub/nft.py

@@ -24,9 +24,12 @@ def prepare_nft_batch(
     ctx: DiffusionLossContext,
     batch: list[dict],
     *,
+    microbatch_id: int,
+    optim_step_idx: int,
     pad_to_len: int | None = None,
 ) -> PreparedBatch:
-    """Corrupt clean x0 at each pair's sigma; CFG-free cond."""
+    """Corrupt clean x0 at each pair's sigma; CFG-free cond. ``microbatch_id`` seeds the noise; ``optim_step_idx`` is
+    part of every prepare hook's signature and unused here."""
     if len(ctx.models) != 1:
         raise ValueError("DiffusionNFT currently supports a single DiT component")
     device = ctx.device
@@ -41,7 +44,7 @@ def prepare_nft_batch(
     pos_cond = config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len)
 
     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)
+        stable_hash("nft_corrupt", int(ctx.args.seed), ctx.rollout_id, microbatch_id, ctx.dp_rank)
     )
     xt = corrupt(x0, t, sample_noise(x0, generator=noise_generator))
     # Rollout trajectories carry only the visual stream.
@@ -111,6 +114,8 @@ def nft_loss_formula(
     old_pred: dict[str, torch.Tensor] | None,
     ref_pred: dict[str, torch.Tensor] | None,
     metrics: MetricBuffer,
+    optim_step_idx: int,
+    train_output_inputs: list,
     write_old_log_prob: bool = False,
     old_log_prob_from_new: bool = False,
 ) -> torch.Tensor: