[diffusion] fix: restore nested condition rows in FlowGRPO actor replay - #534
Open
NancyFyong wants to merge 1 commit into
Open
[diffusion] fix: restore nested condition rows in FlowGRPO actor replay#534NancyFyong wants to merge 1 commit into
NancyFyong wants to merge 1 commit into
Conversation
Ref2VA reference condition rows are padded to a global length and turned into jagged nested tensors by embeds_padding_2_no_padding. The FlowGRPO diffusion engine never restored them, so the MiniMax H3 Ref2VA adapter's [:, :count] slice on condition_video_rows raised 'RuntimeError: slice() not supported for NestedTensor on dim=0' before the first training step. Move _unpad_condition_rows from the DiffusionNFT engine into the shared DiffusersFSDPEngine base and call it from PPODiffusersFSDPEngine.forward path so condition rows reach the adapter as dense padded tensors; the adapter slices the first count rows on dim 1, which is the existing T2VA/ FL2VA behavior. AI assistance (pi coding agent) was used for this change. Co-authored-by: pi coding agent Signed-off-by: NancyFyong <2742092809@qq.com>
NancyFyong
marked this pull request as ready for review
September 4, 2026 12:00
NancyFyong
requested review from
AndyZhou952,
SamitHuang,
knlnguyen1802,
wtomin and
zhtmike
as code owners
September 4, 2026 12:00
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes a
RuntimeError: slice() not supported for NestedTensor on dim=0that blocksMiniMax H3 Ref2VA FlowGRPO training on current
main. T2VA and FL2VA areunaffected.
Ref2VA reference condition rows (
condition_video_rows/condition_audio_rows)are padded to a global length and turned into jagged nested tensors by
embeds_padding_2_no_padding. The FlowGRPO diffusion engine never restored themto dense tensors, so the MiniMax H3 Ref2VA adapter's
[:, :count]slice indiffusers_training_adapter.pycrashed before the first training step.Checklist Before Starting
open work is [3/N][diffusion, rollout, tests] refactor: unify diffusion rollout request contract #480 (rollout request refactor) and [doc, recipe] chore: document MiniMax H3 training support #529 (docs), which are
unrelated to nested condition rows.
[diffusion] fix: ...(validated withtests/special_sanity/check_pr_title.py).Test
CPU regression coverage added in
tests/pipelines/test_minimax_h3_ref2va_flow_grpo_on_cpu.py:test_ref2va_engine_unpads_nested_condition_rows: builds a nestedcondition_video_rows/condition_audio_rows+ mask + count, runs_unpad_condition_rows,and asserts the rows are restored to dense tensors and that
[:, :count]slicing works.test_ref2va_engine_rejects_mismatched_nested_mask: nested rows without a nestedmask are rejected rather than sliced.
Commands run (project venv, CPU):
Note:
test_diffusers_fsdp_engine.pyrequires GPUs/Ray and is skipped; this changedoes not alter the enabled engine paths for T2VA/FL2VA.
Design & Code Changes
_unpad_condition_rowsfromNFTDiffusersFSDPEngineinto the sharedDiffusersFSDPEnginebase so it is available to all diffusion engines.self._unpad_condition_rows(micro_batch)inPPODiffusersFSDPEngine.prepare_model_inputsbefore delegating to the adapter, socondition rows reach the MiniMax H3 Ref2VA adapter as dense padded tensors.
_unpad_condition_rowsdefinition from the NFT engine(it inherits the base version; its existing call site still resolves correctly).
Checklist Before Submitting
AI assistance
AI assistance (pi coding agent) was used for this change.
A human submitter (NancyFyong) reviewed every changed line and verified the
commands/results above.