From 34e0d48503855caca4b254d4952bc69ed20676c4 Mon Sep 17 00:00:00 2001 From: jcarin-sys Date: Mon, 17 Aug 2026 13:53:59 +0000 Subject: [PATCH] Add mhc_split_axis_contraction: contract the mHC rate and embed axes separately The flattened (rate, embed) axis forces GSPMD to all-gather the activation's TP-sharded embed dim. Contracting the axes separately keeps the activation sharded; the weights are created in the consumed layout with values equal to the flat init reshaped. Measured 7.711 to 6.811 s/step (1.13x). Co-authored-by: Sudarsanan Co-authored-by: Armin Co-authored-by: utlz --- .../utils/param_mapping.py | 11 ++- src/maxtext/configs/base.yml | 1 + src/maxtext/configs/types.py | 8 ++ src/maxtext/layers/mhc.py | 81 ++++++++++++++----- src/maxtext/utils/muon_utils.py | 10 +++ tests/unit/mhc_test.py | 54 +++++++++++++ tests/unit/muon_utils_test.py | 13 +++ tests/unit/param_mapping_test.py | 17 ++++ 8 files changed, 172 insertions(+), 23 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 94e96173f1..b7d287c1db 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/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 diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index a4cf6877e3..7761316a06 100644 --- a/src/maxtext/configs/base.yml +++ b/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. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1f8659f7e0..86aa20fbd0 100644 --- a/src/maxtext/configs/types.py +++ b/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""" diff --git a/src/maxtext/layers/mhc.py b/src/maxtext/layers/mhc.py index 817def4dc0..494c2eeb4e 100644 --- a/src/maxtext/layers/mhc.py +++ b/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 :] diff --git a/src/maxtext/utils/muon_utils.py b/src/maxtext/utils/muon_utils.py index ff77c57807..26b51a03a2 100644 --- a/src/maxtext/utils/muon_utils.py +++ b/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); diff --git a/tests/unit/mhc_test.py b/tests/unit/mhc_test.py index 921004ed78..c5e64d4517 100644 --- a/tests/unit/mhc_test.py +++ b/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.""" diff --git a/tests/unit/muon_utils_test.py b/tests/unit/muon_utils_test.py index 58bfadf29a..a5f2b72309 100644 --- a/tests/unit/muon_utils_test.py +++ b/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() diff --git a/tests/unit/param_mapping_test.py b/tests/unit/param_mapping_test.py index cea6485817..f8e004ec4a 100644 --- a/tests/unit/param_mapping_test.py +++ b/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},