[Stacked PR 1/5] Fix upstream pyink formatting and GMM v2 compatibility for Qwen3.5

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py

@@ -1090,7 +1090,7 @@ def validate_inputs(
 
   assert group_offset.shape == (1,)
 
-  size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype)
+  size_lhs_sublane = max(pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype), 16)
   size_lhs_sublane = min(size_lhs_sublane, size_m)
   if fuse_act is not None:
     num_lanes = pltpu.get_tpu_info().num_lanes

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py

@@ -221,6 +221,9 @@ def make_tgmm_configs(
   size_lhs_sublane = min(size_lhs_sublane, size_m)
   size_rhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(rhs.dtype)
   size_rhs_sublane = min(size_rhs_sublane, size_m)
+  common_sublane = min(size_lhs_sublane, size_rhs_sublane)
+  size_lhs_sublane = common_sublane
+  size_rhs_sublane = common_sublane
   assert size_lhs_sublane == size_rhs_sublane, (
       f"size_lhs_sublane should be the same as size_rhs_sublane {lhs.dtype=}," f" {rhs.dtype=}"
   )

src/maxtext/models/qwen3.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Qwen3 family of model decoder layers."""
+
 # pylint: disable=arguments-differ
 # pylint: disable=no-name-in-module
 
@@ -624,9 +625,7 @@ def _explicit_activation_shardings(self, batch: int):
     cp_len = LENGTH if gdn_context_axes(self.config) else None
 
     def _sharding(logical_axes):
-      pspec = logical_to_mesh_axes(
-          logical_axes, mesh=self.mesh, rules=logical_rules
-      )
+      pspec = logical_to_mesh_axes(logical_axes, mesh=self.mesh, rules=logical_rules)
       # Training microbatches can be smaller than the physical batch partition.
       # Only dim 0 is inspected, so the trailing sizes are placeholders.
       shape = (batch,) + (1,) * (len(logical_axes) - 1)
@@ -654,9 +653,7 @@ def __call__(
     # hidden_states: (B, S, E)
     cfg = self.config
     batch, seq_len, _ = hidden_states.shape
-    flat_sharding, head_sharding, state_sharding = (
-        self._explicit_activation_shardings(batch)
-    )
+    flat_sharding, head_sharding, state_sharding = self._explicit_activation_shardings(batch)
 
     active_cache = kv_cache if kv_cache is not None else self.cache
 
@@ -759,13 +756,9 @@ def __call__(
     b_raw, a_raw = jnp.split(mixed_ba, split_indices_ba, axis=3)
 
     # b: (B, S, H_v)
-    b = jnp.reshape(
-        b_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding
-    )
+    b = jnp.reshape(b_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding)
     # a: (B, S, H_v)
-    a = jnp.reshape(
-        a_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding
-    )
+    a = jnp.reshape(a_raw, (batch, seq_len, self.num_v_heads), out_sharding=flat_sharding)
 
     if use_paged_state:
       # =========================================================================
@@ -1051,9 +1044,7 @@ def extract_state(c_in, v_len):
         value = jax.sharding.reshard(value, qkv_pspec)
         g = jax.sharding.reshard(g, g_beta_pspec)
         beta = jax.sharding.reshard(beta, g_beta_pspec)
-        recurrent_state_arg = jax.sharding.reshard(
-            recurrent_state_arg, state_pspec
-        )
+        recurrent_state_arg = jax.sharding.reshard(recurrent_state_arg, state_pspec)
 
       @functools.partial(
           jax.shard_map,
@@ -1127,15 +1118,11 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h):
     # The normalization and gating is applied per-head on the value dimension.
 
     # Apply the norm and gate. Output shape: (B, S, H_v, D_v)
-    gated_output_reshaped = self.norm(
-        core_attn_out, z, out_sharding=head_sharding
-    )
+    gated_output_reshaped = self.norm(core_attn_out, z, out_sharding=head_sharding)
 
     # Reshape back to a single feature dimension for the final projection.
     # Shape from (B, S, H_v, D_v) -> (B, S, value_dim)
-    gated_output = jnp.reshape(
-        gated_output_reshaped, (batch, seq_len, -1), out_sharding=flat_sharding
-    )
+    gated_output = jnp.reshape(gated_output_reshaped, (batch, seq_len, -1), out_sharding=flat_sharding)
 
     # Final output shape: (B, S, E)
     output = self.out_proj(gated_output, out_sharding=out_sharding)
@@ -1692,9 +1679,7 @@ def __init__(
     # Physical shardings used to pin sublayer outputs under ShardMode.EXPLICIT. In
     # ShardMode.AUTO the callees ignore these and let GSPMD infer the layout.
     if cfg.shard_mode == ShardMode.EXPLICIT:
-      self.out_sharding = create_sharding(
-          mesh, self.activation_axis_names, rules=get_logical_axis_rules()
-      )
+      self.out_sharding = create_sharding(mesh, self.activation_axis_names, rules=get_logical_axis_rules())
       self.mlp_intermediate_sharding = create_sharding(
           mesh, self.mlp_activation_axis_names, rules=get_logical_axis_rules()
       )
@@ -1777,15 +1762,11 @@ def __call__(
 
     # First LayerNorm, applied before the attention block.
     hidden_states = self.input_layernorm(inputs, out_sharding=self.out_sharding)
-    hidden_states = self._maybe_shard_with_logical(
-        hidden_states, self.activation_axis_names
-    )
+    hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names)
 
     # Conditionally apply either the Linear Attention or Full Attention block.
     if isinstance(self.attention, Qwen3NextFullAttention):
-      attention_output, new_kv_cache = cast(
-          Qwen3NextFullAttention, self.attention
-      )(
+      attention_output, new_kv_cache = cast(Qwen3NextFullAttention, self.attention)(
           hidden_states,
           decoder_segment_ids,
           decoder_positions,
@@ -1796,9 +1777,7 @@ def __call__(
           out_sharding=self.out_sharding,
       )
     else:
-      attention_output, new_kv_cache = cast(
-          Qwen3NextGatedDeltaNet, self.attention
-      )(
+      attention_output, new_kv_cache = cast(Qwen3NextGatedDeltaNet, self.attention)(
           hidden_states,
           model_mode=model_mode,
           kv_cache=kv_cache,
@@ -1808,24 +1787,16 @@ def __call__(
       )
 
     # First residual connection after attention
-    attention_output = self._maybe_shard_with_logical(
-        attention_output, self.activation_axis_names
-    )
+    attention_output = self._maybe_shard_with_logical(attention_output, self.activation_axis_names)
     hidden_states = residual + attention_output
-    hidden_states = self._maybe_shard_with_logical(
-        hidden_states, self.activation_axis_names
-    )
+    hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names)
 
     # Prepare for the MoE block by capturing the new residual
     residual = hidden_states
 
     # Second LayerNorm, applied before the MoE block.
-    hidden_states = self.post_attention_layernorm(
-        hidden_states, out_sharding=self.out_sharding
-    )
-    hidden_states = self._maybe_shard_with_logical(
-        hidden_states, self.activation_axis_names
-    )
+    hidden_states = self.post_attention_layernorm(hidden_states, out_sharding=self.out_sharding)
+    hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names)
 
     # Instantiate and call our `Qwen3NextSparseMoeBlock`.
     mlp_output, load_balance_loss = self.mlp(
@@ -1841,9 +1812,7 @@ def __call__(
       self.moe_lb_loss = nnx.Intermediate(load_balance_loss)
 
     # Final residual connection (after the MoE block)
-    mlp_output = self._maybe_shard_with_logical(
-        mlp_output, self.activation_axis_names
-    )
+    mlp_output = self._maybe_shard_with_logical(mlp_output, self.activation_axis_names)
     layer_output = residual + mlp_output
     layer_output = self._maybe_shard_with_logical(
         layer_output,

src/maxtext/models/qwen3_5.py

@@ -144,9 +144,7 @@ def __init__(
     # Physical shardings used to pin sublayer outputs under ShardMode.EXPLICIT. In
     # ShardMode.AUTO the callees ignore these and let GSPMD infer the layout.
     if cfg.shard_mode == ShardMode.EXPLICIT:
-      self.out_sharding = create_sharding(
-          mesh, self.activation_axis_names, rules=get_logical_axis_rules()
-      )
+      self.out_sharding = create_sharding(mesh, self.activation_axis_names, rules=get_logical_axis_rules())
       self.mlp_intermediate_sharding = create_sharding(
           mesh, self.mlp_activation_axis_names, rules=get_logical_axis_rules()
       )
@@ -226,15 +224,11 @@ def __call__(
 
     # First LayerNorm, applied before the attention block.
     hidden_states = self.input_layernorm(inputs, out_sharding=self.out_sharding)
-    hidden_states = self._maybe_shard_with_logical(
-        hidden_states, self.activation_axis_names
-    )
+    hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names)
 
     # Conditionally apply either the Linear Attention or Full Attention block.
     if isinstance(self.attention, Qwen3_5FullAttention):
-      attention_output, new_kv_cache = cast(
-          Qwen3_5FullAttention, self.attention
-      )(
+      attention_output, new_kv_cache = cast(Qwen3_5FullAttention, self.attention)(
           hidden_states,
           decoder_segment_ids,
           decoder_positions,
@@ -245,9 +239,7 @@ def __call__(
           out_sharding=self.out_sharding,
       )
     else:
-      attention_output, new_kv_cache = cast(
-          Qwen3_5GatedDeltaNet, self.attention
-      )(
+      attention_output, new_kv_cache = cast(Qwen3_5GatedDeltaNet, self.attention)(
           hidden_states,
           model_mode=model_mode,
           kv_cache=kv_cache,
@@ -257,24 +249,16 @@ def __call__(
       )
 
     # First residual connection after attention
-    attention_output = self._maybe_shard_with_logical(
-        attention_output, self.activation_axis_names
-    )
+    attention_output = self._maybe_shard_with_logical(attention_output, self.activation_axis_names)
     hidden_states = residual + attention_output
-    hidden_states = self._maybe_shard_with_logical(
-        hidden_states, self.activation_axis_names
-    )
+    hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names)
 
     # Prepare for the MoE block by capturing the new residual
     residual = hidden_states
 
     # Second LayerNorm, applied before the MoE block.
-    hidden_states = self.post_attention_layernorm(
-        hidden_states, out_sharding=self.out_sharding
-    )
-    hidden_states = self._maybe_shard_with_logical(
-        hidden_states, self.activation_axis_names
-    )
+    hidden_states = self.post_attention_layernorm(hidden_states, out_sharding=self.out_sharding)
+    hidden_states = self._maybe_shard_with_logical(hidden_states, self.activation_axis_names)
 
     # Instantiate and call our `Qwen3_5SparseMoEBlock`.
     mlp_output, load_balance_loss = self.mlp(
@@ -291,9 +275,7 @@ def __call__(
       self.sow(nnx.Intermediate, "moe_lb_loss", load_balance_loss)
 
     # Final residual connection (after the MoE block)
-    mlp_output = self._maybe_shard_with_logical(
-        mlp_output, self.activation_axis_names
-    )
+    mlp_output = self._maybe_shard_with_logical(mlp_output, self.activation_axis_names)
     layer_output = residual + mlp_output
     layer_output = self._maybe_shard_with_logical(
         layer_output,

tests/integration/train_tests.py

@@ -222,9 +222,7 @@ class TrainTests(unittest.TestCase):
       # The Qwen3.5 model configs default to a HuggingFace tokenizer that is not
       # vendored in the repo; use the checked-in tiktoken asset instead.
       "tokenizer_type=tiktoken",
-      (
-          rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}"
-      ),
+      (rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}"),
   ]
 
   # Same sublayers as Qwen3.5, wired together by Qwen3NextScannableBlock rather than a
@@ -1072,26 +1070,20 @@ def test_tpu_qwen3_hybrid_explicit_sharding_matches_auto(self):
     for decoder_block, model_overrides in self._QWEN3_HYBRID_MODELS.items():
       with self.subTest(decoder_block=decoder_block):
         args = parallelism[decoder_block]
-        auto_losses = self._losses(
-            f"{decoder_block}_auto", model_overrides, args + ["shard_mode=auto"]
-        )
+        auto_losses = self._losses(f"{decoder_block}_auto", model_overrides, args + ["shard_mode=auto"])
         explicit_losses = self._losses(
             f"{decoder_block}_explicit",
             model_overrides,
             args + ["shard_mode=explicit"],
         )
         print(f"[{decoder_block}] auto losses: {auto_losses}", flush=True)
-        print(
-            f"[{decoder_block}] explicit losses: {explicit_losses}", flush=True
-        )
+        print(f"[{decoder_block}] explicit losses: {explicit_losses}", flush=True)
         self.assertTrue(auto_losses, "auto run produced no metrics")
         # `activation_batch` carries the expert axis, so pinning it reassociates the
         # backward reductions: the forward pass is bit-for-bit and the drift only appears
         # once gradients flow. Over 20 steps it stays below 4e-5 relative and changes
         # sign, i.e. it is float noise rather than the two runs pulling apart.
-        np.testing.assert_allclose(
-            explicit_losses, auto_losses, rtol=1e-4, atol=0.0
-        )
+        np.testing.assert_allclose(explicit_losses, auto_losses, rtol=1e-4, atol=0.0)
 
   @pytest.mark.integration_test
   @pytest.mark.tpu_only
@@ -1127,8 +1119,7 @@ def test_tpu_qwen3_hybrid_zero1_gradient_accumulation(self):
         sharded = self._losses(
             f"{decoder_block}_ga_zero1",
             model_overrides,
-            zero1_ga
-            + ["shard_mode=explicit", "shard_optimizer_over_data=True"],
+            zero1_ga + ["shard_mode=explicit", "shard_optimizer_over_data=True"],
         )
         print(f"[{decoder_block}] auto + GA losses: {baseline}", flush=True)
         print(

tests/unit/pyconfig_test.py

@@ -285,14 +285,10 @@ def initialize(decoder_block=decoder_block, **kwargs):
         with self.assertRaisesRegex(Exception, "requires `sparse_matmul=True`"):
           initialize(sparse_matmul=False, megablox=False)
 
-        with self.assertRaisesRegex(
-            Exception, "does not support context parallelism"
-        ):
+        with self.assertRaisesRegex(Exception, "does not support context parallelism"):
           initialize(ici_context_parallelism=2)
 
-        with self.assertRaisesRegex(
-            Exception, "does not support context parallelism"
-        ):
+        with self.assertRaisesRegex(Exception, "does not support context parallelism"):
           initialize(ici_context_usp_ulysses_parallelism=2)
 
   def test_explicit_sharding_mistral_decoder_support(self):