[biswapanda]

Summary

• Register Nemotron-H Omni with Prime-RL's multimodal batching.
• Materialize and validate the image-processor tensors required by Nemotron 3.5 Super.
• Preserve frozen vision weights after their first FSDP forward when packed ranks process different image counts.
• Complete the three-PR stack for native Nemotron 3.5 Super image-conditioned RL training.

Why

Nemotron-H Omni image preprocessing produces tensors and image-token expansion metadata that Prime-RL must materialize and validate before building training batches. Registering its multimodal adapter lets batching materialize those tensors and validate their token alignment. Packed ranks can process different numbers of images; resharding the frozen vision encoder after every forward can therefore cause mismatched all-gather collectives. Retaining its gathered weights prevents those repeated collectives while preserving the existing behavior for other vision encoders.

Changes

• Add NemotronHOmniAdapter for image tensor materialization.
• Validate placeholder expansion lengths and reject mixed processed image shapes.
• Add the model opt-in used to prevent rank-divergent repeated vision all-gathers.
• Add focused multimodal adapter and FSDP behavior tests.

Validation

• After moving architecture metadata into the model layer, the combined stack tree remains identical. Actual text-first and image-second acceptance on the updated heads is pending GPU access.

• Ruff format: passed for all three reconstructed stack layers.

• Ruff check: passed for all three reconstructed stack layers.

• Python compileall: passed for all three reconstructed stack layers.

• The pre-split equivalent passed 64 focused Nemotron model, conversion, multimodal batching, routed-forwarding, and composite-weight tests.

• A three-step Nemotron 3.5 Super image RL run passed with Python dynamo.vllm, TP4 inference, four trainer GPUs, inline routed experts, sampling masks, and NCCL policies v0 through v3. Mismatch KL: 0.0021, 0.0006, 0.0030.

• An eight-step follow-up passed with NCCL policies v0 through v8. Mismatch KL: 0.0023, 0.0022, 0.0022, 0.0095, 0.0005, 0.0107, 0.0035, 0.0030. All 32 rollouts were trainable with zero errors or cancellations.

• Focused pytest collection for the reconstructed stack is blocked on this host because its local environment lacks flash_attn; GitHub CI is the clean test environment.

Stack

1. feat(trainer): add Nemotron-H Omni checkpoint conversion #3730 — Nemotron-H Omni checkpoint conversion.
2. feat(trainer): add Nemotron-H Omni composite model #3731 — Nemotron-H Omni composite trainer model.
3. This PR — multimodal training integration.

This branch is based on commit b514c4894 from #3731. Because the branches originate from a fork, all three upstream PRs target main; upper diffs shrink as lower PRs merge.

vLLM compatibility

This PR does not include the Prime-RL runtime backport of vLLM #56231. Nemotron LoRA validation should use a vLLM release candidate or later release containing that upstream fix.