fix(post_train): Fix Tunix is_update_step signature and Qwix LoRA FSDP mesh sharding

[RexBearIU]

Description

This PR fixes multiple post-training regressions affecting Tunix SFT and Qwix LoRA fine-tuning in MaxText:

1. Tunix is_update_step signature mismatch in MaxTextPeftTrainer:

   • Tunix's PeftTrainer.train() passes (model, optimizer, grad_accumulator) as partial arguments and calls train_step(inputs, is_update_step=...).
   • MaxTextPeftTrainer.create_train_step_fn() had an outdated signature (model, optimizer, inputs, grad_accumulator=None) which caused TypeError: train_step() got an unexpected keyword argument 'is_update_step' across all Tunix SFT workloads (gemma3-4b.sft, gpt-oss-20b.sft, llama3_1_70b.sft).
   • Fixed argument order to (model, optimizer, grad_accumulator, inputs, is_update_step=True, **kwargs) and added support for conditional/accumulated optimizer updates when gradient accumulation is active.

2. Qwix LoRA Sharding IndivisibleError on multi-device FSDP meshes:

   • lora_utils.apply_lora_to_model generated dummy tracing inputs assuming batch size dp_size = mesh.shape['data'] (defaulting to 1 when data=1).
   • On partitioned multi-device meshes (e.g. v5p-128 with fsdp=64), JAX sharding checks failed with IndivisibleError because array dimension 0 (batch size 1) was not evenly divisible by the partition factor (64).
   • Fixed dp_size to compute the product of all data-parallel mesh axes (data, fsdp, fsdp_transpose, expert) and seq_len to compute the product of sequence-parallel axes (tensor_sequence, context).

Tests

• Unit tests added and verified in tests/post_training/unit/train_sft_test.py and tests/post_training/unit/lora_utils_test.py.
• Pre-commit hooks (codespell, pylint, pyink) pass cleanly.

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 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.

[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 36.36364% with 14 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 ⚠️

📢 Thoughts on this report? Let us know!