Fix Tunix state translation to handle Flax NNX rng count structures

[chiajunglien]

Description

Fix Tunix state translation (unstacking) crashes in train_rl.py caused by Flax NNX 1D arrays and scalar RNG states when syncing weights from MaxText to vLLM rollout engine.

Root Cause

Tunix’s transfer_state_directly logic in tunix/generate/utils.py defaults to unstacking layers over scan_axis=1.
This assumption fails critically against Flax NNX's state representations under scan_layers=True, resulting in two distinct IndexError failure modes depending on the model architecture:

1. Llama-3.1-70B (1D Array Slicing IndexError):
   State parameters like dropout.rngs.params.count evaluate as 1D arrays (e.g., shape (80,)) when stacked over layers.
   Tunix attempts to slice this along scan_axis=1, immediately throwing tuple index out of range in _bulk_align_and_unstack because the requested axis does not exist.

   DAG Error Log

   Shortened Error Trace:

   [2026-07-29, 05:02:29 UTC] {xpk.py:365} INFO - [rank0]:   File "/usr/local/lib/python3.12/site-          packages/tunix/generate/utils.py", line 1697, in intersect_trees
   [2026-07-29, 05:02:29 UTC] {xpk.py:365} INFO - [rank0]:     unstacked_cache[cache_key] = _bulk_align_and_unstack(
   [2026-07-29, 05:02:29 UTC] {xpk.py:365} INFO - [rank0]:                                  ^^^^^^^^^^^^^^^^^^^^^^^^
   [2026-07-29, 05:02:29 UTC] {xpk.py:365} INFO - [rank0]:   File "/usr/local/lib/python3.12/site-     packages/tunix/generate/utils.py", line 1257, in _bulk_align_and_unstack
   [2026-07-29, 05:02:29 UTC] {xpk.py:365} INFO - [rank0]:     + (arr.shape[scan_axis],)
   [2026-07-29, 05:02:29 UTC] {xpk.py:365} INFO - [rank0]:        ~~~~~~~~~^^^^^^^^^^^
   [2026-07-29, 05:02:29 UTC] {xpk.py:365} INFO - [rank0]: IndexError: tuple index out of range

2. Gemma3-4B (Scalar Pure-Unstack Indexing limit):
   Due to Gemma's sliding window attention blocks, unroll_gemma_scanned_weights bypasses scalar dropout and rng fields from being explicitly unrolled. When these scalars (shape ()) reach Tunix, _unstack_scanned_param natively
   processes them via the "pure unstacking" code path and returns a tuple of length 1 (i.e. (src_val,)).
   However, vLLM expects values for all 34 independent layers. The subsequent state loading logic attempts to access unstacked_cache[cache_key][layer_idx] (e.g., layer_idx=33), resulting in a critical IndexError traversing the length-1 tuple.

   DAG Error Log

   Shortened Error Trace:

   [2026-07-29, 04:14:06 UTC] {xpk.py:365} INFO - [rank0]:   File "/usr/local/lib/python3.12/site-packages/tunix/generate/utils.py", line 1767, in transfer_state_directly
   [2026-07-29, 04:14:06 UTC] {xpk.py:365} INFO - [rank0]:     final_source, final_spec = intersect_trees(full_source_dict,   full_target_spec)
   [2026-07-29, 04:14:06 UTC] {xpk.py:365} INFO - [rank0]:                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
   [2026-07-29, 04:14:06 UTC] {xpk.py:365} INFO - [rank0]:   File "/usr/local/lib/python3.12/site-packages/tunix/generate/utils.py", line 1702, in intersect_trees
   [2026-07-29, 04:14:06 UTC] {xpk.py:365} INFO - [rank0]:     sliced_val = unstacked_cache[cache_key][layer_idx]
   [2026-07-29, 04:14:06 UTC] {xpk.py:365} INFO - [rank0]:                  ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^
   [2026-07-29, 04:14:06 UTC] {xpk.py:365} INFO - [rank0]: IndexError: tuple index out of range

Solution

Expanded the _tpu_inference_compat_patches context manager in train_rl.py to gracefully handle these NNX state edge cases during tunix.generate.utils operations:
1. Dynamic scan_axis fallback: Intercepted both _bulk_align_and_unstack and _unstack_scanned_param to check array dimensions prior to execution. If an array has insufficient dimensions (len(arr.shape) <= scan_axis), it safely assigns scan_axis=0 to accommodate 1D state arrays smoothly.
2. Broadcasting scalars robustly: Within _compat_unstack, if the unwrapped result yields a length-1 tuple (due to
src_shape == tgt_shape for fully broadcast/scalar parameters), the tuple is replicated (res * 256) to guarantee
subsequent [layer_idx] lookups succeed uniformly across diverse model depths.

Tests

• Evaluated RL-Training on Llama3.1-70b

Log: https://cloudlogging.app.goo.gl/ebVYu3HB79cfkKxQA

xpk workload create-pathways \
  --cluster=v5p-128-bodaborg-europe-west4-b \
  --workload=llama70-rl-$(date +%m%d%H%M) \
  --tpu-type=v5p-128 \
  --num-slices=1 \
  --priority=very-high \
  --project=cloud-tpu-multipod-dev \
  --zone=europe-west4-b \
  --skip-validation \
  --docker-image=[gcr.io/tpu-prod-env-multipod/llama-rl](http://gcr.io/tpu-prod-env-multipod/maxtext_post_training_nightly:28210168847):latest \
  --command="set -xue; \
    export TPU_MIN_LOG_LEVEL=0; \
    export TF_CPP_MIN_LOG_LEVEL=0; \
    export TPU_STDERR_LOG_LEVEL=0; \
    export JAX_PLATFORMS='proxy,cpu'; \
    export JAX_BACKEND_TARGET='grpc://127.0.0.1:29000'; \
    export ENABLE_PATHWAYS_PERSISTENCE='1'; \
    python3 -m maxtext.trainers.post_train.rl.train_rl \
      base_output_directory='gs://cloud-tpu-multipod-dev-emma-data/llama3.1-70b-0729/rl' \
      load_parameters_path='gs://runner-maxtext-logs/llama3.1-70b/to_maxtext/scanned/post-20260729T003944/0/items' \
      tokenizer_path='meta-llama/Llama-3.1-70B-Instruct'  \
      run_name='rl_test_0630' \
      rl.loss_algo='grpo' \
      scan_layers=true \
      num_batches=5 \
      batch_size=1 \
      num_test_batches=5 \
      model_name=llama3.1-70b \
      enable_single_controller=true \
      checkpoint_storage_use_zarr3=False \
      checkpoint_storage_use_ocdbt=False \
      rollout_tensor_parallelism=4 \
      vllm_hf_overrides='{\"architectures\": [\"MaxTextForCausalLM\"]}' \
      vllm_additional_config='{\"maxtext_config\": {\"model_name\": \"llama3.1-70b\", \"log_config\": \"false\"}}'"

• Evaluated end-to-end on Gemma3-4b

Log: https://cloudlogging.app.goo.gl/2YAsHHGbKdfUNerm9

xpk workload create-pathways \
  --cluster=v5p-128-bodaborg-europe-west4-b \
  --workload=gemma3-rl \
  --tpu-type=v5p-128 \
  --num-slices=1 \
  --priority=very-high \
  --project=cloud-tpu-multipod-dev \
  --zone=europe-west4-b \
  --docker-image=[gcr.io/tpu-prod-env-multipod/gemma3-rl](http://gcr.io/tpu-prod-env-multipod/maxtext_post_training_nightly:28210168847):latest \
  --env GCS_OUTPUT=gs://ml-auto-solutions/output/unowned/rl-f8c-v5p-128-2026-07-29-03-29-13/ \
  --skip-validation \
  --command='set -xue; export TPU_MIN_LOG_LEVEL=0 && export TF_CPP_MIN_LOG_LEVEL=0 && export TPU_STDERR_LOG_LEVEL=0 && export JAX_PLATFORMS=proxy,cpu && export JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 && export ENABLE_PATHWAYS_PERSISTENCE=1 && bash tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh post-20260729T030450 true'

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 0% with 17 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/post_train/rl/train_rl.py	0.00%	17 Missing ⚠️

📢 Thoughts on this report? Let us know!