Skip to content

Commit cb23330

Browse files
stanley1208akoumpa
andauthored
fix(moe): export fused expert lora in the corrected peft >= 0.19.1 layout (NVIDIA-NeMo#3626)
* fix(moe): export fused expert lora in the corrected peft >= 0.19.1 layout Fixes NVIDIA-NeMo#3592. peft flipped the in/out interpretation of non-transposed 3-D expert parameters in ParamWrapper in 0.19.1 (huggingface/peft#3165), which transposes every fused expert lora weight it creates. The v5 export used the pre-flip layout, so adapters exported for qwen3_moe, minimax_m2 and nemotron_v3 fail to load with a size mismatch on peft >= 0.19.1, which is what a plain pip install gives today. Following the plan discussed on the issue: - export emits the corrected layout by default; the explicit checkpoint option legacy_paramwrapper_layout keeps the pre-flip layout for consumers pinned to an older peft (separate from v4_compatible) - the peft save stamps the layout into automodel_peft_config.json and the loader resolves it from that metadata first, falling back to shape detection (tensors whose shapes distinguish the generations decide for the whole file, which covers minimax_m2 where hidden == 2x moe inter makes the gate_up pair shape-identical), and fails with a clear error when neither settles it - peft floor raised to 0.19.1; uv.lock and docker/common/uv-pytorch.lock regenerated with uv 0.8.22 The real-peft round-trip suite (loading, merged deltas, forward parity) now runs against the corrected layout with the locked peft, and unit tests cover both layouts for gated and non-gated experts. Signed-off-by: stanley1208 <stanley.mei08@gmail.com> * fix(moe): fail closed on the paramwrapper layout stamp and prove peft parity Follow-up to the review on NVIDIA-NeMo#3626. - Qwen3MoeStateDictAdapter.convert_single_tensor_to_hf now forwards legacy_paramwrapper_layout, so a legacy export is legacy for Qwen3 too instead of modern tensors under a peft-0.18 stamp. - _convert_paramwrapper_to_native accepts only the two known layout ids and raises when the stamp is unknown, when it contradicts unambiguous tensor shapes, or when the tensors disagree among themselves, instead of treating any non-legacy value as modern and letting a stamp mask a mixed file. - The PEFT v5 adapter test compares the delta real peft merges against the delta computed from the native grouped A/B tensors, parameterized over nemotron_v3, minimax_m2 and qwen3_moe. A converter that transposed A/B consistently in both directions passed the old round trip and fails this. The legacy export is checked the same way through peft's pre-flip delta formula, since no pre-0.19.1 peft can be installed next to the test floor. - get_peft_version / is_peft_min_version leave import_utils; the parity test requires peft >= 0.19.1 via importorskip. The unused logger the first commit added to the mixin is gone as well. - finetune.mdx asks for peft >= 0.19.1 and explains when to use checkpoint.legacy_paramwrapper_layout and what it implies. Signed-off-by: stanley1208 <stanley.mei08@gmail.com> --------- Signed-off-by: stanley1208 <stanley.mei08@gmail.com> Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
1 parent 1e56f25 commit cb23330

13 files changed

Lines changed: 684 additions & 100 deletions

File tree

‎docker/common/uv-pytorch.lock‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4362,7 +4362,7 @@ magi = [{ name = "magi-attention", git = "https://github.com/SandAI-org/MagiAtte
43624362
test = [
43634363
{ name = "coverage" },
43644364
{ name = "onnxruntime-gpu", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'", index = "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-13/pypi/simple/" },
4365-
{ name = "peft", specifier = ">=0.18.1" },
4365+
{ name = "peft", specifier = ">=0.19.1" },
43664366
{ name = "pytest" },
43674367
{ name = "pytest-shard" },
43684368
{ name = "ruamel-yaml" },
@@ -5286,7 +5286,7 @@ wheels = [
52865286

52875287
[[package]]
52885288
name = "peft"
5289-
version = "0.18.1"
5289+
version = "0.20.0"
52905290
source = { registry = "https://pypi.org/simple" }
52915291
dependencies = [
52925292
{ name = "accelerate" },
@@ -5301,9 +5301,9 @@ dependencies = [
53015301
{ name = "tqdm" },
53025302
{ name = "transformers" },
53035303
]
5304-
sdist = { url = "https://files.pythonhosted.org/packages/d8/48/147b3ea999560b40a34fd78724c7777aa9d18409c2250bdcaf9c4f2db7fc/peft-0.18.1.tar.gz", hash = "sha256:2dd0d6bfce936d1850e48aaddbd250941c5c02fc8ef3237cd8fd5aac35e0bae2", size = 635030, upload-time = "2026-01-09T13:08:01.136Z" }
5304+
sdist = { url = "https://files.pythonhosted.org/packages/b4/08/02541a2c29be7c78698f73d438bc0e733b214f3f35ec5db79fca8da8fc61/peft-0.20.0.tar.gz", hash = "sha256:4769c8093a4ca145fd6fb3fd4dd50449675f5fe46434ad1e98b285a132d4b1d0", size = 880503, upload-time = "2026-07-28T13:46:01.85Z" }
53055305
wheels = [
5306-
{ url = "https://files.pythonhosted.org/packages/b3/14/b4e3f574acf349ae6f61f9c000a77f97a3b315b4bb6ad03791e79ae4a568/peft-0.18.1-py3-none-any.whl", hash = "sha256:0bf06847a3551e3019fc58c440cffc9a6b73e6e2962c95b52e224f77bbdb50f1", size = 556960, upload-time = "2026-01-09T13:07:55.865Z" },
5306+
{ url = "https://files.pythonhosted.org/packages/28/79/13bcabb8048126422d5c4b880575d40886c726f354db88cfeed4325525bb/peft-0.20.0-py3-none-any.whl", hash = "sha256:0fbba16ffebfad3de96e06f2da6860fd860292324b85b6141909fa1e26ea9233", size = 775777, upload-time = "2026-07-28T13:45:59.809Z" },
53075307
]
53085308

53095309
[[package]]

‎docs/guides/llm/finetune.mdx‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -480,9 +480,11 @@ print(tokenizer.decode(output[0], skip_special_tokens=True))
480480
PEFT adapters must be loaded on top of the base model. Because the base NeMo AutoModel installation does not include the Hugging Face PEFT library, install a tested version in your active environment before running the inference code:
481481

482482
```bash
483-
uv pip install "peft>=0.18.1"
483+
uv pip install "peft>=0.19.1"
484484
```
485485

486+
LoRA adapters for models with fused MoE expert weights (for example Qwen3-MoE, MiniMax-M2, Nemotron 3) are exported in the ParamWrapper layout peft has used since 0.19.1, when peft changed it (huggingface/peft#3165); that is why the step above installs `peft>=0.19.1`. If such an adapter has to be loaded by a consumer pinned to an older peft, set `checkpoint.legacy_paramwrapper_layout: true` in the training config to export the pre-0.19.1 layout instead. An adapter exported that way does not load with peft 0.19.1 or newer.
487+
486488
```python
487489
import torch
488490
from transformers import AutoModelForCausalLM, AutoTokenizer

‎nemo_automodel/components/checkpoint/addons.py‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,12 @@
2828
)
2929
from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter
3030
from nemo_automodel.components.checkpoint.stateful_wrappers import ModelState
31-
from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin
31+
from nemo_automodel.components.moe.state_dict_mixin import (
32+
_PARAMWRAPPER_LAYOUT_LEGACY,
33+
_PARAMWRAPPER_LAYOUT_METADATA_KEY,
34+
_PARAMWRAPPER_LAYOUT_MODERN,
35+
MoESplitExpertsStateDictMixin,
36+
)
3237
from nemo_automodel.shared.parameter_names import canonical_parameter_fqn
3338

3439
if TYPE_CHECKING:
@@ -268,6 +273,12 @@ def pre_save(self, **kwargs) -> None:
268273
process_group = kwargs.get("process_group")
269274
hf_peft_config = _get_hf_peft_config(peft_config, model_state, v4_compatible=v4_compatible)
270275
automodel_peft_metadata = _get_automodel_peft_metadata(peft_config)
276+
adapter = getattr(_unwrap_ddp_model(model_state.model[0]), "state_dict_adapter", None)
277+
layout_stamp = _get_paramwrapper_layout_stamp(
278+
adapter, v4_compatible, kwargs.get("legacy_paramwrapper_layout", False)
279+
)
280+
if layout_stamp is not None:
281+
automodel_peft_metadata[_PARAMWRAPPER_LAYOUT_METADATA_KEY] = layout_stamp
271282
if _is_group_rank_0(process_group):
272283
# if the HF model has custom model code, we need to save it as part of the checkpoint
273284
model_part = model_state.model[0] if model_state is not None else None
@@ -354,6 +365,19 @@ def _get_hf_peft_config(peft_config: "PeftConfig", model_state: ModelState, v4_c
354365
return config
355366

356367

368+
def _get_paramwrapper_layout_stamp(adapter, v4_compatible: bool, legacy_paramwrapper_layout: bool) -> "str | None":
369+
"""Which fused expert LoRA layout this PEFT save uses, for the metadata stamp.
370+
371+
peft flipped the ParamWrapper layout in 0.19.1 (huggingface/peft#3165);
372+
stamping the layout into automodel_peft_config.json lets the loader
373+
resolve it from metadata instead of tensor shapes. Returns None for
374+
adapters that don't emit fused (v5 ParamWrapper) expert LoRA at all.
375+
"""
376+
if v4_compatible or not getattr(adapter, "_v5_peft_target_parameters", ()):
377+
return None
378+
return _PARAMWRAPPER_LAYOUT_LEGACY if legacy_paramwrapper_layout else _PARAMWRAPPER_LAYOUT_MODERN
379+
380+
357381
def _get_automodel_peft_metadata(peft_config: "PeftConfig") -> dict:
358382
"""
359383
Get the PEFT metadata in the format expected by Automodel.

‎nemo_automodel/components/checkpoint/checkpointing.py‎

Lines changed: 41 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -620,6 +620,7 @@ def save_model(
620620
quantization=False,
621621
device_mesh=self.moe_mesh,
622622
v4_compatible=self.config.v4_compatible,
623+
legacy_paramwrapper_layout=self.config.legacy_paramwrapper_layout,
623624
)
624625
# MoE adapters return non-contiguous views; safetensors.save rejects those.
625626
_materialize_to_hf_views_for_save(state_dict)
@@ -646,6 +647,7 @@ def save_model(
646647
fqn_to_dtype_mapping=fqn_to_dtype_mapping,
647648
original_model_path=self._get_original_model_path(model_state),
648649
v4_compatible=self.config.v4_compatible,
650+
legacy_paramwrapper_layout=self.config.legacy_paramwrapper_layout,
649651
process_group=consolidation_process_group,
650652
)
651653
self._maybe_write_offline_consolidation_script(model_dir)
@@ -880,7 +882,10 @@ def load_model(
880882
t_disk = time.monotonic()
881883
if state_dict_from_disk is not None:
882884
state_dict_from_disk = _maybe_adapt_state_dict_from_hf(
883-
model_state.model[0], state_dict_from_disk, moe_mesh=self.moe_mesh
885+
model_state.model[0],
886+
state_dict_from_disk,
887+
moe_mesh=self.moe_mesh,
888+
paramwrapper_layout_hint=_read_paramwrapper_layout_metadata(model_path),
884889
)
885890
else:
886891
state_dict_from_disk = {}
@@ -1052,7 +1057,12 @@ def load_model(
10521057
if compat_tied_lm_head_source_key is not None and isinstance(lm_head_param_name, str):
10531058
state_dict[lm_head_param_name] = state_dict.pop(compat_tied_lm_head_source_key)
10541059

1055-
state_dict = _maybe_adapt_state_dict_from_hf(model_state.model[0], state_dict, moe_mesh=self.moe_mesh)
1060+
state_dict = _maybe_adapt_state_dict_from_hf(
1061+
model_state.model[0],
1062+
state_dict,
1063+
moe_mesh=self.moe_mesh,
1064+
paramwrapper_layout_hint=_read_paramwrapper_layout_metadata(model_path),
1065+
)
10561066
adapter_complete = time.monotonic()
10571067
expected_keys_for_diff = {k for k in expected_keys if not k.endswith("_extra_state")}
10581068
loaded_keys_for_diff = {k for k in state_dict if not k.endswith("_extra_state")}
@@ -2785,14 +2795,41 @@ def _load_hf_bin_checkpoint(model_path: str) -> dict[str, torch.Tensor] | None:
27852795

27862796

27872797
def _maybe_adapt_state_dict_from_hf(
2788-
model_part: nn.Module, state_dict: dict[str, torch.Tensor], moe_mesh: DeviceMesh | None = None
2798+
model_part: nn.Module,
2799+
state_dict: dict[str, torch.Tensor],
2800+
moe_mesh: DeviceMesh | None = None,
2801+
paramwrapper_layout_hint: str | None = None,
27892802
) -> dict[str, torch.Tensor]:
27902803
"""
27912804
Custom models use state dict adapters to convert the state dict from the Hugging Face format to the native format.
2805+
2806+
``paramwrapper_layout_hint`` carries the fused expert LoRA layout recorded in
2807+
the checkpoint's automodel_peft_config.json (see _read_paramwrapper_layout_metadata),
2808+
so the adapter resolves the peft ParamWrapper layout from metadata instead of shapes.
27922809
"""
27932810
adapter = getattr(_unwrap_ddp_model(model_part), "state_dict_adapter", None)
27942811
if adapter:
27952812
ep_mesh_dims = [dim for dim in moe_mesh.mesh_dim_names if dim != "pp"] if moe_mesh is not None else []
27962813
ep_mesh = moe_mesh[tuple(ep_mesh_dims)] if ep_mesh_dims else moe_mesh
2797-
return adapter.from_hf(state_dict, device_mesh=ep_mesh)
2814+
adapter._paramwrapper_layout_hint = paramwrapper_layout_hint
2815+
try:
2816+
return adapter.from_hf(state_dict, device_mesh=ep_mesh)
2817+
finally:
2818+
adapter._paramwrapper_layout_hint = None
27982819
return state_dict
2820+
2821+
2822+
def _read_paramwrapper_layout_metadata(model_path: str | os.PathLike) -> str | None:
2823+
"""Return the fused expert LoRA layout stamped into a PEFT checkpoint, if any.
2824+
2825+
The PEFT save path records which peft ParamWrapper layout the adapter was
2826+
exported in (peft flipped it in 0.19.1, huggingface/peft#3165) inside
2827+
automodel_peft_config.json. Absent or unstamped checkpoints return None and
2828+
the adapter falls back to shape detection.
2829+
"""
2830+
metadata_path = os.path.join(model_path, "automodel_peft_config.json")
2831+
try:
2832+
with open(metadata_path) as f:
2833+
return json.load(f).get("paramwrapper_layout")
2834+
except (FileNotFoundError, NotADirectoryError, json.JSONDecodeError):
2835+
return None

‎nemo_automodel/components/checkpoint/config.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,9 @@ class CheckpointingConfig:
122122
staging_dir: str | None = None
123123
v4_compatible: bool = False # If True, save the original pretrained config.json (with quantization_config removed)
124124
# instead of the in-memory v5 config. Useful when downstream consumers (e.g. vLLM) expect a v4-format config.
125+
legacy_paramwrapper_layout: bool = False # If True, export fused expert LoRA (PEFT v5 ParamWrapper) weights in
126+
# the pre-flip peft <= 0.19.0 layout instead of the corrected peft >= 0.19.1 layout (huggingface/peft#3165).
127+
# Only for consumers pinned to an older peft; unrelated to v4_compatible.
125128
diffusers_compatible: bool = False # If True, use diffusers-compatible index filename
126129
# (diffusion_pytorch_model.safetensors.index.json) so checkpoints are loadable via diffusers from_pretrained().
127130
best_metric_key: str = "default" # Validation metric key used to select the best checkpoint.

‎nemo_automodel/components/models/qwen3_moe/state_dict_adapter.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,9 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
8181
8282
When ``v4_compatible=False`` (the default), LoRA expert tensors are
8383
emitted in PEFT v0.18+ ParamWrapper format so that
84-
``PeftModel.from_pretrained()`` can load them directly. When
84+
``PeftModel.from_pretrained()`` can load them directly; the layout
85+
follows peft >= 0.19.1 unless ``legacy_paramwrapper_layout=True``
86+
asks for the pre-flip layout (huggingface/peft#3165). When
8587
``v4_compatible=True``, the legacy per-expert split is used instead
8688
(via the parent mixin).
8789
@@ -101,7 +103,9 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t
101103
expert_segment = self._expert_path_segment
102104
for suffix in _LORA_EXPERT_SUFFIXES:
103105
if fqn.endswith(f".{suffix}") and f".{expert_segment}.{suffix}" in fqn:
104-
result = self._convert_lora_to_paramwrapper(fqn, tensor)
106+
result = self._convert_lora_to_paramwrapper(
107+
fqn, tensor, legacy_layout=kwargs.get("legacy_paramwrapper_layout", False)
108+
)
105109
if exclude_key_regex:
106110
result = [(k, v) for k, v in result if not re.match(exclude_key_regex, k)]
107111
return result

0 commit comments

Comments
 (0)