perf(unified-model): per-track micro-batch geometry for AR and image

[leviking98z-rgb]

Summary

UnifiedModelTrainStack slices both lineage levels with one shared micro_batch_size, but the two tracks are shaped differently:

• AR rows are variable-length responses — a larger micro pads to the longest row in the micro, so the activation peak is set by the worst row.
• image rows are fixed-shape latents — a larger micro is clean parallelism with no padding waste.

One number cannot suit both, so today the image track is pinned to whatever the AR track can afford.

This adds optional ar_micro_batch_size / image_micro_batch_size. Both default to the shared micro_batch_size, so omitting them is a behavioral no-op. _optimizer_step_slices and prepare_segment now take the size explicitly, which matters for correctness: a replay-anchored track (FlowGRPO under old_logp_source='replay') recomputes its π_old anchor at the exact (mini, micro) slices training will use, so the anchor must be frozen at the same per-track micro size or the on-policy ratio stops being 1.

Second, smaller change: _backward_part weights each micro's metrics by its sample share. The backward pass already uses a sample-share loss_scale, but metrics were averaged unweighted, so a ragged final micro (batch not divisible by the micro size) skewed the reported mean toward its fewer rows. aggregate_numeric_metrics grows a keyword-only weights argument, renormalized per key so a metric present in only some micros still averages correctly. Equal-sized micros reduce to the previous plain average.

Loss-equivalence caveat (documented in code): micro-batching only reconstructs the whole-batch objective when the per-micro sample share is the right weight. That holds for a per-row mean — FlowGRPO's flat latent mean, and GRPO's seq-mean-* modes — but not for GRPO's default token-mean when responses differ in length, where the exact objective weights by token count rather than row count. So ar_micro_batch_size > 1 wants a seq-mean-* loss_agg_mode. The image track (FlowGRPO) has no such constraint. This is why the recipes ship both knobs as null rather than raising the image default here.

Related Issue

N/A

Test Plan

All commands run from the repo root at 133c11c, venv with torch 2.13.0+cu130.

1. Behavior harness (24 assertions; uncommitted per the tests/-removal policy, source quoted below):

$ PYTHONPATH=$PWD python /tmp/verify_pertrack_mbs.py
PASS  unweighted mean unchanged          PASS  omitted arg uses shared mbs
PASS  equal weights == plain mean        PASS  ragged final micro
PASS  ragged weighted mean               PASS  sample shares sum to 1
PASS  ragged differs from unweighted     PASS  ragged share is smaller
PASS  missing key renormalizes           PASS  2 updates -> 2 step lists
PASS  length mismatch raises             PASS  each update 4 micros of 2
PASS  zero weight falls back             PASS  updates are disjoint+ordered
PASS  empty list                         PASS  ar_mbs=0 rejected
PASS  bool+tensor with weights           PASS  ar_mbs=-1 rejected
PASS  defaults fall back to shared       PASS  ar mbs=1 -> 8 micros
PASS  image mbs=2 -> 8 micros            PASS  image slices are width 2
PASS  ar covers all rows                 PASS  image covers all rows

all checks passed

Covers: defaults reproduce the shared value; _optimizer_step_slices omitting the new kwarg is identical to passing the shared size; AR=8@mbs1 → 8 micros while image=16@mbs2 → 8 micros of width 2, each covering all rows; ragged 5@mbs2 → [(0,2),(2,4),(4,5)] with shares summing to 1 and the ragged share strictly smaller; num_updates=2 splits before micro-batching into disjoint ordered mini-batches; non-positive sizes rejected; and for aggregate_numeric_metrics — unweighted path unchanged, equal weights == plain mean, ragged weights give 3.0 where the unweighted mean gives 2.5, per-key renormalization when a key is missing, length mismatch raises, zero total weight falls back, bool/tensor coercion preserved.

2. Recipe keys resolve against the constructor:

$ PYTHONPATH=$PWD python -c "<yaml load + inspect.signature check>"
hi3_vllmomni.yaml: ar=None image=None unknown_keys=[]
hi3_vllmomni_veomni_ep.yaml: ar=None image=None unknown_keys=[]
bagel_trainside_unigrpo.yaml: ar=None image=None unknown_keys=[]
all stack keys map to __init__ params

3. Repo gates:

$ ruff check unirl/train/unified_model_stack.py unirl/utils/misc.py
All checks passed!
$ ruff format --check unirl/train/unified_model_stack.py unirl/utils/misc.py
2 files already formatted
$ python lint/check_recipe_targets.py
check-recipe-targets: 2311 recipe _target_ paths resolve.
$ python lint/check_experimental_boundaries.py
check-experimental-boundaries: ok

Not run: multi-GPU training A/B on this branch. See Reviewer Notes — the measurement I have is from a different base and I am not presenting it as a result for this diff.

Compatibility / Risk

• No behavior change by default. Both knobs default to the shared micro_batch_size; the three shipped recipes set them to null. _optimizer_step_slices / prepare_segment keep their old call form working (the new argument is keyword-only with a None default).
• aggregate_numeric_metrics(weights=...) is keyword-only and defaults to None; the three existing callers (diffusionnft.py, train/stack/base.py ×2) are untouched and unaffected.
• Metrics values do change for a ragged final micro — that is the intended fix. Equal-sized micros are unchanged.
• No config-schema break: the new recipe keys are additive and optional.
• Raising ar_micro_batch_size above 1 shifts the AR objective under loss_agg_mode: token-mean (see Summary). Left at the default, nothing shifts.

Reviewer Notes

Verification: every line is mine to defend; the diff and all commands above were reviewed and run.

Duplicate-work check: scanned the 33 open PRs. Nothing else touches unirl/train/unified_model_stack.py or unirl/utils/misc.py. #299 (share single-stream sampling/replay loops) and #156 (batched-step replay to more models) are in the diffusion-replay area and do not overlap this change; #253 (batch trainside UniGRPO rollout forwards) touches bagel/* model code plus the same recipe file, but a different section — worth a glance for recipe conflict, not for logic.

On the performance claim — please read before merging. I measured image_micro_batch_size=2 at roughly −20% train-step time on 8×H20 HI3 (batch_size 8, 4 diffusion steps / 2 SDE, 512px, 1 step), at the cost of ~5 GB more peak allocated. But that run sat on a different, older base with a different fix for the DP-sharding problem — I had independently hit the "image Part gets replicated to every rank" bug and solved it with a new per-argument DP_SCATTER_INDEPENDENT dispatch mode, before finding that main already fixes it (and better) by passing the whole lineage Sample and tree-sharding it. So:

• That number is not a measurement of this diff. Treat it as motivation, not as a result.
• It was a single run, not repeated. Same-config repeats elsewhere in that matrix spread ~4–9%.
• I no longer hold the GPUs to re-run it on this base.

Happy to re-measure if someone can point me at capacity, and equally happy for this to sit as draft until then. The code is correctness-complete and default-inert regardless; the open question is purely how much the knob buys on current main.

Suggested review order: the prepare_segment anchor-geometry coupling is the only subtle part — if a replay-anchored track's anchor were frozen at a different micro size than training uses, the ratio would silently stop being 1. Everything else is plumbing.