test(peft): add a shared HF PEFT reload suite - #3945
Conversation
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>
yuhezhang-ai
left a comment
There was a problem hiding this comment.
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() |
There was a problem hiding this comment.
[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) |
There was a problem hiding this comment.
[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.
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>
|
/ok to test a279c5e |
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
yuhezhang-ai
left a comment
There was a problem hiding this comment.
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.
|
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 |
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_modelwithis_peft=True, reload withPeftModel.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 adaptedtalker.model.*and left those adapters random.map_peft_target_module_to_hfwas 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.