feat(minimax-h3): vLLM-Omni rollout backend for H3 FlowGRPO at 32-GPU scale

.gitignore

@@ -70,6 +70,9 @@ outputs/
 logs/
 local/
 !unirl/reward/local/
+
+# Reward-model weight cache (ImageBind resolves .checkpoints/ relative to cwd)
+.checkpoints/
 **/rollout_data/
 **/buffer_stats/
 *.out

INSTALL.md

@@ -4,7 +4,7 @@ UniRL ships two mutually exclusive inference engines (`vllm` and `sglang`) — i
 
 | Engine | CUDA | glibc |
 |---|---|---|
-| **vllm-omni** | 12.9 | ≥ 2.28 |
+| **vllm-omni** | 13.0 | ≥ 2.28 |
 | **sglang** | 13.0 | ≥ 2.34 |
 
 ## vllm-omni
@@ -15,6 +15,11 @@ export VLLM_USE_PRECOMPILED=1   # else 30+ min CUDA build
 uv pip install -e ".[vllm,train,infer]"
 ```
 
+The vLLM extra temporarily pins the reviewed commit from
+`nussejzz/vllm-omni` that backs the MiniMax-H3 integration. Replace this
+direct reference with the first upstream vLLM-Omni release containing PR
+`#6351`. This stack targets vLLM 0.27, PyTorch 2.13, and CUDA 13.0.
+
 ## sglang
 
 ```bash
@@ -26,7 +31,7 @@ uv pip install -e ".[sglang,train,infer]" --prerelease=allow
 
 | Extra | Adds | Use when |
 |---|---|---|
-| `vllm` | `vllm`, `vllm-omni`, torch +cu129 stack, PyAV | Running any vllm-omni-based example |
+| `vllm` | `vllm`, pinned H3-capable `vllm-omni`, torch +cu130 stack, PyAV | Running any vllm-omni-based example |
 | `sglang` | `sglang[diffusion]`, `flash-attn-4`, torch +cu130 stack, PyAV | Running VLM/LLM examples or `sd3_sglang_*` |
 | `train` | `wandb`, `aiohttp` | Training runs (almost always wanted) |
 | `infer` | `accelerate` | HunyuanImage3 and similar models |

examples/diffusion/minimax_h3/README.md

@@ -0,0 +1,125 @@
+# MiniMax-H3 T2VA — 32-GPU hybrid-parallel deployment
+
+MiniMax-H3 is a 33B dense omni-modal transformer that denoises video and stereo
+audio jointly in one packed sequence. These recipes run it under vLLM-Omni
+rollout with an FSDP trainer and a resident reward service on the same GPUs.
+
+## Scope of what is verified
+
+This directory documents a **deployment** that has passed capacity, LoRA-sync,
+rollout/replay parity and throughput gates. It is not a converged RL recipe:
+
+- Topology and the three-phase lifecycle are verified and reproducible.
+- Rollout/replay log-prob parity holds at `~4e-5`, well inside the `1e-3` gate.
+- **Reward convergence is still open.** No H3 recipe here has produced a
+  sustained rise in held-out visual reward. A Wan2.1 trainside control on the
+  same trainer, PickScore and GRPO path does rise (`0.6985 → 0.7201` over 20
+  rollouts), which is what localizes the open problem to H3/reward rather than
+  to the shared training stack.
+
+Treat these files as infrastructure and as a correctness/performance baseline.
+
+## Topology
+
+Four nodes × 8 × 96GB GPUs. Three roles time-share the same 32 GPUs by
+switching residency per phase, rather than statically partitioning them.
+
+```mermaid
+flowchart LR
+  Trainer["Trainer world32<br/>HSDP shard8 x replicate4"]
+  Rollout["Rollout DP8<br/>8 replicas x 4 GPUs"]
+  Reward["Reward service DP32"]
+
+  subgraph Replica["One rollout replica (4 GPUs)"]
+    DiT["H3 DiT<br/>TP2 x Ulysses2"]
+    Text["Qwen text encoder<br/>TP4"]
+    VAE["Video VAE VPP1<br/>Audio VAE replicated"]
+  end
+
+  Trainer -->|LoRA IPC sync| Rollout
+  Rollout -->|trajectory + decoded media| Reward
+  Reward -->|rewards| Trainer
+  Rollout --- Replica
+```
+
+**Trainer** — `world_size=32`, FSDP `hybrid` (shard degree 8 inside a node,
+replicated across the 4 nodes), FP32 master weights, BF16 mixed precision,
+activation checkpointing on, LoRA on the attention Q/K/V/out and FFN of the 50
+denoising blocks, `micro_batch_size=1`.
+
+**Rollout** — 8 data-parallel replicas of 4 GPUs each; DiT TP2 × Ulysses2, text
+encoder TP4, video VAE VPP1.
+
+**Reward** — reward service DP32, long-lived process whose model residency
+follows the phase switches.
+
+**Phase order** — sync LoRA to the 8 replicas → rollout wakes and generates
+video, audio and a sparse FP32 trajectory → rollout sleeps, reward scores →
+trainer replays, backwards and steps → next rollout.
+
+The switches that make this fit are `layout: colocate`,
+`transport: colocate_store`, `enable_fsdp_offload: true`,
+`offload_train_during_reward: true` and `rollout.config.enable_sleep_mode: true`.
+
+## Measured throughput
+
+Same K8 conditions: 8 prompts × 8 samples, 256×448×107, 24 transitions,
+eta 0.6, PickScore-only.
+
+| Rollout topology | median generate | peak memory | max parity drift |
+| --- | --- | --- | --- |
+| HSDP4 + UP4 | 112.697 s | 66,876 MiB | 3.73e-5 |
+| TP2 × UP2 | 91.087 s | 77,436 MiB | 3.94e-5 |
+
+TP2 × UP2 is 19.18% faster. K16 generate is ~182 s, roughly linear in sample
+count. A TP4 arm raced a stale GPU keep-alive process and OOM'd, so that number
+was discarded; TP2 × UP2 is therefore the fastest **verified** candidate, not a
+proven global optimum.
+
+## Recipes
+
+```text
+minimax_h3_t2va_vllmomni_32c_quality100_tp2_up2
+  -> minimax_h3_t2va_vllmomni_32c_quality100
+  -> minimax_h3_t2va_vllmomni_32c_8x8_hsdp8x4
+  -> minimax_h3_t2va_vllmomni_32c_8x8
+  -> minimax_h3_t2va_vllmomni_2x4_timeshare
+  -> minimax_h3_t2va_trainside
+```
+
+`minimax_h3_t2va_trainside` is the in-process baseline; the `32c_*` layers add
+the vLLM-Omni rollout, the 32-GPU geometry and the topology overrides.
+`_tp4` and `_hsdp4` variants exist for topology comparison.
+
+## Geometry constraint
+
+The reference canvas resolver starts from a 768-pixel short edge, but that is a
+default rather than a model floor. `MiniMaxH3Geometry.resolve` enforces the
+structural constraints: both axes must be multiples of 32, the aspect ratio must
+stay within 1:4–4:1, the area must not exceed 768×1344, and the duration and
+frame count must round-trip through the VAE. Lower-resolution runs remain
+quality-sensitive and should be qualified independently.
+
+## Known limitations
+
+- **Rollout siblings are serial.** The adapter issues one request per sibling,
+  so K8 → K16 roughly doubles generate time; request batching is not effective
+  yet.
+- **Decode is duplicated.** Under VPP1 every rank in a replica runs video then
+  audio decode and only the output rank returns, instead of splitting the two
+  decoders across ranks.
+- **Trajectories travel through the driver.** Trajectory, reward media and
+  reward rows are orchestrated through Ray and the driver rather than by
+  rollout-worker to reward-worker GPU P2P.
+- **Communication is not yet itemized.** TP all-reduce, Ulysses all-to-all, text
+  TP4 broadcast, LoRA sync, VAE decode and sleep/wake have not been separated;
+  that needs NVTX traces and per-rank timelines.
+- **Reward DP32 is fine-grained.** At PickScore batch size 2 each rank does
+  little work, so launch, scatter and on/offload overhead may dominate.
+- **Trainer replay retains activations.** Sparse replay keeps several steps in
+  one autograd graph even at micro-batch 1; step-wise backward and selective
+  activation checkpointing are unexplored.
+
+Deep performance work should wait until a fixed-eval reward rise is reproducible
+on one H3 recipe — optimizing throughput before then just runs a wrong objective
+faster.

examples/diffusion/minimax_h3/minimax_h3_t2va_vllmomni_2x4_timeshare.yaml

@@ -0,0 +1,54 @@
+# @package _global_
+# PR5 H3 baseline with FSDP8 training, 2 x HSDP4+UP4 rollout, and reward DP8.
+
+defaults:
+  - minimax_h3_t2va_trainside
+  - _self_
+
+num_devices: 8
+devices_per_node: 8
+batch_size: 8
+num_rollouts: 1
+transport: colocate_store
+layout: colocate
+reward_fraction: 0.0
+enable_fsdp_offload: true
+offload_train_during_reward: true
+eval_interval: 0
+
+logging:
+  report_to_wandb: false
+  log_media: false
+  run_name: h3-fsdp8-vllmomni-2x4-rmdp8
+
+rollout:
+  _target_: unirl.rollout.engine.vllm_omni.engine.VLLMOmniRolloutEngine
+  model_config: ${bundle.config}
+  config:
+    _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig
+    model_path: ${bundle.config.pretrained_model_ckpt_path}
+    modality: minimax_h3_t2va
+    enable_sleep_mode: true
+    replica_size: 4
+    tp_size: 4
+
+reward:
+  backend:
+    config:
+      batch_size: 8
+      frame_selection: middle
+      weights:
+        videopickscore: 1.0
+        clap: 0.0
+
+data_source:
+  args:
+    run:
+      data_path: ${oc.env:H3_PROMPT_FILE,examples/diffusion/minimax_h3/smoke_prompts.txt}
+      eval_data_path: ${oc.env:H3_PROMPT_FILE,examples/diffusion/minimax_h3/smoke_prompts.txt}
+
+sampling:
+  # Measured peak with one replay transition is already 92,188 MiB on a
+  # 97,871-MiB H20. Keep the default runnable until replay can backward one
+  # timestep at a time instead of retaining all selected steps in one graph.
+  sde_indices: [6]

examples/diffusion/minimax_h3/minimax_h3_t2va_vllmomni_32c_8x8.yaml

@@ -0,0 +1,47 @@
+# @package _global_
+# Target 32-GPU collocate geometry:
+#   8 root prompts x 8 samples = 64 rows
+#   rollout DP8 x 4 GPUs, resident reward DP32 (2 rows/rank),
+#   trainer world32 (2 rows/rank, micro_batch_size=1).
+# This is the two-step/one-transition gate recipe; the long-run recipe can
+# raise num_inference_steps and sde_indices only after lifecycle gates pass.
+
+defaults:
+  - minimax_h3_t2va_vllmomni_2x4_timeshare
+  - _self_
+
+num_devices: 32
+batch_size: 8
+num_rollouts: 2
+
+bundle:
+  config:
+    use_lora: true
+
+sync:
+  _target_: unirl.distributed.weight_sync.lora.LocalLoraWeightSync
+  param_prefix: "transformer."
+  adapter_name: default
+  # The H3 engine has four sequence-parallel subprocesses per controller.
+  copy: true
+  verify: false
+
+logging:
+  report_to_wandb: true
+  log_media: true
+  media_max_items: 2
+
+reward:
+  backend:
+    config:
+      batch_size: 2
+
+stack:
+  micro_batch_size: 1
+  num_updates_per_batch: 1
+
+sampling:
+  num_inference_steps: 2
+  # One replay transition for the first 8x8 memory/backward gate.
+  sde_indices: [0]
+  samples_per_prompt: 8

examples/diffusion/minimax_h3/minimax_h3_t2va_vllmomni_32c_8x8_fsdp32.yaml

@@ -0,0 +1,14 @@
+# @package _global_
+# Comparison target: one parameter shard group across all 32 GPUs.
+
+defaults:
+  - minimax_h3_t2va_vllmomni_32c_8x8
+  - _self_
+
+logging:
+  run_name: h3-8p8s-vllmomni-fsdp32-rmdp32
+  tags: [minimax-h3, vllm-omni, 8p8s, fsdp32, rm-dp32]
+
+backend:
+  fsdp_cfg:
+    fsdp_mode: full

examples/diffusion/minimax_h3/minimax_h3_t2va_vllmomni_32c_8x8_hsdp8x4.yaml

@@ -0,0 +1,14 @@
+# @package _global_
+# Primary target: node-local shard8 x four replicated nodes.
+
+defaults:
+  - minimax_h3_t2va_vllmomni_32c_8x8
+  - _self_
+
+logging:
+  run_name: h3-8p8s-vllmomni-hsdp8x4-rmdp32
+  tags: [minimax-h3, vllm-omni, 8p8s, hsdp8x4, rm-dp32]
+
+backend:
+  fsdp_cfg:
+    fsdp_mode: hybrid