diff --git a/docs/user-guide/customization.md b/docs/user-guide/customization.md
index 97e33f8a8..3a5f30314 100644
--- a/docs/user-guide/customization.md
+++ b/docs/user-guide/customization.md
@@ -279,12 +279,14 @@ 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.
@@ -292,11 +294,15 @@ above. Builds DiT inputs from train pairs.
 ### `--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`
 
diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py
index ac5b0cf1e..d296992a4 100644
--- a/miles/backends/fsdp_utils/actor.py
+++ b/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,
         )
diff --git a/miles/backends/fsdp_utils/loss_hub/__init__.py b/miles/backends/fsdp_utils/loss_hub/__init__.py
index bbe51d5e5..651570b37 100644
--- a/miles/backends/fsdp_utils/loss_hub/__init__.py
+++ b/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] = {
diff --git a/miles/backends/fsdp_utils/loss_hub/dmd2.py b/miles/backends/fsdp_utils/loss_hub/dmd2.py
index 2e3a56900..2388cd17f 100644
--- a/miles/backends/fsdp_utils/loss_hub/dmd2.py
+++ b/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:
diff --git a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py b/miles/backends/fsdp_utils/loss_hub/flow_grpo.py
index c33a495bc..48eab4052 100644
--- a/miles/backends/fsdp_utils/loss_hub/flow_grpo.py
+++ b/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:
diff --git a/miles/backends/fsdp_utils/loss_hub/nft.py b/miles/backends/fsdp_utils/loss_hub/nft.py
index 4ef42999e..88765cd94 100644
--- a/miles/backends/fsdp_utils/loss_hub/nft.py
+++ b/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:
diff --git a/miles/backends/fsdp_utils/loss_hub/sft.py b/miles/backends/fsdp_utils/loss_hub/sft.py
index b17acfae7..df1ca25a3 100644
--- a/miles/backends/fsdp_utils/loss_hub/sft.py
+++ b/miles/backends/fsdp_utils/loss_hub/sft.py
@@ -12,7 +12,7 @@
 
 
 def select_component(
-    ctx: DiffusionLossContext, noise_schedule: FlowNoiseSchedule, device: torch.device
+    ctx: DiffusionLossContext, noise_schedule: FlowNoiseSchedule, device: torch.device, *, microbatch_id: int
 ) -> tuple[str, nn.Module, torch.Tensor]:
     """Pick one rank-aligned DiT component and the indices of the grid points it serves."""
     config = ctx.train_pipeline_config
@@ -25,7 +25,7 @@ def select_component(
     else:
         # dp_rank stays out of this seed: every rank must forward the same expert.
         expert_generator = torch.Generator().manual_seed(
-            stable_hash("expert", int(ctx.args.seed), ctx.rollout_id, ctx.microbatch_id)
+            stable_hash("expert", int(ctx.args.seed), ctx.rollout_id, microbatch_id)
         )
         drawn = int(torch.randint(len(timesteps), (1,), generator=expert_generator))
         component_name = component_for_timestep(timesteps[drawn])
@@ -38,6 +38,8 @@ def prepare_sft_batch(
     ctx: DiffusionLossContext,
     batch: list[dict],
     *,
+    microbatch_id: int,
+    optim_step_idx: int,
     pad_to_len: int | None = None,
     grid_index_range: tuple[int, int] | None = None,
 ) -> PreparedBatch:
@@ -45,7 +47,8 @@ def prepare_sft_batch(
 
     The visual grid picks the DiT component and the index is drawn uniformly from the points it serves, within
     ``grid_index_range`` (inclusive) when given. Sharing the index pairs the streams' sigmas as inference steps them,
-    one step index read through each stream's shift.
+    one step index read through each stream's shift. ``microbatch_id`` seeds the draws; ``optim_step_idx`` is part of
+    every prepare hook's signature and unused here.
     """
     device = ctx.device
     config = ctx.train_pipeline_config
@@ -58,13 +61,15 @@ def prepare_sft_batch(
         stream_name: FlowNoiseSchedule.shifted(ctx.args.fsdp_flow_shift[stream_name], config.num_train_timesteps)
         for stream_name in stream_names
     }
-    component_name, model, visual_pool = select_component(ctx, noise_schedules["visual"], device)
+    component_name, model, visual_pool = select_component(
+        ctx, noise_schedules["visual"], device, microbatch_id=microbatch_id
+    )
     if grid_index_range is not None:
         visual_pool = visual_pool[(visual_pool >= grid_index_range[0]) & (visual_pool <= grid_index_range[1])]
 
     # The visual stream keeps the single-stream seed, so single-stream runs draw as before; other streams add their
     # name, so adding a stream never changes another stream's noise.
-    seed_parts = ("sample", int(ctx.args.seed), ctx.rollout_id, ctx.microbatch_id, ctx.dp_rank)
+    seed_parts = ("sample", int(ctx.args.seed), ctx.rollout_id, microbatch_id, ctx.dp_rank)
     generators = {
         stream_name: torch.Generator(device=device).manual_seed(
             stable_hash(*seed_parts) if stream_name == "visual" else stable_hash(*seed_parts, stream_name)
@@ -123,6 +128,8 @@ def sft_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:
diff --git a/miles/backends/fsdp_utils/loss_hub/types.py b/miles/backends/fsdp_utils/loss_hub/types.py
index fec342ea8..129d08cad 100644
--- a/miles/backends/fsdp_utils/loss_hub/types.py
+++ b/miles/backends/fsdp_utils/loss_hub/types.py
@@ -10,9 +10,12 @@
 import torch.nn as nn
 
 
-@dataclass
+@dataclass(frozen=True, eq=False)
 class DiffusionLossContext:
-    """Train-side handles for prepare and loss hooks."""
+    """Rollout-level, read-only handles for the loss hooks; what changes per micro-batch travels as arguments.
+
+    Compared by identity: nothing hashes a context, and a synthesized hash would fail on the dict fields.
+    """
 
     models: dict[str, torch.nn.Module]
     train_pipeline_config: Any
@@ -21,14 +24,10 @@ class DiffusionLossContext:
     forward_dtype: torch.dtype
     device: torch.device
     rollout_id: int = 0
-    microbatch_id: int = 0
-    optim_step_idx: int = 0
     dp_rank: int = 0
     reference_models: dict[str, torch.nn.Module] = field(default_factory=dict)
     teacher_models: dict[str, torch.nn.Module] = field(default_factory=dict)
     external_data: dict[str, Any] = field(default_factory=dict)
-    # Filled by loss hooks during one train call and read by the train-output hook.
-    train_output_inputs: list[Any] = field(default_factory=list)
 
 
 @dataclass(frozen=True)
diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py
index 740280fd0..2b26fdac3 100644
--- a/miles/utils/arguments.py
+++ b/miles/utils/arguments.py
@@ -960,7 +960,7 @@ def add_algo_arguments(parser):
                 default=None,
                 help=(
                     "Dotted path to prepare hook "
-                    "`fn(ctx, batch, *, pad_to_len=None) -> PreparedBatch`. "
+                    "`fn(ctx, batch, *, microbatch_id, optim_step_idx, pad_to_len=None) -> PreparedBatch`. "
                     "Builds DiT inputs; --loss-type custom_loss only, default is Flow-GRPO SDE-pair stacking."
                 ),
             )
@@ -991,8 +991,8 @@ def add_algo_arguments(parser):
                 type=str,
                 default=None,
                 help=(
-                    "Dotted path to `fn(ctx, train_pairs) -> Any`, run without grad after the train steps of "
-                    "a rollout. Its return value reaches the paired rank of every trainer group that trains "
+                    "Dotted path to `fn(ctx, train_pairs, train_output_inputs) -> Any`, run without grad after the "
+                    "train steps of a rollout with the list the loss formula filled. Its return value reaches the paired rank of every trainer group that trains "
                     "later on the same rollout as `ctx.external_data[<role>]`. --loss-type custom_loss only."
                 ),
             )
@@ -2038,6 +2038,8 @@ def validate_dmd2_args(args) -> None:
         )
     if args.ref_mode != "none":
         raise ValueError("--loss-type dmd2 uses no reference model; drop --ref-mode and --diffusion-kl-beta")
+    if args.diffusion_recompute_old_log_prob:
+        raise ValueError("--diffusion-recompute-old-log-prob is only supported for policy_loss / Flow-GRPO, not DMD2")
     if args.num_steps_per_rollout != 1:
         raise ValueError(
             f"--loss-type dmd2 takes one actor step per rollout, got --num-steps-per-rollout "
diff --git a/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py b/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py
index f84fe1a57..6fc39746d 100644
--- a/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py
+++ b/tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py
@@ -92,7 +92,15 @@ def check_roles(mesh, adapter_paths, reference_has_adapter, cpu_offload):
             target = reference_oracle(inputs)
             torch.testing.assert_close(teacher[component](inputs), teacher_oracle(inputs))
         harness, observed = make_actor_forward_harness(actor, reference, component, inputs, "ref")
-        load_actor_forward_method()(harness, None, [], metrics=None).backward()
+        load_actor_forward_method()(
+            harness,
+            None,
+            [],
+            microbatch_id=0,
+            optim_step_idx=0,
+            train_output_inputs=[],
+            metrics=None,
+        ).backward()
         expected_actor_output = oracle(inputs)
         (expected_actor_output - target).square().mean().backward()
         torch.testing.assert_close(observed["ref_pred"], target)
@@ -116,7 +124,15 @@ def check_roles(mesh, adapter_paths, reference_has_adapter, cpu_offload):
             base_target = base_oracle(inputs)
         forward_count = actor[component].base_model.model.forward_count.item()
         harness, observed = make_actor_forward_harness(actor, {}, component, inputs, "lora_base")
-        load_actor_forward_method()(harness, None, [], metrics=None).backward()
+        load_actor_forward_method()(
+            harness,
+            None,
+            [],
+            microbatch_id=0,
+            optim_step_idx=0,
+            train_output_inputs=[],
+            metrics=None,
+        ).backward()
         expected_actor_output = oracle(inputs)
         (expected_actor_output - base_target).square().mean().backward()
         torch.testing.assert_close(observed["ref_pred"], base_target)
diff --git a/tests/fast/backends/fsdp_utils/test_h3_streams.py b/tests/fast/backends/fsdp_utils/test_h3_streams.py
index a47a9e47a..0f56dd4b0 100644
--- a/tests/fast/backends/fsdp_utils/test_h3_streams.py
+++ b/tests/fast/backends/fsdp_utils/test_h3_streams.py
@@ -104,7 +104,7 @@ def test_t2va_prediction_matches_the_sft_target_and_skips_padding_rows():
         device=torch.device("cpu"),
     )
     batch = [{"latent": {"visual": torch.ones(3, 96)}, "cond_kwargs": cond}]
-    prepared = prepare_sft_batch(context, batch)
+    prepared = prepare_sft_batch(context, batch, microbatch_id=0, optim_step_idx=0)
     prediction = config.compute_noise_pred(
         model=context.models["transformer"],
         latents_input=prepared.latents,
@@ -125,7 +125,15 @@ def test_t2va_prediction_matches_the_sft_target_and_skips_padding_rows():
     assert inputs["timestep_indices"].shape == (7,)
     metrics = new_metric_buffer(None, torch.device("cpu"), ["transformer"], sigma_buckets=5)
     loss = sft_loss_formula(
-        context, batch, prepared, new_pred=prediction, old_pred=None, ref_pred=None, metrics=metrics
+        context,
+        batch,
+        prepared,
+        new_pred=prediction,
+        old_pred=None,
+        ref_pred=None,
+        metrics=metrics,
+        optim_step_idx=0,
+        train_output_inputs=[],
     )
     assert torch.isfinite(loss) and loss.item() > 0
 
@@ -405,7 +413,7 @@ def test_real_tiny_diffusers_h3_joint_sft_forward_and_backward(forward_dtype):
         forward_dtype=forward_dtype,
         device=torch.device("cpu"),
     )
-    prepared = prepare_sft_batch(context, batch)
+    prepared = prepare_sft_batch(context, batch, microbatch_id=0, optim_step_idx=0)
     autocast_enabled = forward_dtype != torch.float32 and config.enable_autocast
     with torch.autocast("cpu", dtype=forward_dtype, enabled=autocast_enabled):
         prediction = config.compute_noise_pred(
@@ -426,7 +434,15 @@ def test_real_tiny_diffusers_h3_joint_sft_forward_and_backward(forward_dtype):
     }
     metrics = new_metric_buffer(None, torch.device("cpu"), ["transformer_ref"], sigma_buckets=5)
     loss = sft_loss_formula(
-        context, batch, prepared, new_pred=prediction, old_pred=None, ref_pred=None, metrics=metrics
+        context,
+        batch,
+        prepared,
+        new_pred=prediction,
+        old_pred=None,
+        ref_pred=None,
+        metrics=metrics,
+        optim_step_idx=0,
+        train_output_inputs=[],
     )
     assert torch.isfinite(loss) and loss.item() > 0
     assert prepared.extras["target"]["visual"].dtype == torch.float32
diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py
index eca56850e..c05447b30 100644
--- a/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py
+++ b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py
@@ -48,6 +48,7 @@
 
 register_cpu_ci(est_time=10, suite="stage-a-cpu", labels=[])
 
+import dataclasses
 from argparse import Namespace
 
 import pytest
@@ -136,14 +137,14 @@ def _ctx(teacher_guidance_scale=3.0, critic_scale=1.0, teacher_scale=1.0):
     )
 
 
-def _critic_step(ctx, pairs, optim_step_idx, metrics=None):
+def _critic_step(ctx, pairs, optim_step_idx, train_output_inputs, metrics=None):
     """One critic optimizer step as _train_core runs it: prepare, the critic's training forward, the loss hook."""
-    ctx.optim_step_idx = optim_step_idx
     steps = []
     for start in range(0, len(pairs), ctx.args.micro_batch_size):
         batch = pairs[start : start + ctx.args.micro_batch_size]
-        ctx.microbatch_id = optim_step_idx * 10 + start
-        prepared = dmd2.prepare_critic_batch(ctx, batch)
+        prepared = dmd2.prepare_critic_batch(
+            ctx, batch, microbatch_id=optim_step_idx * 10 + start, optim_step_idx=optim_step_idx
+        )
         new_pred = ctx.train_pipeline_config.compute_noise_pred(
             model=prepared.model,
             latents_input=prepared.latents,
@@ -157,7 +158,15 @@ def _critic_step(ctx, pairs, optim_step_idx, metrics=None):
             true_cfg_scale=None,
         )
         loss = dmd2.critic_loss_formula(
-            ctx, batch, prepared, new_pred=new_pred, old_pred=None, ref_pred=None, metrics=metrics or _Metrics()
+            ctx,
+            batch,
+            prepared,
+            new_pred=new_pred,
+            old_pred=None,
+            ref_pred=None,
+            metrics=metrics or _Metrics(),
+            optim_step_idx=optim_step_idx,
+            train_output_inputs=train_output_inputs,
         )
         steps.append((batch, prepared, new_pred, loss))
     return steps
@@ -179,9 +188,8 @@ def test_critic_batch_is_the_sft_batch_on_the_student_x0():
     pairs = _pairs(3)
 
     for optim_step_idx in range(CRITIC_STEPS - 1):
-        ctx.optim_step_idx = optim_step_idx
-        critic_batch = dmd2.prepare_critic_batch(ctx, pairs)
-        sft_batch = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs))
+        critic_batch = dmd2.prepare_critic_batch(ctx, pairs, microbatch_id=0, optim_step_idx=optim_step_idx)
+        sft_batch = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs), microbatch_id=0, optim_step_idx=optim_step_idx)
 
         assert critic_batch.model is ctx.models["transformer"]
         for field in ("latents", "timesteps", "timesteps_for_model"):
@@ -194,10 +202,12 @@ def test_critic_batch_is_the_sft_batch_on_the_student_x0():
 def test_last_step_draws_from_the_scoring_range():
     ctx = _ctx()
     pairs = _pairs(256)
-    ctx.optim_step_idx = CRITIC_STEPS - 1
+    last_step = CRITIC_STEPS - 1
 
-    critic_batch = dmd2.prepare_critic_batch(ctx, pairs)
-    sft_batch = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs), grid_index_range=(20, 980))
+    critic_batch = dmd2.prepare_critic_batch(ctx, pairs, microbatch_id=0, optim_step_idx=last_step)
+    sft_batch = prepare_sft_batch(
+        ctx, _sft_pairs(ctx, pairs), microbatch_id=0, optim_step_idx=last_step, grid_index_range=(20, 980)
+    )
 
     assert torch.equal(critic_batch.latents["visual"], sft_batch.latents["visual"])
     assert torch.equal(critic_batch.extras["target"]["visual"], sft_batch.extras["target"]["visual"])
@@ -205,20 +215,21 @@ def test_last_step_draws_from_the_scoring_range():
     sigmas = critic_batch.extras["sigmas"]["visual"]
     assert set(sigmas.tolist()) <= set(grid[20:981].tolist())
     # The same draw without the range leaves it, so the restriction is exercised.
-    assert not set(prepare_sft_batch(ctx, _sft_pairs(ctx, pairs)).extras["sigmas"]["visual"].tolist()) <= set(
-        grid[20:981].tolist()
-    )
+    unrestricted = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs), microbatch_id=0, optim_step_idx=last_step)
+    assert not set(unrestricted.extras["sigmas"]["visual"].tolist()) <= set(grid[20:981].tolist())
 
 
 def test_score_sigma_range_bounds_the_unshifted_sigma():
     ctx = _ctx()
     ctx.args.dmd2_score_sigma_range = [0.1, 0.5]
     pairs = _pairs(256)
-    ctx.optim_step_idx = CRITIC_STEPS - 1
+    last_step = CRITIC_STEPS - 1
 
-    critic_batch = dmd2.prepare_critic_batch(ctx, pairs)
+    critic_batch = dmd2.prepare_critic_batch(ctx, pairs, microbatch_id=0, optim_step_idx=last_step)
     # Grid index i carries the unshifted sigma 1 - i / 1000, so [0.1, 0.5] is grid points [500, 900].
-    sft_batch = prepare_sft_batch(ctx, _sft_pairs(ctx, pairs), grid_index_range=(500, 900))
+    sft_batch = prepare_sft_batch(
+        ctx, _sft_pairs(ctx, pairs), microbatch_id=0, optim_step_idx=last_step, grid_index_range=(500, 900)
+    )
 
     sigmas = critic_batch.extras["sigmas"]["visual"]
     assert torch.equal(sigmas, sft_batch.extras["sigmas"]["visual"])
@@ -231,9 +242,17 @@ def test_critic_loss_is_the_sft_loss():
         ctx = _ctx()
         metrics = _Metrics()
         sft_metrics = _Metrics()
-        for batch, prepared, new_pred, loss in _critic_step(ctx, _pairs(3), optim_step_idx, metrics):
+        for batch, prepared, new_pred, loss in _critic_step(ctx, _pairs(3), optim_step_idx, [], metrics):
             sft_loss = sft_loss_formula(
-                ctx, batch, prepared, new_pred=new_pred, old_pred=None, ref_pred=None, metrics=sft_metrics
+                ctx,
+                batch,
+                prepared,
+                new_pred=new_pred,
+                old_pred=None,
+                ref_pred=None,
+                metrics=sft_metrics,
+                optim_step_idx=optim_step_idx,
+                train_output_inputs=[],
             )
             assert torch.equal(loss, sft_loss)
         assert metrics.seen == sft_metrics.seen
@@ -244,7 +263,7 @@ def test_x0_critic_loss_weights_the_velocity_mse():
         ctx = _ctx()
         ctx.models["transformer"] = _NoisyInputVelocity()
         ctx.args.dmd2_critic_loss_space = "x0"
-        for batch, prepared, new_pred, loss in _critic_step(ctx, _pairs(3), optim_step_idx):
+        for batch, prepared, new_pred, loss in _critic_step(ctx, _pairs(3), optim_step_idx, []):
             x_s = prepared.latents["visual"]
             sigma = prepared.extras["sigmas"]["visual"].view(-1, 1, 1, 1)
             x0 = torch.stack([pair["x0"] for pair in batch])
@@ -258,12 +277,14 @@ def test_only_the_last_step_keeps_scoring_inputs():
     ctx = _ctx()
     pairs = _pairs(9)
 
+    inputs = []
     for optim_step_idx in range(CRITIC_STEPS - 1):
-        _critic_step(ctx, pairs[optim_step_idx * 3 : optim_step_idx * 3 + 3], optim_step_idx)
-    assert ctx.train_output_inputs == []
+        _critic_step(ctx, pairs[optim_step_idx * 3 : optim_step_idx * 3 + 3], optim_step_idx, inputs)
+    assert inputs == []
 
-    _critic_step(ctx, pairs[6:], CRITIC_STEPS - 1)
-    scores = dmd2.score_last_step_samples(ctx, pairs)
+    _critic_step(ctx, pairs[6:], CRITIC_STEPS - 1, inputs)
+    assert len(inputs) == 2  # one entry per micro-batch of the last step
+    scores = dmd2.score_last_step_samples(ctx, pairs, inputs)
 
     assert list(scores) == [106, 107, 108]
 
@@ -273,8 +294,9 @@ def test_scores_match_separate_forwards_at_the_kept_input():
     ctx.models["transformer"] = _NoisyInputVelocity()
     pairs = _pairs(5)
 
-    steps = _critic_step(ctx, pairs, CRITIC_STEPS - 1)
-    scores = dmd2.score_last_step_samples(ctx, pairs)
+    inputs = []
+    steps = _critic_step(ctx, pairs, CRITIC_STEPS - 1, inputs)
+    scores = dmd2.score_last_step_samples(ctx, pairs, inputs)
 
     assert list(scores) == [100, 101, 102, 103, 104]
     for batch, prepared, _, _ in steps:
@@ -292,8 +314,9 @@ def test_teacher_runs_cfg_and_the_critic_runs_unguided():
     ctx = _ctx(teacher_guidance_scale=3.0, teacher_scale=1.0, critic_scale=2.0)
     pairs = _pairs(2)
 
-    [(_, prepared, _, _)] = _critic_step(ctx, pairs, CRITIC_STEPS - 1)
-    scores = dmd2.score_last_step_samples(ctx, pairs)
+    inputs = []
+    [(_, prepared, _, _)] = _critic_step(ctx, pairs, CRITIC_STEPS - 1, inputs)
+    scores = dmd2.score_last_step_samples(ctx, pairs, inputs)
 
     x_s = prepared.latents["visual"]
     sigma = prepared.extras["sigmas"]["visual"].view(-1, 1, 1, 1)
@@ -307,22 +330,25 @@ def test_teacher_runs_cfg_and_the_critic_runs_unguided():
 
 def test_teacher_cfg_needs_the_rollout_negative_cond():
     ctx = _ctx(teacher_guidance_scale=3.0)
-    _critic_step(ctx, _pairs(2, with_neg=False), CRITIC_STEPS - 1)
+    inputs = []
+    _critic_step(ctx, _pairs(2, with_neg=False), CRITIC_STEPS - 1, inputs)
     with pytest.raises(ValueError, match="needs the rollout's negative cond"):
-        dmd2.score_last_step_samples(ctx, _pairs(2, with_neg=False))
+        dmd2.score_last_step_samples(ctx, _pairs(2, with_neg=False), inputs)
     # Unguided, the teacher reads only the positive cond.
     ctx = _ctx(teacher_guidance_scale=1.0)
-    _critic_step(ctx, _pairs(2, with_neg=False), CRITIC_STEPS - 1)
-    dmd2.score_last_step_samples(ctx, _pairs(2, with_neg=False))
+    inputs = []
+    _critic_step(ctx, _pairs(2, with_neg=False), CRITIC_STEPS - 1, inputs)
+    dmd2.score_last_step_samples(ctx, _pairs(2, with_neg=False), inputs)
 
 
 def test_scoring_reshards_the_teacher(monkeypatch):
     resharded = []
     monkeypatch.setattr(dmd2, "reshard_model", resharded.append)
     ctx = _ctx()
-    _critic_step(ctx, _pairs(2), CRITIC_STEPS - 1)
+    inputs = []
+    _critic_step(ctx, _pairs(2), CRITIC_STEPS - 1, inputs)
 
-    dmd2.score_last_step_samples(ctx, _pairs(2))
+    dmd2.score_last_step_samples(ctx, _pairs(2), inputs)
 
     assert resharded == [ctx.teacher_models["transformer"]]
 
@@ -333,11 +359,12 @@ def test_actor_keeps_the_critic_last_step_pairs():
     schedule = build_microbatch_schedule(
         num_pairs_per_optim_step=3, num_optim_steps_per_rollout=CRITIC_STEPS, micro_batch_size=3
     )
+    inputs = []
     for optim_step_idx, [(pair_lo, pair_hi)] in enumerate(schedule):
-        _critic_step(ctx, pairs[pair_lo:pair_hi], optim_step_idx)
+        _critic_step(ctx, pairs[pair_lo:pair_hi], optim_step_idx, inputs)
 
     actor_ctx = _ctx()
-    actor_ctx.external_data["critic"] = dmd2.score_last_step_samples(ctx, pairs)
+    actor_ctx.external_data["critic"] = dmd2.score_last_step_samples(ctx, pairs, inputs)
 
     assert dmd2.select_scored_pairs(actor_ctx, pairs) == pairs[6:]
 
@@ -371,12 +398,22 @@ def _metrics():
 
 def _loss(ctx, pairs, x0):
     """dmd2_loss_formula with the student predicting ``x0`` from each pair's rollout input."""
-    prepared = dmd2.prepare_generator_batch(ctx, pairs)
+    prepared = dmd2.prepare_generator_batch(ctx, pairs, microbatch_id=0, optim_step_idx=0)
     sigmas = prepared.extras["sigmas"].view(-1, 1, 1, 1)
     new_pred = {"visual": (prepared.latents["visual"] - x0) / sigmas}
     metrics = _metrics()
     return (
-        dmd2.dmd2_loss_formula(ctx, pairs, prepared, new_pred=new_pred, old_pred=None, ref_pred=None, metrics=metrics),
+        dmd2.dmd2_loss_formula(
+            ctx,
+            pairs,
+            prepared,
+            new_pred=new_pred,
+            old_pred=None,
+            ref_pred=None,
+            metrics=metrics,
+            optim_step_idx=0,
+            train_output_inputs=[],
+        ),
         metrics,
     )
 
@@ -392,7 +429,7 @@ def test_generator_batch_is_the_student_step_on_its_input():
     ctx = _ctx()
     pairs = _generator_pairs(2)
 
-    prepared = dmd2.prepare_generator_batch(ctx, pairs)
+    prepared = dmd2.prepare_generator_batch(ctx, pairs, microbatch_id=0, optim_step_idx=0)
 
     assert prepared.model is ctx.models["transformer"]
     assert torch.equal(prepared.latents["visual"], torch.stack([pair["latent"] for pair in pairs]).float())
@@ -406,7 +443,7 @@ def test_generator_needs_a_single_component():
     ctx = _ctx()
     ctx.models["transformer_2"] = ctx.models["transformer"]
     with pytest.raises(ValueError, match="single DiT component"):
-        dmd2.prepare_generator_batch(ctx, _generator_pairs(1))
+        dmd2.prepare_generator_batch(ctx, _generator_pairs(1), microbatch_id=0, optim_step_idx=0)
 
 
 def test_surrogate_gradient_matches_the_reference():
@@ -414,7 +451,7 @@ def test_surrogate_gradient_matches_the_reference():
     pairs = _generator_pairs(2)
     x0 = torch.randn(2, *LATENT_SHAPE, requires_grad=True)
     x0_real, x0_fake = torch.randn(2, *LATENT_SHAPE), torch.randn(2, *LATENT_SHAPE)
-    ctx.external_data = {"critic": _scores(pairs, x0_real, x0_fake)}
+    ctx = dataclasses.replace(ctx, external_data={"critic": _scores(pairs, x0_real, x0_fake)})
 
     loss_sum, _ = _loss(ctx, pairs, x0)
     (loss_sum / len(pairs)).backward()
@@ -434,7 +471,7 @@ def test_critic_matching_the_teacher_gives_no_gradient():
     pairs = _generator_pairs(2)
     x0 = torch.randn(2, *LATENT_SHAPE, requires_grad=True)
     x0_pred = torch.randn(2, *LATENT_SHAPE)
-    ctx.external_data = {"critic": _scores(pairs, x0_pred, x0_pred)}
+    ctx = dataclasses.replace(ctx, external_data={"critic": _scores(pairs, x0_pred, x0_pred)})
 
     loss_sum, metrics = _loss(ctx, pairs, x0)
     loss_sum.backward()
@@ -460,13 +497,8 @@ def posterior_mean(mu, x_s, a, b):
         x0 = student_mu + torch.randn(2, *LATENT_SHAPE, generator=generator)
         a, b = 0.5, 0.5
         x_s = a * x0.detach() + b * torch.randn(x0.shape, generator=generator)
-        ctx.external_data = {
-            "critic": _scores(
-                pairs,
-                posterior_mean(teacher_mu, x_s, a, b),
-                posterior_mean(float(student_mu), x_s, a, b),
-            )
-        }
+        scores = _scores(pairs, posterior_mean(teacher_mu, x_s, a, b), posterior_mean(float(student_mu), x_s, a, b))
+        ctx = dataclasses.replace(ctx, external_data={"critic": scores})
         loss_sum, _ = _loss(ctx, pairs, x0)
         (loss_sum / len(pairs)).backward()
         optimizer.step()
@@ -480,7 +512,7 @@ def test_trainer_and_rollout_x0_diff_is_recorded():
     ctx = _ctx()
     pairs = _generator_pairs(2)
     x0 = torch.stack([pair["x0"] for pair in pairs]) + 0.25
-    ctx.external_data = {"critic": _scores(pairs, x0, x0)}
+    ctx = dataclasses.replace(ctx, external_data={"critic": _scores(pairs, x0, x0)})
 
     _, metrics = _loss(ctx, pairs, x0)
 
@@ -494,7 +526,10 @@ def test_health_metrics_are_recorded_overall_and_per_step():
     pairs = _generator_pairs(2)
     # Pair n stopped at step n: x0 = 0.5 + n against a teacher x0 of 0 and a critic x0 of 0.25.
     x0 = torch.stack([pair["x0"] for pair in pairs]) + torch.tensor([0.5, 1.5]).view(2, 1, 1, 1)
-    ctx.external_data = {"critic": _scores(pairs, torch.zeros(2, *LATENT_SHAPE), torch.full((2, *LATENT_SHAPE), 0.25))}
+    ctx = dataclasses.replace(
+        ctx,
+        external_data={"critic": _scores(pairs, torch.zeros(2, *LATENT_SHAPE), torch.full((2, *LATENT_SHAPE), 0.25))},
+    )
 
     _, metrics = _loss(ctx, pairs, x0)
 
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 c3e596099..c711e54a6 100644
--- a/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py
+++ b/tests/fast/backends/fsdp_utils/test_loss_hub_nft.py
@@ -89,6 +89,8 @@ def loss(kl_beta):
                 old_pred=old_pred,
                 ref_pred={"visual": new_pred["visual"] - 1.0},
                 metrics=metrics,
+                optim_step_idx=0,
+                train_output_inputs=[],
             )
 
         torch.testing.assert_close(loss(0.5) - loss(0.0), torch.tensor(0.5 * 2))
@@ -240,7 +242,7 @@ class _Env:
         pos_cond_kwargs = None
         neg_cond_kwargs = None
 
-    def _ctx(self, config, microbatch_id=0):
+    def _ctx(self, config):
         return DiffusionLossContext(
             models={"transformer": torch.nn.Identity()},
             train_pipeline_config=config,
@@ -248,7 +250,6 @@ def _ctx(self, config, microbatch_id=0):
             args=Namespace(seed=42),
             forward_dtype=torch.float32,
             device=torch.device("cpu"),
-            microbatch_id=microbatch_id,
         )
 
     def _batch(self):
@@ -271,18 +272,18 @@ def process_sigma_as_timesteps_input(self, sigmas):
                 seen["sigmas"] = sigmas.clone()
                 return sigmas
 
-        prepare_nft_batch(self._ctx(_Recording()), self._batch())
+        prepare_nft_batch(self._ctx(_Recording()), self._batch(), microbatch_id=0, optim_step_idx=0)
         assert torch.equal(seen["sigmas"], torch.tensor(self.SIGMAS))
 
     def test_sd3_style_family_gets_sigma_times_num_train_timesteps(self):
-        prepared = prepare_nft_batch(self._ctx(_Sd3StyleConfig()), self._batch())
+        prepared = prepare_nft_batch(self._ctx(_Sd3StyleConfig()), self._batch(), microbatch_id=0, optim_step_idx=0)
         assert torch.equal(prepared.timesteps["visual"], torch.tensor(self.SIGMAS))
         assert torch.equal(
             prepared.timesteps_for_model["visual"], torch.tensor(self.SIGMAS) * float(_StubConfig.num_train_timesteps)
         )
 
     def test_qwen_style_family_gets_the_sigma_bit_exactly(self):
-        prepared = prepare_nft_batch(self._ctx(_QwenStyleConfig()), self._batch())
+        prepared = prepare_nft_batch(self._ctx(_QwenStyleConfig()), self._batch(), microbatch_id=0, optim_step_idx=0)
         # equal, not allclose: the wrong hook would still pass allclose.
         assert torch.equal(prepared.timesteps_for_model["visual"], torch.tensor(self.SIGMAS))
 
@@ -354,12 +355,14 @@ def test_each_sample_draws_its_own_permutation(self):
     def test_corruption_noise_reproduces(self):
         ctx = TestPrepareNftBatch()._ctx(_Sd3StyleConfig())
         batch = TestPrepareNftBatch()._batch()
-        first = prepare_nft_batch(ctx, batch)
-        second = prepare_nft_batch(ctx, batch)
+        first = prepare_nft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
+        second = prepare_nft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         assert torch.equal(first.latents["visual"], second.latents["visual"])
 
     def test_a_different_microbatch_draws_different_noise(self):
         harness = TestPrepareNftBatch()
-        first = prepare_nft_batch(harness._ctx(_Sd3StyleConfig()), harness._batch())
-        second = prepare_nft_batch(harness._ctx(_Sd3StyleConfig(), microbatch_id=1), harness._batch())
+        first = prepare_nft_batch(harness._ctx(_Sd3StyleConfig()), harness._batch(), microbatch_id=0, optim_step_idx=0)
+        second = prepare_nft_batch(
+            harness._ctx(_Sd3StyleConfig()), harness._batch(), microbatch_id=1, optim_step_idx=0
+        )
         assert not torch.equal(first.latents["visual"], second.latents["visual"])
diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_selection.py b/tests/fast/backends/fsdp_utils/test_loss_hub_selection.py
index f04403a15..0f1019464 100644
--- a/tests/fast/backends/fsdp_utils/test_loss_hub_selection.py
+++ b/tests/fast/backends/fsdp_utils/test_loss_hub_selection.py
@@ -80,7 +80,7 @@ def test_the_table_has_no_other_entries():
 
 
 def test_custom_loss_loads_every_path():
-    def prepare(ctx, batch, *, pad_to_len=None):
+    def prepare(ctx, batch, *, microbatch_id, optim_step_idx, pad_to_len=None):
         pass
 
     def formula(ctx, batch, prepared, **kwargs):
@@ -89,7 +89,7 @@ def formula(ctx, batch, prepared, **kwargs):
     def train_input(ctx, train_pairs):
         pass
 
-    def train_output(ctx, train_pairs):
+    def train_output(ctx, train_pairs, train_output_inputs):
         pass
 
     args = _custom_args(
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 d3b7e3afb..bd63f01d2 100644
--- a/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py
+++ b/tests/fast/backends/fsdp_utils/test_loss_hub_sft.py
@@ -4,7 +4,7 @@
                 │
     prepare_sft_batch:
       grid       ──► FlowNoiseSchedule.shifted(--fsdp-flow-shift[name])
-      select_component(visual grid), once, rank-aligned ──► DiT component + the visual grid points it serves
+      select_component(visual grid, microbatch_id), once, rank-aligned ──► DiT component + the visual grid points it serves
       each stream gets its own seeded generator
                        (visual: the single-stream seed; other streams add their name to it)
       grid index ──► one per pair, uniform over the component's points, drawn by the visual generator
@@ -36,6 +36,7 @@
 
 register_cpu_ci(est_time=30, suite="stage-a-cpu", labels=[])
 
+import dataclasses
 from argparse import Namespace
 
 import pytest
@@ -77,7 +78,7 @@ def _noise_schedule():
     return FlowNoiseSchedule(sigmas * NUM_TRAIN_TIMESTEPS, torch.cat([sigmas, torch.zeros(1)]))
 
 
-def _ctx(models, rollout_id=3, microbatch_id=0, dp_rank=0, config=None):
+def _ctx(models, rollout_id=3, dp_rank=0, config=None):
     return DiffusionLossContext(
         models=models,
         train_pipeline_config=config if config is not None else _Config(),
@@ -91,7 +92,6 @@ def _ctx(models, rollout_id=3, microbatch_id=0, dp_rank=0, config=None):
         forward_dtype=torch.float32,
         device=torch.device("cpu"),
         rollout_id=rollout_id,
-        microbatch_id=microbatch_id,
         dp_rank=dp_rank,
     )
 
@@ -119,12 +119,18 @@ def emit_mean(self, key, *, total, count):
         self.seen[key] = (prev_total + float(total), prev_count + count)
 
 
+def test_the_context_is_read_only():
+    ctx = _ctx({"transformer": nn.Identity()})
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        ctx.rollout_id = 4
+
+
 class TestPrepareSftBatch:
     def test_corruption_and_target_identity(self):
         torch.manual_seed(0)
         ctx = _ctx({"transformer": nn.Identity()})
         batch = _batch()
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
 
         x0 = torch.stack([pair["latent"]["visual"] for pair in batch]).float()
         sigma = (prepared.timesteps["visual"] / NUM_TRAIN_TIMESTEPS).view(-1, 1, 1, 1, 1)
@@ -136,7 +142,7 @@ def test_corruption_and_target_identity(self):
 
     def test_single_model_serves_the_whole_grid(self):
         ctx = _ctx({"transformer": nn.Identity()}, config=_SingleConfig())
-        _, _, pool = select_component(ctx, _noise_schedule(), torch.device("cpu"))
+        _, _, pool = select_component(ctx, _noise_schedule(), torch.device("cpu"), microbatch_id=0)
         assert torch.equal(pool, torch.arange(NUM_GRID))
 
     def test_single_wan_expert_serves_only_its_timesteps(self):
@@ -144,7 +150,7 @@ def test_single_wan_expert_serves_only_its_timesteps(self):
         timesteps = _noise_schedule().timesteps
         for component_name in ("transformer", "transformer_2"):
             ctx = _ctx({component_name: nn.Identity()}, config=config)
-            name, _, pool = select_component(ctx, _noise_schedule(), torch.device("cpu"))
+            name, _, pool = select_component(ctx, _noise_schedule(), torch.device("cpu"), microbatch_id=0)
             expected = [
                 i
                 for i, timestep in enumerate(timesteps)
@@ -160,8 +166,7 @@ def test_dual_expert_micro_batch_is_phase_pure(self):
         noise_schedule = _noise_schedule()
         picked = set()
         for call in range(20):
-            ctx.microbatch_id = call
-            name, model, pool = select_component(ctx, noise_schedule, torch.device("cpu"))
+            name, model, pool = select_component(ctx, noise_schedule, torch.device("cpu"), microbatch_id=call)
             picked.add(name)
             assert model is models[name]
             for i in pool.tolist():
@@ -174,11 +179,11 @@ def test_prepare_is_independent_of_global_rng_state(self):
         torch.manual_seed(0)
         global_state = torch.get_rng_state()
 
-        first = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+        first = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch, microbatch_id=0, optim_step_idx=0)
         assert torch.equal(torch.get_rng_state(), global_state)
 
         torch.rand(1000)
-        second = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+        second = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch, microbatch_id=0, optim_step_idx=0)
         assert torch.equal(first.timesteps["visual"], second.timesteps["visual"])
         assert torch.equal(first.latents["visual"], second.latents["visual"])
         assert torch.equal(_visual_target(first), _visual_target(second))
@@ -188,7 +193,7 @@ def test_visual_stream_draws_what_the_single_stream_seed_draws(self):
         # the draws single-stream SFT made before latents were keyed by stream, rebuilt by hand.
         ctx = _ctx({"transformer": nn.Identity()})
         batch = _batch()
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
 
         generator = torch.Generator().manual_seed(stable_hash("sample", 42, 3, 0, 0))
         noise_schedule = FlowNoiseSchedule.shifted(FLOW_SHIFT, NUM_TRAIN_TIMESTEPS)
@@ -201,17 +206,21 @@ def test_visual_stream_draws_what_the_single_stream_seed_draws(self):
 
     def test_samples_within_microbatch_use_different_noise(self):
         batch = _batch(bsz=4)
-        prepared = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+        prepared = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch, microbatch_id=0, optim_step_idx=0)
         x0 = torch.stack([pair["latent"]["visual"] for pair in batch]).float()
         noise = _visual_target(prepared) + 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)
+        base = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch, microbatch_id=0, optim_step_idx=0)
+        other_rollout = prepare_sft_batch(
+            _ctx({"transformer": nn.Identity()}, rollout_id=4), batch, microbatch_id=0, optim_step_idx=0
+        )
+        other_slot = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch, microbatch_id=1, optim_step_idx=0)
+        other_dp_rank = prepare_sft_batch(
+            _ctx({"transformer": nn.Identity()}, dp_rank=1), batch, microbatch_id=0, optim_step_idx=0
+        )
         for other in (other_rollout, other_slot, other_dp_rank):
             assert not torch.equal(_visual_target(base), _visual_target(other))
 
@@ -219,8 +228,8 @@ 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())
+            a = prepare_sft_batch(_ctx(models, dp_rank=0), _batch(), microbatch_id=slot, optim_step_idx=0)
+            b = prepare_sft_batch(_ctx(models, dp_rank=1), _batch(), microbatch_id=slot, optim_step_idx=0)
             assert a.component_name == b.component_name
             assert not torch.equal(_visual_target(a), _visual_target(b))
 
@@ -229,7 +238,9 @@ def test_prepare_draws_from_the_flow_shift_grid(self):
         noise_schedule = FlowNoiseSchedule.shifted(3.0, 4)
         torch.testing.assert_close(noise_schedule.sigmas, torch.tensor([1.0, 0.9, 0.75, 0.5, 0.0]))
         torch.testing.assert_close(noise_schedule.timesteps, torch.tensor([4.0, 3.6, 3.0, 2.0]))
-        prepared = prepare_sft_batch(_ctx({"transformer": nn.Identity()}), _batch(bsz=16))
+        prepared = prepare_sft_batch(
+            _ctx({"transformer": nn.Identity()}), _batch(bsz=16), microbatch_id=0, optim_step_idx=0
+        )
         grid = FlowNoiseSchedule.shifted(FLOW_SHIFT, NUM_TRAIN_TIMESTEPS)
         assert set(prepared.extras["sigmas"]["visual"].tolist()) <= set(grid.sigmas[:-1].tolist())
         assert set(prepared.timesteps["visual"].tolist()) <= set(grid.timesteps.tolist())
@@ -240,10 +251,18 @@ def _loss(self, prediction_offset, metrics, log_loss_sigma_bucket=5):
         ctx = _ctx({"transformer": nn.Identity()})
         ctx.args.log_loss_sigma_bucket = log_loss_sigma_bucket
         batch = _batch()
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         new_pred = {"visual": _visual_target(prepared) + prediction_offset}
         return prepared, sft_loss_formula(
-            ctx, batch, prepared, new_pred=new_pred, old_pred=None, ref_pred=None, metrics=metrics
+            ctx,
+            batch,
+            prepared,
+            new_pred=new_pred,
+            old_pred=None,
+            ref_pred=None,
+            metrics=metrics,
+            optim_step_idx=0,
+            train_output_inputs=[],
         )
 
     def test_zero_loss_on_exact_velocity(self):
@@ -295,7 +314,9 @@ class TestJointStreams:
     def test_each_stream_draws_from_its_own_flow_shift_grid(self):
         # --fsdp-flow-shift visual=3,audio=1: each stream's sigmas lie on its own grid, not on the other's.
         batch = _joint_batch(16)
-        prepared = prepare_sft_batch(_ctx({"transformer": nn.Identity()}, config=_SingleConfig()), batch)
+        prepared = prepare_sft_batch(
+            _ctx({"transformer": nn.Identity()}, config=_SingleConfig()), batch, microbatch_id=0, optim_step_idx=0
+        )
         assert set(prepared.latents) == {"visual", "audio"}
         grids = {
             name: set(FlowNoiseSchedule.shifted(FLOW_SHIFTS[name], NUM_TRAIN_TIMESTEPS).sigmas[:-1].tolist())
@@ -316,14 +337,14 @@ def test_streams_share_the_grid_index_but_not_the_noise(self):
         batch = _joint_batch(32)
         for pair in batch:
             pair["latent"]["audio"] = pair["latent"]["visual"].clone()
-        first = prepare_sft_batch(ctx, batch)
+        first = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         grids = {name: FlowNoiseSchedule.shifted(FLOW_SHIFTS[name], NUM_TRAIN_TIMESTEPS) for name in FLOW_SHIFTS}
         idx = torch.searchsorted(-grids["visual"].sigmas.float(), -first.extras["sigmas"]["visual"])
         assert torch.equal(first.extras["sigmas"]["audio"], grids["audio"].sigmas.float()[idx])
         assert not torch.equal(first.extras["target"]["visual"], first.extras["target"]["audio"])
         for pair in batch:
             pair["latent"]["action"] = torch.ones(10, 7)
-        second = prepare_sft_batch(ctx, batch)
+        second = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         for name in ("visual", "audio"):
             assert torch.equal(first.extras["target"][name], second.extras["target"][name])
             assert torch.equal(first.timesteps[name], second.timesteps[name])
@@ -332,14 +353,22 @@ def test_loss_averages_each_stream_over_its_own_size(self):
         # visual (96 elements) off by 1, audio (12 elements) off by 2: per pair 1 + 4, not a size-weighted mix.
         ctx = _ctx({"transformer": nn.Identity()}, config=_SingleConfig())
         batch = _joint_batch(2)
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         predictions = {
             "visual": (prepared.extras["target"]["visual"] + 1.0).requires_grad_(),
             "audio": (prepared.extras["target"]["audio"] + 2.0).requires_grad_(),
         }
         metrics = _Metrics()
         loss = sft_loss_formula(
-            ctx, batch, prepared, new_pred=predictions, old_pred=None, ref_pred=None, metrics=metrics
+            ctx,
+            batch,
+            prepared,
+            new_pred=predictions,
+            old_pred=None,
+            ref_pred=None,
+            metrics=metrics,
+            optim_step_idx=0,
+            train_output_inputs=[],
         )
         assert loss.item() == pytest.approx(10.0)
         loss.backward()
@@ -354,13 +383,23 @@ def test_joint_metrics_split_by_stream_and_bucket_by_own_sigma(self):
 
         ctx = _ctx({"transformer": nn.Identity()}, config=_SingleConfig())
         batch = _joint_batch(2)
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         predictions = {
             "visual": prepared.extras["target"]["visual"] + 1.0,
             "audio": prepared.extras["target"]["audio"] + 2.0,
         }
         metrics = _Metrics()
-        sft_loss_formula(ctx, batch, prepared, new_pred=predictions, old_pred=None, ref_pred=None, metrics=metrics)
+        sft_loss_formula(
+            ctx,
+            batch,
+            prepared,
+            new_pred=predictions,
+            old_pred=None,
+            ref_pred=None,
+            metrics=metrics,
+            optim_step_idx=0,
+            train_output_inputs=[],
+        )
         assert metrics.seen["loss"] == (10.0, 2)
         assert metrics.seen["loss_visual"] == (2.0, 2)
         assert metrics.seen["loss_audio"] == (8.0, 2)
@@ -380,7 +419,7 @@ def test_a_stream_the_config_does_not_predict_adds_no_loss(self):
         # H3 Ref2VA on a silent target: the audio stream is noised and fed to the DiT, but only visual is predicted.
         ctx = _ctx({"transformer": nn.Identity()}, config=_SingleConfig())
         batch = _joint_batch(2)
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         loss = sft_loss_formula(
             ctx,
             batch,
@@ -389,6 +428,8 @@ def test_a_stream_the_config_does_not_predict_adds_no_loss(self):
             old_pred=None,
             ref_pred=None,
             metrics=_Metrics(),
+            optim_step_idx=0,
+            train_output_inputs=[],
         )
         assert loss.item() == pytest.approx(2.0)
 
@@ -397,7 +438,7 @@ def test_a_stream_left_out_of_supervised_streams_adds_no_loss(self):
         ctx = _ctx({"transformer": nn.Identity()}, config=_SingleConfig())
         ctx.args.fsdp_supervised_streams = ["visual"]
         batch = _joint_batch(2)
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         metrics = _Metrics()
         loss = sft_loss_formula(
             ctx,
@@ -410,6 +451,8 @@ def test_a_stream_left_out_of_supervised_streams_adds_no_loss(self):
             old_pred=None,
             ref_pred=None,
             metrics=metrics,
+            optim_step_idx=0,
+            train_output_inputs=[],
         )
         assert loss.item() == pytest.approx(2.0)
         # One supervised stream: its loss is the total, and only its own buckets appear.
@@ -420,8 +463,10 @@ def test_visual_stream_selects_a_rank_aligned_expert(self):
         models = {"transformer": nn.Identity(), "transformer_2": nn.Identity()}
         batch = _joint_batch(4)
         for slot in range(5):
-            contexts = [_ctx(models, microbatch_id=slot, dp_rank=rank) for rank in (0, 1)]
-            first, second = [prepare_sft_batch(context, batch) for context in contexts]
+            contexts = [_ctx(models, dp_rank=rank) for rank in (0, 1)]
+            first, second = [
+                prepare_sft_batch(context, batch, microbatch_id=slot, optim_step_idx=0) for context in contexts
+            ]
             assert first.component_name == second.component_name
             for timestep in first.timesteps["visual"]:
                 component = contexts[0].train_pipeline_config.component_for_timestep(float(timestep))
@@ -430,16 +475,24 @@ def test_visual_stream_selects_a_rank_aligned_expert(self):
     def test_rejects_misshaped_predictions(self):
         ctx = _ctx({"transformer": nn.Identity()}, config=_SingleConfig())
         batch = _joint_batch(1)
-        prepared = prepare_sft_batch(ctx, batch)
+        prepared = prepare_sft_batch(ctx, batch, microbatch_id=0, optim_step_idx=0)
         predictions = dict(prepared.extras["target"])
         predictions["audio"] = predictions["audio"][0]
         with pytest.raises(ValueError, match="prediction shape"):
             sft_loss_formula(
-                ctx, batch, prepared, new_pred=predictions, old_pred=None, ref_pred=None, metrics=_Metrics()
+                ctx,
+                batch,
+                prepared,
+                new_pred=predictions,
+                old_pred=None,
+                ref_pred=None,
+                metrics=_Metrics(),
+                optim_step_idx=0,
+                train_output_inputs=[],
             )
 
     def test_rejects_different_stream_sets_in_one_microbatch(self):
         batch = _joint_batch(2)
         del batch[1]["latent"]["audio"]
         with pytest.raises(ValueError, match="latent streams"):
-            prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch)
+            prepare_sft_batch(_ctx({"transformer": nn.Identity()}), batch, microbatch_id=0, optim_step_idx=0)
diff --git a/tests/fast/backends/fsdp_utils/test_reference_models.py b/tests/fast/backends/fsdp_utils/test_reference_models.py
index 8f6f2ad99..f57aaadfb 100644
--- a/tests/fast/backends/fsdp_utils/test_reference_models.py
+++ b/tests/fast/backends/fsdp_utils/test_reference_models.py
@@ -158,9 +158,14 @@ def make_actor_forward_harness(
         true_cfg_scale=None,
     )
 
+    def prepare(ctx, batch, **kwargs):
+        observed["prepare_kwargs"] = kwargs
+        return prepared
+
     def loss_formula(ctx, batch, prepared, *, new_pred, old_pred, ref_pred, **kwargs):
         predictions = {"new_pred": new_pred, "old_pred": old_pred, "ref_pred": ref_pred}
         observed.update({name: pred["visual"] for name, pred in predictions.items() if pred is not None})
+        observed["formula_kwargs"] = kwargs
         return (new_pred["visual"] - ref_pred["visual"]).square().mean()
 
     harness = SimpleNamespace(
@@ -171,7 +176,7 @@ def loss_formula(ctx, batch, prepared, *, new_pred, old_pred, ref_pred, **kwargs
         reference_models=reference_models,
         train_pipeline_config=TinyPipelineConfig(),
         _forward_dtype=torch.float32,
-        loss_hooks=SimpleNamespace(prepare=lambda *args, **kwargs: prepared, formula=loss_formula),
+        loss_hooks=SimpleNamespace(prepare=prepare, formula=loss_formula),
     )
     return harness, observed
 
@@ -259,7 +264,15 @@ def test_loaded_reference_and_teacher_are_independent_and_frozen(
             teacher_output = teacher[component](inputs)
         assert not torch.equal(target, teacher_output)
         harness, observed = make_actor_forward_harness(actor, reference, component, inputs, "ref")
-        load_actor_forward_method()(harness, None, [], metrics=None).backward()
+        load_actor_forward_method()(
+            harness,
+            None,
+            [],
+            microbatch_id=0,
+            optim_step_idx=0,
+            train_output_inputs=[],
+            metrics=None,
+        ).backward()
         oracle_output = oracle(inputs)
         (oracle_output - target).square().mean().backward()
         torch.testing.assert_close(observed["new_pred"], oracle_output)
@@ -299,7 +312,15 @@ def test_lora_base_forward_restores_adapters_and_actor_gradients(tmp_path, unsha
     with torch.no_grad():
         target = base(inputs)
     harness, observed = make_actor_forward_harness(actor, {}, component, inputs, "lora_base")
-    load_actor_forward_method()(harness, None, [], metrics=None).backward()
+    load_actor_forward_method()(
+        harness,
+        None,
+        [],
+        microbatch_id=0,
+        optim_step_idx=0,
+        train_output_inputs=[],
+        metrics=None,
+    ).backward()
     oracle_output = oracle(inputs)
     (oracle_output - target).square().mean().backward()
     torch.testing.assert_close(observed["ref_pred"], target)
@@ -353,7 +374,15 @@ def test_lora_base_reference_keeps_adapter_gradients_on_fsdp_shards():
             harness, _ = make_actor_forward_harness(
                 {"transformer": model}, {}, "transformer", torch.ones(2, 3), "lora_base"
             )
-            forward(harness, None, [], metrics=None).backward()
+            forward(
+                harness,
+                None,
+                [],
+                microbatch_id=0,
+                optim_step_idx=0,
+                train_output_inputs=[],
+                metrics=None,
+            ).backward()
             for name, parameter in adapter_parameters.items():
                 assert parameter.requires_grad and parameter.grad is not None, (step, name)
             model.zero_grad(set_to_none=True)
@@ -390,7 +419,15 @@ def test_nft_gets_ema_as_old_and_lora_base_as_kl_reference(tmp_path, unsharded_m
     harness, observed = make_actor_forward_harness(
         actor, {}, component, inputs, "lora_base", loss_type="nft", ema_optimizer=ema_optimizer
     )
-    load_actor_forward_method()(harness, None, [], metrics=None)
+    load_actor_forward_method()(
+        harness,
+        None,
+        [],
+        microbatch_id=0,
+        optim_step_idx=0,
+        train_output_inputs=[],
+        metrics=None,
+    )
     torch.testing.assert_close(observed["new_pred"], expected_new)
     torch.testing.assert_close(observed["old_pred"], expected_old)
     torch.testing.assert_close(observed["ref_pred"], expected_ref)
@@ -452,3 +489,25 @@ def gather_object(*, obj, object_gather_list, **kwargs):
     namespace[method.name](harness, [("block.proj.lora_A.weight", torch.ones(2, 3))], "transformer")
     assert len(sent) == 1
     assert sent[0]["lora_rank"] == 2 and sent[0]["lora_alpha"] == 4
+
+
+def test_the_trunk_hands_the_hooks_their_own_counters_and_the_one_list():
+    model = torch.nn.Linear(3, 3)
+    harness, observed = make_actor_forward_harness(
+        {"block": model}, {"block": model}, "block", torch.ones(2, 3), "ref"
+    )
+    train_output_inputs = []
+
+    load_actor_forward_method()(
+        harness,
+        None,
+        [],
+        microbatch_id=7,
+        optim_step_idx=2,
+        train_output_inputs=train_output_inputs,
+        metrics=None,
+    )
+
+    assert observed["prepare_kwargs"] == {"microbatch_id": 7, "optim_step_idx": 2, "pad_to_len": None}
+    assert observed["formula_kwargs"]["optim_step_idx"] == 2
+    assert observed["formula_kwargs"]["train_output_inputs"] is train_output_inputs
diff --git a/tests/fast/backends/fsdp_utils/test_train_input_hook.py b/tests/fast/backends/fsdp_utils/test_train_input_hook.py
index ee94596ff..ccb573cc9 100644
--- a/tests/fast/backends/fsdp_utils/test_train_input_hook.py
+++ b/tests/fast/backends/fsdp_utils/test_train_input_hook.py
@@ -3,11 +3,14 @@
     train pairs (the converted rollout output) + ctx (external_data from earlier groups)
         --train-input hook (filter / reorder / rewrite)--> preprocessed pairs
     preprocessed pairs --> micro-batch schedule --> training forwards, in order
-    preprocessed pairs --> train-output hook (the same ctx)
+    preprocessed pairs --> train-output hook (the same ctx, the list the forwards filled)
     no hook            --> the converted train pairs, as before
 
+The trunk numbers each forward: ``microbatch_id`` counts the rollout's micro-batches, ``optim_step_idx`` is the
+window; the old-log-prob recompute of windows 1.. reuses their training forwards' numbers.
+
 The production _train_core runs on CPU with the forward, metrics and rank checks replaced; the
-optimizer step is skipped, so only the pairs and ctx each stage receives are observed.
+optimizer step is skipped, so only the pairs, ctx and numbers each stage receives are observed.
 """
 
 from tests.ci.ci_register import register_cpu_ci
@@ -51,8 +54,10 @@ def _load_train_core():
     return namespace["_train_core"]
 
 
-def _actor(train_input, train_output, forwarded):
+def _actor(train_input, train_output, forwarded, **args):
+    lists = []  # the train_output_inputs object each forward received
     return SimpleNamespace(
+        lists=lists,
         models={},
         reference_models={},
         teacher_models={},
@@ -62,10 +67,10 @@ def _actor(train_input, train_output, forwarded):
             num_steps_per_rollout=2,
             micro_batch_size=1,
             loss_type="policy_loss",
-            diffusion_recompute_old_log_prob=False,
             debug_skip_optimizer_step=True,
             log_loss_sigma_bucket=0,
             diffusion_num_steps=4,
+            **{"diffusion_recompute_old_log_prob": False, **args},
         ),
         _forward_dtype=torch.float32,
         parallel_state=SimpleNamespace(dp_rank=0, dp_group=None),
@@ -73,7 +78,13 @@ def _actor(train_input, train_output, forwarded):
         optimizer=SimpleNamespace(zero_grad=lambda set_to_none: None),
         global_step=0,
         _maybe_legacy_window_pad_len=lambda train_pairs, microbatch_ranges: None,
-        _forward_train_pair_batch=lambda ctx, chunk, **kwargs: forwarded.append((ctx.microbatch_id, chunk)),
+        _forward_train_pair_batch=lambda ctx, chunk, **kwargs: (
+            forwarded.append(
+                (kwargs["microbatch_id"], kwargs["optim_step_idx"], kwargs.get("write_old_log_prob", False), chunk)
+            ),
+            kwargs["train_output_inputs"].append(kwargs["microbatch_id"]),
+            lists.append(kwargs["train_output_inputs"]),
+        ),
         _log_metrics=lambda rollout_id, metrics, step: None,
     )
 
@@ -96,25 +107,46 @@ def train_input(ctx, train_pairs):
         ]
 
     forwarded = []
-    actor = _actor(train_input, lambda ctx, train_pairs: (ctx, train_pairs), forwarded)
-    ctx, output_pairs = _load_train_core()(
+    actor = _actor(train_input, lambda ctx, train_pairs, inputs: (ctx, train_pairs, inputs), forwarded)
+    ctx, output_pairs, train_output_inputs = _load_train_core()(
         actor, rollout_id=3, rollout_data={"train_data": pairs}, external_data={"critic": scores}
     )
 
     preprocessed = [{"sample_index": index, "score": scores[index]} for index in (1, 2, 4, 5)]
     assert seen_by_input == [(ctx, pairs)]
     assert (ctx.rollout_id, ctx.external_data) == (3, {"critic": scores})
-    assert forwarded == [(microbatch_id, [pair]) for microbatch_id, pair in enumerate(preprocessed)]
+    # Two optimizer steps of two micro-batches each: the id counts on, the step index follows the window.
+    assert forwarded == [
+        (0, 0, False, [preprocessed[0]]),
+        (1, 0, False, [preprocessed[1]]),
+        (2, 1, False, [preprocessed[2]]),
+        (3, 1, False, [preprocessed[3]]),
+    ]
     assert output_pairs == preprocessed
+    # Every forward appended to the one list the train-output hook received.
+    assert train_output_inputs == [0, 1, 2, 3]
+    assert all(received is train_output_inputs for received in actor.lists)
     assert actor.global_step == 2
 
 
 def test_without_the_hook_the_converted_pairs_train():
     pairs = [{"sample_index": index} for index in range(4)]
     forwarded = []
-    actor = _actor(None, lambda ctx, train_pairs: train_pairs, forwarded)
+    actor = _actor(None, lambda ctx, train_pairs, inputs: train_pairs, forwarded)
 
     output_pairs = _load_train_core()(actor, rollout_id=0, rollout_data={"train_data": pairs}, external_data={})
 
-    assert forwarded == [(index, [pair]) for index, pair in enumerate(pairs)]
+    assert forwarded == [(index, index // 2, False, [pair]) for index, pair in enumerate(pairs)]
     assert output_pairs == pairs
+
+
+def test_the_recompute_passes_the_numbers_of_the_training_forwards_of_windows_1_on():
+    pairs = [{"sample_index": index} for index in range(4)]
+    forwarded = []
+    actor = _actor(None, lambda ctx, train_pairs, inputs: None, forwarded, diffusion_recompute_old_log_prob=True)
+
+    _load_train_core()(actor, rollout_id=0, rollout_data={"train_data": pairs}, external_data={})
+
+    recompute = [(2, 1, True, [pairs[2]]), (3, 1, True, [pairs[3]])]
+    training = [(index, index // 2, False, [pair]) for index, pair in enumerate(pairs)]
+    assert forwarded == recompute + training
diff --git a/tests/fast/utils/test_dmd2_arguments.py b/tests/fast/utils/test_dmd2_arguments.py
index c4752cd9f..6427344aa 100644
--- a/tests/fast/utils/test_dmd2_arguments.py
+++ b/tests/fast/utils/test_dmd2_arguments.py
@@ -7,7 +7,7 @@
                       takes --dmd2-critic-steps-per-rollout steps over each DP shard
     other losses  ──► --critic-load / --critic-lr have no consumer and are rejected
     teacher score ──► CFG --dmd2-teacher-guidance-scale >= 1
-    actor         ──► unguided rollouts, one step per rollout
+    actor         ──► unguided rollouts, one step per rollout, no old-log-prob recompute (no ratio)
 """
 
 from tests.ci.ci_register import register_cpu_ci
@@ -36,6 +36,7 @@ def _args(**overrides):
         diffusion_true_cfg_scale=None,
         use_lora=False,
         ref_mode="none",
+        diffusion_recompute_old_log_prob=False,
         num_steps_per_rollout=1,
         global_batch_size=40,
         micro_batch_size=1,
@@ -127,6 +128,7 @@ def test_valid_dmd2_args_pass_unchanged():
         ({"use_lora": True}, "drop --use-lora"),
         ({"teacher_lora_adapter_paths": ["teacher-adapter"]}, "drop --use-lora and --teacher-lora-adapter-path"),
         ({"ref_mode": "ref"}, "uses no reference model"),
+        ({"diffusion_recompute_old_log_prob": True}, "not DMD2"),
         ({"num_steps_per_rollout": 2}, "one actor step per rollout"),
         ({"dmd2_critic_steps_per_rollout": 0}, "must be >= 1"),
         # 40 samples over 2 DP shards = 20 per shard, which 3 critic steps cannot split.