Skip to content

refactor(checkpoint): make the target-module rename hook a StateDictAdapter method - #3591

Open
stanley1208 wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
stanley1208:feat/target-module-interface
Open

refactor(checkpoint): make the target-module rename hook a StateDictAdapter method#3591
stanley1208 wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
stanley1208:feat/target-module-interface

Conversation

@stanley1208

@stanley1208 stanley1208 commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

@HuiyingLi follow up to #3435.

made the rename hook a real method on StateDictAdapter instead of a getattr. so now the base class has the method, it just returns the name unchanged, and kimi_k3 overrides it with its own renames.

some older adapters (llama, qwen2/3, kimivl) don't inherit from StateDictAdapter, so I added an isinstance check to keep their saves working like before.

…dapter method

Follow-up to NVIDIA-NeMo#3435. The rename hook for adapter_config.json's
target_modules was discovered with getattr in _extract_target_modules,
which is easy to miss when writing a new adapter. It's now a regular
method on the StateDictAdapter base that returns the name unchanged, so
adapters whose to_hf renames modules just override it (kimi_k3 already
does) and addons.py calls it through an isinstance check instead.

Legacy adapters that don't subclass StateDictAdapter (llama, qwen2/3,
kimivl, and the mixin-only test adapters) skip the hook like they did
with the old getattr guard, so their PEFT saves keep working.

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

copy-pr-bot Bot commented Aug 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.

@HuiyingLi

Copy link
Copy Markdown
Contributor

Thank you I appreciate it!

@HuiyingLi

Copy link
Copy Markdown
Contributor

/claude review

@HuiyingLi

Copy link
Copy Markdown
Contributor

/ok to test 4cc45eb

@claude

claude Bot commented Aug 20, 2026

Copy link
Copy Markdown
Contributor

LGTM

Reviewed all 3 changed files. The change swaps the duck-typed getattr(adapter, "map_peft_target_module_to_hf", None) probe for a typed isinstance(adapter, StateDictAdapter) check plus an identity-default base method — a clean move from runtime attribute probing to a typed contract.

Verified:

  • KimiK3StateDictAdapter (the only real overrider) subclasses StateDictAdapter, so the isinstance gate preserves its rename behavior.
  • No adapter defines map_peft_target_module_to_hf without subclassing StateDictAdapter, so nothing regresses; legacy duck-typed adapters correctly skip the hook (covered by the new test).
  • EncoderStateDictAdapter subclasses the base and inherits the identity default (no-op), and adapter=None is safely False.
  • New base method has a Google-style docstring; no tensor inputs, no circular import.
  • Tests cover the identity default, the overriding adapter, and the bare non-subclass adapter.

@HuiyingLi

Copy link
Copy Markdown
Contributor

/ok to test 35e1eba

@akoumpa

akoumpa commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

/ok to test 0556d59

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.

3 participants