Fix QLoRA NNX issues#4480
Draft
hsuan-lun-chiang wants to merge 2 commits into
Draft
Conversation
…oading
Fix the Errors:
1. DeepSeek Unscanned Checkpoint Mismatch:
When restoring unscanned DeepSeek checkpoints, loading failed with:
'ValueError: Checkpoint structure mismatch: 374 of 377 model parameter paths were not found... Example missing paths: decoder.dense_layer_0...'
Cause: _init_sequential_deepseek registered layers as singular 'dense_layer' and 'moe_layer', producing PyTree paths like 'decoder.dense_layer_0'. However, Linen checkpoints and MaxText converter scripts format layer keys as plural ('decoder.dense_layers_0', 'decoder.moe_layers_0').
2. Generic Unscanned Model Checkpoint Mismatch (GPT-OSS, Llama4, etc.):
When restoring unscanned generic models, loading failed with:
'ValueError: Checkpoint structure mismatch: 456 of 459 model parameter paths were not found... Example missing paths: decoder.layers.0...'
Cause: _init_sequential_generic used _create_and_register_layer which appended to self.layers initialized as nnx.List([]). In Flax NNX, tracking submodules in an nnx.List container forces state PyTree extraction to use list-index paths ('decoder.layers.0') instead of named attribute paths ('decoder.layers_0'), breaking compatibility with on-disk Linen checkpoints.
Summary of Fixes:
- Updated _init_sequential_deepseek to use plural prefixes ('dense_layers', 'moe_layers').
- Used _create_and_register_named_layer in sequential initialization (_init_sequential_deepseek, _init_sequential_generic) so submodules register exclusively under named string attributes ('layers_0', 'layers_1').
- Replaced nnx.List with a plain Python list for self.layers and updated the forward pass to retrieve named attributes via getattr(self, f'layers_{lyr}').
- Sequential Layer Naming & Flax 0.12.0 Compatibility (nnx_decoders.py):
- Fixed DeepSeek sequential layer retrieval in _apply_sequential_layers to look up 'dense_layers_{i}' and 'moe_layers_{i}' attributes when scan_layers=False.
- Removed self.layers.append(layer) in _init_gemma4_small_layers to prevent Flax 0.12.0 from raising a static container attribute ValueError on unannotated Python lists.
- MaxEngine KV-Cache Stacking & Sharding (maxengine.py):
- Updated self.prefill_kv_cache_shardings, _maybe_stack_prefill_result_cache, and _maybe_unstack_prefill_result_cache to handle named per-layer keys ('layers_0', 'layers_1') in addition to stacked 'layers' keys when scan_layers=False. Unstacks KV cache matching decode_state['cache'] keys to prevent PyTree structure mismatches during engine insert.
- LoRA Module Regex Pattern Matching (lora_utils.py & lora_utils_test.py):
- Updated _get_lora_module_path regex replacement to r'layers(?:_[0-9]+|/[0-9]+)?/' so LoRA target verification matches named attributes ('decoder/layers_0/...') as well as indexed ('decoder/layers/0/...') and scanned ('decoder/layers/...') paths.
- Updated expected regex string assertions in test_get_lora_module_path.
hsuan-lun-chiang
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
abhinavclemson,
aireenmei,
bvandermoon,
darisoy,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
igorts-git,
jiangjy1982,
khatwanimohit,
parambole,
richjames0,
shralex,
shuningjin,
suexu1025,
vipannalla and
xibinliu
as code owners
hsuan-lun-chiang
changed the title
hsuan-lun-chiang
force-pushed
the
fix-qlora-nnx
branch
2 times, most recently
from
2c06f55 to
c718861
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
- Increase model size in lora_utils_test.py to avoid IndivisibleError on FSDP. - Monkeypatch qwix LoraProvider.dot_general to find parameter names before quantization. - Refactor NNXDecoder to support dynamic graph init with scan_layers=False. - Add assertions to verify quantization in tests. TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
hsuan-lun-chiang
force-pushed
the
fix-qlora-nnx
branch
from
c718861 to
bb9bcef
Compare
mesakhcienet
force-pushed
the
xibin/nnx_layers
branch
from
e00c9be to
63e125d
Compare
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
2 participants
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
Fix QLoRA IndivisibleError and parameter missing issues for NNX.
Root Cause Analysis & Resolution
See analysis in the PR description or referenced commits:
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.TAG=agy
CONV=a499f437-b38f-44a6-99d3-11a38fb2e032