DSV4: Add mhc_split_axis_contraction to contract the mHC rate and embed axes separately

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -4094,13 +4094,16 @@ def reshape_transpose_o_a(input_tensor, target_shape=None):
 
   # Functions for mHC split
   def mhc_split_fn_pre(input_tensor, target_shape=None):
-    return np.transpose(input_tensor[0:4, :])
+    tensor = np.transpose(input_tensor[0:4, :])
+    return tensor.reshape(target_shape) if target_shape is not None else tensor
 
   def mhc_split_fn_post(input_tensor, target_shape=None):
-    return np.transpose(input_tensor[4:8, :])
+    tensor = np.transpose(input_tensor[4:8, :])
+    return tensor.reshape(target_shape) if target_shape is not None else tensor
 
   def mhc_split_fn_res(input_tensor, target_shape=None):
-    return np.transpose(input_tensor[8:24, :])
+    tensor = np.transpose(input_tensor[8:24, :])
+    return tensor.reshape(target_shape) if target_shape is not None else tensor
 
   def mhc_split_base_pre(input_tensor, target_shape=None):
     return input_tensor[0:4]
@@ -4169,7 +4172,7 @@ def mhc_split_scale_res(input_tensor, target_shape=None):
     def mhc_concat_fn(input_tensors, target_shape=None):
       if len(input_tensors) != 3:
         raise ValueError(f"mhc_concat_fn expected 3 tensors (pre, post, res), got {len(input_tensors)}")
-      tensors = [np.asarray(t) for t in input_tensors]
+      tensors = [np.asarray(t).reshape((-1, t.shape[-1])) for t in input_tensors]
       res = np.transpose(np.concatenate(tensors, axis=1))
       return res.reshape(target_shape) if target_shape is not None else res
 

src/maxtext/configs/base.yml

@@ -1348,6 +1348,7 @@ mhc_pallas_kernel_fwd_block_size: 256
 # TPU v7x if we go higher. For TPU v6 (Trillium), there is more VMEM, so 256
 # works and provides better results.
 mhc_pallas_kernel_bwd_block_size: 128
+mhc_split_axis_contraction: False
 
 ################################## DeepSeek Engram ##################################
 # Indices of transformer layers where Engram are integrated; leave empty [] to disable.

src/maxtext/configs/types.py

@@ -1775,6 +1775,14 @@ def validate_mhc_kernel(self) -> "ManifoldConstrainedHyperConnections":
       raise ValueError("use_mhc_pallas_kernel=True requires enable_mhc_lite=True.")
     return self
 
+  mhc_split_axis_contraction: bool = Field(
+      False,
+      description=(
+          "Whether to contract the mHC rate and embed axes separately instead of flattening them, "
+          "so the activation's TP-sharded embed dim is never all-gathered."
+      ),
+  )
+
 
 class DilocoParams(BaseModel):
   """Diloco Hyperparameters"""

src/maxtext/layers/mhc.py

@@ -75,6 +75,25 @@ def sinkhorn(t, iters=20):
   return t.astype(initial_dtype)
 
 
+class _SplitAxesRMSNorm(nnx.Module):
+  """RMS normalization across the mHC rate and embedding axes."""
+
+  def __init__(self, rate: int, dim: int, dtype, weight_dtype, epsilon: float, rngs: nnx.Rngs):
+    self.dtype = dtype
+    self.epsilon = epsilon
+    self.scale = nnx.Param(
+        jax.nn.initializers.ones(rngs.params(), (rate, dim), weight_dtype),
+        out_sharding=(None, "activation_embed"),
+    )
+
+  def __call__(self, x: Array) -> Array:
+    x = jnp.asarray(x, jnp.float32)
+    mean2 = jnp.mean(jax.lax.square(x), axis=(-2, -1), keepdims=True)
+    y = jnp.asarray(x * jax.lax.rsqrt(mean2 + self.epsilon), self.dtype)
+    scale = jnp.asarray(self.scale.get_value(), self.dtype)
+    return y * scale
+
+
 class ManifoldConstrainedHyperConnections(nnx.Module):
   """Implements Manifold-Constrained Hyper-Connections (mHC).
 
@@ -107,15 +126,24 @@ def __init__(
     if getattr(self.config, "use_mhc_pallas_kernel", False) and not self.config.enable_mhc_lite:
       raise ValueError("use_mhc_pallas_kernel=True requires enable_mhc_lite=True.")
 
-    # Norm layer
-    self.mhc_norm = RMSNorm(
-        num_features=self.k * self.dim,
-        dtype=self.config.dtype,
-        weight_dtype=self.weight_dtype,
-        kernel_axes=("norm",),
-        epsilon=self.config.normalization_layer_epsilon,
-        rngs=self.rngs,
-    )
+    if self.config.mhc_split_axis_contraction:
+      self.mhc_norm = _SplitAxesRMSNorm(
+          rate=self.k,
+          dim=self.dim,
+          dtype=self.config.dtype,
+          weight_dtype=self.weight_dtype,
+          epsilon=self.config.normalization_layer_epsilon,
+          rngs=self.rngs,
+      )
+    else:
+      self.mhc_norm = RMSNorm(
+          num_features=self.k * self.dim,
+          dtype=self.config.dtype,
+          weight_dtype=self.weight_dtype,
+          kernel_axes=("norm",),
+          epsilon=self.config.normalization_layer_epsilon,
+          rngs=self.rngs,
+      )
 
     # Scalars
     self.res_alpha_scale = nnx.Param(
@@ -142,15 +170,23 @@ def __init__(
       res_beta_shape = (self.k, self.k)
       res_beta_sharding = (None, None)
 
-    # Weight matrices
     scale_init = nd_dense_init(1.0, "fan_in", "normal")
-    in_axis = 0
-    out_axis = 1
-    weight_sharding_axis_name = ("activation_embed", None)
+    if self.config.mhc_split_axis_contraction:
+      in_axis = (0, 1)
+      out_axis = 2
+      weight_sharding_axis_name = (None, "activation_embed", None)
+      res_alpha_shape = (self.k, self.dim, res_out_dim)
+      alpha_shape = (self.k, self.dim, self.k)
+    else:
+      in_axis = 0
+      out_axis = 1
+      weight_sharding_axis_name = ("activation_embed", None)
+      res_alpha_shape = (self.k * self.dim, res_out_dim)
+      alpha_shape = (self.k * self.dim, self.k)
     self.res_alpha = nnx.Param(
         scale_init(
             self.rngs.params(),
-            (self.k * self.dim, res_out_dim),
+            res_alpha_shape,
             self.weight_dtype,
             in_axis=in_axis,
             out_axis=out_axis,
@@ -160,7 +196,7 @@ def __init__(
     self.pre_alpha = nnx.Param(
         scale_init(
             self.rngs.params(),
-            (self.k * self.dim, self.k),
+            alpha_shape,
             self.weight_dtype,
             in_axis=in_axis,
             out_axis=out_axis,
@@ -170,7 +206,7 @@ def __init__(
     self.post_alpha = nnx.Param(
         scale_init(
             self.rngs.params(),
-            (self.k * self.dim, self.k),
+            alpha_shape,
             self.weight_dtype,
             in_axis=in_axis,
             out_axis=out_axis,
@@ -286,8 +322,11 @@ def __call__(
       )
     else:
       with jax.named_scope("mhc_norm"):
-        # 1. Flatten the tensor, and RMS normalization
-        norm_x = self.mhc_norm(jnp.reshape(x, (b, s, k * d)))
+        if self.config.mhc_split_axis_contraction:
+          norm_x = self.mhc_norm(x)
+        else:
+          # 1. Flatten the tensor, and RMS normalization
+          norm_x = self.mhc_norm(jnp.reshape(x, (b, s, k * d)))
 
       # Fused Projections
       pre_alpha = jnp.asarray(self.pre_alpha[...], self.dtype)
@@ -297,7 +336,11 @@ def __call__(
       alpha_concat = jnp.concatenate([pre_alpha, post_alpha, res_alpha], axis=-1)
 
       # MatMul on normalized input
-      h_concat = jnp.einsum("bsm,mn -> bsn", norm_x, alpha_concat, precision=self.matmul_precision)
+      if self.config.mhc_split_axis_contraction:
+        # Keeping the TP-sharded embed axis separate avoids gathering it before contraction.
+        h_concat = jnp.einsum("bskd,kdn -> bsn", norm_x, alpha_concat, precision=self.matmul_precision)
+      else:
+        h_concat = jnp.einsum("bsm,mn -> bsn", norm_x, alpha_concat, precision=self.matmul_precision)
       h_pre = h_concat[..., : self.k]
       h_post = h_concat[..., self.k : 2 * self.k]
       h_res = h_concat[..., 2 * self.k :]

src/maxtext/utils/muon_utils.py

@@ -140,6 +140,16 @@ def get_muon_weight_dimension_numbers(model, config, verbose=False):
     def apply_transform_nnx(path: Tuple[jax.tree_util.KeyEntry, ...], leaf):
       # Convert jax.tree_util.KeyEntry path to Tuple[str, ...]
       path_strings = tuple(p.key for p in path if isinstance(p, jax.tree_util.DictKey))
+      if (
+          config is not None
+          and getattr(config, "mhc_split_axis_contraction", False)
+          and _is_path_contain_any(("mhc_attention", "mhc_mlp"), path_strings)
+          and _is_path_contain_any(("pre_alpha", "post_alpha", "res_alpha"), path_strings)
+      ):
+        ndim = len(leaf.shape)
+        scan_axis = getattr(config, "param_scan_axis", 1) % ndim if ndim == 4 else None
+        reduction_axes = tuple(axis for axis in range(ndim - 1) if axis != scan_axis)
+        return mdn(reduction_axes, (-1,))
       return transform_logic(path_strings)
 
     # NNX abstract_param is an nnx.State (not Linen's dict of LogicallyPartitioned leaves);

tests/unit/mhc_test.py

@@ -106,6 +106,7 @@ def _setup_mhc(
       sequence_length=7,
       per_device_batch_size=None,
       dtype=None,
+      mhc_split_axis_contraction=False,
   ):
     """Sets up the common configurations and modules for MHC testing."""
     self.dim = dim
@@ -128,6 +129,7 @@ def _setup_mhc(
         "base_emb_dim": self.dim,
         "mhc_expansion_rate": rate,
         "enable_mhc_lite": enable_mhc_lite,
+        "mhc_split_axis_contraction": mhc_split_axis_contraction,
         "use_mhc_pallas_kernel": use_mhc_pallas_kernel,
         "decoder_block": "deepseek",
         "num_experts": 4,
@@ -614,6 +616,58 @@ def forward_kernel(x):
           atol=5e-2,
       )
 
+  @parameterized.named_parameters(("Sinkhorn", False), ("Lite", True))
+  def test_split_axis_contraction_matches_baseline(self, enable_mhc_lite):
+    outputs = {}
+    for split in (False, True):
+      self._setup_mhc(4, enable_mhc_lite=enable_mhc_lite, mhc_split_axis_contraction=split)
+      with nn_partitioning.axis_rules(self.config.logical_axis_rules):
+        module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs)
+        norm_scale = jnp.arange(self.config.mhc_expansion_rate * self.dim, dtype=self.config.weight_dtype) / 32 + 0.5
+        module.mhc_norm.scale[...] = jnp.reshape(norm_scale, module.mhc_norm.scale.shape)
+        layer = linears.MlpBlock(
+            config=self.config,
+            mesh=self.mesh,
+            in_features=self.config.emb_dim,
+            intermediate_dim=self.config.moe_mlp_dim,
+            activations=self.config.mlp_activations,
+            intermediate_dropout_rate=self.config.dropout_rate,
+            dtype=self.config.dtype,
+            weight_dtype=self.config.weight_dtype,
+            model_mode=self.config.model_call_mode,
+            rngs=self.rngs,
+        )
+        out, _ = module(self.pre_norm, layer, x=self.x, mhc_type=HyperConnectionType.MLP_DENSE)
+        outputs[split] = np.asarray(out, dtype=np.float32)
+
+    np.testing.assert_allclose(outputs[False], outputs[True], rtol=1e-5, atol=1e-5)
+
+  def test_split_axis_param_layout(self):
+    params = {}
+    for split in (False, True):
+      self._setup_mhc(4, mhc_split_axis_contraction=split)
+      with nn_partitioning.axis_rules(self.config.logical_axis_rules):
+        module = mhc.ManifoldConstrainedHyperConnections(self.config, self.dim, self.mesh, self.rngs)
+      params[split] = {
+          "norm": np.asarray(module.mhc_norm.scale[...]),
+          "res": np.asarray(module.res_alpha[...]),
+          "pre": np.asarray(module.pre_alpha[...]),
+          "post": np.asarray(module.post_alpha[...]),
+      }
+
+    k, d = self.config.mhc_expansion_rate, self.dim
+    self.assertEqual(params[False]["norm"].shape, (k * d,))
+    self.assertEqual(params[False]["res"].shape, (k * d, k * k))
+    self.assertEqual(params[False]["pre"].shape, (k * d, k))
+    self.assertEqual(params[False]["post"].shape, (k * d, k))
+    self.assertEqual(params[True]["norm"].shape, (k, d))
+    self.assertEqual(params[True]["res"].shape, (k, d, k * k))
+    self.assertEqual(params[True]["pre"].shape, (k, d, k))
+    self.assertEqual(params[True]["post"].shape, (k, d, k))
+
+    for name in ("norm", "res", "pre", "post"):
+      np.testing.assert_array_equal(params[True][name], params[False][name].reshape(params[True][name].shape))
+
 
 def _get_permutation_matrices(k: int) -> jax.Array:
   """Generates all permutation matrices for k streams."""

tests/unit/muon_utils_test.py

@@ -173,6 +173,19 @@ def test_nnx_model_dispatches_to_tree_map_with_path(self):
     # 'w_standard' does not trigger any special rule → standard mdn.
     self.assertEqual(result["w_standard"], mdn((0,), (-1,)))
 
+  def test_split_mhc_alphas_reduce_rate_and_embed_axes(self):
+    model = nnx.Module()
+    model.mhc_attention = nnx.Module()
+    model.mhc_attention.pre_alpha = nnx.Param(jnp.ones((4, 8, 4)))
+    model.mhc_mlp = nnx.Module()
+    model.mhc_mlp.res_alpha = nnx.Param(jnp.ones((4, 2, 8, 16)))
+    config = mock.Mock(mhc_split_axis_contraction=True, param_scan_axis=1)
+
+    result = muon_utils.get_muon_weight_dimension_numbers(model, config)
+
+    self.assertEqual(result["mhc_attention"]["pre_alpha"], mdn((0, 1), (-1,)))
+    self.assertEqual(result["mhc_mlp"]["res_alpha"], mdn((0, 2), (-1,)))
+
   def test_nnx_verbose_path_executes_print_debug(self):
     """verbose=True should also execute _print_structure_debug without raising."""
     buf = io.StringIO()

tests/unit/param_mapping_test.py

@@ -245,6 +245,23 @@ def getter(name):
         np.testing.assert_array_equal(out[f"b{b}_l{l}"], value_of(b, l))
 
   # Specific tests with assertions
+  def test_deepseek4_mhc_reshaped_hooks(self):
+    config = {"num_hidden_layers": 1, "n_routed_experts": 2}
+    maxtext_config = mock.Mock()
+    prefix = "params-decoder-layers_0-mhc_attention"
+    keys = tuple(f"{prefix}-{name}_alpha" for name in ("pre", "post", "res"))
+    hf_weight = np.arange(24 * 8, dtype=np.float32).reshape(24, 8)
+
+    from_hf = param_mapping.DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config)
+    split_weights = (
+        from_hf[keys[0]](hf_weight, (4, 2, 4)),
+        from_hf[keys[1]](hf_weight, (4, 2, 4)),
+        from_hf[keys[2]](hf_weight, (4, 2, 16)),
+    )
+
+    to_hf = param_mapping.DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, saving_to_hf=True)
+    np.testing.assert_array_equal(to_hf[keys](split_weights, hf_weight.shape), hf_weight)
+
   def test_reshape_kernel_hook(self):
     config = {
         "text_config": {"num_hidden_layers": 2, "hidden_size": 256},