Skip to content

save_adapters/load_adapters do not round-trip: subset saves fail to load, and names are matched by substring #16102

Description

@udsy19

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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions