FIX safe_merge leaving the base layer modified, and merging several adapters at once on torchao by yupengtang · Pull Request #3779 · huggingface/peft

yupengtang · 2026-09-21T14:34:08Z

safe_merge could leave the base layer modified when it raised (#3710)

safe_merge=True is supposed to raise before anything is written to the base layer. Several layers either validated after writing or did not validate at all:

BOFT Conv2d had no finiteness check in its safe merge branch.
AdaMSS and RoAd wrote the merged weight before validating the merged bias; the RoAd bnb layers never validated the bias.
The LoRA bnb layers also wrote the weight first, and their bias check was missing .all(), so it raised RuntimeError: Boolean value of Tensor with more than one value is ambiguous for any real bias.
LoRA + torchao ran the check on the dequantized base weight before the delta was added, so a broken adapter passed the check.

While checking the bnb path I also found that the LoRA bnb layers merge and unmerge lora_B.bias without the adapter scaling. #2489 fixed that for the non-quantized layers only; with lora_alpha=16, r=8 the merged 8bit logits were off by 6.89 (0.078 after the fix).

Merging several adapters at once on torchao (#3728)

TorchaoLoraLinear.merge read base_layer.weight once before the loop and deleted the local at the end of each iteration, so merge_adapter(adapter_names=["a", "b"]) raised UnboundLocalError on the second adapter. The weight is now read inside the loop, as unmerge already does.

Tests

tests/testing_common.py: two common tests run over all test_custom_models.py cases. One makes the adapter weights non-finite, the other makes the base layer bias non-finite so the weight check passes first. Both check that safe_merge=True raises and leaves the base layer unchanged.
tests/test_common_gpu.py / tests/test_gpu_examples.py: bnb and torchao versions of the same check, plus the bnb bias scaling.
tests/test_quantization.py: merging two adapters in one call on torchao gives the same base weights as merging them one at a time.

Results:

CPU, pytest tests/test_custom_models.py -k safe_merge: 948 passed, 54 skipped. With the new tests on main's source: 17 failed (BOFT Conv2d, and AdaMSS / RoAd with a non-finite bias).
CPU, pytest tests/test_quantization.py -k torchao: 55 passed, 2 skipped. The new multi-adapter test raises the UnboundLocalError from LoRA + torchao: merging several adapters in one call raises UnboundLocalError #3728 without the torchao change.
ruff check / ruff format --check clean with ruff 0.16.4.
GPU (L40S, torch 2.11 + cu130, bitsandbytes 0.50.2, torchao 0.18.0), the safe merge, bnb bias and torchao merge tests from tests/test_common_gpu.py, tests/test_gpu_examples.py and tests/test_quantization.py: 14 passed. On main's source with the same tests: 8 failed (missing .all() in the LoRA bnb bias check, RoAd bnb not raising, unscaled lora_B.bias in 8bit and 4bit, torchao safe merge not raising, and the LoRA + torchao: merging several adapters in one call raises UnboundLocalError #3728 UnboundLocalError), and the 6 pre-existing merge tests pass on both.
GPU, pytest tests/test_quantization.py -k torchao: 55 passed, 2 skipped.

AI assistance was used for this PR (Claude Code), mainly to write the audit scripts that compared the safe_merge implementations across tuners and to draft tests and this description. I reviewed every changed line and ran the tests listed above.

One small thing noticed while checking the approval format: scripts/triage_prs.py matches /peft-triage approved, while .ai/AGENTS.md says @peft-triage approved and the header comment in .github/workflows/triage_prs.yml says /peft-triage: approved. Following either document would not match the automation. Happy to send a separate PR for that if useful.

BenjaminBossan

Thanks for this PR with the pooled changes. I only skimmed the fixes so far and focused on the testing first. Please check my comments.

BenjaminBossan · 2026-09-23T12:36:44Z


+    @pytest.mark.single_gpu_tests
+    def test_torchao_safe_merge_broken_adapter_leaves_base_layer_unchanged(self):
+        from transformers import TorchAoConfig


Use global imports.

Done, TorchAoConfig is now imported at the top of the file.

BenjaminBossan · 2026-09-23T12:39:15Z

+        bias = bias.detach().clone() if bias is not None else None
+        return weight, bias
+
+    def _assert_base_layer_unchanged(self, layer, snapshot):


Just pass weight and bias directly instead of snapshot. Call them orig_weight and orig_bias.

Done, the helper now takes orig_weight and orig_bias directly and the snapshot helper is gone.

BenjaminBossan · 2026-09-23T12:40:26Z

+
+        with hub_online_once(model_id):
+            model = self.transformers_class.from_pretrained(model_id)
+            config = config_cls(base_model_name_or_path=model_id, **config_kwargs)


Not needed, right?

Suggested change

config = config_cls(base_model_name_or_path=model_id, **config_kwargs)

config = config_cls(**config_kwargs)

Right, not needed. Applied.

BenjaminBossan · 2026-09-23T12:46:49Z

+        if not layers:
+            pytest.skip("No mergeable tuner layers in this model.")
+
+        for layer in layers:


Not a fan of looping through the layers and merging them individually. This is not what a user would typically do. Please poison an individual weight, create a copy of the tensors, call model.merge_adapter(), then check if the tensors are preserved. Same idea for the other test.

Changed both tests to do this: poison one value, keep copies of the base weight and bias, call model.merge_adapter(safe_merge=True) and check the copies are preserved. The poisoned layer is the first tuner layer in module order, which is the first one merge_adapter visits, so the merge raises before any layer is merged.

BenjaminBossan · 2026-09-23T12:50:22Z

+                continue
+            with torch.no_grad():
+                for param in adapter_params:
+                    param.fill_(float("nan"))


Let's just change a single value to nan or inf, not the whole tensor. Also, I'm a bit surprised that this works, as torch.equal should return False for nan.

Now a single value. On the torch.equal question, the comparison never sees a NaN. In the broken adapter test the NaN goes into the adapter parameters and what gets compared is the base weight and bias, which stay finite. In the bias test the base bias is poisoned with inf rather than NaN, and torch.equal treats inf as equal to inf, so the check is meaningful there as well.

BenjaminBossan · 2026-09-23T13:00:01Z

        check_outputs_similar(out_before, out_unloaded)
+
+    @pytest.mark.skipif(not is_torchao_available(), reason="torchao is not installed")
+    def test_torchao_lora_merge_several_adapters_in_one_call(self):


test_quantization.py is for generic quantization tests. For specific tests like this, move it to test_gpu_examples.py.

Moved it to test_gpu_examples.py, into TestPeftTorchao next to the other torchao merge test.

BOFT Conv2d had no finiteness check in its safe merge branch. AdaMSS and RoAd wrote the merged weight before validating the merged bias, and the RoAd bnb layers never validated the bias at all. The LoRA bnb layers also wrote the weight first, and their bias check was missing .all(), so it raised a RuntimeError for any real bias. LoRA + torchao ran the check on the dequantized base weight before the delta was added. The LoRA bnb layers also merged and unmerged lora_B.bias without scaling. huggingface#2489 fixed that for the non-quantized layers only. Add common tests that make the adapter or the base bias non-finite and check that safe_merge raises without touching the base layer, plus bnb and torchao tests for the quantized code paths.

TorchaoLoraLinear.merge read base_layer.weight once before the loop and deleted the local at the end of every iteration, so merging a second adapter in the same call raised UnboundLocalError. Read the weight inside the loop, as unmerge already does.

* Poison a single value instead of whole tensors and go through model.merge_adapter(safe_merge=True) instead of merging layers one by one. * The bnb tests now load a full model and merge through it. * Move the torchao several-adapters test to test_gpu_examples.py and use the global TorchAoConfig import. * Pass orig_weight/orig_bias directly and explain the weight guard.

yupengtang · 2026-09-26T15:41:24Z

Thanks for the review. All comments are addressed, and I rebased on current main.

Reran everything after the rework, each time against main's source as well so the new tests are known to catch the bugs:

CPU, pytest tests/test_custom_models.py -k safe_merge: 957 passed, 54 skipped. On main's source: 17 failed.
GPU (L40S), TestBnbMergeBias in test_common_gpu.py: 6 passed. On main's source: 6 failed.
GPU, the two torchao merge tests in test_gpu_examples.py: 2 passed. On main's source: 2 failed.
GPU, pytest tests/test_quantization.py -k torchao: 54 passed, 2 skipped (one fewer than before, since the several-adapters test moved to test_gpu_examples.py).
ruff check and ruff format --check clean.

BenjaminBossan

Thanks for the updates, I did another review. Please check my comments.

Moreover, I only just now noticed that this PR and #3757 partially fix the same problem, namely the missing application of scaling when LoRA bias is enabled and when the safe merge check is applied. It only covers the LoRA bnb part though, the rest of this PR is not covered. Sorry for not noticing this earlier. I think the best course of action is to remove that part of the change from this PR, as this PR is already quite big as is. I hope that the other PR can be merged very soon, then this PR can be updated to take the change into account.

BenjaminBossan · 2026-09-28T10:39:16Z

+        # poison the first layer that merge_adapter visits, so the merge fails before any layer is merged; one
+        # non-finite value per adapter parameter is enough to make the delta weight non-finite


The comment states that the first layer is poisoned, but below, we poison every layer, so either change the comment or the code.

Also, IIUC, if we poison the last layer, the first layer will still be in merged state before the safe merge check fails, right? So in a sense, this problem is only partially resolved. That is fine for now, maybe just add a comment to that effect.

Only layers[0] is poisoned; the loop goes over that single layer's own adapter parameters (lora_A and lora_B for LoRA), which I can see reads like every layer. I reworded the comment to make that explicit and added the note about earlier layers staying merged when a later layer is the one that fails.

BenjaminBossan · 2026-09-28T10:41:12Z

+        orig_weight = base_layer.weight.detach().clone()
+        orig_bias = base_layer.bias.detach().clone()
+
+        try:


Not a big fan of try ... except ... else here. I'd rather see with pytest.raises in combination with a sentinel value for the else case.

Before changing this I checked which cases actually raise here. Of the 315 cases this test runs in test_custom_models.py, 39 raise (GLoRA, LoRA with lora_bias, IA3 without feedforward_modules, BEFT, RoAd, AdaMSS) and the other 276 merge nothing into the base bias, so for most methods the no-error path is the expected one. That makes the outcome depend on the method, so I was not sure what the sentinel should be based on. Two options I can see:

an explicit flag in the test cases saying whether the method merges into the base bias, then pytest.raises(...) or nullcontext() depending on it

deciding it at runtime by merging a copy of the model without safe_merge and checking whether the base bias moves

Happy to do either, or something else if you had another pattern in mind.

BenjaminBossan · 2026-09-28T10:43:44Z

+
+    @pytest.mark.single_gpu_tests
+    def test_torchao_lora_merge_several_adapters_in_one_call(self):
+        # Merging several adapters in one call should give the same weights as merging them one at a time


Let's reference the original issue for this.

Added a link to #3728.

BenjaminBossan · 2026-09-28T10:47:07Z

+        model_one_call = get_model_with_two_adapters()
+        model_one_call.merge_adapter(adapter_names=["default", "other"])
+
+        model_one_by_one = get_model_with_two_adapters()


Would it be possible to deep-copy the first model instead of having a function to create it twice?

Done. The model is loaded once and the second one is a deepcopy of it, so both now start from exactly the same adapters.

BenjaminBossan · 2026-09-28T10:47:37Z

+        def get_model_with_two_adapters():
+            torch.manual_seed(0)
+            quantization_config = TorchAoConfig(quant_type=self.get_quant_type("int8_weight_only"))
+            model = AutoModelForCausalLM.from_pretrained(


Let's use the with hub_online_once context.

… cleanups - Drop the LoRA bnb bias scaling and its test, which huggingface#3757 covers; the bnb merge now only computes and checks the bias before writing the weight. - Clarify which layer the broken adapter test poisons and note that earlier layers stay merged if a later one fails. - torchao several adapters test: link huggingface#3728, load the model once inside hub_online_once and deepcopy it.

yupengtang · 2026-09-30T18:28:36Z

Thanks for the second pass, and for spotting the overlap with #3757.

I took the LoRA bnb bias scaling out of this PR, together with its test, so that part is left to #3757. What remains in lora/bnb.py is only the safe merge fix: the merged bias is computed and checked before the quantized weight is written, with .all() on the check. unmerge is back to what it is on main. Once #3757 is merged I will rebase; the only overlap left is the order of those lines in merge.

The inline comments are addressed, except the try ... except ... else one, where I have a question (see there).

Reran after the changes, against main's source as well:

CPU, pytest tests/test_custom_models.py -k safe_merge: 957 passed, 54 skipped. On main's source: 17 failed.
GPU (A100), the bnb and torchao merge tests in test_common_gpu.py and test_gpu_examples.py: 12 passed. On main's source: the 6 new ones fail (4 TestBnbMergeBias, 2 torchao), the rest pass.
GPU, pytest tests/test_quantization.py -k torchao: 54 passed, 2 skipped, same as on main.
ruff check and ruff format --check clean.

Since the LoRA bnb error message is back to main's wording, the bnb bias test now matches on NaNs detected only.

BenjaminBossan · 2026-10-15T23:54:02Z

@yupengtang As there has been no update on #3757 for several weeks now, I would like to give this PR priority for the LoRA BNB part. Could you please work this into the PR? Thanks!

BenjaminBossan self-assigned this Sep 22, 2026

BenjaminBossan requested changes Sep 23, 2026

View reviewed changes

github-actions Bot added the triaged label Sep 24, 2026

yupengtang added 3 commits September 26, 2026 11:12

yupengtang force-pushed the fix-safe-merge-leaves-base-layer-untouched branch from 2bbb183 to 6a48146 Compare September 26, 2026 15:40

BenjaminBossan requested changes Sep 28, 2026

View reviewed changes

	config = config_cls(base_model_name_or_path=model_id, **config_kwargs)
	config = config_cls(**config_kwargs)

		# poison the first layer that merge_adapter visits, so the merge fails before any layer is merged; one
		# non-finite value per adapter parameter is enough to make the delta weight non-finite

Conversation

yupengtang commented Sep 21, 2026

safe_merge could leave the base layer modified when it raised (#3710)

Merging several adapters at once on torchao (#3728)

Tests

Uh oh!

BenjaminBossan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yupengtang commented Sep 26, 2026

Uh oh!

BenjaminBossan left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yupengtang commented Sep 30, 2026

Uh oh!

BenjaminBossan commented Oct 15, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants