Add rename_weight_keys helper and adopt it in pure-rename models - #334
Merged
Conversation
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR refactors duplicated preprocess_weights() key-renaming loops by introducing a shared helper (rename_weight_keys) in mobius._weight_utils, then adopting it in several models that only require pure substring renames. It also adds focused unit tests to lock down replacement ordering/cascading semantics, collision detection, and tensor value sharing.
Changes:
- Added
rename_weight_keys(state_dict, replacements)helper (ordered substring replacements; raises on key collisions). - Replaced hand-written rename loops in 5 models (
phi,hunyuan_dit,cogvideox,hunyuan_v1,chatglm) with the helper. - Added unit tests covering behavior (ordering/cascade, identity, collisions, shared tensor values, empty input).
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/_weight_utils.py | Adds rename_weight_keys utility with ordered replacements and collision detection. |
| src/mobius/_weight_utils_test.py | Adds unit tests validating rename semantics, collision behavior, and value sharing. |
| src/mobius/models/phi.py | Switches Phi weight renames to use the shared helper. |
| src/mobius/models/hunyuan_v1.py | Replaces dict-comprehension rename logic with rename_weight_keys. |
| src/mobius/models/hunyuan_dit.py | Refactors DiT rename loop to rename_weight_keys. |
| src/mobius/models/cogvideox.py | Refactors CogVideoX rename loop to rename_weight_keys. |
| src/mobius/models/chatglm.py | Replaces in-place pop/replace logic with rename_weight_keys (pure substring renames). |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Many model `preprocess_weights` implementations hand-write the same
rename loop:
new = {}
for name, tensor in state_dict.items():
name = name.replace(old1, new1)
name = name.replace(old2, new2)
new[name] = tensor
Add `rename_weight_keys(state_dict, replacements)` to `_weight_utils.py`
to centralise this. It applies an ordered sequence of (old, new)
substring replacements to every key (replacements cascade, matching the
hand-written loops), shares tensor values without cloning, and raises
ValueError on a key collision (two source keys mapping to the same
renamed key) so a silently-dropped weight surfaces as a hard error.
Adopt it in the five models whose preprocess_weights is a pure
substring-rename: phi, hunyuan_dit, cogvideox, hunyuan_v1, chatglm. The
transformations are behaviour-preserving; models with interleaved tensor
transforms (reshape/split/filter) are intentionally left as-is.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
Address PR review: when two source keys rename to the same target, the ValueError now identifies the first producer key (not just "another key"), which makes the collision actionable when it fires in a large checkpoint. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/weight-rename-helper
branch
from
June 5, 2026 15:39
66d5165 to
3bf0bef
Compare
justinchuby
added a commit
that referenced
this pull request
Jun 5, 2026
Reflect the public-API and DRY changes shipped in PRs #333/#334/#336: - build_from_gguf is now a top-level export; update the import examples in docs/api/build_from_gguf.md and docs/getting-started.md to `from mobius import build_from_gguf`. - weight-name-alignment skill: add a 'Shared helpers' section documenting the _weight_utils rename helpers, including the new rename_weight_keys and vlm_vision_weights, so future model work reuses them instead of hand-written rename loops. - multimodal-models skill: point at the shared vlm_* weight helpers. - moe-models skill: note Qwen35MoEBlock subclasses Qwen2MoELayer and fix the class file paths (models/qwen.py -> models/qwen35.py). Depends on #333 (build_from_gguf export), #334 (rename_weight_keys) and #336 (vlm_vision_weights) landing first. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Jun 5, 2026
Updates outdated docs and agent skills to match the public-API + DRY refactor shipped across the themed PRs. ## Changes - **`build_from_gguf` import** — now a top-level export, so `docs/api/build_from_gguf.md` and `docs/getting-started.md` use `from mobius import build_from_gguf` instead of the internal `mobius.integrations.gguf` path. - **weight-name-alignment skill** — new *Shared helpers* section documenting the `_weight_utils` rename helpers, including the new `rename_weight_keys` and `vlm_vision_weights`, so future model work reuses them instead of hand-written rename loops. - **multimodal-models skill** — points at the shared `vlm_*` weight helpers. - **moe-models skill** — notes `Qwen35MoEBlock` now subclasses `Qwen2MoELayer`, and fixes stale class file paths (`models/qwen.py` → `models/qwen35.py`). ## Dependencies Depends on the API PRs landing first (the helpers/exports documented here only exist on those branches): - #333 — `build_from_gguf` top-level export - #334 — `rename_weight_keys` - #336 — `vlm_vision_weights` Docs-only; no code or tests affected. --------- Signed-off-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Justin Chu <11205048+justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
DRY refactor: centralise the duplicated weight-key rename loop.
Many model
preprocess_weightsimplementations hand-write the same idiom:Changes
rename_weight_keys(state_dict, replacements)to_weight_utils.py. Applies an ordered sequence of(old, new)substring replacements to every key (replacements cascade, matching the hand-written loops it replaces), shares tensor values (no clone), and raisesValueErroron a key collision (two source keys mapping to the same renamed key) so a silently-dropped weight becomes a hard error instead of a subtle bug.phi,hunyuan_dit,cogvideox,hunyuan_v1,chatglm. Behaviour-preserving.Out of scope (deliberately left as-is)
Models with interleaved tensor transforms (reshape/split/filter + rename) keep their custom loops — forcing them through the helper would obscure control flow. The vision-tower
fc1/fc2rename triple-duplication (gemma3/llava/mllama) is a separate, filtering-based pattern best handled in the multimodal-task DRY PR.Verification
ruff check+ruff format --checkclean on all changed files.