perf(minimax-h3): shard the video VAE decode across the SP group

[leviking98z-rgb]

Draft. Opening early for direction-level feedback. The measurement below is on a
fork branch that also carries native Ulysses SP for MiniMax-H3 (not yet upstream); on
main today this code path is inert because H3 has no SP installer, so sp_size is
always 1 and the stock decode is taken unchanged. See "Applicability" below.

Summary

Under Ulysses SP the MiniMax-H3 video VAE decode runs sp_size times for one kept result.

SP shards only the trainable transformer, and the SP hooks gather at norm_out, so every
rank finishes the denoise loop holding identical full-sequence latents and then decodes
the whole video independently. Dispatch.DP_SCATTER's collect
(unirl/distributed/group/dispatch.py) keeps only sp_rank == 0, so the other ranks'
decodes are discarded.

Stage timing on a real Flow-GRPO rollout (768×768×124, sp_size=8, forward_batch_size=1)
shows this dominates the step — not the 33B DiT:

stage	per sample	share
embed (conditioner)	0.07s	0.1%
denoise (DiT, 10 steps)	18.20s	26.1%
video decode (VAE)	51.41s	73.7%
audio decode	0.08s	0.1%

All 8 ranks log that same 51.4s decode at the same wall clock — 1439s wasted per 32-sample
rollout, 41% of step time.

This PR spreads the decoder's independent temporal _decode_clip calls (7 at this geometry)
over the SP group and all-gathers them. The blend/trim/concat post-loop is byte-for-byte the
vendor logic, so only the expensive clip forwards move. Opt-in via
bundle.config.video_decode_shard_across_sp (default false); at sp_size == 1 the stock
vae.decode(...) path is taken unchanged, and unirl/models/minimax_h3/vendor/ is untouched
so a re-vendor stays a recopy.

Why not shard by sample (the more obvious fix): SP attention does an all-to-all in every
block, so the group must enter each denoise in lockstep. A rank that skipped its decode would
simply block in the next collective and save nothing.

Related Issue

N/A.

Test Plan

Hardware 8×H20; MiniMax-H3 33B; diffusion/minimax_h3/minimax_h3_t2va_trainside with
+backend.fsdp_cfg.sp_size=8, num_devices=8 batch_size=8 samples_per_prompt=4 (32 samples),
768×768×124, 10 denoise steps, sde_indices=[0,3,6], real VideoPickScore+CLAP reward,
num_updates_per_batch=2, +offload_train_during_reward=true, seed 42.

End-to-end A/B, one rollout per arm, flag off vs on:

metric	off	on
perf/step_time_s	3536.2	2160.7	−38.9%
perf/generate_time_s	2513.3	1128.1	−55.1%
video decode / sample	51.47s	8.22s	6.26×
denoise / sample	19.40s	19.39s	unchanged
peak GPU memory	82171 MiB	82171 MiB	unchanged
rollout/reward_mean	0.4394542	0.4394542	identical

Both arms exit_code=0, train/optimizer_updates=2, no NaN. Sanity: decode saved
32 × (51.47 − 8.22) = 1384s against a measured step delta of 1375.5s, so the step
improvement is fully explained by the decode and nothing else silently changed.

Bit-exactness — real VAE weights, GPU, fp32, same process, stock vae.decode(...) vs the
sharded path, 3 seeds:

max_abs_diff 0.0   max_ulp 0   mismatched 0 / 219,414,528

reward_mean alone is not a sufficient check here (VideoPickScore scores only the middle
frame), so each arm also hashed every decoded video: the 32 SHA-256 values match exactly
across arms, and all 8 ranks agree within each arm.

Algorithm equivalence across world sizes (CPU, gloo): 3 geometries (37/22/12 latent
frames) × world 1/2/4/8 = 12 cases, all max|Δ| == 0, including chunk counts that do not
divide evenly by the world size.

Fail-closed under injected faults (CPU, gloo, 8 ranks), fault on rank 3:

decoder raises      -> 8/8 ranks raise; no hang
ragged clip extent  -> 8/8 ranks raise; no hang

An earlier iteration failed the ragged case (1/8 raised, 7 blocked in the gather — an 1800s
NCCL timeout in a real job); the shape agreement is now derived from the gathered probes so
every rank raises together.

Not run: pre-commit/repo lint (no local environment for it on the cluster where this was
measured); multi-node SP.

Compatibility / Risk

• New config field video_decode_shard_across_sp, default false — existing recipes and
  checkpoints unaffected, decode byte-identical when off.
• No effect at sp_size == 1, which is every configuration on main today.
• No checkpoint, data-format or API change. Peak memory measured unchanged; the gather is one
  video per sample over NVLink.
• Requires uniform clip shapes across the SP group. Verified for 17n+5 frame counts at
  768×768; a geometry producing ragged clips raises a clear collective error on every rank
  rather than hanging.

Reviewer Notes

• Applicability: this is only reachable once MiniMax-H3 has an SP installer. It was
  developed against a fork branch that adds native Ulysses SP for H3. If that is not a
  direction you want, this PR is moot and I'll close it — that is the main thing I'd like
  feedback on before polishing further.
• _decode_clip, _blend and the tiling attributes are read from the vendored autoencoder.
  They are private-by-name; a re-vendor that renames them would need this updated. Tiling
  itself cannot be disabled — the released frames are the blended-tile ones — which is why the
  geometry is reproduced rather than bypassed.
• Only temporal chunks are sharded. Each clip also contains a 4×4 spatial tile grid (112
  decoder calls per sample); sharding those too adds ~0.8pp by projection and was left out as
  not worth the extra stitching complexity.
• At this geometry there are 7 chunks for 8 ranks, so one rank idles — measured 6.26× rather
  than 7×. Longer clips divide better.
• AI assistance: developed and measured with Claude Code. Duplicate-work check:
  unirl/models/wan21/wan_video_vae.py has a similarly-shaped tiled_parallel_decode, but it
  is unreachable (sp_group=None default, no caller) and shards spatial tiles of a different
  VAE; nothing existing covers H3.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.

[leviking98z-rgb]

150-rollout SP1/SP8 stability follow-up

I ran the longer real Flow-GRPO check requested after the original one-rollout A/B. Both arms used the same fixed source and inputs, 8×H20 each, 192×192×124, 10 denoise steps, seed 42, four samples per prompt, two updates per batch, and video_decode_shard_across_sp=true.

arm	topology	result	reward mean	train rows
SP1	SP1 / DP8	150/150, exit 0	0.37178	300
SP8	SP8 / DP1	150/150, exit 0	0.35804	300

Validation checks:

• no traceback, OOM, NaN/Inf, NCCL error, or timeout in either log;
• 150 reward rows and 300 finite loss/ratio/grad-norm rows per arm;
• zero sample-count or group-count mismatches across all 150 paired steps;
• mean paired difference SP8−SP1 -0.01373; moving-block bootstrap 95% interval [-0.02204, -0.00585];
• first-20 / last-20 mean difference -0.01961 / -0.00345.

W&B: SP1, SP8. Each cloud run has 150 unique rollout steps; the offline-to-online sync introduced no duplicates.

Interpretation: this is a long-horizon liveness/stability check, including the SP1 fallback fixed in 363d5bac and the SP8 sharded path. It changes DP topology at the same time (DP8 versus DP1), so the reward gap cannot be attributed to VAE decode sharding. The decoder-specific causal check remains the within-SP8 flag-off/flag-on result in the PR body: byte-identical decoded videos and matching SHA-256 values.