fix: resolve Gemma 3/4 RL rollout gibberish by unrolling scanned weights for vLLM adapter

[mesakhcienet]

Description

Problem

When running train_rl.py RL rollouts for Gemma 3 and Gemma 4 with scan_layers=true, the vLLM integration generates pure gibberish (e.g., "enenenen...").
The root cause is a structural mismatch: Gemma 3/4's cyclic attention forces MaxText to split scanned layers into nested blocks (layers.layers_0 and layers_remainder). Because Gemma relies on the MaxTextForCausalLM adapter (which expects flat, integer-keyed layers like standard Hugging Face models), the weight mapping silently failed. As a result, vLLM initialized with random weights, leading to the gibberish output.

Solution

1. Unroll Scanned Weights: fix previously introduced unroll_gemma_scanned_weights in maxtext_vllm_rollout.py. This function detects the layers.layers_0 structure and dynamically unrolls it into a flat, integer-keyed dictionary, ensuring exact compatibility with the MaxTextForCausalLM vLLM adapter.
2. Type Safety Enforcement: The unrolling logic strictly enforces int keys instead of str (e.g., 0 instead of "0") to comply with the adapter's nnx.List target structure.
3. End-to-End Test Update: Updated tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh to explicitly include the required use_standalone_converter=true flag.

Why use_standalone_converter=true is required for Gemma 3/4

Unlike Llama 3 or Qwen 3, Gemma 3 and 4 are bleeding-edge architectures that are not yet natively supported by upstream vLLM. To run them, MaxText overrides the vLLM model registry using a custom Python adapter (MaxTextForCausalLM).

Because of this custom adapter and Gemma's complex cyclic attention (layers_0 and layers_remainder), the default tunix fast-mapping path (transfer_state_with_mappings) completely fails to match the tensor structures. When tunix fails, it silently falls back to random weights (resulting in gibberish).

Setting use_standalone_converter=true explicitly bypasses tunix and routes the weights through our custom conversion script—triggering the new unroll_gemma_scanned_weights logic and ensuring the adapter receives a properly formatted dictionary.

If you look at the exact configuration passed to the vLLM rollout adapter inside src/maxtext/trainers/post_train/rl/train_rl.py (around line 499), you will see this explicitly defined
flag:

 rollout_vllm_additional_config=rollout_additional_config,
 rollout_vllm_init_with_random_weights=True, # <--- THIS LINE
 rollout_vllm_enable_dp_attention=trainer_config.enable_dp_attention,

1. The Optimization Hack: To save memory and time during the RL loop, MaxText does not ask vLLM to load real weights from disk. Because MaxText is about to stream the Actor's trained weights directly from its own memory over to vLLM, it tells vLLM to just boot up instantly using jax.random (pure garbage initialization). This is what rollout_vllm_init_with_random_weights=True does.

2. The tunix Failure Mode: tunix maps weights by matching dictionary keys. If a key from MaxText (e.g. layers.layers_0) doesn't perfectly match a key in the vLLM model registry (e.g. layers.0), tunix simply skips that tensor. It does not crash or throw a fatal error.

3. The Resulting State: Because tunix skipped those tensors, it never overwrites the garbage memory. vLLM proceeds to execute the decoding rollout using those skipped, randomly initialized weights.

When you feed random weights through a Transformer's attention layers, the math collapses into a repeating feedback loop of garbage tokens, which is exactly why it spit out "enenenenenenenen".

This single line of configuration (rollout_vllm_init_with_random_weights=True) combined with tunix's silent failure mode is the definitive proof of why it generates gibberish!

Tests

gemma3-4b unscan

• Script: here
• Log: here

gemma3-4b scan

use_standalone_converter and scan_layers should be set to true

• Script: here,
• Log: here

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.

[codecov]

Codecov Report

❌ Patch coverage is 5.12821% with 37 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...c/maxtext/integration/vllm/maxtext_vllm_rollout.py	2.63%	37 Missing ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

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

[hengtaoguo]

Could you also take a look at the failing CPU tests and see if relevant? Thanks!

[mesakhcienet]

Could you also take a look at the failing CPU tests and see if relevant? Thanks!

Looks like it was from main branch, i have rebased our branch and the old failing CPU tests are passed. Thank you!