feat: QLoRA support for Dense/MoE models across Pathways and McJAX

[RexBearIU]

Description

This PR introduces comprehensive, out-of-the-box support for QLoRA (Quantized Low-Rank Adaptation) and fixes critical initialization crashes that occurred when running LoRA/QLoRA Fine-Tuning in multi-host TPU environments. It
unifies the array distribution logic so that both Dense models (like Llama 3) and Mixture-of-Experts models (like Gemma 4) work seamlessly on both Pathways and standard McJAX backends.

What was added/changed:

1. QLoRA Configuration Support: Added lora_weight_qtype (e.g., nf4) and lora_tile_size to configuration types, passing them to the underlying qwix library to enable quantized base model training. Includes a new
   lora_llama3_demo.ipynb notebook.
2. Native JAX Resharding (jax.make_array_from_callback): Added a native _safe_reshard function to lora_utils.py. Instead of attempting cross-host streaming (which crashes McJAX) or relying on external utilities, each host now
   gracefully constructs the global array by extracting its required physical slice directly from its local CPU copy.
3. MoE Partition Safety: Added strict type and partition spec checks to safely ignore dynamic routing states (which yield None partition specs) preventing AttributeError crashes in MoE models.
4. Documentation Updates: Updated the lora_on_multi_host.md tutorial with the correct gemma4-26b configurations, safe default memory limits, and explicit Pathways initialization flags.

Tests

Validated initialization, compilation, and training on a 2-slice v6e-8 cluster:

• Llama 3.1 8B on McJAX
• Llama 3.1 8B on Pathways
• Gemma 4 26B on McJAX (Successfully bypassed NoneType MoE partition crashes)
• Gemma 4 26B on Pathways (Successfully completed full training loop)

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 68.13187% with 29 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	53.84%	5 Missing and 1 partial ⚠️
src/maxtext/input_pipeline/input_pipeline_utils.py	75.00%	1 Missing and 2 partials ⚠️
src/maxtext/layers/nnx_decoders.py	90.00%	0 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[RexBearIU]

Closing this PR in favor of splitting the work into three smaller, focused PRs: #3968 , #3969 , and #3970.