Fix DeepSeek4HyperHead compilation error

[olufiyin19]

Fix DeepSeek4HyperHead compilation error

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Code Review

This pull request introduces the DeepSeek4HyperHead module to support DeepSeek V4 hyper heads, which reduces the expansion rate dimension and applies a series of dense layers. This head is integrated into the decoder initialization and call paths when mhc_expansion_rate > 1 and the decoder block is DEEPSEEK4. Feedback on these changes highlights two main issues: first, accessing config.mhc_expansion_rate directly can raise an AttributeError on configurations where it is not defined, so using getattr is recommended; second, the DenseGeneral layers in DeepSeek4HyperHead should be configured with proper sharding, precision, and offloading parameters to prevent potential out-of-memory (OOM) errors and ensure consistency.

[gemini-code-assist]

Accessing config.mhc_expansion_rate directly will raise an AttributeError for any model configuration that does not define mhc_expansion_rate (such as standard LLaMA or Gemma models). Using getattr(config, "mhc_expansion_rate", 1) prevents this crash and safely defaults to 1 when the attribute is missing.

Suggested change

	if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4:
	if getattr(config, "mhc_expansion_rate", 1) > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4:

[gemini-code-assist]

The DenseGeneral layers inside DeepSeek4HyperHead are initialized without specifying logical kernel_axes, shard_mode, matmul_precision, and parameter_memory_host_offload. Without these, the weight matrices will not be sharded (replicated instead), which can lead to high memory usage and potential Out-Of-Memory (OOM) errors during large-scale training. Additionally, the layers will not respect the user's configuration for sharding, precision, and offloading. Specifying these parameters ensures proper FSDP sharding and consistency with the rest of the model.

            linears.DenseGeneral(
                in_features_shape=self.dim,
                out_features_shape=self.dim,
                dtype=self.dtype,
                weight_dtype=self.weight_dtype,
                kernel_axes=("embed", None),
                shard_mode=config.shard_mode,
                matmul_precision=config.matmul_precision,
                parameter_memory_host_offload=config.parameter_memory_host_offload,
                rngs=self.rngs,
            )

[codecov]

Codecov Report

❌ Patch coverage is 23.80952% with 16 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/mhc.py	31.25%	11 Missing ⚠️
src/maxtext/layers/nnx_decoders.py	0.00%	4 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!