fix: robust layer-count parsing for nested checkpoint prefixes (complements #336) - #339
HEETMEHTA18 wants to merge 4 commits into
Conversation
Fixes ValueError: invalid literal for int() with base 10: 'layers' when a checkpoint's LM weights are nested under a wrapper module (e.g. language_model.model.layers.N...). Mirrors the .find()+.isdigit() guard already used for MoE expert-key parsing in airllm_base.py. Complements lyogavin#336, which fixes only the layer_names branch with a fixed split-index shift; this covers both branches for arbitrary nesting depth.
|
@Tradunsky Thanks for the suggestion! I’ll add tests covering different model/checkpoint prefix patterns to make sure the layer-count parsing handles various nesting depths and model architectures correctly. I’ll make the changes and push the updated tests to the same branch. |
Extract _count_layers to a module-level helper so it is directly unit-testable, and add test_layer_count.py exercising a matrix of the prefixes seen across supported architectures (model.layers, transformer.h, transformer.encoder.layers, language_model.model.layers) plus the nested VLM / ge_model shapes that triggered lyogavin#335. Confirms non-layer keys, empty indexes, non-numeric segments and duplicate indices are handled safely. Existing test_kimi_k3_split suite still passes unchanged.
Use .find() + walrus in a single guard to drop the layer-count parser from 19 lines to 13 without changing behavior.
| nums = set() | ||
| for k in index_keys: | ||
| pos = k.find(marker) | ||
| if pos != -1 and (head := k[pos + len(marker):].split('.', 1)[0]).isdigit(): |
There was a problem hiding this comment.
Although this version is shorter, it has higher cognitive load to understand and debug.
I would suggest keeping the previous version.
There was a problem hiding this comment.
Agreed — readability wins over a couple of saved lines here. Reverted to the explicit loop version in 1e15a5a. Thanks for the review.
| def test_kimi_k3_language_model_prefix(self): | ||
| # Kimi K3 nests the decoder under language_model (airllm_kimi_k3.py) | ||
| prefix = 'language_model.model.layers' | ||
| keys = [f'{prefix}.{i}.block_sparse_moe.experts.0.w1.weight' for i in range(3)] |
There was a problem hiding this comment.
| keys = [f'{prefix}.{i}.block_sparse_moe.experts.0.w1.weight' for i in range(3)] | |
| keys = layer_keys(prefix, 3, 'block_sparse_moe.experts.0.w1.weight') |
There was a problem hiding this comment.
Applied — now uses layer_keys(prefix, 3, 'block_sparse_moe.experts.0.w1.weight') in 1e15a5a. All 21 tests still pass.
Tradunsky
left a comment
Tradunsky
left a comment
There was a problem hiding this comment.
A few suggestions. Overall looks good! Thank you for adding tests!
Per review feedback: the walrus-operator one-liner was more compact but harder to read and debug, so keep the explicit loop. Also use the layer_keys() test helper for the Kimi K3 prefix per the suggestion.
|
Hey @lyogavin @Tradunsky — this PR is ready to merge. It fixes #335 by replacing the fixed-position layer-index parsing in |
Tradunsky
left a comment
Tradunsky
left a comment
There was a problem hiding this comment.
Thank you for considering my suggestions.
Ready to fly! 🚢
|
@Tradunsky thank you!! for providing more suggestions to the changes that i have committed!! |
This branch has not been deployed
2 participants
Fixes #335 (both branches -- see comment on #336 for why that PR's
fix is scoped to only one branch).
utils.py's layer-counting logic assumed every key matchingmodel.layersstarts with that prefix, so it always read the layerindex from a fixed split position. Checkpoints where the LM sits
under a wrapper module (e.g.
language_model.model.layers.0...,common in VLM/multimodal architectures) still match the substring
filter but break the positional assumption, causing:
This applies the same
.find()+.isdigit()guard already usedfor per-expert MoE key parsing in
airllm_base.py, to both thelayer_names is Noneand custom-layer_namesbranches, so it'scorrect for arbitrary nesting depth rather than one fixed position.
Verified against both a normal flat checkpoint and a nested
wrapper-style checkpoint (see test snippet in code comments / commit
message).
Why not just the #336 approach
#336 moves
split('.')[1]tosplit('.')[2]in the custom-layer_namesbranch only. Two gaps:
layer_names is Nonebranch (default architecture path) is nottouched and crashes on the same nested-key shape reproduced in ValueError: invalid literal for int() with base 10: 'layers' #335.
layer_namesbranch, a fixed split index is the sameclass of bug at a different depth. For Key K3's
language_model.model.layers.0.self_attn...keys the post-prefixsegment is
.0.self_attn..., where[1]is the layer index but[2]is
'self_attn'— i.e. the[2]shift would also raise the exactbug on the K3 layout exercised by
air_llm/tests/test_kimi_k3_split.py.The
.find()+.isdigit()approach handles both branches at anynesting depth.
air_llm/tests/test_kimi_k3_split.pypasses unchangedwith this change.