Skip to content

FIX AdaMSS: register adamss_newB in other_param_names (delete_adapter leak)#3367

Closed
DaoyuanLi2816 wants to merge 1 commit into
huggingface:mainfrom
DaoyuanLi2816:fix/adamss-register-newB-buffer
Closed

FIX AdaMSS: register adamss_newB in other_param_names (delete_adapter leak)#3367
DaoyuanLi2816 wants to merge 1 commit into
huggingface:mainfrom
DaoyuanLi2816:fix/adamss-register-newB-buffer

Conversation

@DaoyuanLi2816

Copy link
Copy Markdown
Contributor

What does this PR do?

AdamssLayer.adamss_newB is a per-adapter frozen projection buffer (BufferDict(persistent=True)), but it is listed in neither adapter_layer_names nor other_param_names. The generic BaseTunerLayer machinery iterates adapter_layer_names + other_param_names, so it never touches adamss_newB:

  • delete_adapter(name) removes name from adamss_A / adamss_B but leaves it in adamss_newB — the deleted adapter's projection buffer leaks.
  • per-adapter device moves (_move_adapter_to_device_of_base_layer) skip it. (In the common get_peft_model(...).to(device) flow nn.Module.to() still moves the registered submodule, so that only bites exotic multi-device-base setups; the delete leak is the concrete issue.)

Frozen per-adapter buffers belong in other_param_names — this is exactly what VeRA does with its vera_A / vera_B BufferDicts. The existing comment shows the author correctly kept adamss_newB out of adapter_layer_names (it isn't trainable) but missed that it still needs registering in other_param_names.

How it was verified

  • Repro: with two AdaMSS adapters, delete_adapter("a1") leaves "a1" in every layer's adamss_newB on main (while adamss_A is correctly cleaned); with the fix it's removed.
  • Added test_delete_adapter_removes_newB_buffer — it fails on main and passes with the fix.
  • delete_adapter still works with no crash (the adamss_ prefix means the module-level delete already filters these keys), and forward / merge_and_unload are unaffected. tests/test_adamss_asa.py passes (9 passed); ruff check / ruff format --check clean.

AI assistance

AI assistance was used to prepare this change. I reviewed every changed line, verified the behavior as described above, and take responsibility for the patch.

adamss_newB is a per-adapter frozen projection buffer (BufferDict) but was
listed in neither adapter_layer_names nor other_param_names, so the generic
BaseTunerLayer machinery (which iterates adapter_layer_names + other_param_names)
never touched it: delete_adapter() left the deleted adapter in adamss_newB (leak)
and per-adapter device moves skipped it. Frozen per-adapter buffers belong in
other_param_names, like vera_A/vera_B in VeRA.
@DaoyuanLi2816

Copy link
Copy Markdown
Contributor Author

Superseded by #3376, which now registers adamss_newB alongside the same fix for the other tuners (per the reviewer's request to pool all cases into one PR).

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