Repository navigation
fix(peft): save shared-outer MoE LoRA adapters in PEFT format - #6076
huvunvidia merged 3 commits into
Conversation
``save_hf_adapter`` failed for shared-outer MoE LoRA (``experts_shared_outer_loras``) with ``ValueError: Incomplete adapter export for target parameter``. The default export emits the shared factor once as a ``[1, ...]`` tensor under the expert-agnostic HF name and the per-expert factor as numbered 2D slices, so ``convert_adapter_weights_to_peft_state`` saw a lone 3D tensor, treated it as a packed ``target_parameters`` entry and found no partner to pack with. The ``stack_3d_moe`` layout (both sides 3D with different expert dims) hit the same packer with mismatched shapes. Recognise both shapes of a shared-outer export and keep them as exported module weights instead of trying to fold them into a PEFT parameter target; a lone 3D side without any per-expert sibling still raises, so a genuinely incomplete export is not silently saved. Based on #33 (@Arist12), re-applied on the current export layout without changing what the streamer emits. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
yaoyu-33
added
area:peft
| for base_name in shared_outer_targets: | ||
| for lora_suffix, tensor in parameter_weights[base_name].items(): | ||
| name = f"{base_name}{lora_suffix}" | ||
| adapter_state[f"base_model.model.{name}"] = tensor |
There was a problem hiding this comment.
save_hf_adapter promises a directory loadable by PeftModel.from_pretrained, but this removes the tensor from target_parameters and writes an expert-agnostic 3D key as a normal module weight. Standard PEFT has no matching module for that key and may ignore it while loading the per-expert partner. The new tests only inspect the rewritten dictionary. Please add a real PeftModel.from_pretrained load plus numerical adapter parity, as the existing packed-expert test does, or make this a clearly serving-specific export path rather than PEFT format.
There was a problem hiding this comment.
Agreed — fixed in 1e4eaef by taking the second option: this layout is a serving-only export, not PEFT format, and the code now says so.
| lora_a = weights.get(".lora_A.weight") | ||
| lora_b = weights.get(".lora_B.weight") | ||
| if lora_a is not None and lora_b is not None: | ||
| return lora_a.shape[0] != lora_b.shape[0] |
There was a problem hiding this comment.
This classifies every unequal pair of leading dimensions as shared-outer, including malformed ordinary exports such as expert dimensions 8 and 4. That turns a shape mismatch into a silently accepted checkpoint. Please require exactly one side to have expert_dim=1 and validate the expected shared side/layout before bypassing target-parameter packing.
There was a problem hiding this comment.
Fixed in 1e4eaef. _is_shared_outer_expert_export is replaced by _classify_shared_outer_expert_export, which returns the shared side only for a well-formed shared-outer pair and raises for everything else
…g layout Review feedback on the shared-outer adapter save path: * The classifier no longer treats every unequal pair of leading dims as shared-outer. ``_classify_shared_outer_expert_export`` requires exactly one side with expert dim 1, checks that the shared side is the one shared-outer actually shares (lora_A on gate/up, lora_B on down; the vLLM 3D-MoE stems ``experts.base_layer`` / ``experts`` map the same way), checks that the per-expert partner exists (numbered 2D slices or an ``[E, ...]`` stack) and that both sides agree on the rank. Anything else -- 8 vs 4 expert dims, a shared factor on the wrong side, a rank mismatch, a shared factor without a partner -- raises instead of being written as a checkpoint. * The shared-outer layout is now an explicit serving-only export. ``convert_adapter_weights_to_peft_state`` and ``save_hf_adapter`` take ``allow_serving_layout``; by default a shared-outer export raises with a message pointing at ``expand_shared_outer=True`` (the PEFT-loadable form) or ``allow_serving_layout=True`` (the layout SGLang's ``experts_shared_outer_loras`` reader consumes, not loadable by ``PeftModel.from_pretrained``). Docstrings say so instead of claiming a PEFT shared layout. Tests: PEFT default refuses shared-outer; serving opt-in keeps the shared factor as a module weight for both the numbered-2D and the 3D pair layouts; 3D-MoE stems and unknown projections; rejection of unequal expert dims, wrong shared side, rank mismatch and missing partner; save_hf_adapter end-to-end (refuses by default, writes the serving layout with no target_parameters when asked). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>
|
/ok to test 1e4eaef |
yaoyu-33
left a comment
yaoyu-33
left a comment
There was a problem hiding this comment.
LGTM. The latest changes address the review comments and add coverage for the affected boundary.
yaoyu-33
added
ready-to-merge
|
/ok to test 4d99c12 |
This branch was successfully deployed
4 participants
What does this PR do ?
Makes
save_hf_adapterhandle shared-outer MoE LoRA (experts_shared_outer_loras=True), which today fails with the opaqueValueError: Incomplete adapter export for target parameter: model.layers.N.mlp.experts.<proj>. Shared-outer has no PEFTParamWrapperrepresentation, so the layout is exposed as an explicit serving-only export (allow_serving_layout=True) and refused with a clear message in the default PEFT mode.Changelog
convert_adapter_weights_to_peft_state(..., allow_serving_layout=False):_classify_shared_outer_expert_exportrecognises a shared-outer export only when exactly one side has expert dim 1, that side is the one shared-outer shares (lora_A on gate/up incl. the 3D-MoEexperts.base_layerstem, lora_B on down incl. the bareexpertsstem), the per-expert partner exists (numbered 2Dexperts.N.<proj>slices or an[E, ...]stack under the same name) and both sides agree on the rank. Unequal expert dims without a shared side (e.g. 8 vs 4), a shared factor on the wrong side, a rank mismatch or a missing partner raise instead of being written.expand_shared_outer=True(replicates the shared factor into the PEFT-loadable per-expert layout) orallow_serving_layout=True. With the opt-in both sides are written exactly as exported (shared factor once as[1, ...]under the expert-agnostic name), which is the layout SGLang'sexperts_shared_outer_lorasreader consumes; it is not loadable byPeftModel.from_pretrained, and the docstrings say so.save_hf_adapter(..., allow_serving_layout=False)forwards the flag; itsexpand_shared_outerdocstring no longer claims a "PEFT shared[1, ...]layout".save_hf_adapterend-to-end); serving opt-in keeps the shared factor as a module weight for the numbered-2D and the 3D-pair layouts, accepts the 3D-MoE stems and unknown projections; rejection of unequal expert dims, wrong shared side, rank mismatch and missing partner; incomplete export still raises.The streamer's output layouts are unchanged; only the PEFT on-disk conversion learns to validate the layout it already produces and to name it for what it is.
GitHub Actions CI
See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.
Before your PR is "Ready for review"
Pre checks:
Additional Information
--experts-shared-outer-lorason 8× MI350X and the saved adapter was served by SGLang's shared-outer loader. Re-applied here on top of feat(peft): expand_shared_outer export flag for shared-outer MoE LoRA #5588 / Support 3D-MoE LoRA export for shared-outer routed-expert adapters #5939 without changing whatstream_adapter_weights_megatron_to_hfemits.Validation
tests/unit_tests/peft+test_model_bridge.py+test_model_bridge_lora.py+test_adapter_export.py: 627 passed, 7 skipped (baselineupstream/main@2eae3c971on the same box: 443 + 181 = 624). New tests PASSED by name:TestSaveHfAdapter::test_convert_shared_outer_adapter_keeps_shared_factor_as_module_weight,TestSaveHfAdapter::test_convert_shared_outer_stacked_pair_keeps_both_sides_as_module_weights,TestSaveHfAdapter::test_convert_incomplete_target_parameter_still_raises.rx devbox2× H200 (h200-sci-k8s), imagenvcr.io/nvidia/pytorch:26.08-py3(the upstream CI base image), TE 2.18.0+27486e03 (matches upstreamuv.lock), Megatron-Core at.main.commitc9b53d0a8, Megatron-Bridge editable. Command mirrors CI'sLaunch_Unit_Tests_Core.sh:CUDA_VISIBLE_DEVICES=0,1 pytest -m "not pleasefixme".ruff check,ruff check --select I,ruff format --check(ruff 0.9.9) pass on the changed files.Revision 2 (review: PEFT-format promise; classifier too permissive)
allow_serving_layout), refused by default with guidance; the classifier validates exactly-one-shared-side, expected side per projection, partner presence and rank agreement.nvcr.io/nvidia/pytorch:26.08-py3, TE 2.18.0+27486e03, Megatron-Corec9b53d0a8): fail-before with revision 1's sources + the new tests →9 failed(7 ×unexpected keyword argument 'allow_serving_layout', 2 ×DID NOT RAISE ValueErrorfor the PEFT-mode refusal in the converter and insave_hf_adapter); pass-after →21 passed(the 9 new/changed tests plus the 12 existing shared-outer/expand tests selected by name);test_adapter_export.py+test_model_bridge_lora.py+test_auto_bridge.py+tests/unit_tests/peft→ 691 passed, 7 skipped (baselineupstream/main681 passed / 7 skipped, no new failures).🤖 Generated with Claude Code