feat: add lora warm-start by pstjohn · Pull Request #3874 · NVIDIA-NeMo/RL

pstjohn · 2026-08-27T16:03:52Z

What does this PR do?

Adds LoRA warm-start support so a new training run can initialize its adapters from an existing checkpoint—for example, continuing from an SFT LoRA checkpoint with GRPO.

This supports both training backends:

DTensor/AutoModel: loads a PEFT adapter directory containing adapter_model.safetensors and adapter_config.json.
Megatron Core: loads adapters from a native Megatron iter_* checkpoint or checkpoint root.

The implementation:

Adds restore_from to the AutoModel LoRA and Megatron PEFT configurations.
Validates donor paths, adapter configuration, and adapter tensor coverage before loading.
Supports custom AutoModel state-dict adapters and expert-parallel MoE models.
Loads only adapter weights; optimizer, RNG, train, and rerun state start fresh.
Gives normal resume checkpoints precedence over restore_from.
Anchors the KL reference model to the warm-started adapters.
Adds defaults to example and reference configurations.
Documents the workflow and backend-specific checkpoint formats.

Issues

N/A

Usage

DTensor/AutoModel

policy:
  dtensor_cfg:
    _v2: true
    lora_cfg:
      enabled: true
      restore_from: /path/to/sft-run/step_100/policy/weights

restore_from may point directly to the adapter directory or its model/ subdirectory.

Megatron Core

policy:
  megatron_cfg:
    peft:
      enabled: true
      restore_from: /path/to/sft-run/step_100/policy/weights

restore_from may point to an iter_XXXXXXX directory or a checkpoint root that resolves to one.

The donor LoRA configuration must be compatible with the new run. In particular, adapter rank and scaling must match.

Testing

AutoModel adapter validation and loading: 11 passed
Megatron PEFT warm-start suite: 24 passed
Configuration validation tests: 14 passed
Targeted pre-commit checks: passed
git diff --check: passed

The backend-specific tests were run in the prebuilt DTensor and Megatron container runtimes.

Before your PR is "Ready for review"

Read and followed the contributor guidelines
Added unit tests covering both backends, invalid configurations, resume behavior, relative paths, expert parallelism, and failure cleanup
Ran the relevant unit tests locally
Added documentation and updated example configurations
Full functional test suite will run in CI

Additional Information

Warm-starting is distinct from resuming:

restore_from initializes a new run from donor adapter weights.
Optimizer and training progress are not restored.
If a normal NeMo RL resume checkpoint exists, its policy weights take precedence.
For RL algorithms using a reference policy, the reference remains anchored to the warm-started initial policy.

Enable GRPO to load a pre-trained LoRA adapter in order to continue LoRA fine-tuning from a SFT LoRA checkpoint. Signed-off-by: Peter St. John <pstjohn@nvidia.com>

copy-pr-bot · 2026-08-27T16:03:55Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

pstjohn · 2026-08-27T21:02:48Z

+    checkpoints. Optimizer state is not loaded: warm starts begin with a
+    fresh optimizer.
+    """
+    adapter_dir = _resolve_lora_adapter_dir(restore_from)


1 action item. Please fix this in this PR. setup.py:725

This is PR-introduced. With a relative restore_from, the symlink stores a relative target inside the temporary directory, so it resolves beneath that directory and dangles. A relative directory named model also skips staging, while AutoModel selects PEFT loading using the literal "/model" in path check.

AI-1

Normalize the resolved directory before basename/symlink handling and cover both relative-path cases.

Suggested change

adapter_dir = _resolve_lora_adapter_dir(restore_from)

adapter_dir = os.path.abspath(_resolve_lora_adapter_dir(restore_from))

Resolved in cc2a58a.

pstjohn · 2026-08-27T21:02:48Z

+    stripped_donor_keys = {
+        key[len(prefix) :] if key.startswith(prefix) else key for key in donor_keys
+    }
+    expected_keys = {


1 action item. Please fix this in this PR. setup.py:691

This is PR-introduced. For Qwen3-MoE with EP=1, donor_keys are HF-format names while expected_keys are native parameter names. AutoModel maps native lora_gate_and_up_A to PEFT base_layer.lora_B.weight, then performs the reverse mapping during load. This validator runs first and raises the new “donor adapter key mismatch” error for a valid donor.

AI-1

Compare keys in the same post-conversion namespace and add a Qwen3-MoE EP=1 save/warm-start round-trip test.

Resolved in cc2a58a.

pstjohn · 2026-08-27T21:02:48Z

+    ``gate_up_linear.weight0``) that do not appear in ``named_parameters()``,
+    so their coverage cannot be checked against parameter names.
+    """
+    if _has_expert_parallelism(model):


1 action item. Please fix this in this PR. setup.py:674

This is PR-introduced. With EP>1, a QKV-only donor and a current QKV+FC configuration pass here because validation returns immediately. AutoModel’s EP loader then continues past missing or unmatched donor keys and only logs counts, leaving the missing FC adapters freshly initialized while warm-start reports success.

AI-1

Validate normalized complete adapter-key sets for EP checkpoints and add the QKV-only donor versus QKV+FC regression case.

Resolved in cc2a58a.

pstjohn · 2026-08-27T21:02:48Z

+        ckpt_cfg.load = restore_dir
+        # finetune=False is required for the adapter-only filter, but
+        # optimizer/RNG state must not come from the donor run.
+        ckpt_cfg.finetune = False


1 action item. Please fix this in this PR. setup.py:629

This is PR-introduced. Disabling optimizer and RNG restore does not make this donor load weights-only: Megatron-Bridge independently restores the singleton rerun state, and its PEFT filtering preserves non-model metadata. MCore then overwrites mode, iteration, state, and pending-rerun flags; resetting TrainState does not undo that.

AI-1

Preserve/reset the current RerunStateMachine around donor loading, or use a true weights-only path, and add a donor-with-pending-rerun regression test.

Resolved in cc2a58a.

pstjohn · 2026-08-27T21:02:48Z

+            "megatron.bridge.training.checkpointing",
+            reason="requires the mcore extra (Megatron-Bridge)",
+        )
+        src = inspect.getsource(checkpointing)


1 action item. Please fix this in this PR. test_megatron_setup.py:3997

This test gap is PR-introduced. The five searched strings occur together in both the upstream adapter-filtering predicate and non-strict-load predicate. Removing either required block leaves every string elsewhere, so this tripwire still passes; the other warm-start tests mock the loader boundary.

AI-1

Replace this module-wide source search with a focused loader/storage-mocked behavior test that proves both adapter filtering and non-strict loading remain active.

Resolved in cc2a58a.

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

copy-pr-bot · 2026-08-27T21:58:15Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

pstjohn · 2026-08-27T21:58:43Z

/ok to test cc2a58a

feat: add lora warm-start

88f4e40

Enable GRPO to load a pre-trained LoRA adapter in order to continue LoRA fine-tuning from a SFT LoRA checkpoint. Signed-off-by: Peter St. John <pstjohn@nvidia.com>

addressing local review

b5fc6d6

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

github-actions Bot added the Documentation Improvements or additions to documentation label Aug 27, 2026

Merge branch 'main' into feat-lora-warm-start

cd5a3d9

pstjohn commented Aug 27, 2026

View reviewed changes

fix: respond to reviewer comments

cc2a58a

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

pstjohn marked this pull request as ready for review August 27, 2026 21:58

pstjohn requested review from a team as code owners August 27, 2026 21:58

Navigation Menu

Uh oh!

feat: add lora warm-start - #3874

feat: add lora warm-start#3874
pstjohn wants to merge 4 commits into
NVIDIA-NeMo:mainfrom
pstjohn:feat-lora-warm-start

pstjohn commented Aug 27, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

pstjohn Aug 27, 2026

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

pstjohn commented Aug 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	adapter_dir = _resolve_lora_adapter_dir(restore_from)
	adapter_dir = os.path.abspath(_resolve_lora_adapter_dir(restore_from))

Uh oh!

Conversation

pstjohn commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Issues

Usage

DTensor/AutoModel

Megatron Core

Testing

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

Choose a reason for hiding this comment

AI-1

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

AI-1

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

AI-1

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

AI-1

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

AI-1

Uh oh!

Choose a reason for hiding this comment

Uh oh!

copy-pr-bot Bot commented Aug 27, 2026

Uh oh!

pstjohn commented Aug 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

pstjohn commented Aug 27, 2026 •

edited

Loading