DSV4: Add mhc_split_axis_contraction to contract the mHC rate and embed axes separately - #4912
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the mhc_split_axis_contraction configuration option, which allows contracting the mHC rate and embedding axes separately instead of flattening them, thereby avoiding all-gathering the TP-sharded embedding dimension. To support this, a new _SplitAxesRMSNorm layer is added, and the parameter shapes, initializations, and matmuls in ManifoldConstrainedHyperConnections are adjusted accordingly. Additionally, Muon optimizer utilities and checkpoint conversion hooks are updated, and comprehensive unit tests are introduced. The review feedback suggests converting input tensors to numpy arrays before accessing .shape in mhc_concat_fn to prevent potential AttributeErrors, and using the standard [...] syntax instead of get_value() on nnx.Param for consistency and compatibility.
| 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] |
There was a problem hiding this comment.
Accessing t.shape directly on the elements of input_tensors before converting them to numpy arrays can raise an AttributeError if any element is a list or other non-array-like object. It is safer and more robust to convert to a numpy array first and then access .shape on the converted array.
| tensors = [np.asarray(t).reshape((-1, t.shape[-1])) for t in input_tensors] | |
| tensors = [arr.reshape((-1, arr.shape[-1])) for arr in map(np.asarray, input_tensors)] |
| 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) |
There was a problem hiding this comment.
Using self.scale.get_value() is inconsistent with the rest of the file (and the MaxText codebase), which uses the standard [...] syntax (e.g., self.pre_alpha[...], self.mhc_norm.scale[...]) to access nnx.Param values. Additionally, get_value() is non-standard in Flax NNX and may cause compatibility issues or runtime errors depending on the Flax version. Please use self.scale[...] instead.
| scale = jnp.asarray(self.scale.get_value(), self.dtype) | |
| scale = jnp.asarray(self.scale[...], self.dtype) |
…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 <[email protected]> Co-authored-by: Armin <[email protected]> Co-authored-by: utlz <[email protected]>
d936232 to
34e0d48
Compare
2 participants
Description
Adds
mhc_split_axis_contraction, an opt-in path that keeps the mHC rate and TP-sharded embedding axes separate through RMS normalization and the fused alpha projection. The split contractions replace flatten-induced activation gathers with reductions over the sharded embedding axis.When enabled, the checkpoint layout of the mHC parameters changes: the norm scale uses
(k, dim)instead of(k * dim,), and the alphas use(k, dim, n)instead of(k * dim, n). Their initialized values equal the flat initialization reshaped, so converting a flag-off checkpoint is a pure reshape.The flag defaults to false, preserving the existing RMSNorm module, parameter shapes, RNG draw order, and contraction path. DeepSeek-V4 Hugging Face conversion reshapes alpha weights for either layout, and Muon contracts both split input axes.
Performance
Step time was 7.711 s/step with
mhc_split_axis_contractiondisabled and 6.811 s/step with it enabled, a 1.13x result.Reproduction setup: v6e-128, DeepSeek-V4-Flash 284B, LoRA fine-tuning,
ici_expert_parallelism=16,ici_tensor_parallelism=8,ici_fsdp_parallelism=1,max_target_length=16384,per_device_batch_size=0.125(global batch 16),LIBTPU_INIT_ARGS=--xla_tpu_scoped_vmem_limit_kib=98304, dynamic-splash attention path enabled, withmhc_split_axis_contractiontoggled.Tests
JAX_PLATFORMS=cpu pytest -q tests/unit/mhc_test.py -k split_axis— 3 passed; the split path matches Sinkhorn and mHC-lite outputs, and its initialized parameters are exact reshapes of the flat layout.JAX_PLATFORMS=cpu pytest -q tests/unit/mhc_test.py— 27 passed, 2 skipped.JAX_PLATFORMS=cpu pytest -q tests/unit/param_mapping_test.py tests/unit/muon_utils_test.py— 53 passed; split-layout Hugging Face conversion round-trips and Muon uses the intended contraction axes.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.