Skip to content

test(peft): add a shared HF PEFT reload suite - #3945

Merged
yuhezhang-ai merged 7 commits into
NVIDIA-NeMo:mainfrom
stanley1208:feat/peft-reload-suite
Sep 21, 2026
Merged

yuhezhang-ai merged 7 commits into
NVIDIA-NeMo:mainfrom
stanley1208:feat/peft-reload-suite

Conversation

@stanley1208

Copy link
Copy Markdown
Contributor

First deliverable for #3867. Fixes #3944.

the suite

one registry of families, each a tiny transformers model plus the automodel adapter that owns its naming. per family: save through Checkpointer.save_model with is_peft=True, reload with PeftModel.from_pretrained, check tensors and logits match. a second test checks bulk and per-tensor exports agree.

covers llama (dense control), nemotron_v3, qwen3_moe, minimax_m2, qwen3_omni_moe.

i attach the adapter to the transformers model rather than building the native class, since the native moe classes call torch.cuda.current_device() unconditionally and cannot be built on cpu. both routes give identical output.

the bug it found

qwen3_omni_moe exported tensors under thinker. but target_modules without it, so peft also adapted talker.model.* and left those adapters random. map_peft_target_module_to_hf was never overridden; qwen2_5_omni already does this. detail in #3944.

reverting the override fails only omni. reverting the nemotron fix from #3866 fails only nemotron.

not covered

fused expert lora, since only nn.Linear gets wrapped here. that stays with test_moe_peft_v5_state_dict_adapters.py.

First deliverable for NVIDIA-NeMo#3867.

Export defects keep recurring because the outer prefix, the model-owned
renames, the fused expert layout and the target metadata are produced in
separate places, and the existing tests mostly check that AutoModel can
read back what AutoModel wrote. That passes even when the converter is
wrong in both directions, which is how the artifact ends up unloadable in
real PEFT.

Each family is one entry in a registry: a tiny Transformers model plus
the AutoModel adapter that owns its naming. Every entry saves through
Checkpointer.save_model with is_peft, reloads with
PeftModel.from_pretrained, and must match tensor for tensor and in its
forward output. A second test requires the bulk and per-tensor exports to
agree, since streaming consumers use the latter.

The model comes from Transformers rather than the native AutoModel class
because several native MoE classes build rope buffers on
torch.cuda.current_device() and cannot be constructed on CPU. The adapter
is the component under test either way.

Covers llama (dense control, no adapter), nemotron_v3, qwen3_moe and
minimax_m2. Reverting the Nemotron PEFT namespace selection, the
AMINT-330 defect fixed in NVIDIA-NeMo#3866, fails only that family's case.

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>
Adds qwen3_omni_moe. The family needs two things the text families do not:
the source model and the reload target are different classes, and the
comparison runs against the adapted submodule rather than the wrapper.

AutoModel trains the Omni thinker on its own, so its keys carry no
thinker. segment and the adapter adds one on export. The artifact
therefore targets the full Omni model, and the suite now lets a family
name a separate reload target plus the attribute that corresponds to the
source.

The case is marked xfail: the exported tensors carry the thinker.
segment but target_modules does not, so PEFT's suffix match also adapts
talker.model.* on the full model. Those adapters have no counterpart in
the file and are initialized randomly, which silently adapts the talker.
Tensor names and target metadata are produced by different code paths,
which is the asymmetry NVIDIA-NeMo#3867 is about.

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>
Found by the reload suite added in this branch.

to_hf and convert_single_tensor_to_hf both put exported adapter tensors
under thinker., but target_modules went out unchanged because the
adapter never overrode map_peft_target_module_to_hf. PEFT suffix-matches
that list against the receiving model, so on the full Omni model an
entry like model.layers.0.self_attn.q_proj also matched
talker.model.layers.0.self_attn.q_proj. The talker got adapter modules
the checkpoint has no weights for, so they stayed randomly initialized:
a model silently carrying untrained adapters on a tower the user never
fine-tuned.

Override the hook the same way kimi_k3 does, whose docstring already
states the rule: target_modules entries need the same renames the
state-dict keys get. Both export formats keep the namespace, since the
omni checkpoint has no layout that addresses thinker modules without it.

The suite's xfail machinery is now declarative rather than an imperative
pytest.xfail() call, which aborted before the body ran and would have
let a fixed export keep its marker unnoticed.

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>
@stanley1208
stanley1208 requested a review from a team as a code owner September 19, 2026 06:02
@copy-pr-bot

copy-pr-bot Bot commented Sep 19, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@stanley1208 stanley1208 changed the title test(peft): add a shared HF PEFT reload suite, and fix the qwen3_omni_moe target_modules it caught test(peft): add a shared HF PEFT reload suite Sep 19, 2026

@yuhezhang-ai yuhezhang-ai left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, @stanley1208, for adding the shared reload suite and catching the Omni namespace mismatch. Testing the exported artifact with the real PEFT consumer is valuable coverage.

I found two issues in this revision, detailed inline. The focused CPU run of the new suite and existing Qwen3 Omni prefix tests reported 4 failed, 13 passed, 1 skipped with Transformers 5.15.1, PEFT 0.19.1, and TorchAO 0.16.0.

assert exported, f"{family.id}: the export wrote no adapter tensors"
assert len(set(exported)) == len(exported), f"{family.id}: duplicate keys in the export"

loaded = PeftModel.from_pretrained(reference, str(adapter_dir), key_mapping={}, autocast_adapter_dtype=False).eval()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] Resolve the PEFT/Transformers reload incompatibility

Four new reload cases fail with the repository's declared Transformers 5.15.1 and minimum PEFT 0.19.1 versions: nemotron_v3, qwen3_moe, minimax_m2, and qwen3_omni_moe. PEFT's automatic v5 conversion passes distributed_operation to WeightConverter, whose constructor no longer accepts it, raising TypeError: WeightConverter.__init__() got an unexpected keyword argument 'distributed_operation' before the assertions. key_mapping={} does not disable this conversion.

Please resolve the dependency compatibility before adding these as unconditional unit tests. The suite passes with the older Transformers 5.12.1 / PEFT 0.18.1 environment, but that does not validate the declared dependency combination.

Returns:
Target-module name in the HF omni layout.
"""
return self._add_thinker_prefix(name)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Honor the detected thinker namespace

from_hf() supports standalone-thinker checkpoints and sets _uses_thinker_prefix=False when their expert keys omit thinker.. Both to_hf() and convert_single_tensor_to_hf() then preserve that layout, but this hook adds thinker. unconditionally. Consequently, adapter_config.json targets thinker.model.layers... while the exported tensor keys remain base_model.model.model.layers..., and real PEFT rejects the standalone-thinker reload with Target modules ... not found.

Please honor _uses_thinker_prefix here, matching both tensor exporters, and add a standalone-thinker reload case. I reproduced the failure through Checkpointer.save_model; restoring only the inherited pre-PR hook allowed reload with exact adapter tensor equality. That control used Transformers 5.12.1 / PEFT 0.18.1 to isolate this regression from the separate dependency incompatibility above.

@svcnvidia-nemo-ci svcnvidia-nemo-ci added the waiting-on-customer Waiting on the original author to respond label Sep 19, 2026
map_peft_target_module_to_hf added the thinker. namespace every time. to_hf
and convert_single_tensor_to_hf both check _uses_thinker_prefix first, so on
a thinker-only base the config said thinker.model.layers... while the tensors
stayed base_model.model.model.layers..., and peft refused the reload with
"Target modules ... not found".

The repo already relies on that layout: two existing tests assert the flag
goes False for a thinker-less checkpoint. The full omni case is unchanged,
since a fresh export and a full omni load both leave the flag True.

Adds a standalone-thinker case to the reload suite and a unit test tying
target_modules to the exported tensor names.

Also raises the peft floor to 0.20.0. 0.19.1 passes distributed_operation
into WeightConverter, which has never accepted it in any transformers v5,
so any MoE lora reload dies with a TypeError before it reads a weight.
It reproduces the same on transformers 5.12.1 and 5.15.1, and it already
breaks three merged tests in test_moe_peft_v5_state_dict_adapters.py.
0.20.0 assigns those attributes after construction instead.

Signed-off-by: stanley1208 <stanley.mei08@gmail.com>
@stanley1208
stanley1208 requested a review from a team as a code owner September 20, 2026 10:07
@yuhezhang-ai

Copy link
Copy Markdown
Contributor

/ok to test a279c5e

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

@yuhezhang-ai yuhezhang-ai left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for addressing both review findings. The target-module mapping now follows the detected full-Omni or standalone-thinker layout, and PEFT 0.20.0 resolves the Transformers 5.15.1 reload incompatibility.

I reran the shared reload suite, Omni prefix tests, and existing fused-MoE PEFT tests after merging current main: 28 passed, 7 skipped (six existing GPU-only cases and the inapplicable dense-model adapter comparison). Ruff format/check also pass. The generated dependency files have been refreshed to resolve the CI installation failures. Approving the code; full CI remains the merge gate.

@yuhezhang-ai

Copy link
Copy Markdown
Contributor

Approved after verifying both fixes. I merged current main and refreshed the generated dependency files on this branch; the dependency check now passes. The focused tests still report 28 passed and 7 expected skips.

Started maintainer CI proxy #3965 at the identical source commit e7a0008c03ea64121e0ea81fe93202cbbde3b264: CI run. This PR remains the merge target.

@yuhezhang-ai
yuhezhang-ai merged commit eb11d83 into NVIDIA-NeMo:main Sep 21, 2026
87 checks passed

This branch was successfully deployed

3 active deployments
public e7a0008c Deployed Sep 21, 2026 by copy-pr-bot[bot] via release / finalize / notify #4720
test e7a0008c Deployed Sep 21, 2026 by copy-pr-bot[bot] via cicd-wait-in-queue #10891
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

qwen3_omni_moe adapter_config target_modules miss the thinker. namespace, so peft also adapts the talker

4 participants