Skip to content

fix(vllm): load Super-Omni RADIO final LayerNorm - #4036

Open
DanialTaheri wants to merge 2 commits into
super-v3.5-posttrainingfrom
smohsenitahe/fix-radio-final-layernorm-super35
Open

fix(vllm): load Super-Omni RADIO final LayerNorm#4036
DanialTaheri wants to merge 2 commits into
super-v3.5-posttrainingfrom
smohsenitahe/fix-radio-final-layernorm-super35

Conversation

@DanialTaheri

Copy link
Copy Markdown
Contributor

What does this PR do ?

Aligns vLLM's Nemotron Super-Omni RADIO feature path with Megatron/HF by adding the checkpoint-backed final LayerNorm that is missing from vLLM 0.25.1.

The pinned vLLM implementation omitted the LayerNorm module, its forward application, and checkpoint/refit name mapping. With NeMo-RL's dummy startup load followed by refit, the missing module and mapping prevent the learned affine parameters from reaching the rollout model, causing visual features—and therefore rollout token log probabilities—to diverge from the Megatron policy.

This change:

  • Creates the final RADIO LayerNorm for Super-Omni MTP configurations.
  • Applies it in FP32 to both dynamic and fixed/chunked image/video feature paths before projection.
  • Accepts both vision_final_layernorm.* and vision_projector.vision_final_layernorm.* checkpoint names.
  • Includes the parameters in the multimodal connector/refit path and loads them with default_weight_loader.
  • Applies the vLLM source patch atomically and idempotently.

Issues

None.

Usage

No new recipe option is required. The patch is selected from the model configuration and applied automatically during vLLM worker initialization. Successful activation emits:

Enabled checkpoint-backed RADIO final LayerNorm
Loaded RADIO final LayerNorm affine parameters

Before your PR is "Ready for review"

Pre checks:

  • Read and followed the contributor guidelines.
  • Added focused coverage for forward paths, checkpoint mapping, syntax, and idempotency.
  • Ran the focused vLLM-backed tests in the Super-Omni container: 2 passed.
  • No documentation update is required because this is an automatic model-parity fix with no user-facing configuration.

Additional Information

Target branch: super-v3.5-posttraining.

Focused test command:

python -m pytest -q --vllm-only \
  tests/unit/models/generation/test_vllm_patches.py::test_radio_final_layernorm_patch_covers_model_forward_and_loader \
  tests/unit/models/generation/test_vllm_patches.py::test_radio_final_layernorm_patch_is_idempotent

Signed-off-by: Seyed Taheri <smohsenitahe@aws-cmh-slurm-1-vscode-01.cm.cluster>
@DanialTaheri
DanialTaheri requested review from a team as code owners September 7, 2026 04:26
@copy-pr-bot

copy-pr-bot Bot commented Sep 7, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@DanialTaheri

Copy link
Copy Markdown
Contributor Author

/ok to test 194f391

Signed-off-by: Seyed Taheri <smohsenitahe@aws-cmh-slurm-1-vscode-01.cm.cluster>
@DanialTaheri
DanialTaheri requested review from a team as code owners September 7, 2026 04:51
@DanialTaheri DanialTaheri added the CI:L2 Run doctests, unit tests, functional tests, and convergence tests label Sep 7, 2026
@DanialTaheri

Copy link
Copy Markdown
Contributor Author

/ok to test 4500a7a

@DanialTaheri

Copy link
Copy Markdown
Contributor Author

CI follow-up for 4500a7acb2c9a3fd9c8f49c0b704e6913e1ab073:\n\n- Applied Ruff 0.9.9 formatting to the PR-owned files.\n- Updated stale test-only expectations/helper support for the existing RADIO ls1/ls2 patch; runtime LayerScale code is unchanged.\n- Fixed the target branch\u0027s repository-wide Ruff/Pyrefly blockers found by the --all-files lint gate.\n- Local validation: repository-wide Ruff and Pyrefly pass; patch tests pass in normal mode (13/13) and vLLM-only mode (8/8); affected rollout test passes (2/2).\n- Applied CI:L2 and retriggered CI with the full SHA.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L2 Run doctests, unit tests, functional tests, and convergence tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant