feat: Ultra Megatron LoRA Support

[vinhngx]

PR Notes: Ultra Megatron LoRA Support

Summary

This PR integrates Megatron LoRA support for Nemotron 3 Ultra on top of
ultra-v3, adds the Ultra 6-node LoRA starter recipe, fixes the LoRA-to-vLLM
refit path, and documents/validates post-training HF PEFT adapter export and
vLLM LoRA serving.

Provenance

The Megatron LoRA enablement was cherry-picked/adapted from Virginia's NeMo RL
LoRA PR work: NVIDIA-NeMo/RL#1889.
The relevant upstream changes provide the core Megatron PEFT plumbing: policy config schema, Megatron-Bridge LoRA construction, pre-wrap
hooks, adapter-only trainable parameter setup, and Megatron LoRA functional test
coverage.

This Ultra branch layers the Ultra-specific work on top: the 6n4g Ultra recipe,
the conditional merged-weight export required by vLLM rollout refit, the Ultra
adapter DCP-to-HF PEFT converter, Ultra documentation, and GB200 Slurm runtime
validation.

Why Ultra-Specific Work Was Needed

Virginia's PR provides the core Megatron LoRA training support, but Ultra needed
additional work for three reasons:

1. vLLM rollout refit expects plain model weights. During RL training, the
   non-colocated vLLM workers are initialized as the plain Ultra model, not as a
   PEFT-wrapped model with dynamically loaded adapters. Raw Megatron-Bridge LoRA
   exports can include PEFT wrapper names such as .base_layer, .adapter, or
   .lora_, which do not match vLLM's base-model state dict. Ultra LoRA rollout
   therefore needs LoRA deltas merged into the exported base weights for refit.

2. Ultra uses fused attention projections. The Ultra Megatron checkpoint
   layout stores attention LoRA weights under fused linear_qkv modules, while
   HF PEFT/vLLM adapter serving validates against unfused q_proj, k_proj,
   and v_proj module names. Post-training adapter serving therefore needs an
   Ultra-aware DCP-to-HF converter that splits fused QKV LoRA-B tensors.

3. The Ultra recipe has a different cluster/topology shape. The validated
   workflow uses 6 GB200 nodes with 4 training nodes and 2 generation nodes. The
   recipe and docs need to encode this shape, plus the scheduler resume override
   observed during resume-to-step-10 validation.

In short: Virginia's PR enables Megatron LoRA training, while the Ultra-specific
work makes that support usable for Ultra's non-colocated vLLM rollout path,
post-training PEFT adapter serving, and GB200 Slurm recipe.

What Changed

• Cherry-picked/adapted Virginia's Megatron PEFT/LoRA policy support:
  • PEFT config schema.
  • Megatron-Bridge LoRA adapter setup.
  • PEFT pre-wrap hook wiring.
  • Adapter-only trainable parameter setup.
• Added Ultra LoRA recipe:
  • examples/configs/recipes/llm/grpo-ultra-6n4g-megatron-lora.yaml
  • 6 GB200 nodes x 4 GPUs: 4 training nodes + 2 non-colocated vLLM generation nodes.
• Fixed vLLM rollout refit for Megatron LoRA:
  • full-weight refit keeps merge_adapter_weights=False;
  • LoRA refit uses merge_adapter_weights=True;
  • LoRA refit raises if unmerged PEFT wrapper keys are exported.
• Added unit coverage:
  • tests/unit/models/policy/test_megatron_lora_refit_export.py
• Added Ultra adapter conversion utility:
  • examples/converters/convert_megatron_lora_dcp_to_hf_ultra.py
  • converts Megatron LoRA DCP adapter checkpoints to HF PEFT adapter format;
  • splits Ultra fused-QKV LoRA-B tensors into PEFT q/k/v tensors.
• Updated Ultra docs:
  • docs/guides/nemotron-3-ultra.md
  • explains training-time merged-weight vLLM refit versus post-training PEFT adapter serving.

Validation

Unit / Static

python -m py_compile \
  nemo_rl/models/policy/workers/megatron_policy_worker.py \
  tests/unit/models/policy/test_megatron_lora_refit_export.py

python -m pytest --confcutdir=tests/unit/models/policy \
  tests/unit/models/policy/test_megatron_lora_refit_export.py -q

Result: 5 passed.

The Ultra adapter converter also passed syntax validation:

python -m py_compile examples/converters/convert_megatron_lora_dcp_to_hf_ultra.py

Slurm Runtime: LoRA Training Smoke

Validated a 6-node GB200 run with 4 Megatron training nodes and 2 non-colocated
vLLM generation nodes.

Evidence from the successful run:

PEFT Statistics:
  Total parameters: 51,007,391,744
  Trainable parameters: 3,194,880
  Trainable percentage: 0.01%

prepare_for_generation/transfer_and_update_weights end elapsed=9.9008s

The run completed rollout generation, logprob inference, policy training, and a
step checkpoint.

W&B: https://wandb.ai/hwinf_dcm/dapo-ultra-lora/runs/17sdkhbv

Slurm Runtime: Resume to Step 10

Resumed a LoRA checkpoint and trained through step 10 after enabling scheduler
resume override:

policy:
  megatron_cfg:
    scheduler:
      override_opt_param_scheduler: true

Checkpoints produced:

step_1
step_5
step_10

W&B: https://wandb.ai/hwinf_dcm/dapo-ultra-lora/runs/wpjkl47d

Adapter Export and vLLM Serving

Converted the step-10 Megatron LoRA DCP adapter checkpoint to HF PEFT format.

Result:

Wrote 96 tensors
Skipped 6 non-LoRA/unmapped entries

Served the converted adapter with vLLM --enable-lora.

Validation:

• /v1/models listed both nvidia/nemotron-3-ultra and dapo-lora.
• Chat smoke prompt Solve: What is 12 + 25? Give only the answer. returned 37.

Notes

During RL training, vLLM does not serve a separate PEFT adapter. NeMo RL exports
Megatron policy weights with LoRA deltas merged into the base tensors, then
refits the plain vLLM model. The HF PEFT adapter conversion script is for
post-training adapter serving.

[copy-pr-bot]

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.

[terrykong]

@yfw to review

[yfw]

examples/converters/convert_megatron_lora_dcp_to_hf_ultra.py:81

🔴 Blocker — silent QKV scramble for GQA models. This splits the fused QKV LoRA-B contiguously ([:q_dim] | [q_dim:q_dim+kv] | [rest]), assuming layout [all-Q | all-K | all-V]. But Megatron's fused linear_qkv is interleaved per KV-group:

# q1 q2 k1 v1 | q3 q4 k2 v2 | q5 q6 k3 v3 | ...

(Megatron-LM attention.py:1399). The LoRA adapter's linear_out adds its delta before the q/k/v split, so its rows follow that same interleaved order — a contiguous slice is correct only for num_query_groups == 1 (MQA).

Nemotron-3-Ultra is GQA — examples/converters/ultra/configuration_nemotron_h.py:142 sets num_attention_heads=32, num_key_value_heads=8 (8 query groups). So the exported HF/vLLM adapter is silently corrupt on the only model this converter targets (it still passes the size check; q-rows leak into k_proj/v_proj).

Fix: reshape to [ng, (q_per_group + 2), head_dim] and gather q/k/v per group before concatenating, and derive q_dim/head_dim from the model config rather than the CLI default. This only affects the offline post-training serving converter (not training or in-training refit), so it's recoverable, but it's a merge-blocker for the adapter-export feature.

[yfw]

Typo, suggest:

Suggested change

	lora_dtype: None
	lora_dtype: null

[yfw]

nemo_rl/models/megatron/setup.py:797

Typo megtatron_cfg → megatron_cfg in this user-facing error string. Same typo appears at L801, L975, and L979.

Suggested change

	"If megtatron_cfg.peft.enabled is True, dim must be set in peft_cfg"
	"If megatron_cfg.peft.enabled is True, dim must be set in peft_cfg"

[yfw]

examples/configs/recipes/llm/grpo-ultra-6n4g-megatron-lora.yaml:175

This comment is now stale: the PR removes the ValueError guard in setup.py that blocked PEFT + freeze_moe_router. The removal is safe (the router weight is a base param, frozen anyway under PEFT, and the freeze hook runs before LoRA wrapping), so the code is fine — but the comment should be updated/removed to avoid implying an unsupported combo.

[yfw]

tests/unit/models/policy/test_megatron_lora_refit_export.py:25

This path is cwd-relative and resolved at import time, so the whole module hits a FileNotFoundError collection error when run as cd tests && uv run pytest unit/... (the invocation the testing skill documents). It passes from the repo root (the L0 runner's cwd). Anchoring on __file__ makes it cwd-independent:

Suggested change

	source = Path("nemo_rl/models/policy/workers/megatron_policy_worker.py")
	source = (
	Path(__file__).resolve().parents[4]
	/ "nemo_rl/models/policy/workers/megatron_policy_worker.py"
	)

[yfw]

tests/test_suites/nightly.txt:79

Nit: the new recipe uses lowercase qwen3-8b (which matches the testing skill's examples), but the existing sibling is grpo-qwen3-8B-base-1n8g-fsdp2-lora. Worth aligning the family to lowercase 8b for consistency.

[yfw]

nemo_rl/models/policy/workers/megatron_policy_worker.py:636

This swaps reference/policy weights via a manual copy_ loop instead of load_state_dict(..., strict=True). It's defensible — PyTorch's load_state_dict also does in-place copy_ under no_grad, so no param-identity benefit is lost, and the bare reference_state_dict[k] index still fails loudly on a missing key. Two caveats: it drops the symmetric key validation, and it skips non-tensor _extra_state (FP8 scale/amax) — harmless for LoRA (policy and reference share identical frozen base weights) but a latent gap for FP8 full-finetune reference swaps.

Worth flagging that this reference-model PEFT support already exists on main, where it's implemented more robustly via _apply_state_dict_to_model(raise_if_key_missing=True) plus FP8 _extra_state handling and a sampling_params reset. This PR re-introduces the capability on ultra-v3 with the older bare-copy_ approach, so it isn't new functionality versus main — it brings ultra-v3 to parity, but via a path that will conflict when ultra-v3 rebases/merges onto main. Consider adopting main's _apply_state_dict_to_model here (or at minimum a one-time key-set assertion + a comment on why copy_ over load_state_dict) to minimize that reconciliation.

[yfw]

@vinhngx can you also address the CI failing checks?

[vinhngx]

addressed the Lora checkpoint conversion issue and all CI issues. Looks like some CI issues are inherited from the ultra-v3 branch (wonder if those CI checks also failed on the ultra-v3 branch).

[vinhngx]

Revision Notes

Addressed reviewer feedback for the Ultra Megatron LoRA PR.

Code Updates

• Fixed Ultra HF LoRA export for grouped-query QKV layout by deriving Q/K/V dimensions from the base model config instead of using a contiguous QKV split.
• Added regression coverage for the Ultra LoRA converter, including grouped QKV split behavior.
• Switched reference model state restoration to strict load_state_dict application.
• Fixed megatron_cfg typo in Megatron setup error messages.
• Changed YAML lora_dtype: None to lora_dtype: null.
• Removed stale MoE router freeze comment from the Ultra LoRA recipe.
• Made the Megatron LoRA refit-export unit test cwd-independent.
• Normalized Qwen3 8B nightly suite naming to lowercase.
• Fixed copyright/signoff issues.

Validation

• Local py_compile passed for touched Python files.
• Local focused pytest passed, including running the refit-export test from both repo root and tests/.
• Live Slurm validation completed on job 3312576:
  • Container unit tests: 8 passed in 2.20s
  • One-step Ultra Megatron LoRA training completed and saved checkpoint
  • HF LoRA converter wrote 96 adapter tensors, skipping only optimizer state
  • vLLM loaded the exported adapter as dapo-lora
  • LoRA chat smoke returned 37

Commit

• Updated revision commit: ac173f0c2 fix: address Ultra LoRA PR feedback
• Includes Signed-off-by trailer.