Describe the bug
AdapterModelPTMixin.save_adapters() / load_adapters() do not round-trip. There are two
independent defects, both in nemo/core/classes/mixins/adapter_mixins.py.
(1) A file saved for a subset of adapters cannot be loaded with the documented default
name=None.
save_adapters writes state-dict entries only for the adapters it was asked to save
(:891-901) but stores the whole model adapter config as __cfg__ (:929).
load_adapters(name=None) then takes its list of adapters to restore from that config rather
than from the state dict (:972-973), and raises KeyError for every adapter the file never
contained.
(2) Adapter state-dict entries are selected and stripped by substring, so adapters with
overlapping names contaminate each other.
module.adapter_layer is an nn.ModuleDict holding every adapter of that module, so its
state_dict() keys are <adapter name>.<...> for all of them. save_adapters filters those
with if adapter_name in k (:923), so saving an4 also serialises all of an4_v2's tensors
into a file advertised as holding only an4. load_adapters then strips the name with
k.replace(f"{adapter_name}.", "") (:1043), which leaves the foreign keys intact, and the
per-module load_state_dict(..., strict=True) fails.
This is the workflow documented at docs/source/core/adapters/intro.rst:132 ("storing just the
adapter module(s) separately from the Model, so that you can use the same base model, and share
just the Adapter modules") and taught in tutorials/02_NeMo_Adapters.ipynb.
Steps/Code to reproduce bug
Both reproducers use the fixtures already in
tests/core/mixins/adapters/test_adapter_model_mixin.py. CPU only, no GPU or pretrained model.
import os, tempfile
os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1" # see load_adapters() docstring, adapter_mixins.py:963
import torch
from tests.core.mixins.adapters.test_adapter_model_mixin import (
DefaultAdapterModel, get_adapter_cfg, get_model_config,
)
# (1) save one of two adapters, then load with the documented default
m = DefaultAdapterModel(cfg=get_model_config(in_features=50))
m.add_adapter("encoder:en", cfg=get_adapter_cfg())
m.add_adapter("encoder:de", cfg=get_adapter_cfg())
with tempfile.TemporaryDirectory() as d:
fp = os.path.join(d, "en.pt")
m.save_adapters(fp, name="encoder:en")
DefaultAdapterModel(cfg=get_model_config(in_features=50)).load_adapters(fp)
KeyError: "Requested to load adapter with name `encoder:de`, but could not the adapter in the
state dict. \nAvailable adapter names in state dict are: ['encoder:en']"
# (2) two adapters where one name is a prefix of the other
m = DefaultAdapterModel(cfg=get_model_config(in_features=50))
m.add_adapter("encoder:an4", cfg=get_adapter_cfg())
m.add_adapter("encoder:an4_v2", cfg=get_adapter_cfg())
with tempfile.TemporaryDirectory() as d:
fp = os.path.join(d, "an4.pt")
m.save_adapters(fp, name="encoder:an4")
print(sorted(torch.load(fp, weights_only=False)["encoder:an4"][0].keys()))
DefaultAdapterModel(cfg=get_model_config(in_features=50)).load_adapters(fp, name="encoder:an4")
['an4.module.0.bias', 'an4.module.0.weight', 'an4.module.1.weight', 'an4.module.3.weight',
'an4_v2.module.0.bias', 'an4_v2.module.0.weight', 'an4_v2.module.1.weight', 'an4_v2.module.3.weight']
RuntimeError: Error(s) in loading state_dict for LinearAdapter:
Unexpected key(s) in state_dict: "an4_v2.module.0.weight", "an4_v2.module.0.bias",
"an4_v2.module.1.weight", "an4_v2.module.3.weight".
Note the first line: the "small shareable adapter file" silently contains a second adapter's
weights.
Controls, to show neither result is an artefact of the reproducer:
- In (1),
load_adapters(fp, name="encoder:en") succeeds, and
save_adapters(fp, name=None) -> load_adapters(fp) succeeds.
- In (2), renaming
an4_v2 to other makes the same script succeed and produce a file with only
the four an4.* keys.
Expected behavior
load_adapters(path) restores exactly the adapters the file contains.
- A file saved for adapter
X contains only X's parameters, regardless of whether another
adapter's name has X as a prefix.
Environment overview
- Environment location: bare-metal virtualenv
- Method of install: source checkout of
main @ d0a07c056
Environment details
- OS: macOS 26.5.1 (darwin, arm64)
- PyTorch 2.12.0
- Python 3.10.18
Additional context
Both are prefix-vs-substring / wrong-source-of-truth issues:
k.startswith(f"{adapter_name}.") and k[len(prefix):] instead of in / .replace(), and
name = list(state_dict.keys()) instead of list(config.keys()).
Unrelated but adjacent, and offered only as context: the duplicate-adapter-name guard at
adapter_mixins.py:741 reads
if hasattr(self, '_restoring_adapters') and self._restoring_adapters is not True:.
_restoring_adapters exists only during restore (:687 sets it, :693 deletes it), so on the
normal add_adapter path hasattr() is False and the check never executes. Same-module
duplicates are still rejected by AdapterModuleMixin.add_adapter (:260-263), so only the
cross-module case (encoder:tuning then decoder:tuning) slips through. Enabling the guard
would be a behaviour change rather than a bug fix, so I have deliberately left it alone.
I have a branch with the fix for (1) and (2) plus regression tests and am happy to open a PR.
Describe the bug
AdapterModelPTMixin.save_adapters()/load_adapters()do not round-trip. There are twoindependent defects, both in
nemo/core/classes/mixins/adapter_mixins.py.(1) A file saved for a subset of adapters cannot be loaded with the documented default
name=None.save_adapterswrites state-dict entries only for the adapters it was asked to save(
:891-901) but stores the whole model adapter config as__cfg__(:929).load_adapters(name=None)then takes its list of adapters to restore from that config ratherthan from the state dict (
:972-973), and raisesKeyErrorfor every adapter the file nevercontained.
(2) Adapter state-dict entries are selected and stripped by substring, so adapters with
overlapping names contaminate each other.
module.adapter_layeris annn.ModuleDictholding every adapter of that module, so itsstate_dict()keys are<adapter name>.<...>for all of them.save_adaptersfilters thosewith
if adapter_name in k(:923), so savingan4also serialises all ofan4_v2's tensorsinto a file advertised as holding only
an4.load_adaptersthen strips the name withk.replace(f"{adapter_name}.", "")(:1043), which leaves the foreign keys intact, and theper-module
load_state_dict(..., strict=True)fails.This is the workflow documented at
docs/source/core/adapters/intro.rst:132("storing just theadapter module(s) separately from the Model, so that you can use the same base model, and share
just the Adapter modules") and taught in
tutorials/02_NeMo_Adapters.ipynb.Steps/Code to reproduce bug
Both reproducers use the fixtures already in
tests/core/mixins/adapters/test_adapter_model_mixin.py. CPU only, no GPU or pretrained model.Note the first line: the "small shareable adapter file" silently contains a second adapter's
weights.
Controls, to show neither result is an artefact of the reproducer:
load_adapters(fp, name="encoder:en")succeeds, andsave_adapters(fp, name=None)->load_adapters(fp)succeeds.an4_v2toothermakes the same script succeed and produce a file with onlythe four
an4.*keys.Expected behavior
load_adapters(path)restores exactly the adapters the file contains.Xcontains onlyX's parameters, regardless of whether anotheradapter's name has
Xas a prefix.Environment overview
main@d0a07c056Environment details
Additional context
Both are prefix-vs-substring / wrong-source-of-truth issues:
k.startswith(f"{adapter_name}.")andk[len(prefix):]instead ofin/.replace(), andname = list(state_dict.keys())instead oflist(config.keys()).Unrelated but adjacent, and offered only as context: the duplicate-adapter-name guard at
adapter_mixins.py:741readsif hasattr(self, '_restoring_adapters') and self._restoring_adapters is not True:._restoring_adaptersexists only during restore (:687sets it,:693deletes it), so on thenormal
add_adapterpathhasattr()isFalseand the check never executes. Same-moduleduplicates are still rejected by
AdapterModuleMixin.add_adapter(:260-263), so only thecross-module case (
encoder:tuningthendecoder:tuning) slips through. Enabling the guardwould be a behaviour change rather than a bug fix, so I have deliberately left it alone.
I have a branch with the fix for (1) and (2) plus regression tests and am happy to open a PR.