Fix Qwen 3.5 35B RL gibberish output issue and show end-to-end RL convergence by AntonyMei · Pull Request #4770 · AI-Hypercomputer/maxtext

AntonyMei · 2026-08-07T00:14:53Z

Description

This PR fixes the Qwen 3.5 RL gibberish output issue and shows end-to-end RL training convergence with Qwen 3.5 35B + v5p-64 on GSM8K. This is a clean up of the fixes in yixuanm-dev-35b-new branch.

FIXES: b/521604343
FIXES: b/542769108

Tests

On yixuanm-dev-35b-new, Qwen 3.5 35B + v5p-64 on GSM8K shows end-to-end convergence:
gs://yixuanm-maxtext-logs/ConvergenceTesting/ym-qw35-35/tensorboard

Checklist

Before submitting this PR, please make sure (put X in square brackets):

[ x] I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
[ x] I have necessary comments in my code, particularly in hard-to-understand areas.
[ x] I have run end-to-end tests tests and provided workload links above if applicable.
[ x] 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 adds support for Qwen3.5 post-training reinforcement learning (RL) using GRPO on TPU hardware. Key enhancements include a hybrid cache layout helper for Qwen's Gated DeltaNet (GDN) to prevent silent quantization, scanned-weight unrolling for direct MaxText-to-MaxText weight synchronization, support for partial multi-dimensional rotary position embeddings (MRoPE), and dynamic replication of batch dimensions in the Mixture of Experts (MoE) layer to avoid sharding mismatches. Additionally, answer extraction utilities have been updated to handle prefilled reasoning start tokens. The code review feedback suggests several robustness improvements, including converting all path segments to strings during weight unrolling, adding defensive checks for null values in the MoE layer to prevent AttributeError crashes, and making the answer extraction regex robust to empty reasoning blocks.

gemini-code-assist · 2026-08-07T00:16:54Z

+    if container_idx == -1 or "dropout" in key or "rngs" in key:
+      continue
+    if not hasattr(value, "shape") or len(value.shape) <= scan_axis:
+      raise ValueError(f"Qwen scanned parameter {'.'.join(key)} has no scan axis {scan_axis}: {value!r}")


Using '.'.join(key) can raise a TypeError if any of the path segments in key are not strings (for example, integers representing sequence indices, which are common in Flax/NNX parameter paths). To ensure robustness and consistency with other parts of the codebase (such as line 255), we should convert all path segments to strings using map(str, key).

Suggested change

raise ValueError(f"Qwen scanned parameter {'.'.join(key)} has no scan axis {scan_axis}: {value!r}")

raise ValueError(f"Qwen scanned parameter {'.'.join(map(str, key))} has no scan axis {scan_axis}: {value!r}")

gemini-code-assist · 2026-08-07T00:16:55Z

+      if pre_bias_logits_pspec is not None:
+        pre_bias_logits_pspec = remove_incompatible_mesh_axes_from_partition_spec(
+            pre_bias_logits_pspec,
+            pre_bias_logits.shape,
+            self.mesh,
+            dims=(0,),
+        )


If pre_bias_logits is None (which is the case for non-DeepSeek3/4 models), attempting to access pre_bias_logits.shape will raise an AttributeError. We should defensively check that pre_bias_logits is not None before accessing its shape.

Suggested change

if pre_bias_logits_pspec is not None:

pre_bias_logits_pspec = remove_incompatible_mesh_axes_from_partition_spec(

pre_bias_logits_pspec,

pre_bias_logits.shape,

self.mesh,

dims=(0,),

)

if pre_bias_logits_pspec is not None and pre_bias_logits is not None:

pre_bias_logits_pspec = remove_incompatible_mesh_axes_from_partition_spec(

pre_bias_logits_pspec,

pre_bias_logits.shape,

self.mesh,

dims=(0,),

)

gemini-code-assist · 2026-08-07T00:16:55Z

+      if decoder_tokens_pspec is not None:
+        decoder_tokens_pspec = remove_incompatible_mesh_axes_from_partition_spec(
+            decoder_tokens_pspec,
+            input_ids.shape,
+            self.mesh,
+            dims=(0,),
+        )


If input_ids is None, attempting to access input_ids.shape will raise an AttributeError. To prevent potential runtime crashes, we should defensively check that input_ids is not None before accessing its properties.

Suggested change

if decoder_tokens_pspec is not None:

decoder_tokens_pspec = remove_incompatible_mesh_axes_from_partition_spec(

decoder_tokens_pspec,

input_ids.shape,

self.mesh,

dims=(0,),

)

if decoder_tokens_pspec is not None and input_ids is not None:

decoder_tokens_pspec = remove_incompatible_mesh_axes_from_partition_spec(

decoder_tokens_pspec,

input_ids.shape,

self.mesh,

dims=(0,),

)

gemini-code-assist · 2026-08-07T00:16:55Z

      (
-          rf"{tmvp_config.reasoning_start_token}.+{tmvp_config.reasoning_end_token}.*?"
-          rf"{tmvp_config.solution_start_token}(.+?){tmvp_config.solution_end_token}"
+          rf"{reasoning_start}.+{re.escape(tmvp_config.reasoning_end_token)}.*?"


Using .+ requires at least one character before the reasoning_end_token. If reasoning_start_token_in_prompt is True and the model immediately outputs the closing reasoning tag (i.e., an empty reasoning block), the regex will fail to match. Changing .+ to .*? (non-greedy, 0 or more characters) makes the pattern more robust and prevents matching failures on empty reasoning blocks.

Suggested change

rf"{reasoning_start}.+{re.escape(tmvp_config.reasoning_end_token)}.*?"

rf"{reasoning_start}.*?{re.escape(tmvp_config.reasoning_end_token)}.*??"

NuojCheng · 2026-08-07T23:13:07Z

  return jax.sharding.PartitionSpec(*new_spec)


+def remove_incompatible_mesh_axes_from_partition_spec(pspec, shape, mesh, dims=None):


I am worried this function gonna silently fallback sharding configs instead of shouting out and throwing errors. I suggested adding a flag allow_remove_axes=False and this fallback happens only when the flag is true.

NuojCheng · 2026-08-07T23:13:25Z

+    maybe_shard_with_pspec,
+    remove_expert_from_partition_spec,
+    remove_incompatible_mesh_axes_from_partition_spec,
+)


nice clean up, thank you

NuojCheng · 2026-08-07T23:16:35Z

+          inputs.shape,
+          self.mesh,
+          dims=(0,),
+      )


could you move all logics from 1664-1698 to a separate small function, maybe named maybe_remove_incompatible_mesh_axes_from_partition_spec, so that readers don't need to read this logic unless necessary

NuojCheng

Sharding and MoE components LGTM if my comments are resolved. Thank you for the code cleanup Yixuan!

khatwanimohit · 2026-08-07T23:22:25Z

+  return -1, -1
+
+
+def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Optional[int] = None):


can you use

maxtext/src/maxtext/utils/max_utils.py

Line 1064 in 1555e1f

def unscan_train_state_params(params, sharding, mesh, scan_axis, layer_groups):

for unrolling params ? unless there is something unique about qwen3 scanned weights which I don't know about

NuojCheng · 2026-08-07T23:23:17Z

+    raise ValueError(f"PartitionSpec rank {len(pspec)} exceeds tensor rank {len(shape)}")
+
+  dims_to_check = None if dims is None else set(dims)
+  compatible_pspec = pspec


what if dims < 0?

A9isha · 2026-08-07T22:24:11Z

@@ -0,0 +1,159 @@
+# Copyright 2026 Google LLC


we can skip this file.

A9isha · 2026-08-07T22:33:14Z

+      # Qwen GDN keeps its short convolution history in BF16, but recurrence is
+      # accumulated and persisted in FP32. Declaring both caches as the model
+      # dtype silently quantizes the recurrent state after every generated token.
+      mamba_shapes, mamba_dtypes, unpadded_mamba_page_size = build_qwen_gdn_cache_layout(


but this file is general across all models. You probably want to make this build_qwen_gdn_cache_layout conditional on the model

A9isha · 2026-08-07T22:40:28Z

+from typing import Any
+
+
+def build_qwen_gdn_cache_layout(cfg: Any, torch_module: Any):


is this qwen specific? could other models also need such cache utilities?

A9isha · 2026-08-07T22:57:33Z

      min_timescale: Start of the geometric index (typically 1).
      max_timescale: End of the geometric index (rope_theta, e.g., 1000000).
-      embedding_dims: Dimension of the embedding (head_dim).
+      embedding_dims: Dimension of the attention head.


now that #4764 is merged, please rebase to remove those change from this PR

A9isha · 2026-08-07T23:34:32Z

+    )
+    response = "We calculate 6 * 7.</think>\nThe final answer is 42. do do do"
+
+    self.assertEqual(utils_rl.extract_answer(response, config), "42")


I know the native final answer helps when the model did utter the answer and it is frustrating to not give it the reward. but it was not within answer tags, but if we allow it to put the answer this way it will not learn the format well, I know in the easy gsm8k dataset it is learning the format finally after ~150 training steps

AntonyMei added 11 commits August 6, 2026 11:12

[Fix] Apply partial MRoPE to Qwen3.5 attention

68085ac

Fix scanned Qwen direct weight synchronization

2322620

Fix tensor-parallel MoE direct-sync layout

e80dd3a

Make direct weight-sync validation fail closed

7db5ad9

Preserve mixed precision in Qwen GDN cache

15753ce

Truncate Qwen GDN metadata to active requests

cd7a01e

Disable unsafe prefix caching for recurrent rollouts

c69ce36

Support small batches in Qwen GDN and MoE sharding

8ddbe6e

Handle prefilled reasoning markers in RL rewards

2e0d825

Add clean Qwen3.5 35B GSM8K RL recipe

d61e31e

Forward rollout options through the MaxText adapter

bb10ffe

AntonyMei requested review from NuojCheng, RissyRan, abhinavclemson, aireenmei, bvandermoon, darisoy, gagika, gobbleturk, huytransformer, igorts-git, jiangjy1982, khatwanimohit, michelle-yooh, parambole, richjames0, shralex, shuningjin, vipannalla and xibinliu as code owners August 7, 2026 00:14

AntonyMei requested review from A9isha, SurbhiJainUSC, dipannita08 and hengtaoguo as code owners August 7, 2026 00:14

gemini-code-assist Bot reviewed Aug 7, 2026

View reviewed changes

NuojCheng reviewed Aug 7, 2026

View reviewed changes

NuojCheng approved these changes Aug 7, 2026

View reviewed changes

khatwanimohit reviewed Aug 7, 2026

View reviewed changes

NuojCheng reviewed Aug 7, 2026

View reviewed changes

A9isha reviewed Aug 7, 2026

View reviewed changes

YixuanWang-99 approved these changes Aug 8, 2026

View reviewed changes

RissyRan added the gemini-review label Aug 8, 2026

[Fix] Add final compatibility fixes.

92c2eca

AntonyMei force-pushed the yixuanm-qwen35-rl-fixes-clean branch from 6567fec to 92c2eca Compare August 8, 2026 02:25

	raise ValueError(f"Qwen scanned parameter {'.'.join(key)} has no scan axis {scan_axis}: {value!r}")
	raise ValueError(f"Qwen scanned parameter {'.'.join(map(str, key))} has no scan axis {scan_axis}: {value!r}")

	rf"{reasoning_start}.+{re.escape(tmvp_config.reasoning_end_token)}.*?"
	rf"{reasoning_start}.?{re.escape(tmvp_config.reasoning_end_token)}.??"

		return jax.sharding.PartitionSpec(*new_spec)


		def remove_incompatible_mesh_axes_from_partition_spec(pspec, shape, mesh, dims=None):

		return -1, -1


		def unroll_qwen_scanned_weights(weights, scan_axis: int = 1, pattern_length: Optional[int] = None):

		from typing import Any


		def build_qwen_gdn_cache_layout(cfg: Any, torch_module: Any):

Uh oh!

Conversation

AntonyMei commented Aug 7, 2026

Description

Tests

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng Aug 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

NuojCheng Aug 7, 2026 •

edited

Loading

NuojCheng Aug 7, 2026 •

edited

Loading

NuojCheng left a comment •

edited

Loading