feat(nnx): support native Flax NNX PEFT/LoRA training loop - #4501
Open
RexBearIU wants to merge 2 commits into
Open
feat(nnx): support native Flax NNX PEFT/LoRA training loop#4501RexBearIU wants to merge 2 commits into
RexBearIU wants to merge 2 commits into
Conversation
RexBearIU
requested review from
A9isha,
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
July 16, 2026 01:29
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
RexBearIU
force-pushed
the
nnx-lora-support
branch
from
July 16, 2026 07:44
7ee8331 to
d54e2a3
Compare
RexBearIU
force-pushed
the
nnx-lora-support
branch
from
July 17, 2026 09:45
d54e2a3 to
108f993
Compare
SurbhiJainUSC
approved these changes
Jul 17, 2026
RexBearIU
force-pushed
the
nnx-lora-support
branch
2 times, most recently
from
July 22, 2026 15:41
282a533 to
558f0f2
Compare
RexBearIU
force-pushed
the
nnx-lora-support
branch
27 times, most recently
from
July 23, 2026 21:38
7f11e55 to
70c1883
Compare
YixuanWang-99
approved these changes
Jul 28, 2026
- Resolve pjit sharding metadata mismatches by re-evaluating state_mesh_shardings after parameter restoration and LoRA/QLoRA module injection. - Update l2norm_pytree to skip non-inexact leaves (e.g. packed void arrays or WithAux wrappers). - Implement normalized key translation (0 <-> layers_0) in checkpointing and model creation utilities for seamless parameter merging across NNX and Linen models. - Add mesh axis shape sanitization for non-divisible dimensions. - Include unit tests for l2norm_pytree void quantization, key translation, and upfront LoRA state restoration.
Three unit-test failures on a multi-device mesh: - get_nnx_var_named_sharding_with_scan_axis applied adjust_pspec_for_indivisible_shapes to every resolved spec, silently dropping a physical axis whenever a dimension was not divisible by the mesh axis size. On a single device every axis has size 1 so this was a no-op, which is why it only showed up on the 2x2 pathways mesh. - The composite-pytree and pre-resolved-NamedSharding branches returned v.replace(sharding=...), which sets metadata and leaves the value as the raw arrays. nnx_construct_named_sharding must replace the leaf value with a NamedSharding. - nnx_remove_scan_axis dropped its inconsistent-metadata guard and instead popped or truncated the sharding names, which is what corrupted DeepSeek scan-axis sharding before. Also drops TestLeafSqSum, which tests a qvalue-unwrapping l2norm_pytree that no longer exists in max_utils.
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
4 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.
Description
This PR adds native parameter-efficient fine-tuning (PEFT), LoRA, and QLoRA training support and complete checkpoint lifecycle resume/restoration within Flax NNX training in MaxText (
train.py,train_sft.py,maxtext_utils.py,checkpointing.py,train_state_nnx.py,lora_utils.py,sharding.py).Key Changes:
Dynamic Parameter Filtering (
wrt):nnx.Optimizer(wrt=...)across MaxText (train.py,train_sft.py,train_compile.py,generate_param_only_checkpoint.py,standalone_checkpointer.py,train_utils.py) to targetnnx.LoRAParamwhen LoRA is enabled, ensuring gradient updates affect adapter weights.Mesh Sharding & Activation Protection:
adjust_pspec_for_indivisible_shapesinsrc/maxtext/utils/sharding.pystrictly to parameter variables (isinstance(inputs, nnx.Variable)), preserving activation partition specs insideshard_mapacross multi-host TPU topologies.Checkpoint Lifecycle & Resume:
src/maxtext/common/checkpointing.py, supporting both Linen'params'wrappers and root NNX state dictionaries.TrainStateNNX.split_for_checkpointintrain_state_nnx.pyto routennx.Paramandnnx.LoRAParaminto parameter collections for complete PyTree state tracking.QWiX Quantization & Scanned Layer Execution:
getattr(self.kernel, "value", self.kernel)) insrc/maxtext/layers/linears.pyfor QWiX quantized parameter wrappers.nnx.split/nnx.mergefor QWiX layers insrc/maxtext/layers/nnx_decoders.pyand filteredRngStateandIntermediatevariables during layer scanning.SFT Mesh Placement & Imports:
train_sft.pywithinwith jax.set_mesh(mesh)to ensure multi-device mesh placement.train_sftin test helpers and decorated SFT test methods with@pytest.mark.post_training.Tests
pyink,pylint,codespell,yamllint).tests/integration/pipeline_parallelism_test.pypassed (test_full_train_circular,test_full_train_non_circular,test_subset_layers,test_full_train_fp8,test_full_train_nanoo_fp8).tests/unit/maxtext_utils_nnx_test.pyandtests/unit/lora_utils_nnx_test.pypassed 100%.test_qlora_e2e_gemma4_sft_nf4,test_qlora_e2e_qwen3_sft_nf4,test_qlora_e2e_gptoss_unscanned_nf4) passed 100%.Checklist
gemini-reviewlabel.