fix(post_train): Fix Tunix SFT signatures, Qwix LoRA mesh sharding, and scale Qwen/LLaMA parallelism

[RexBearIU]

Description

Motivation & Background

This PR addresses multiple blocking regressions across post-training workflows (Tunix SFT, Qwix LoRA fine-tuning, and RL rollouts), MoE checkpoint conversions, and hardware scalability in MaxText. These issues caused persistent failures on Cloud Composer TPU end-to-end testing pipelines.

---

Detailed Changes

1. Gemma 4 MoE per_expert_scale Forward Pass Fix

• Issue: In Linen, RoutedMoELegacy scaled wo_kernel during non-fused forward passes (wo_kernel * per_expert_scale). During the Flax NNX migration, wo was converted to linears.DenseGeneral, and per_expert_scale multiplication was accidentally omitted for non-fused/training calls in RoutedMoE.
• Fix: In src/maxtext/layers/moe.py, restored per_expert_scale multiplication on top_k_weights inside RoutedMoE.gate() for selected top-k experts when fuse_expert_scales is false.

2. Tunix is_update_step Signature Alignment in MaxTextPeftTrainer

• Issue: The upstream Tunix PeftTrainer.train() passes (model, optimizer, grad_accumulator) as partial arguments and invokes train_step(inputs, is_update_step=...). MaxTextPeftTrainer.create_train_step_fn() had an outdated signature (model, optimizer, inputs, grad_accumulator=None), causing TypeError: train_step() got an unexpected keyword argument 'is_update_step' across all Tunix SFT runs (gemma3-4b.sft, gpt-oss-20b.sft, llama3_1_70b.sft).
• Fix:
  • Standardized argument order to (model, optimizer, grad_accumulator, inputs, is_update_step=True, **kwargs).
  • Added conditional optimizer updates when gradient accumulation is active (is_update_step=True), accumulating gradients across micro-batches otherwise.

3. Qwix LoRA Sharding IndivisibleError on Multi-Device FSDP Meshes

• Issue: lora_utils.apply_lora_to_model created dummy tracing inputs using dp_size = mesh.shape['data'] (defaulting to 1 when data=1). On partitioned multi-device meshes (e.g. v5p-128 with fsdp=64), JAX threw IndivisibleError because batch dimension 1 was not divisible by the partition factor (64).
• Fix:
  • Updated dp_size calculation to compute the product of all data-parallel mesh axes:
    dp_size = mesh.shape.get('data', 1) * mesh.shape.get('fsdp', 1) * mesh.shape.get('fsdp_transpose', 1) * mesh.shape.get('expert', 1)
  • Updated seq_len to compute the product of sequence-parallel axes (tensor_sequence, context).

4. MoE Checkpoint Conversion & Qwen3 Hardware Scalability

• Expert Weight Transposition: Added reshape_expert_kernel in param_mapping.py to preserve expert axis 0 (transpose(0, 2, 1)) for 3D MoE expert weights, separating 2D gate/router hooks from 3D expert hooks.
• Mesh Scalability (v5p-32): Scaled Qwen3-30B SFT and RL test scripts from 128 cores ($4 \times 4 \times 4$) to 32 cores (v5p-32 / 16 devices) with TP=4, FSDP=2, EP=2 (test_qwen3_sft.sh, test_qwen3_rl.sh), reducing resource footprint by 75%.
• CPU Logit Verification: Added dtype=float32 and per_device_batch_size=1 to test_qwen3_to_mt.sh forward-pass logit checker with --max_kl_div=0.05.

5. Memory Protection & Model Stability

• LLaMA 3.1 70B Rematerialization: Added remat_policy=full across SFT and RL test scripts to prevent Out-Of-Memory (OOM) failures on 70B parameter footprints.
• GPT-OSS 20B Dynamic Sharding: Removed hardcoded ici_fsdp_parallelism=64 from test_gpt_oss.sh to allow dynamic auto-sharding on v5p-8, and set rollout TP=1 for single-host execution.
• Dependencies: Added torchvision install to test_qwen3_multimodal_sft.sh for multimodal image processing.

---

Tests

• Unit Tests: Pass cleanly (tests/post_training/unit/train_sft_test.py, tests/post_training/unit/lora_utils_test.py).
• Pre-commit Checks: codespell, pylint, pyink, and yamllint all pass with 0 errors.
• E2E Cloud Composer: Test dashboard: Cloud Composer Airflow UI

---

Checklist

• 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 unit tests and verified pre-commit hooks.

[gemini-code-assist]

Code Review

This pull request introduces support for gradient accumulation and conditional optimizer updates in the SFT trainer, and updates LoRA initialization to dynamically compute dummy input shapes based on the mesh topology. A critical issue was identified in the gradient accumulation logic: checking the truthiness of grad_accumulator.grads can evaluate to False on the first step if it is empty or uninitialized, which would permanently bypass gradient accumulation. It is recommended to simplify this check to only verify the presence of the add method.

[gemini-code-assist]

Checking bool(getattr(grad_accumulator, "grads", None)) can cause gradient accumulation to be completely bypassed. If grad_accumulator.grads is initialized to None or is empty on the first step, this condition evaluates to False. As a result, the code will fall back to the else block, directly updating the optimizer and never calling grad_accumulator.add(grads). This means gradient accumulation will be permanently disabled. To fix this, simplify the condition to only check if grad_accumulator is not None and has the add method.

      if grad_accumulator is not None and hasattr(grad_accumulator, "add"):

[codecov]

Codecov Report

❌ Patch coverage is 45.94595% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/post_train/sft/train_sft.py	6.66%	14 Missing ⚠️
...xtext/checkpoint_conversion/utils/param_mapping.py	50.00%	6 Missing ⚠️

📢 Thoughts on this report? Let us know!