Fix JAX 0.11.0 sharding compatibility and SFT trainer signature for Tunix integration

[RexBearIU]

Description

This PR resolves CI notebook test failures (maxtext_jupyter_notebooks / run) under JAX 0.11.0+ and updates SFT trainer interfaces for Tunix integration.

Summary of Changes:

1. src/maxtext/integration/tunix/tunix_adapter.py:
   • Added _compat_wsc and _compat_top_k shims around jax.lax.with_sharding_constraint and jax.lax.top_k.
   • Fixes AssertionError: with_sharding_constraint acts as an assert when all axes of mesh are of type Explicit and ShardingTypeError on JAX 0.11.0+ when resharding sharded logits during vLLM decoding.
2. src/maxtext/trainers/post_train/sft/train_sft.py:
   • Updated train_step signature to accept grad_accumulator: Any = None.
   • Fixes TypeError: train_step() got an unexpected keyword argument 'grad_accumulator' when called by PeftTrainer.train.
3. src/maxtext/examples/lora_llama3_demo.ipynb:
   • Set MODEL_NAME = "llama3.1-8b" to satisfy MaxText pyconfig Pydantic model name validation schema.
   • Retained %%capture in Cell 18.
   • Initialized AutoTokenizer and fallback chat template in Cell 21 when RUN_EVALUATION is enabled.

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

• Tested end-to-end execution of lora_llama3_demo.ipynb on TPU v6e-8 using Papermill with both RUN_EVALUATION="false" and RUN_EVALUATION="true".
• Verified 100% clean exit code 0 and Pre/Post SFT evaluation metrics.
• Ran local pre-commit run --files (codespell, pylint, pyink) - 100% passed cleanly.

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 40.74074% with 16 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/integration/tunix/tunix_adapter.py	42.30%	15 Missing ⚠️
src/maxtext/trainers/post_train/sft/train_sft.py	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[SurbhiJainUSC]

Why do we need to use non-instruct version?

[RexBearIU]

We fixed this in PR #4417. We accidentally removed Instruct while testing the new dependencies.