Skip to content

[Megatron] Fix reference adapter loading for LoRA RLHF - #9934

Open
taking-lying-flat wants to merge 2 commits into
modelscope:mainfrom
taking-lying-flat:agent/fix-megatron-ref-adapter-loading
Open

[Megatron] Fix reference adapter loading for LoRA RLHF#9934
taking-lying-flat wants to merge 2 commits into
modelscope:mainfrom
taking-lying-flat:agent/fix-megatron-ref-adapter-loading

Conversation

@taking-lying-flat

@taking-lying-flat taking-lying-flat commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Summary

  • load --ref_adapters into the policy model chunks that are actually used for adapter-based reference forwards
  • load --mcore_ref_adapter into ref_adapter by mapping source default lookup keys to the target reference slot
  • support safetensors adapter loading for lora_llm in both policy and reference slots
  • normalize one newly added Qwen f-string for the Python 3.10 lint runner

Root cause

MegatronRLHFTrainer.prepare_model() initialized self.ref_models as an empty list for LoRA training, but passed that list to bridge.load_weights() for --ref_adapters. The bridge therefore opened the checkpoint but converted zero model chunks. Reference forwards later switched self.unwrapped_models to ref_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_name and retained every PEFT key without remapping its checkpoint lookup key. A normal policy LoRA checkpoint therefore matched and overwrote the current model default adapter instead of initializing ref_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_llm no longer silently skips safetensors policy or reference adapters.

@taking-lying-flat
taking-lying-flat force-pushed the agent/fix-megatron-ref-adapter-loading branch 2 times, most recently from 2548291 to 59f23e4 Compare August 17, 2026 11:35
@taking-lying-flat
taking-lying-flat marked this pull request as ready for review August 17, 2026 11:39
@taking-lying-flat
taking-lying-flat force-pushed the agent/fix-megatron-ref-adapter-loading branch from 59f23e4 to f106742 Compare August 17, 2026 11:44
@taking-lying-flat taking-lying-flat changed the title [Megatron] Fix reference adapter loading for LoRA RLHF [Megatron] Fix RLHF reference adapters and checkpoint RNG state Aug 17, 2026
@taking-lying-flat
taking-lying-flat force-pushed the agent/fix-megatron-ref-adapter-loading branch from 26ced08 to f106742 Compare August 17, 2026 12:09
@taking-lying-flat taking-lying-flat changed the title [Megatron] Fix RLHF reference adapters and checkpoint RNG state [Megatron] Fix reference adapter loading for LoRA RLHF Aug 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant