Support sync_ref_model with PEFT by syncing the reference adapter - #6975
Support sync_ref_model with PEFT by syncing the reference adapter#6975behroozazarkhalili wants to merge 13 commits into
Conversation
`sync_ref_model=True` raised NotImplementedError on every PEFT model in GRPO, DPO, KTO and RLOO, on the grounds that PEFT keeps no separate `ref_model` to synchronize. All four trainers already build what a sync needs: a frozen "ref" adapter copied from "default", created when an already-PEFT model is passed in. The guard rejected that case too, so it was broader than the limitation it described. SyncRefModelCallback now handles it. The existing `_sync_target_model` zips the parameters of two modules, which cannot express a reference that lives inside the policy model, so the PEFT branch pairs each "default" parameter with its "ref" counterpart by name, reusing the mapping that initializes the adapter. The ZeRO-3 gather wrapper mirrors the one already beside it. The trainers also create the "ref" adapter when the caller passes `peft_config` and asks for a synced reference. Without that, the same config would work or raise depending only on how the model was constructed. A caller who leaves `sync_ref_model=False` still gets no extra adapter, so nobody pays for one they did not ask for. The guard now fires only when no reference adapter could be created at all, which is `peft<0.20.0` with a LoRA config using `target_parameters`. Resolves #3108
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Jokasa7
left a comment
Jokasa7
left a comment
There was a problem hiding this comment.
I traced the new EMA path through PEFT 0.20.0's actual parameter layouts and all four duplicated trainer branches. Default LoRA matrices follow the intended path, but two supported LoRA configurations remain either silently unsynchronized or fail while constructing the reference adapter; details inline. I used AI assistance for source navigation and verified both findings against the exact PR head and exact upstream methods with executable reproductions.
| # `_sync_target_model`'s parameter-wise zip of two modules does not apply. Pair each `"default"` parameter with | ||
| # its `"ref"` counterpart by name instead; this is the same mapping used to initialize the `"ref"` adapter. | ||
| for name, param in model.named_parameters(): | ||
| if ".default." in name: |
There was a problem hiding this comment.
[P2] Map PEFT state structurally instead of matching .default.
This predicate misses valid parameters from LoraConfig(trainable_token_indices=...). PEFT 0.20.0 stores those deltas in a ParameterDict under names ending in ...trainable_tokens_delta.default, so there is no trailing dot after the adapter name. The identical initialization loops in the four trainers and this callback therefore never copy or EMA-update the corresponding .ref delta. Executing the exact-head method with a default delta of 7 and a ref delta of 0 left ref at 0 even with alpha=1.0, while reference forwards explicitly select "ref". The same hard-coding also copies default rather than model.active_adapter when an already-PEFT model is training a custom adapter. Please map through PEFT adapter state dicts, or otherwise handle terminal adapter keys and derive the source adapter explicitly; a trainable_token_indices regression test would catch the silent stale-reference case.
There was a problem hiding this comment.
Confirmed. Fixed in be926d5.
On a tiny Qwen2 with LoraConfig(trainable_token_indices=[0, 1]) the parameter is base_model.model.model.embed_tokens.token_adapter.trainable_tokens_delta.default. It ends in .default with nothing after it, so ".default." in name is false and one of five adapter parameter pairs was skipped. Setting the policy delta to 7.0 and the reference to 0.0, then running the callback with alpha=1.0, left the reference at 0.0.
Getting the replacement right took three attempts, and the two wrong ones are worth recording because they fail in opposite directions:
- Matching
"default"as a path component fixes the terminal case but matches a base-model parameter of that name. PEFT reserves neitherdefaultnorref, so a model owningbase_model.model.defaultproducedbase_model.model.ref, which nothing provides, andget_parameterraised during construction. - Taking the last such component fixes that but breaks
modules_to_save. PEFT wraps the saved module in aModuleDictkeyed by adapter, so a saved module with its owndefaultchild reads...block.modules_to_save.default.default.weight, where the adapter key is the first component. The last one belongs to the base module, so the parameter was skipped and its reference never moved, reintroducing exactly the defect you reported.
What actually distinguishes an adapter parameter is the container: PEFT keys them by adapter name inside a ModuleDict (LoRA matrices, modules_to_save) or a ParameterDict (the token deltas). The scan now walks the default components from the end and takes the first whose container also holds a ref key, which additionally declines a base ModuleDict that happens to have a default entry and no counterpart.
Five tests in TestSyncRefModelCallbackAdapterPairing cover the terminal token delta, a base module and parameter named default, a modules_to_save child of the same name, a base ModuleDict with no ref key, and the EMA equation itself. Each of the four ways to get the rule wrong fails one of them: reverting to the last-occurrence rule, dropping the ref membership test, dropping the container-type test, and replacing the EMA update with a plain copy.
| # model is the reference. That equivalence only holds while the reference stays fixed: with | ||
| # `sync_ref_model=True` the reference has to track the policy, which requires parameters of its own to | ||
| # move. So in that case create the "ref" adapter here as well. | ||
| needs_ref_adapter = args.sync_ref_model and ref_model is None |
There was a problem hiding this comment.
[P2] Handle bias-bearing LoRA before creating the ref adapter
This new peft_config path also reaches model.add_adapter("ref", default_config) for valid LoraConfig(bias="all") and bias="lora_only" configurations. PEFT 0.20.0 rejects that second adapter in BaseTuner._check_new_adapter_config with supports only 1 adapter with bias, which I reproduced using the exact upstream method. All four trainers fail here before the new "ref" not in peft_config guard, so the guard's message that only old target_parameters configurations lack a sync target is incomplete. Since these bias parameters are shared rather than adapter-local, please either reject this configuration explicitly before adding "ref" with an accurate explanation, or use a synchronization design with an independent target for it.
There was a problem hiding this comment.
Confirmed. Fixed in be926d5.
LoraConfig(bias="all") and bias="lora_only" both raise ValueError: LoraModel supports only 1 adapter with bias from add_adapter("ref", ...), while bias="none" succeeds. The call sat after the guard, so those configurations failed while the trainer was still being constructed.
The guard now skips the reference adapter for them, in the same shape as the existing target_parameters skip. With sync_ref_model=True that lands on the existing check for a missing "ref" adapter, so the configuration is rejected with an explanation rather than a PEFT stack trace. Skipping at this layer rather than raising here is deliberate: needs_ref_adapter also covers a pretrained PEFT model without synchronization, where a fixed base-model reference is supported and raising would break a working setup.
You were right that the guard's message was incomplete, and it was worse than incomplete once bias could reach it. It named target_parameters and peft<0.20.0 as the only cause, which would have been an actively misleading diagnosis for a bias config. It now names both causes and both remedies.
On your point about rejecting versus falling back: the reference log probs do come from the base model, which is fixed and cannot track the policy, so asking for both a bias-bearing adapter and sync_ref_model=True is refused rather than silently trained against a frozen reference. The regression test asserts the raise rather than a successful run.
…ring LoRA
The "default" to "ref" pairing matched the adapter name as the substring
".default.", which requires a path component after it. `trainable_token_indices`
stores its deltas in a `ParameterDict` keyed by adapter, so those parameter names
end in ".default" with nothing following. The predicate skipped them, and the
reference copy of the token deltas stayed frozen while the policy moved. Matching
"default" as a path component covers both placements: on a tiny Qwen2 with
`trainable_token_indices=[0, 1]`, one of five adapter parameter pairs went
unsynced before, none after.
A LoRA config with `bias` other than "none" trains bias terms shared with the
base model, and PEFT permits one such adapter per model: `add_adapter("ref", ...)`
raises `LoraModel supports only 1 adapter with bias` for both "all" and
"lora_only", while "none" succeeds. That call sat after the guard, so those
configurations failed during trainer construction. The guard now skips the
reference adapter for them, which routes into the existing `sync_ref_model`
rejection rather than a PEFT stack trace.
That rejection named `target_parameters` and `peft<0.20.0` as the only cause,
which is wrong once bias reaches it, so the message now names both.
All five call sites move together, as the trainers require.
…ameters The eight tests added in the previous commit build a `LoraConfig`, which is imported only when PEFT is installed, and they carried no `@require_peft` while every neighbouring PEFT test does. All eight failed the no-optional-dependency CI lane with `NameError: name 'LoraConfig' is not defined`. A local run cannot show this, because PEFT is installed there. Pairing "default" with "ref" on the path component alone was also wrong in both directions. PEFT reserves neither name, so a base model owning a parameter called `default` produced a reference name nothing provides: `base_model.model.default -> base_model.model.ref` raised `AttributeError` during construction. Taking the last such component instead was wrong the other way: `modules_to_save` wraps the saved module in a `ModuleDict` keyed by adapter, so a saved module with its own `default` child reads `...modules_to_save.default.default.weight`, whose adapter key is the first component. The last one belongs to the base module, so the parameter was skipped and its reference never moved, which is the defect this branch set out to fix. Adapter parameters are the ones PEFT keys by adapter name inside a `ModuleDict` or a `ParameterDict`, so the scan now walks the candidates from the end and takes the first whose container also holds a "ref" key. That covers the LoRA matrices, the `trainable_token_indices` deltas and `modules_to_save` alike, and it declines a base container that happens to hold a `default` key with no counterpart. Five callback tests cover the cases: a base module and parameter named `default`, a `modules_to_save` child of the same name, a base `ModuleDict` without a "ref" key, the terminal token delta, and the EMA equation itself. Each of the four ways to get the rule wrong fails one of them.
…f the vLLM sync, and correct the docs A LoRA config with `bias` other than `"none"` trains bias terms that live in the base model. The previous commit turned PEFT's crash on the second adapter into a warning saying the reference would come from the base model with adapters disabled. That reference is not fixed: with `bias="all"` six base biases train, and the adapter-disabled logits move when they move (PEFT warns about the same thing). All four trainers now raise a `ValueError` at construction, with or without `sync_ref_model`, and the `NotImplementedError` for a missing "ref" adapter no longer lists bias as a cause. DPO and KTO also name the case it had left out: a standalone `ref_model` passed alongside a PEFT policy, whose parameters do not pair with the adapter's. The vLLM weight sync skipped the "ref" adapter's LoRA matrices through the tuner prefix but not its `modules_to_save` copy: `lm_head.modules_to_save.ref.weight` was streamed to vLLM under that name. The four sync loops (server and colocate, in `VLLMGeneration` and OnlineDPO) now skip every `modules_to_save` copy except the "default" one. Reachable since the previous commit through GRPO and RLOO with `peft_config`, `sync_ref_model=True`, `use_vllm=True` and `modules_to_save`. RLOO created the "ref" adapter for `beta=0.0` and then raised that there is nothing to synchronize; its guard now carries GRPO's `beta != 0.0` clause. The DPO and KTO sync tests checked only that the adapter's parameters moved. They now also compute the reference log probs through the trainer's own path after training and require them to differ from the base model's, taken with adapters disabled at the same moment: a fresh "ref" adapter reproduces the base model exactly, so only a synced one can differ. A copy of the tree whose reference path ignores the "ref" adapter fails that assertion in both trainers. (A before/after comparison was tried first and rejected: enabling gradient checkpointing during `train()` shifts the reference forward by 1e-3 on identical weights, which a same-moment comparison does not see.) The callbacks fixture said it owned a base parameter named "default" but registered "default_bias"; it now owns one, and the test asserts the sync leaves it alone. The DPO and KTO docs still said `sync_ref_model=True` is unsupported with PEFT models; they now describe the "ref" adapter copy and the three cases that are rejected.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit 2d45dc4. Configure here.
…keep trainable-token state out of the vLLM sync The bias check ran only when a "ref" adapter was about to be created, so with sync_ref_model off a LoRA config with bias="all" was still accepted although the docs and the error message said otherwise. The reason the message gives holds on both paths: the trained biases live in the base model, so disabling the adapter does not give a fixed reference whether or not the reference is synced. The check now runs whenever the trainer takes its reference from the PEFT model, on the peft_config path and on the pretrained-adapter path, in all four trainers. The bias tests cover sync on and off. With trainable_token_indices the vLLM sync pushed three names vLLM does not own: the wrapped embedding weight under its token_adapter prefix and the per-adapter deltas for "default" and "ref". Only the "ref" copy is new in this PR; the other two mean the feature never worked with vLLM. PEFT merges the deltas into the embedding before the sync, so the four loops now skip trainable_tokens_delta and strip the token_adapter prefix, and the merged weight lands under its base name. Measured on a tiny Qwen2 with a fake process group: three leaked names before, none after.
…ns and wrappers The reference-adapter path accessed LoraConfig.target_parameters unconditionally, which older PEFT releases do not define; the version window is now checked before the attribute. AdaLoRA allows a single trainable adapter, so sync_ref_model with an AdaLoraConfig raises a clear ValueError instead of failing inside PEFT. Trainable-token deltas are skipped and the token_adapter prefix stripped in every vLLM weight loop, including the FSDP1 loop and the three Online DPO loops that had not received the modules_to_save filter. The name-pairing predicate now requires the container's owner to be a PEFT tuner layer or ModulesToSaveWrapper, so a base-model ModuleDict keyed "default"/"ref" is left alone. DPO and KTO config text no longer says PEFT is unsupported. Tests: reference log-probs must differ from the adapter-disabled base (GRPO) and the logged KL must match the ref adapter within float noise and differ from the base (RLOO); target_parameters absent, AdaLoRA refused, merged embedding emitted, FSDP1 filter, base ModuleDict left alone, config help text.
2 participants
What
sync_ref_model=Truenow works with PEFT inGRPOTrainer,DPOTrainer,KTOTrainerandRLOOTrainer. It previously raisedNotImplementedError.Why the guard was too broad
Every one of the four trainers already creates the EMA target a sync needs: a frozen
"ref"adapter copied from"default", atgrpo_trainer.py:471,dpo_trainer.py:676,kto_trainer.py:724andrloo_trainer.py:384. The reference forward already selects it viause_adapter(model, adapter_name="ref" if "ref" in model.peft_config else None). The guard rejected that case anyway.Measured on
mainbefore this change, with a LoRA config andbeta != 0:Path B has the target and was still refused.
The three changes
SyncRefModelCallback._sync_target_modeldoeszip(target_model.parameters(), model.parameters(), strict=True), which assumes two modules. With PEFT the reference is a second adapter inside the policy model, so there is nothing to zip. The new branch pairs.default.parameters with their.ref.counterparts by name, which is the same mapping the trainers use to initialize the adapter.sync_ref_adaptermirrors the existingsync_target_modelZeRO-3GatheredParameterswrapper.The four trainers. Both PEFT entry paths now reach one copy of the create-and-copy block through a
needs_ref_adapterflag. On thepeft_configpath the adapter is created only when a synced reference is requested, because a freshly wrapped adapter is zero-initialized and the base model already serves as a fixed reference; that equivalence breaks as soon as the reference has to move. Each trainer uses its own predicate for "a PEFT reference is in use here", taken from its ownelif:args.beta != 0.0in GRPO,ref_model is Nonein DPO and KTO, unconditional in RLOO.The guard. It now fires only when no reference adapter could be created:
peft<0.20.0with a LoRA config usingtarget_parameters(peft#3340), or a standaloneref_modelpassed alongside a PEFT policy, whose parameters do not pair with the adapter's. In both cases the reference cannot track the policy, so there is nothing to sync, and the message says which case applies instead of describing PEFT as unsupported.Bias-bearing LoRA is rejected, not worked around. A
biasother than"none"trains bias terms that live in the base model, so disabling the adapter does not recover a fixed reference: measured on the tiny Qwen2 model,bias="all"makes 6 base biases trainable and the adapter-disabled logits move when they move (PEFT warns about the same thing). PEFT also refuses a second bias-bearing adapter, so no frozen copy is possible. All four trainers raise aValueErrorat construction, with or withoutsync_ref_model; an earlier revision of this PR warned and fell back to the base model, which would have trained against a drifting reference.vLLM weight sync. The sync loops skipped the "ref" adapter's LoRA matrices through the tuner prefix but not its
modules_to_savecopy:lm_head.modules_to_save.ref.weightwas streamed to vLLM under that name. The four loops (server and colocate, inVLLMGenerationand OnlineDPO) now skip everymodules_to_savecopy except"default". Reachable through GRPO and RLOO withpeft_config,sync_ref_model=True,use_vllm=Trueandmodules_to_save; verified against the PEFT state dict of the tiny model rather than a live vLLM server.Verification
Red then green, per trainer. Reverting only the source and keeping the four new tests makes all four fail with
NotImplementedError; with the source they pass.test_train_with_sync_ref_model_and_peftx4, with fixNotImplementedErrortest_train_with_sync_ref_modelx4bias="all"rejected at construction x4End to end on a real two-step GRPO run with
ref_model_mixup_alpha=0.6: withsync_ref_model=Truethe eight.ref.parameters move by 59.73 in total; with it off there are no.ref.parameters and nothing moves.The callback change alone is measurable in isolation: at
alpha=1.0the summed absolute difference between the"default"and"ref"LoRA B matrices goes from 6144.0 to 0.0, meaning the reference becomes the policy exactly. Before the change the same call left it at 6144.0.Memory
sync=Falseon thepeft_configpath still yields['default']in all four trainers, so a LoRA user who does not ask for a synced reference allocates nothing extra. Asking for one costs an adapter, not a second copy of the model.Resolves #3108
Note
Medium Risk
Changes reference-model behavior during RLHF/DPO training and vLLM parameter streaming; incorrect adapter pairing or weight filtering could skew KL/reference terms or break generation sync.
Overview
sync_ref_model=Truenow works with PEFT in DPO, KTO, GRPO, and RLOO. Instead of a separateref_model, trainers add a frozen"ref"adapter (when sync is requested) andSyncRefModelCallbackEMA-updates it from"default"using name-aware pairing—not a naive.default.→.ref.replace—sotrainable_token_indices,modules_to_save, and ambiguous"default"names in the base model are handled correctly.Validation and docs reject incompatible setups: LoRA
bias != "none", AdaLoRA with sync,peft<0.20.0+target_parameters(no ref adapter to sync), and standaloneref_model+ PEFT. Config help text and DPO/KTO compatibility sections describe PEFT behavior.vLLM / Online DPO weight sync skips non-
defaultmodules_to_savecopies (including the ref adapter) andtrainable_tokens_deltaentries, and normalizes PEFT names so only the policy weights vLLM expects are pushed.Reviewed by Cursor Bugbot for commit 6ae5900. Bugbot is set up for automated code reviews on this repo. Configure here.
Update: PEFT version window, AdaLoRA, vLLM loops, pairing predicate
A further adversarial pass found seven defects; all were confirmed by execution and are fixed here.
LoraConfig.target_parametersbefore checking the version. PEFT below 0.17 has no such attribute, so the trainer raisedAttributeErroron a supported release. The version window (0.17 <= peft < 0.20) is now tested first.AdaLoraModelallows a single trainable adapter, soadd_adapter("ref")failed inside PEFT withsync_ref_model=True. The four trainers now refuse anAdaLoraConfigup front with aValueErrorthat names the alternative. Building the frozen copy withinference_mode=Truewas tried and rejected: it passes construction but breaks gradient-checkpointed training with atorch.utils.checkpointrecompute-metadata mismatch on the DPO and KTO sync tests.trainable_token_indicesdeltas were still streamed under names vLLM does not have, and thetoken_adapter.prefix was not stripped from the merged embedding. Every loop now skips the deltas and strips the prefix. The FSDP1 loop and the three Online DPO loops also gain themodules_to_savefilter the other loops already had, so the frozen"ref"copy is never sent.ModuleDictthat happens to hold"default"and"ref"keys was treated as an adapter and its weights overwritten by the sync. The predicate now also requires the container's owner to be a PEFT tuner layer orModulesToSaveWrapper, in the callback and in the four trainers.sync_ref_model=Trueis not compatible with PEFT; they now describe the frozen"ref"adapter.Tests added for each:
target_parameterspatched out ofLoraConfig; AdaLoRA refused; merged embedding emitted with no delta names; FSDP1 filter on a fake wrapped tree; baseModuleDictleft untouched by the sync; config help text. The GRPO sync test now asserts the generated reference log-probs differ from the adapter-disabled base; the RLOO sync test asserts the logged KL matches a recomputation on the"ref"adapter within float32 noise (the trainer scores in chunks, the test in one batch) and differs from the base by more than that noise.On a compute node (job 58040371): 27 targeted tests green after the RLOO tolerance fix (verified locally, 27 s); restoring
callbacks.py,vllm_generation.py,dpo_trainer.py,dpo_config.pyorrloo_trainer.pyfrommainfails the matching tests; a GRPO mutant whose reference forward runs with adapters disabled fails the GRPO oracle, and the same mutant on RLOO fails the KL oracle (0.0017 against 5.6e-5). Thegrpo_trainer.py,kto_trainer.pyandonline_dpo_trainer.pyarms have no test of their own for these lines; the DPO tests cover the shared block and the vLLM tests cover the shared loop shape. ruff 0.13.3 and the pinned doc-builder gate are clean.