[Maxtext][Tunix] Fix Weight Conversion for Qwen3.5-35B

[YixuanWang-99]

Description

This is the new weight converter between maxtext and tunix vllm rollout engine, tests passed.

Key changes

• Fix weight conversion errors for qwen3.5-35b model, for both direct conversion and torchax conversion
• Improve the standalone validate_converter.py using the production weight loading
• Move get_rollout_kwargs_for_parallelism from train_rl to model_creation_utils, works for both single-device and multi-device TP/DP/EP setups

Tests

We have two conversion modes:

1. for maxtext trainer to vllm maxtext rollout (mode 1), using either transfer_stete_directly() and new weightConverter, standalone tests passed
2. for maxtext trainer to vllm native rollout (mode 2), using standalone converter Qwen35MaxTextToVLLMConverter, standalone test passed.

Then we have e2e test for post_train/train_rl 10 step smoke tests also passed:

1. Mode 2, with vllm native HF layout, logs
2. Mode 1, with vlllm maxtext layout, logs

To vllm maxtext layout, Standalone tests passed, E2E test is ongoing:

python -m maxtext.integration.vllm.validate_converter \
      src/maxtext/configs/post_train/rl.yml \
      model_name=qwen3.5-35b-a3b \
      tokenizer_type=huggingface \
      tokenizer_path=Qwen/Qwen3.5-35B-A3B \
      load_parameters_path='gs://hengtaoguo-maxtext-logs/checkpoints/qwen3.5-35b-a3b/scanned/2026-06-11-10-27/0/items' \
      run_name=validate_qwen35_maxtext_integration \
      vllm_hf_overrides=\"{'architectures': ['MaxTextForCausalLM']}\" \
      vllm_additional_config='{\"maxtext_config\": {\"model_name\": \"qwen3.5-35b-a3b\", \"model_call_mode\": \"inference\", \"enable_dp_attention\": false, \"allow_split_physical_axes\": true, \"log_config\": false, \"weight_dtype\": \"bfloat16\", \"prefuse_moe_weights\": true}, \"use_weight_converter\": false}' \
      use_standalone_converter=False \
      max_prefill_predict_length=128 \
      max_target_length=512 \
      scan_layers=true \
      skip_jax_distributed_system=true \
      weight_dtype=bfloat16 \
      rollout_tensor_parallelism=1 \
      hbm_utilization_vllm=0.85 \
      use_chat_template=False \
      hf_access_token=<token>\
      prompt=\"'Translate the following sentence into French: The quick brown fox jumps over the lazy dog.'\""

To vllm native HF layout, Standalone tests passed, E2E test is ongoing:

python -m maxtext.integration.vllm.validate_converter \
      src/maxtext/configs/post_train/rl.yml \
      model_name=qwen3.5-35b-a3b \
      tokenizer_type=huggingface \
      tokenizer_path=Qwen/Qwen3.5-35B-A3B \
      load_parameters_path='gs://hengtaoguo-maxtext-logs/checkpoints/qwen3.5-35b-a3b/scanned/2026-06-11-10-27/0/items' \
      run_name=validate_qwen35_standalone_integration \
      use_standalone_converter=True \
      max_prefill_predict_length=128 \
      max_target_length=512 \
      scan_layers=true \
      skip_jax_distributed_system=true \
      weight_dtype=bfloat16 \
      rollout_tensor_parallelism=1 \
      hbm_utilization_vllm=0.85 \
      use_chat_template=False \
      hf_access_token=<token>\
      prompt='Translate the following sentence into French - The quick brown fox jumps over the lazy dog.' "

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]

Code Review

This pull request integrates MaxText with Tunix and Raiden weight synchronization for RL post-training. Key changes include re-exporting Tunix's shared trainer data types to avoid type mismatches, refactoring MaxTextTrainingEngine to support dynamic/static batch splitting for JIT compilation, and implementing chunked Raiden weight synchronization to prevent host memory OOMs. Additionally, the weight conversion and validation pipelines are updated to align with production rollout behavior. Review feedback suggests treating nnx.Param as a leaf during tree_map in unscan_layers to avoid double-wrapping, and converting nnx.State to a standard dictionary before flattening in _split_into_chunks to ensure robustness against future Flax NNX refactorings.

[codecov]

Codecov Report

❌ Patch coverage is 2.46305% with 198 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/integration/vllm/validate_converter.py	0.00%	154 Missing ⚠️
src/maxtext/utils/model_creation_utils.py	3.84%	25 Missing ⚠️
...t/integration/vllm/torchax_converter/qwen35_moe.py	0.00%	10 Missing ⚠️
...c/maxtext/integration/vllm/maxtext_vllm_rollout.py	25.00%	6 Missing ⚠️
src/maxtext/integration/vllm/convert_utils.py	0.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!