[Rockdu]

Stacked on #367 (the hook-registry PR), which sits on #343: the diff here is the 4 commits above it.

What

• DiffusionLossContext stops carrying loop state. microbatch_id and optim_step_idx become required keyword arguments of every prepare hook (and optim_step_idx of every formula), passed by the two trunk loops; the list the DMD2 critic formula fills for the train-output hook is owned by _train_core and passed to the formulas and to train_output(ctx, train_pairs, train_output_inputs).
• The context is frozen (@dataclass(frozen=True, eq=False)): a rollout-level, read-only view; nothing writes it after construction.
• --loss-type dmd2 rejects --diffusion-recompute-old-log-prob (no log-prob ratio to recompute; with a real step index the recompute loop would run the critic's scoring window twice), like sft_loss and nft already do.

Why

The planned micro-batch prefetch (prepare of micro-batch n+1 while the loss of n runs) races on any per-micro-batch field shared through the context, and a hook-to-hook channel through the context is invisible in the call graph. With the counters and the list as arguments, every per-call fact is in the signature and the context can be shared freely. Numbering is unchanged: training counts from 0; the old-log-prob recompute of windows 1.. reuses their training forwards' numbers.

Not in this PR: moving the DMD2 teacher forward into the trunk and deleting the train-output hook (left to the DMD2 stack's author).

Validation

CPU, on this head:

check	result
tests/fast (minus the Mac-only distributed / ray-socket cases)	634 passed, 3 skipped
one-shot comparison against the hook-registry PR's tree (ctx-field version): sft / nft / dmd2-critic prepares over 4 micro-batch ids x 2 step indices, the critic formula over 2 steps x 2 micro-batches, score_last_step_samples	28 checks bitwise equal
test_train_input_hook.py	the training loop and the recompute loop pass the expected (microbatch_id, optim_step_idx, write_old_log_prob) per forward; the train-output hook receives the same list object every forward appended to
test_reference_models.py	the real _forward_train_pair_batch hands prepare {microbatch_id, optim_step_idx, pad_to_len} and the formula optim_step_idx and the same list
pre-commit run --all-files	passes

GPU (4x H200, sglang-miles head), before the GPU runs were paused, on the same miles/ code except two argparse help strings: test_sd3_ocr_grpo_2xGPU 8/8, test_sd3_nft_pickscore_3xGPU 10/10, test_h3_t2va_sft_2xGPU 2/2 metric series bitwise.

Behaviour ledger (intentional changes): a custom prepare with the old signature prepare(ctx, batch, *, pad_to_len=None) fails at the first micro-batch with TypeError: unexpected keyword 'microbatch_id'; a custom train-output hook with the old two-argument signature fails the same way; a custom formula written as documented (with **kwargs) keeps working; reading ctx.microbatch_id / ctx.optim_step_idx / ctx.train_output_inputs is an AttributeError; writing a context field is a FrozenInstanceError. Everything else is bitwise.

Files

• miles/backends/fsdp_utils/loss_hub/types.py: three fields removed; context frozen
• miles/backends/fsdp_utils/actor.py: the two loops pass the counters; one list per train call
• miles/backends/fsdp_utils/loss_hub/{sft,nft,flow_grpo,dmd2}.py: hook signatures
• miles/backends/fsdp_utils/loss_hub/__init__.py: LossHooks contract
• miles/utils/arguments.py: dmd2 rejects the recompute flag; help strings
• docs/user-guide/customization.md: hook signatures and the numbering rule
• tests/fast/...: call sites updated; new wiring, accumulator-identity and frozen-context tests; tests/fast-gpu/backends/fsdp_utils/_reference_models_worker.py call sites updated

Checklist

• pre-commit run --all-files passes
• Added/updated tests for new behaviour
• pytest -x is green (tests/fast on CPU; tests/fast-gpu was updated but not run on a GPU)
• If launch flags changed, python3 train.py --help still parses (train_diffusion.py --help)
• If a public flag was added, it appears in the CLI reference docs (no flag added)
• If an example was added, it has a real walkthrough (no example added)