Skip to content

[diffusion] fix: restore nested condition rows in FlowGRPO actor replay - #534

Open
NancyFyong wants to merge 1 commit into
verl-project:mainfrom
NancyFyong:fix/ref2va-nestedtensor
Open

[diffusion] fix: restore nested condition rows in FlowGRPO actor replay#534
NancyFyong wants to merge 1 commit into
verl-project:mainfrom
NancyFyong:fix/ref2va-nestedtensor

Conversation

@NancyFyong

Copy link
Copy Markdown
Collaborator

What does this PR do?

Fixes a RuntimeError: slice() not supported for NestedTensor on dim=0 that blocks
MiniMax H3 Ref2VA FlowGRPO training on current main. T2VA and FL2VA are
unaffected.

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 them
to dense tensors, so the MiniMax H3 Ref2VA adapter's [:, :count] slice in
diffusers_training_adapter.py crashed before the first training step.

Checklist Before Starting

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 nested
    condition_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 nested
    mask are rejected rather than sliced.

Commands run (project venv, CPU):

python -m pytest tests/pipelines/test_minimax_h3_ref2va_flow_grpo_on_cpu.py -q          # 15 passed
python -m pytest tests/pipelines/test_minimax_h3_flow_grpo_on_cpu.py \
  tests/pipelines/test_minimax_h3_diffusion_nft_on_cpu.py \
  tests/pipelines/test_minimax_h3_ref2va_flow_grpo_on_cpu.py -q                          # 69 passed
python -m pytest tests/workers/test_diffusers_dtype_on_cpu.py -q                         # 4 passed
pre-commit run --files verl_omni/workers/engine/fsdp/diffusers_impl.py \
  tests/pipelines/test_minimax_h3_ref2va_flow_grpo_on_cpu.py                             # Passed

Note: test_diffusers_fsdp_engine.py requires GPUs/Ray and is skipped; this change
does not alter the enabled engine paths for T2VA/FL2VA.

Design & Code Changes

  • Hoist _unpad_condition_rows from NFTDiffusersFSDPEngine into the shared
    DiffusersFSDPEngine base so it is available to all diffusion engines.
  • Call self._unpad_condition_rows(micro_batch) in
    PPODiffusersFSDPEngine.prepare_model_inputs before delegating to the adapter, so
    condition rows reach the MiniMax H3 Ref2VA adapter as dense padded tensors.
  • Remove the now-duplicated _unpad_condition_rows definition from the NFT engine
    (it inherits the base version; its existing call site still resolves correctly).

Checklist Before Submitting

  • Read the Contribute Guide.
  • Apply pre-commit checks (see Test section).
  • Added CPU regression tests for the restored-rows and mismatched-mask paths.

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.

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>
Copilot AI lite review requested due to automatic review settings September 4, 2026 08:23

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@NancyFyong
NancyFyong marked this pull request as ready for review September 4, 2026 12:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants