feat(minimax-h3): carry FL2VA keyframe condition rows through rollout, replay and NFT

[MinhaoLi0318]

Summary

Step 1a of #391 item 4 (FL2VA: first, last and first+last keyframe conditioning for MiniMax-H3). Rollout stores the noised keyframe condition rows as data; replay and the DiffusionNFT forward reuse them and never recompute them.

Why carry the rows: diffusers 0.40.0 and vLLM-Omni 0.28 draw the condition noise and resize the second keyframe differently, so train-side recomputed rows would silently break the replay log-prob ratio. Carried rows are byte-identical to the rollout rows, as @CjhHa1 asked for step 1 on #391.

Scope: t2va stays bit-identical to main 6c828f5, and to main 0397ddb once merged (Test Plan 5). No code path sets keyframes yet: this PR is the consumer, and step 1b adds the producer (keyframe encode, conditioner prompt, data path), per the 1a / 1b split I proposed on #391; I can fold 1b into this branch if you prefer one PR. Condition rows still run at the target video timestep until step 2.

Changes (6 files in unirl/models/minimax_h3/, +125/-23):

• Conditions. Optional keyframes slot (MiniMaxH3KeyframeCondition) on MiniMaxH3Conditions, one entry per sample: fp32 rows [K * rows_per_frame, 96], anchors, unpadded text-token tags. t2va samples in a keyframe batch carry an empty entry. __post_init__ checks alignment, and a Part.concat that mixes samples with and without keyframes raises in fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582's Batch.concat check instead of shifting rows.
• Layout and forward. build_t2va_layout → build_layout(geometry, conditions) with fail-closed checks. generate, replay and predict_noise_at_step share one predict_noise that prepends the carried rows and returns target-only velocity. trim_text_padding keeps the other slots.
• Docs. New minimax_h3/README.md Gotchas; vendor/__init__.py re-exports MINIMAX_H3_VIDEO_TAG (vendored files untouched).

Related Issue

Part of #391 (item 4, step 1a).

Prior art: verl-omni #401 (MiniMax H3 FL2VA DiffusionNFT, merged 2026-08-24) and #448 (FL2VA FlowGRPO, merged 2026-08-31) also carry the noised keyframe rows from rollout into the training forward and score only the target rows; #401's trainer already runs condition rows at max(t, 0.999), which is step 2 here.

Test Plan

Commits 8976f4b (f3b8b8e rebased, git range-diff 1: f3b8b8e = 1: 8976f4b) and 6dde41e (drops the binding; see Reviewer Notes) on 6c828f5. macOS arm64, Python 3.12.14, torch 2.13.0 CPU, diffusers 0.40.0. CPU plus one CPU→MPS→CPU round trip; no GPU.

SKIP=no-commit-to-branch pre-commit run --from-ref 6c828f5 --to-ref HEAD   # 15 Passed, 4 Skipped
git diff --check 6c828f5 HEAD                                               # clean

One-off CPU harness, not committed (AGENTS.md §5), available on request: real MiniMaxH3Bundle and stage (CPSSDE, eta 0.7), tiny random-config vendored transformer, 64x96x124 canvas; every transformer kwarg, log-prob and output compared with torch.equal plus dtype and shape. All steps exit 0:

mkdir -p main branch m500 out && git fetch upstream refs/pull/500/head
git archive 6c828f5 unirl | tar -x -C main && git archive 6dde41e unirl | tar -x -C branch
PYTHONPATH=main python run_t2va.py out/main.pt && PYTHONPATH=branch python run_t2va.py out/branch.pt
python cmp_t2va.py out/main.pt out/branch.pt && python cmp_t2va_negctl.py out/branch.pt
PYTHONPATH=branch python run_shim.py . run_fl2va_582.py && PYTHONPATH=branch python layout_xref.py
TREE=$(git merge-tree --write-tree 6dde41e d3bdb23) && git archive "$TREE" unirl | tar -x -C m500   # #500
PYTHONPATH=m500 python pr500_bytes.py && PYTHONPATH=m500 python run_shim.py . run_fl2va_582.py

run_fl2va_582.py is the earlier run_fl2va.py with the binding checks replaced by checks against #582's Batch.concat. run_shim.py re-creates _align_track_to_model, which #592 folded into BaseTrainStack._align_track_inputs (same body).

1. t2va vs main → 791 leaves equal (dtype+shape+bytes), 0 differ (rollout, replay, predict_noise_at_step, padded micro-batch, a t2va sample inside mixed t2va/fl2va Parts). Comparator negative control: 3/3 injected changes flagged.
2. FL2VA rollout vs replay → 207 passed, 0 failed: first, last, first+last through Part.concat → pickle → TensorTransport → CPU→MPS→CPU → slice/select → replay → predict_noise_at_step; sparse SDE window; bf16; mixed batches; fail-closed inputs. Parity negative controls (re-drawn noise, fp16 round trip, all-text tags, swapped anchors, main's trim_text_padding) each break parity; the smallest |Δlogp| among them is 1.5e-13, so the check compares bytes. Every Part.concat that mixes samples with and without keyframes (missing first, middle or last; padding path; DP gather shape) raises in fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582's check. With the __post_init__ guard off, its 4 construction and slice cases stop raising and the 4 typed-concat cases still raise in fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582's check; with fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582's check off, [None, f1, f2, f3] silently shifts rows. Before fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582, f3b8b8e on 0fe4f99 gave 210 passed, 0 failed, including the binding's own checks.
3. Layout cross-check → 68 passed, 0 failed against the vendored builder, diffusers 0.40.0, vLLM-Omni v0.28.0 and vLLM-Omni main 0ad67f1d8, four canvases.
4. Merged with perf(minimax-h3): batch replay timesteps #500 → pr500_bytes.py 10 passed, 0 failed (|Δlogp| ≤ 1.03e-13); run_fl2va_582.py 207 passed, 0 failed.
5. Re-check on main 0397ddb (2026-10-16) → its 17 new commits touch nothing in minimax_h3/ or types/conditions/; fix(distributed,reward): access typed handles, batches, and scorer specs directly #604 rewrote batch.py but kept fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582's missing-value check (batch.py:144-146). git merge-tree with 6dde41e is clean. On the merge result: t2va vs 0397ddb 791 equal / 0 differ, negative control 3/3; FL2VA 205 passed, 0 failed (fix(distributed,reward): access typed handles, batches, and scorer specs directly #604 removed Batch.repeat_interleave, so its 2 checks are skipped; the shim follows fix(distributed,reward): access typed handles, batches, and scorer specs directly #604's to_device body); layout 68 passed, 0 failed. With perf(minimax-h3): batch replay timesteps #500 on top: pr500_bytes.py 10/0, FL2VA 205/0. With fix(models): read typed bundle/pipeline contracts directly and drop falsy fallbacks #605 on top: t2va 791/0 vs 0397ddb + fix(models): read typed bundle/pipeline contracts directly and drop falsy fallbacks #605, FL2VA 205/0, layout 68/0.

Parity uses synthetic noised rows (1a has no producer); 1b re-runs it through MiniMaxH3Pipeline.generate. The DiffusionNFT loss itself was not run.

• Not run; reason: no GPU, and nothing reaches the keyframe path in 1a → GPU/FSDP run, incl. empty t2va-entry rows over CUDA IPC (step 1b).
• Not run; reason: no recipe or config change → Hydra --cfg job --resolve.

Compatibility / Risk

• API: build_t2va_layout renamed to build_layout (no other caller); new export MiniMaxH3KeyframeCondition.
• t2va: bit-identical; no config, recipe, checkpoint or data-format change. feat(minimax-h3): vLLM-Omni 0.28 rollout backend for H3 FlowGRPO #378's {"text": ...} dict still works.
• Producers with keyframes emit an entry for every sample or for none. Part.conditions can come from MiniMaxH3Conditions.to_dict or be built by hand.
• FL2VA not usable yet: condition rows run at 1 - sigma, the released model at max(1 - sigma, 0.999) (diffusers 0.40.0 before_denoise.py). The README forbids enabling fl2va before step 2, which must fix both main's row_timestep_plan and perf(minimax-h3): batch replay timesteps #500's batched plan.
• Cost with keyframes: one torch.stack and one torch.cat per forward.
• Open PRs: perf(minimax-h3): batch replay timesteps #500 merges cleanly (verified above). feat(minimax-h3): DiffusionNFT on the vLLM-Omni rollout backend at 32-GPU scale #420 and draft feat(minimax-h3): train original block weights with frozen peripherals #593 also add minimax_h3/README.md (add/add, no contradiction); whichever lands second merges the READMEs. fix(models): read typed bundle/pipeline contracts directly and drop falsy fallbacks #605 also edits minimax_h3/diffusion.py, packing.py and pipeline.py and merges cleanly (Test Plan 5). fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582 is in main.

Duplicate-Work Check

Re-checked 2026-10-11 on the public GitHub pages and with git ls-remote upstream:

• [Roadmap] MiniMax-H3 critical path #391 item 4 is assigned to me (2026-10-02); no [Roadmap] MiniMax-H3 critical path #391 activity since my 2026-10-03 split proposal.
• PRs matching FL2VA: feat(minimax-h3): rework the vLLM-Omni H3 rollout onto 0.27.0rc1 #408 only (closed, vLLM-Omni 0.27.0rc1 rollout). PRs matching keyframe: none. Issues matching fl2va or keyframe: [Roadmap] MiniMax-H3 critical path #391 only.
• Open PRs matching minimax (7), heads from refs/pull/<N>/head, merged with this commit by git merge-tree: perf(minimax-h3): batch replay timesteps #500 (d3bdb23) clean; feat(minimax-h3): DiffusionNFT on the vLLM-Omni rollout backend at 32-GPU scale #420 (521c217) adds only the README add/add; feat(minimax-h3): vLLM-Omni 0.28 rollout backend for H3 FlowGRPO #378 (b21fca7), perf(minimax-h3): shard the video VAE decode across the SP group #469 (363d5ba), feat(trainer): opt-in row-level reward dispatch for diffusion #517 (fca5a55), feat(flux2_klein): support multiple reference images end to end #351 (acabb46), perf(minimax-h3): enable expandable allocator segments for trainside training #583 (81a011d) touch none of these files and gain no new conflict.
• PRs docs(reward): refresh scorer registry and deployment guidance #542-perf(minimax-h3): enable expandable allocator segments for trainside training #583: none touches minimax_h3/ or types/conditions/, or mentions keyframe/fl2va in its commit messages. fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582 (07dec54) changes distributed/tensor/batch.py; see Reviewer Notes.

Re-checked 2026-10-13 after the rebase: PRs matching keyframe or FL2VA are still only #408 and this one; #391 has no new comments. New since 2026-10-11: draft #593 (da8d285, MiniMax-H3 block training) adds minimax_h3/README.md, an add/add with this PR. #582 and #583 are now in main.

Re-checked 2026-10-16 against main 0397ddb: PRs matching keyframe or FL2VA are still only #408 (closed) and this one; issues matching fl2va or keyframe are #391 and the Q2 roadmap #25 (last updated 2026-09-12); #391 has no new comments. Besides this one, 8 open PRs match minimax; the only new one is #605 (5545bcb), which edits minimax_h3/{diffusion,packing,pipeline}.py and merges cleanly with this head, and the other heads are unchanged. Open PRs #589-#609 mention neither keyframe nor fl2va in their commit messages. git merge-tree of this head with each of the other 64 open PRs: only #420 and #593 gain a conflict (the README add/add).

Reviewer Notes

1. @CjhHa1: which PR changes the condition-row timestep? To be discussed with @leviking98z-rgb (perf(minimax-h3): batch replay timesteps #500). 1a already keeps velocity, log-probs and the NFT target on target rows, so I propose step 2 = condition rows at max(1 - sigma, 0.999) in both main's and perf(minimax-h3): batch replay timesteps #500's plan, assertions, and an engine-side row-count check. Alternative: fold it into 1b with perf(minimax-h3): batch replay timesteps #500.
2. Tags live in the keyframes slot (they exist only with keyframes). If you prefer the text side, now is the cheapest time to move them.

• fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582 and the dropped binding: on 0fe4f99 the generic dict concat dropped a missing key, so [None, f1, f2, f3] would give 4 text rows and 3 entries and shift every later sample. The first commit guarded this with a binding (MiniMaxH3KeyframedText plus a from_dict type check). fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582 (790b5ec) makes the generic concat raise on a value missing in a non-empty item, for every model, so 6dde41e drops the binding (2 files, +7/-37) and to_dict / from_dict are main's again. __post_init__'s alignment check stays: it also covers direct construction and a keyframes count that differs from text.
• Review: repo pr-review skill; independent reviewer Sonnet 5.5 (claude-sonnet-5-5), fresh context, 3 rounds; author Claude Opus 5.5. Last round: no P0/P1/R. Since then: README wording, the rebases onto 0fe4f99 and 6c828f5 (git range-diff: 1: 8f4bd75 = 1: f3b8b8e, then 1: f3b8b8e = 1: 8976f4b), and 6dde41e, which removes the binding after fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582 merged; a fourth independent round (Sonnet 5.5, fresh context) on 6dde41e reported no P0/P1/R. Fixed: mixed-batch alignment (the binding and its from_dict unwrap, now replaced by fix(tensor): fail closed when batch concat or DP dispatch sees inconsistent shards #582's check), tuple-only anchors. Kept: target-row velocity slice, README despite feat(minimax-h3): DiffusionNFT on the vLLM-Omni rollout backend at 32-GPU scale #420, generic _concat_value unchanged in this PR.
• For 1b: producer, an entry for every sample, embed-cache and feat(minimax_h3): offline text-embed store and skip-TE load #513 keys with keyframe content, harness via generate, GPU check.
• AI assistance: Claude Code helped design, implement and test this change; I reviewed the full diff and harness results.

Checklist

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

AI assistance: drafted with Claude Code; I reviewed it before posting.

[MinhaoLi0318]

#582 is in main (790b5ec), so I rebased onto 6c828f5 and pushed 6dde41e, which removes the transitional MiniMaxH3KeyframedText binding as planned; a Part.concat that mixes samples with and without keyframes now raises in #582's Batch.concat check. On the new head, t2va is still bit-identical to main (791 leaves equal, 0 differ), the FL2VA harness passes 207/207 and pre-commit passes; the body is updated to match.

AI assistance: drafted with Claude Code; I reviewed it before posting.