[denis-mil]

Note: this PR also fixed #5682 (see comment #5179 (comment) for details)

Bug report: #5195

Description

Motivation:
The reference path had its own mapping helper that cast alpha_scale and
beta down to self.dtype before computing sigmoid(alpha_scale * h + beta),
while the Pallas path called common.compute_sigmoid_gate, which keeps the
scale and bias in float32. With dtype: bfloat16 the two branches therefore
gated on different numbers, so toggling use_mhc_pallas_kernel changed the
layer's output beyond accumulation noise.

Implementation:
Drop ManifoldConstrainedHyperConnections.mapping and call the shared
compute_sigmoid_gate from both pre- and post-mapping, casting the float32
result back to self.dtype so the following GEMM is unchanged. Export the
helper from maxtext.kernels.mhc so the layer can reach it.

Tests:

• test_pallas_kernel_matches_reference_path runs the same weights through
  both branches and asserts they agree; the _tpu variant repeats it against
  the Mosaic-compiled kernel instead of interpret mode.
• test_sigmoid_gate_computes_in_float32 pins the helper's precision directly,
  since the end-to-end comparison cannot resolve a bfloat16 regression there.
• Fold the four copies of the interpret-mode kernel patching into the
  _interpreted_mhc_kernel context manager, and the repeated module/MLP
  construction into _build_mhc_and_mlp.

Tests

Unit tested on CPU and TPU v6e-1

Checklist

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

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• 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 refactors the Manifold-Constrained Hyper Connections (mHC) implementation by sharing the compute_sigmoid_gate helper function between the layer and the Pallas kernel, ensuring identical gating behavior across both paths. The local mapping method in src/maxtext/layers/mhc.py has been removed in favor of this shared helper, which computes in float32 and is subsequently cast back to the activation dtype. Additionally, the unit tests in tests/unit/mhc_test.py have been refactored to reduce duplication using helper methods and a new context manager for kernel interpretation, and new tests have been added to verify that the Pallas kernel matches the reference path and that the sigmoid gate computes correctly in float32. There are no review comments, so no further feedback is provided.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[matrawybarte-lgtm]

هل تم انتهاء المشكله

[dandragona]

LGTM

[denis-mil]

@dandragona thanks for the review.

CI job failed

https://github.com/AI-Hypercomputer/maxtext/actions/runs/36363999320/job/109071898923?pr=5179

Root Cause

• The PR adds the tpu_only test (mhc_test.py:537-540) which runs the real Mosaic kernel with interpret=False. The interpret-mode twin never lowers to Mosaic, so it avoids this issue entirely.
• The test's mesh (mhc_test.py:191-192) covers every device, which evaluates to 4 chips on a v6e-4 environment.
• The call to mhc_kernel.post(...) at mhc.py:342 ultimately ends in a bare pl.pallas_call (mhc_kernels_fwd.py:181) without a shard_map anywhere in the execution path.
• Because layer_out comes from the MLP (which applies sharding constraints over the 4-device mesh), it reaches the kernel already sharded. JAX subsequently refuses to auto-partition the Mosaic custom call.

Note: The test is new, but the problem is not. Enabling use_mhc_pallas_kernel=True would break in the exact same way during real multi-chip training. This new test simply exposes the existing limitation.

Fix Options

1. The Real Fix: Wrap the pre/post kernel calls in jax.shard_map over the mesh, sharding on the batch/sequence dimensions and replicating on the weights. I will create a separate issue to track this implementation.
2. Immediate Action: Remove the TPU test for now, as resolving the sharding limitation is out of scope for this specific PR.

[denis-mil]

Hi @RissyRan @dandragona @bvandermoon,

I've pushed new changes that fix the tests that were failing during the recent CI runs. Could you please take a look at the updates and approve the pending CI workflows so the tests can execute?

Thanks!

Description

Lets the mHC Pallas kernels (use_mhc_pallas_kernel=True) run on a multi-device mesh. Before this change the kernel path only worked on a single device. Each kernel call now runs inside a token-sharded jax.shard_map.

Problem

A bare pallas_call lowers to a Mosaic custom call. On any mesh with more than one device, the mHC kernel path fails with:

Mosaic kernels cannot be automatically partitioned

Solution

Every mHC operation (the RMS norm, the projections and the gates) works within a single token, reducing only over streams * embedding. The kernels can therefore run independently on each batch/length shard, with the weights replicated and the embedding dimension kept whole. shard_map makes this partitioning explicit.

Making the kernels run correctly inside shard_map(check_vma=True) required three supporting changes:

1. Permutation matrices become a compile-time constant. They are no longer a kernel argument. common.permutation_matrices(streams) returns a cached, read-only NumPy array that the kernels build internally, so it is never a traced operand under jit / shard_map / custom_vjp. The layer's reference (non-kernel) path uses the same function, so there is one source for the matrices.
2. Output vma. With check_vma=True, pallas_call requires every out_shape to state the manual mesh axes it varies on. The new common.token_out_shape(shape, dtype, like) copies that information from the token-major input. Outside a shard_map it is a plain ShapeDtypeStruct.
3. Weight-gradient reduction. Each shard produces a weight gradient that is only a partial sum over its own tokens. A custom-VJP rule must return cotangents that vary on the same axes as the primal. The new common.psum_to_primal_vma(grad, primal) therefore psums each weight gradient over the extra axes after the coefficient backward kernel. Outside a shard_map it does nothing.

Implementation

• layers/mhc.py:
  • _kernel_token_axes() maps ("activation_batch", "activation_length") to mesh axes through the logical axis rules. It returns None when there is no mesh or the tokens are not sharded.
  • _sharded_kernel_call() falls back to a direct call in that case. Otherwise it wraps the kernel in jax.shard_map: activations get P(*token_axes, None, ...) and weights get P().
  • _kernel_pre() / _kernel_post() wrap mhc_kernel.pre / post. They split MhcContext into its arrays (x, h_post, residual) to cross the shard_map boundary and rebuild it on the other side.
• kernels/mhc/mhc_kernels_fwd.py and mhc_kernels_bwd.py:
  • All pallas_call output shapes use token_out_shape.
  • The permutations parameter is removed throughout. _pre_op's nondiff_argnums changes from (0, 2) to (0,).
  • _coeff_bwd reduces the weight gradients with psum_to_primal_vma.
• kernels/mhc/common.py: adds token_out_shape, psum_to_primal_vma and permutation_matrices. validate_inputs no longer takes or checks permutations_shape.
• kernels/mhc/api.py: pre() no longer takes permutations and no longer calls stop_gradient on them.

Tests

New and updated tests in tests/unit/mhc_test.py:

• TestMhcKernelsBwd.test_token_sharded_backward_matches_unsharded_tpu (new, tpu_only):
  • Needs an even number of devices, at least 2, and builds a ("data", "fsdp") mesh.
  • Runs pre and post each in its own shard_map, the same way _sharded_kernel_call does: tokens sharded over both axes, weights replicated.
  • Checks the gradients of x and every weight leaf against an unsharded run, with atol = 2% × max|expected|.
• TestMHC.test_pallas_kernel_matches_reference_path_tpu (kept, tpu_only): checks the Mosaic-compiled kernel against the reference layer path.
• The existing forward/backward kernel tests are updated to the new pre() signature.