QLoRA Support + NNX Decoder Sharding Fixes

[RexBearIU]

Description

This PR introduces core support for QLoRA and implements robust sharding metadata synchronization for NNX decoders.

Currently, applying LoRA adapters alongside quantization and complex multi-host sharding setups can lead to PartitionSpec mismatches and cross-backend device_put issues during lax.scan.

This PR solves these issues by:

• Extending the LoRA configuration to support quantization types (lora_weight_qtype) and block sizes (lora_tile_size) for QLoRA.
• Adding generic NNX metadata synchronization helpers (nnx_remove_scan_axis, nnx_add_scan_axis, nnx_sync_moveaxis) to safely manipulate PartitionSpec metadata during lax.scan operations in NNXDecoder.
• Implementing a _safe_reshard workaround using jax.make_array_from_callback in lora_utils.py to natively construct globally sharded arrays, bypassing backend-specific device_put issues on Pathways/McJAX.

Future improvements will include removing the _safe_reshard workaround once the underlying JAX/Qwix parameter materialization issues are fully resolved upstream.

Tests

• Unit tests added for NNX metadata synchronization (tests/utils/test_maxtext_utils_nnx.py).
  • Existing unit tests for lora_utils updated (note: QLoRA specific tests are temporarily skipped pending an upstream qwix fix).
  • Tested via multi-host mock execution to trigger the resharding callback logic.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 59.49367% with 32 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/maxtext_utils_nnx.py	66.07%	8 Missing and 11 partials ⚠️
src/maxtext/utils/lora_utils.py	7.69%	12 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	90.00%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[SurbhiJainUSC]

Why are we skipping QLoRA tests? Does that mean QLoRA won't work even after merging this PR?

[RexBearIU]

Yes, we need a new qwix release for NNX model QLoRA support.

My two CLs addressing the underlying parameter materialization issues have already been merged upstream in qwix. I will sync with the team to ask for their next release plan. Once the new release comes out, we just need to update the dependency version here and unskip these tests.

Should we wait and update the dependency/unskip tests in this PR, or would you prefer to merge this PR as-is and do the update in a separate follow-up PR once the release is published?

[SurbhiJainUSC]

I would say let's wait for Qwix fix, update dependency in MAxText and then merge this PR.

[RexBearIU]

Will do. I'll leave this PR open and update the qwix dependency as soon as the new release is published.