@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1291bdb85e

[EPD][Model Loader] Skip language-model checkpoint shards for `--mm-encoder-only` (#58086)

Signed-off-by: grYe99 <guorongye99@gmail.com> Co-authored-by: grYe99 <guorongye99@gmail.com> Co-authored-by: Cursor Agent <cursoragent@cursor.com>

GuorongYe committed Sep 24, 2026 at 15:41 UTC 1291bdb85edc3cf0a1232be408b7e866b9d3f10a
5 files changed +483 -14
tests/model_executor/model_loader/test_mm_encoder_only_weight_filter.py new
+243
@@ -0,0 +1,243 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for --mm-encoder-only safetensors shard filtering."""
4 +
5 +import json
6 +import os
7 +import tempfile
8 +
9 +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
10 +
11 +from vllm.model_executor.model_loader.weight_utils import (
12 + filter_mm_encoder_only_safetensors_files,
13 + resolve_mm_encoder_only_lm_prefixes,
14 +)
15 +from vllm.model_executor.models.utils import WeightsMapper
16 +
17 +
18 +def _write_index(folder: str, weight_map: dict[str, str]) -> None:
19 + with open(os.path.join(folder, SAFE_WEIGHTS_INDEX_NAME), "w") as f:
20 + json.dump({"weight_map": weight_map}, f)
21 +
22 +
23 +def _qwen25_mapper() -> WeightsMapper:
24 + """Qwen2.5-VL dual layout mapper (nested + flat catch-all)."""
25 + return WeightsMapper(
26 + orig_to_new_prefix={
27 + "model.language_model.": "language_model.model.",
28 + "model.visual.": "visual.",
29 + "lm_head.": "language_model.lm_head.",
30 + "model.": "language_model.model.",
31 + }
32 + )
33 +
34 +
35 +def _qwen3_mapper() -> WeightsMapper:
36 + return WeightsMapper(
37 + orig_to_new_prefix={
38 + "model.visual.": "visual.",
39 + "lm_head.": "language_model.lm_head.",
40 + "model.language_model.": "language_model.model.",
41 + }
42 + )
43 +
44 +
45 +def _molmo_mapper() -> WeightsMapper:
46 + return WeightsMapper(
47 + orig_to_new_prefix={
48 + "model.vision_backbone.": "vision_backbone.",
49 + "model.transformer.blocks.": "model.layers.",
50 + "model.transformer.ln_f.": "model.norm.",
51 + "model.transformer.mlp.down_proj.": "lm_head.",
52 + }
53 + )
54 +
55 +
56 +def test_keeps_vision_shards_drops_language_only():
57 + with tempfile.TemporaryDirectory() as folder:
58 + files = [
59 + os.path.join(folder, "model-00001-of-000003.safetensors"),
60 + os.path.join(folder, "model-00002-of-000003.safetensors"),
61 + os.path.join(folder, "model-00003-of-000003.safetensors"),
62 + ]
63 + for path in files:
64 + open(path, "wb").close()
65 +
66 + _write_index(
67 + folder,
68 + {
69 + "language_model.layers.0.weight": "model-00001-of-000003.safetensors",
70 + "language_model.layers.1.weight": "model-00002-of-000003.safetensors",
71 + "vision_tower.blocks.0.weight": "model-00003-of-000003.safetensors",
72 + "mm_projector.weight": "model-00003-of-000003.safetensors",
73 + },
74 + )
75 +
76 + kept = filter_mm_encoder_only_safetensors_files(
77 + files,
78 + folder,
79 + SAFE_WEIGHTS_INDEX_NAME,
80 + ("language_model.",),
81 + )
82 + assert kept == [files[2]]
83 +
84 +
85 +def test_keeps_mixed_shard_with_vision_and_lm():
86 + with tempfile.TemporaryDirectory() as folder:
87 + mixed = os.path.join(folder, "model-00001-of-000001.safetensors")
88 + open(mixed, "wb").close()
89 + _write_index(
90 + folder,
91 + {
92 + "language_model.embed.weight": "model-00001-of-000001.safetensors",
93 + "vision_tower.patch.weight": "model-00001-of-000001.safetensors",
94 + },
95 + )
96 +
97 + kept = filter_mm_encoder_only_safetensors_files(
98 + [mixed],
99 + folder,
100 + SAFE_WEIGHTS_INDEX_NAME,
101 + ("language_model.",),
102 + )
103 + assert kept == [mixed]
104 +
105 +
106 +def test_no_index_returns_unchanged():
107 + with tempfile.TemporaryDirectory() as folder:
108 + files = [os.path.join(folder, "model.safetensors")]
109 + open(files[0], "wb").close()
110 + kept = filter_mm_encoder_only_safetensors_files(
111 + files,
112 + folder,
113 + SAFE_WEIGHTS_INDEX_NAME,
114 + ("language_model.",),
115 + )
116 + assert kept == files
117 +
118 +
119 +def test_resolve_returns_module_prefixes_without_hf_hardcode():
120 + """No global Qwen HF nest list — just vLLM module prefixes."""
121 + prefixes = resolve_mm_encoder_only_lm_prefixes(["language_model"])
122 + assert prefixes == ("language_model.",)
123 + # With Qwen mapper, still module prefixes (keys classified via mapper).
124 + assert resolve_mm_encoder_only_lm_prefixes(
125 + ["language_model"], weights_mapper=_qwen25_mapper()
126 + ) == ("language_model.",)
127 +
128 +
129 +def test_resolve_empty_language_model_names_fail_closed():
130 + """Missing/empty ``_language_model_names`` → None (no language_model. default)."""
131 + assert resolve_mm_encoder_only_lm_prefixes(None) is None
132 + assert resolve_mm_encoder_only_lm_prefixes([]) is None
133 + assert resolve_mm_encoder_only_lm_prefixes(()) is None
134 +
135 +
136 +def test_resolve_skips_filter_for_shared_hf_root():
137 + """Shared HF root (Molmo-shaped mapper) → fail-closed, not magic name list."""
138 + assert (
139 + resolve_mm_encoder_only_lm_prefixes(["model"], weights_mapper=_molmo_mapper())
140 + is None
141 + )
142 + # Without a mapper, identity layout: module prefix is returned as-is.
143 + assert resolve_mm_encoder_only_lm_prefixes(["model"]) == ("model.",)
144 + assert resolve_mm_encoder_only_lm_prefixes(["model."]) == ("model.",)
145 +
146 +
147 +def test_resolve_keeps_llm_prefix_without_model_broadening():
148 + prefixes = resolve_mm_encoder_only_lm_prefixes(["llm"])
149 + assert prefixes == ("llm.",)
150 +
151 +
152 +def test_qwen3_nested_index_skips_lm_keeps_visual():
153 + with tempfile.TemporaryDirectory() as folder:
154 + files = [
155 + os.path.join(folder, "model-00001-of-000002.safetensors"),
156 + os.path.join(folder, "model-00002-of-000002.safetensors"),
157 + ]
158 + for path in files:
159 + open(path, "wb").close()
160 + _write_index(
161 + folder,
162 + {
163 + "model.language_model.layers.0.weight": (
164 + "model-00001-of-000002.safetensors"
165 + ),
166 + "lm_head.weight": "model-00001-of-000002.safetensors",
167 + "model.visual.blocks.0.weight": "model-00002-of-000002.safetensors",
168 + },
169 + )
170 + prefixes = resolve_mm_encoder_only_lm_prefixes(
171 + ["language_model"], weights_mapper=_qwen3_mapper()
172 + )
173 + assert prefixes is not None
174 + kept = filter_mm_encoder_only_safetensors_files(
175 + files,
176 + folder,
177 + SAFE_WEIGHTS_INDEX_NAME,
178 + prefixes,
179 + weights_mapper=_qwen3_mapper(),
180 + )
181 + assert kept == [files[1]]
182 +
183 +
184 +def test_qwen25_flat_index_skips_lm_keeps_visual():
185 + with tempfile.TemporaryDirectory() as folder:
186 + files = [
187 + os.path.join(folder, "model-00001-of-000002.safetensors"),
188 + os.path.join(folder, "model-00002-of-000002.safetensors"),
189 + ]
190 + for path in files:
191 + open(path, "wb").close()
192 + _write_index(
193 + folder,
194 + {
195 + "model.layers.0.self_attn.q_proj.weight": (
196 + "model-00001-of-000002.safetensors"
197 + ),
198 + "model.embed_tokens.weight": "model-00001-of-000002.safetensors",
199 + "model.norm.weight": "model-00001-of-000002.safetensors",
200 + "lm_head.weight": "model-00001-of-000002.safetensors",
201 + "visual.blocks.0.weight": "model-00002-of-000002.safetensors",
202 + },
203 + )
204 + prefixes = resolve_mm_encoder_only_lm_prefixes(
205 + ["language_model"], weights_mapper=_qwen25_mapper()
206 + )
207 + assert prefixes is not None
208 + kept = filter_mm_encoder_only_safetensors_files(
209 + files,
210 + folder,
211 + SAFE_WEIGHTS_INDEX_NAME,
212 + prefixes,
213 + weights_mapper=_qwen25_mapper(),
214 + )
215 + assert kept == [files[1]]
216 +
217 +
218 +def test_molmo_shaped_index_not_filtered_when_deny_disabled():
219 + """Shared-root LM attr → resolve None; callers leave the file list alone."""
220 + assert (
221 + resolve_mm_encoder_only_lm_prefixes(["model"], weights_mapper=_molmo_mapper())
222 + is None
223 + )
224 + with tempfile.TemporaryDirectory() as folder:
225 + files = [
226 + os.path.join(folder, "vision.safetensors"),
227 + os.path.join(folder, "lm.safetensors"),
228 + ]
229 + for path in files:
230 + open(path, "wb").close()
231 + _write_index(
232 + folder,
233 + {
234 + "model.vision_backbone.patch.weight": "vision.safetensors",
235 + "model.transformer.blocks.0.weight": "lm.safetensors",
236 + },
237 + )
238 + # Document why raw HF startswith("model.") is unsafe without fail-closed:
239 + # whole-shard deny would also drop vision.
240 + if_applied_naively = filter_mm_encoder_only_safetensors_files(
241 + files, folder, SAFE_WEIGHTS_INDEX_NAME, ("model.",)
242 + )
243 + assert if_applied_naively == []
tests/model_executor/model_loader/test_registry.py
+1 -1
@@ -79,7 +79,7 @@ def test_default_loader_hf_still_falls_back_to_pt(tmp_path):
79 # Control: load_format="hf" still picks up .pt weights via fallback.
80 (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00")
81 loader = DefaultModelLoader(LoadConfig(load_format="hf"))
82 - _, files, use_safetensors = loader._prepare_weights(
82 + _, files, use_safetensors, _ = loader._prepare_weights(
83 str(tmp_path),
84 None,
85 None,
vllm/model_executor/model_loader/default_loader.py
+64 -10
@@ -1,11 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from __future__ import annotations
4 +
5 import dataclasses
6 import glob
7 import os
8 import time
9 from collections.abc import Generator, Iterable
8 -from typing import cast
10 +from typing import TYPE_CHECKING, cast
11
12 import torch
13 from torch import nn
@@ -25,6 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import (
27 fastsafetensors_weights_iterator,
28 filter_duplicate_safetensors_files,
29 filter_files_not_needed_for_inference,
30 + filter_mm_encoder_only_safetensors_files,
31 get_quant_config,
32 instanttensor_weights_iterator,
33 maybe_download_from_modelscope,
@@ -32,11 +35,15 @@ from vllm.model_executor.model_loader.weight_utils import (
35 multi_thread_safetensors_weights_iterator,
36 np_cache_weights_iterator,
37 pt_weights_iterator,
38 + resolve_mm_encoder_only_lm_prefixes,
39 safetensors_weights_iterator,
40 )
41 from vllm.tracing import instrument
42 from vllm.transformers_utils.repo_utils import list_filtered_repo_files
43
44 +if TYPE_CHECKING:
45 + from vllm.model_executor.models.utils import WeightsMapper
46 +
47 logger = init_logger(__name__)
48
49
@@ -74,6 +81,9 @@ class DefaultModelLoader(BaseModelLoader):
81 def __init__(self, load_config: LoadConfig):
82 super().__init__(load_config)
83 self.local_expert_ids: set[int] | None = None
84 + # Set in load_weights when --mm-encoder-only; used to drop LM-only shards.
85 + self._encoder_only_lm_prefixes: tuple[str, ...] | None = None
86 + self._encoder_only_weights_mapper: WeightsMapper | None = None
87
88 extra_config = load_config.model_loader_extra_config
89 if not isinstance(extra_config, dict):
@@ -132,7 +142,7 @@ class DefaultModelLoader(BaseModelLoader):
142 revision: str | None,
143 fall_back_to_pt: bool,
144 allow_patterns_overrides: list[str] | None,
135 - ) -> tuple[str, list[str], bool]:
145 + ) -> tuple[str, list[str], bool, str]:
146 """Prepare weights for the model.
147
148 If the model is not local, it will be downloaded."""
@@ -239,20 +249,40 @@ class DefaultModelLoader(BaseModelLoader):
249 f"Cannot find any model weights with `{model_name_or_path}`"
250 )
251
242 - return hf_folder, hf_weights_files, use_safetensors
252 + return hf_folder, hf_weights_files, use_safetensors, index_file
253
254 def _get_weights_iterator(
245 - self, source: "Source"
255 + self, source: Source
256 ) -> Generator[tuple[str, torch.Tensor], None, None]:
257 """Get an iterator for the model weights based on the load format."""
258 extra_config = self.load_config.model_loader_extra_config
249 - hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
250 - source.model_or_path,
251 - source.subfolder,
252 - source.revision,
253 - source.fall_back_to_pt,
254 - source.allow_patterns_overrides,
259 + hf_folder, hf_weights_files, use_safetensors, index_file = (
260 + self._prepare_weights(
261 + source.model_or_path,
262 + source.subfolder,
263 + source.revision,
264 + source.fall_back_to_pt,
265 + source.allow_patterns_overrides,
266 + )
267 )
268 + if (
269 + self._encoder_only_lm_prefixes is not None
270 + and use_safetensors
271 + and hf_weights_files
272 + ):
273 + hf_weights_files = filter_mm_encoder_only_safetensors_files(
274 + hf_weights_files,
275 + hf_folder,
276 + index_file,
277 + self._encoder_only_lm_prefixes,
278 + weights_mapper=self._encoder_only_weights_mapper,
279 + )
280 + if not hf_weights_files:
281 + raise RuntimeError(
282 + "mm-encoder-only shard filter removed all weight files for "
283 + f"`{source.model_or_path}`; check language_model prefixes "
284 + f"{self._encoder_only_lm_prefixes}"
285 + )
286 if self.load_config.load_format == "npcache":
287 # Currently np_cache only support *.bin checkpoints
288 assert use_safetensors is False
@@ -348,6 +378,29 @@ class DefaultModelLoader(BaseModelLoader):
378 allow_patterns_overrides=None,
379 )
380
381 + def _init_mm_encoder_only_weight_filter(
382 + self, model: nn.Module, model_config: ModelConfig
383 + ) -> None:
384 + """Skip pure language-model safetensors shards under --mm-encoder-only."""
385 + mm_config = model_config.multimodal_config
386 + if mm_config is None or not mm_config.mm_encoder_only:
387 + self._encoder_only_lm_prefixes = None
388 + self._encoder_only_weights_mapper = None
389 + return
390 +
391 + # Derive from _language_model_names; fail-closed on shared HF roots
392 + # (Molmo/Phi-4-MM/Muse). Qwen nested/flat keys classified via mapper.
393 + weights_mapper = cast(
394 + "WeightsMapper | None", getattr(model, "hf_to_vllm_mapper", None)
395 + )
396 + self._encoder_only_lm_prefixes = resolve_mm_encoder_only_lm_prefixes(
397 + getattr(model, "_language_model_names", None),
398 + weights_mapper=weights_mapper,
399 + )
400 + self._encoder_only_weights_mapper = (
401 + weights_mapper if self._encoder_only_lm_prefixes is not None else None
402 + )
403 +
404 def _init_ep_weight_filter(self, model_config: ModelConfig) -> None:
405 """Compute local expert ids for EP weight filtering.
406
@@ -423,6 +476,7 @@ class DefaultModelLoader(BaseModelLoader):
476 self.load_config.safetensors_load_strategy = "torchao"
477
478 self._init_ep_weight_filter(model_config)
479 + self._init_mm_encoder_only_weight_filter(model, model_config)
480
481 loaded_weights = model.load_weights(self.get_all_weights(model_config, model))
482
vllm/model_executor/model_loader/weight_utils.py
+169 -1
@@ -15,7 +15,7 @@ import time
15 from collections.abc import Callable, Generator, Iterable
16 from contextlib import contextmanager
17 from pathlib import Path
18 -from typing import IO, Any
18 +from typing import IO, TYPE_CHECKING, Any
19
20 import filelock
21 import huggingface_hub.constants
@@ -61,6 +61,9 @@ except ImportError:
61
62 from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least
63
64 +if TYPE_CHECKING:
65 + from vllm.model_executor.models.utils import WeightsMapper
66 +
67 logger = init_logger(__name__)
68
69 # use system-level temp directory for file locks, so that multiple users
@@ -605,6 +608,171 @@ def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[s
608 return hf_weights_files
609
610
611 +def _normalize_module_prefixes(
612 + language_model_names: Iterable[str] | None,
613 +) -> tuple[str, ...]:
614 + names = tuple(language_model_names or ())
615 + return tuple(n if n.endswith(".") else f"{n}." for n in names)
616 +
617 +
618 +def _mapped_weight_name(weights_mapper: "WeightsMapper | None", key: str) -> str | None:
619 + """Apply ``WeightsMapper.map_name`` when present; else identity."""
620 + if weights_mapper is None:
621 + return key
622 + return weights_mapper.map_name(key)
623 +
624 +
625 +def _hf_prefix_maps_to_lm(
626 + new_name: str | None, lm_module_prefixes: tuple[str, ...]
627 +) -> bool:
628 + if not new_name:
629 + return False
630 + new_p = new_name if new_name.endswith(".") else f"{new_name}."
631 + return any(new_p.startswith(p) or p.startswith(new_p) for p in lm_module_prefixes)
632 +
633 +
634 +def _shared_hf_root_module_prefixes(
635 + lm_module_prefixes: tuple[str, ...],
636 + weights_mapper: "WeightsMapper | None",
637 +) -> tuple[str, ...]:
638 + """Module prefixes that are also an HF root for non-LM weights.
639 +
640 + E.g. Molmo/Phi-4-MM/Muse mark LM as ``model`` while vision HF keys stay
641 + under ``model.vision_*`` / ``model.embed_tokens_extend.*``. Detected via
642 + ``WeightsMapper.orig_to_new_prefix`` (no hard-coded name list).
643 + """
644 + if weights_mapper is None:
645 + return ()
646 + mapping = weights_mapper.orig_to_new_prefix
647 + if not mapping:
648 + return ()
649 +
650 + non_lm_origs: list[str] = []
651 + for orig, new in mapping.items():
652 + orig_p = orig if orig.endswith(".") else f"{orig}."
653 + if not _hf_prefix_maps_to_lm(new, lm_module_prefixes):
654 + non_lm_origs.append(orig_p)
655 +
656 + if not non_lm_origs:
657 + return ()
658 +
659 + shared = [
660 + p
661 + for p in lm_module_prefixes
662 + if any(orig.startswith(p) for orig in non_lm_origs)
663 + ]
664 + return tuple(shared)
665 +
666 +
667 +def resolve_mm_encoder_only_lm_prefixes(
668 + language_model_names: Iterable[str] | None,
669 + *,
670 + weights_mapper: "WeightsMapper | None" = None,
671 +) -> tuple[str, ...] | None:
672 + """Resolve vLLM LM *module* prefixes for ``--mm-encoder-only`` shard skip.
673 +
674 + Prefixes come from ``_language_model_names``. Classification of HF index
675 + keys is done later via optional ``weights_mapper`` (see
676 + ``filter_mm_encoder_only_safetensors_files``), so this helper does **not**
677 + hard-code Qwen/HF nest lists.
678 +
679 + Returns ``None`` (leave the safetensors file list unchanged) when
680 + ``_language_model_names`` is empty/missing, or when a module prefix is a
681 + shared HF checkpoint root for both LM and non-LM weights — fail-closed so
682 + Molmo / Phi-4-MM / Muse cannot under-load the encoder.
683 + """
684 + prefixes = _normalize_module_prefixes(language_model_names)
685 + if not prefixes:
686 + return None
687 + shared = _shared_hf_root_module_prefixes(prefixes, weights_mapper)
688 + if shared:
689 + logger.warning_once(
690 + "mm-encoder-only whole-shard filter not applied: LM module "
691 + "prefix(es) %s share an HF checkpoint root with non-LM weights "
692 + "(via WeightsMapper). Encoder-only load keeps the full "
693 + "safetensors file list; finer HF prefixes or key-level "
694 + "filtering can enable skip later.",
695 + shared,
696 + )
697 + return None
698 + return prefixes
699 +
700 +
701 +def filter_mm_encoder_only_safetensors_files(
702 + hf_weights_files: list[str],
703 + hf_folder: str,
704 + index_file: str,
705 + language_model_prefixes: Iterable[str],
706 + *,
707 + weights_mapper: "WeightsMapper | None" = None,
708 +) -> list[str]:
709 + """Drop safetensors shards that only contain language-model weights.
710 +
711 + Used with ``--mm-encoder-only`` so Encoder-only EPD instances avoid reading
712 + pure LM shards from disk/DRAM.
713 +
714 + Each HF index key is classified by mapping through ``weights_mapper`` (when
715 + provided) and testing the *vLLM* name against ``language_model_prefixes``.
716 + Without a mapper, HF names are compared directly (identity checkpoint
717 + layout, e.g. Kimi ``language_model.*``).
718 +
719 + A shard is kept if it contains any non-LM key. Without an index file,
720 + returns ``hf_weights_files`` unchanged.
721 + """
722 + prefixes = tuple(language_model_prefixes)
723 + if not prefixes:
724 + return hf_weights_files
725 +
726 + index_path = os.path.join(hf_folder, index_file)
727 + if not os.path.isfile(index_path):
728 + logger.warning_once(
729 + "mm-encoder-only shard filter skipped: index file %s not found",
730 + index_path,
731 + )
732 + return hf_weights_files
733 +
734 + with open(index_path) as f:
735 + weight_map: dict[str, str] = json.load(f)["weight_map"]
736 +
737 + keys_by_file: dict[str, list[str]] = {}
738 + for weight_name, weight_file in weight_map.items():
739 + keys_by_file.setdefault(weight_file, []).append(weight_name)
740 +
741 + def _is_lm_key(key: str) -> bool:
742 + mapped = _mapped_weight_name(weights_mapper, key)
743 + if mapped is None:
744 + return False
745 + return any(mapped.startswith(p) for p in prefixes)
746 +
747 + def _is_lm_only(weight_file: str) -> bool:
748 + keys = keys_by_file.get(weight_file)
749 + if not keys:
750 + # Not listed in the index (e.g. extra matched glob) — keep.
751 + return False
752 + return all(_is_lm_key(key) for key in keys)
753 +
754 + kept: list[str] = []
755 + skipped = 0
756 + for path in hf_weights_files:
757 + rel = os.path.relpath(path, hf_folder)
758 + # Index weight_map values are usually basenames; tolerate either form.
759 + basename = os.path.basename(path)
760 + if _is_lm_only(rel) or _is_lm_only(basename):
761 + skipped += 1
762 + continue
763 + kept.append(path)
764 +
765 + if skipped:
766 + logger.info_once(
767 + "mm-encoder-only: skipped %d/%d safetensors shard(s) that only "
768 + "contain language-model weights (prefixes=%s)",
769 + skipped,
770 + skipped + len(kept),
771 + prefixes,
772 + )
773 + return kept
774 +
775 +
776 # explicitly use pure text format, with a newline at the end
777 # this makes it impossible to see the animation in the progress bar
778 # but will avoid messing up with ray or multiprocessing, which wraps
vllm/model_executor/models/utils.py
+6 -2
@@ -78,11 +78,15 @@ class WeightsMapper:
78 orig_to_new_suffix={**self.orig_to_new_suffix, **other.orig_to_new_suffix},
79 )
80
81 - def _map_name(self, key: str) -> str | None:
82 - """Map a weight name (backward-compatible wrapper that discards shard_id)."""
81 + def map_name(self, key: str) -> str | None:
82 + """Map a weight name; returns ``None`` if the weight should be ignored."""
83 result = self._map_name_with_shard(key)
84 return result[0] if result is not None else None
85
86 + def _map_name(self, key: str) -> str | None:
87 + """Backward-compatible alias for :meth:`map_name`."""
88 + return self.map_name(key)
89 +
90 def _map_name_with_shard(self, key: str) -> tuple[str, ShardIds | None] | None:
91 """Map a weight name and extract any shard_id metadata.
92