Fix QLoRA and Gemma4-small decoder issues for NNX#4504
Draft
hsuan-lun-chiang wants to merge 12 commits into
Draft
Fix QLoRA and Gemma4-small decoder issues for NNX#4504hsuan-lun-chiang wants to merge 12 commits into
hsuan-lun-chiang wants to merge 12 commits into
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
4 tasks
…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.
- 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
TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
When scan_layers=False, NNXDecoder updates its layers in a loop.
It was checking `hasattr(self, "layers")` to decide whether to update `self.layers[lyr]` or use `setattr(self, f"layers_{lyr}", new_layer)`.
Because of the newly added `__getattr__` fallback for read-only access to `self.layers`, `hasattr(self, "layers")` was evaluating to True, causing the code to try to mutate the dynamically constructed list in-place. This mutation was lost as it only modified the temporary list returned by `__getattr__`, leaving the actual `layers_{lyr}` attributes unchanged and discarding newly initialized parameters (like LoRA parameters).
This fix checks `hasattr(self, f"layers_{lyr}")` first to prioritize updating the individual layer attributes directly.
TAG=agy
CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
The monkeypatch to `qwix.LoraProvider.dot_general` is not needed because `qwix`'s `find_param` has a shape-based fallback (Approach 2) that successfully matches parameters when ID-based matching fails (which occurs because `PtqProvider` replaces parameters with `WithAux` wrappers during quantization in NNX). The original failure where `LoRAParam` was missing was entirely caused by the sequential mode update bug in `nnx_decoders.py` (which discarded updated layers), not by `find_param` failure. TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
TAG=agy CONV=a499f437-b38f-44a6-99d3-11a38fb2e032
TAG=agy CONV=b1cd074b-f03b-4919-9868-597278555171
Move model_prefix definition to the start of get_data to avoid UnboundLocalError when use_multimodal=False. TAG=agy CONV=b1cd074b-f03b-4919-9868-597278555171
hsuan-lun-chiang
force-pushed
the
fix/nnx_layers
branch
from
July 17, 2026 08:31
b945b4e to
5d8b8c5
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.
This PR fixes several issues with the NNX decoders introduced in PR #4464:
NNXDecoderwhere it was mutating a temporary list returned by__getattr__instead of updating the individuallayers_iattributes. This was causing newly initialized parameters (like LoRA parameters) to be discarded.qwixmonkeypatch sincefind_paramsuccessfully falls back to shape-based matching.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.TAG=agy
CONV=a499f437-b38f-44a6-99d3-11a38fb2e032