[Megatron] Fix reference adapter loading for LoRA RLHF - #9934
Open
taking-lying-flat wants to merge 2 commits into
Open
[Megatron] Fix reference adapter loading for LoRA RLHF#9934taking-lying-flat wants to merge 2 commits into
taking-lying-flat wants to merge 2 commits into
Conversation
taking-lying-flat
force-pushed
the
agent/fix-megatron-ref-adapter-loading
branch
2 times, most recently
from
August 17, 2026 11:35
2548291 to
59f23e4
Compare
taking-lying-flat
marked this pull request as ready for review
August 17, 2026 11:39
taking-lying-flat
force-pushed
the
agent/fix-megatron-ref-adapter-loading
branch
from
August 17, 2026 11:44
59f23e4 to
f106742
Compare
taking-lying-flat
force-pushed
the
agent/fix-megatron-ref-adapter-loading
branch
from
August 17, 2026 12:09
26ced08 to
f106742
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
--ref_adaptersinto the policy model chunks that are actually used for adapter-based reference forwards--mcore_ref_adapterintoref_adapterby mapping sourcedefaultlookup keys to the target reference slotlora_llmin both policy and reference slotsRoot cause
MegatronRLHFTrainer.prepare_model()initializedself.ref_modelsas an empty list for LoRA training, but passed that list tobridge.load_weights()for--ref_adapters. The bridge therefore opened the checkpoint but converted zero model chunks. Reference forwards later switchedself.unwrapped_modelstoref_adapter, so DPO/KTO/GRPO used the freshly initialized adapter instead of the requested checkpoint without raising an error.The MCore checkpoint path separately used the default
adapter_nameand retained every PEFT key without remapping its checkpoint lookup key. A normal policy LoRA checkpoint therefore matched and overwrote the current modeldefaultadapter instead of initializingref_adapter.Impact
Megatron RLHF runs using LoRA plus either reference-adapter format now load the requested weights into the adapter used for reference log-probability computation without overwriting the policy adapter.
lora_llmno longer silently skips safetensors policy or reference adapters.