Fix DeepSeek4HyperHead missing from maxtext.layers.mhc - #4772
Conversation
|
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[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 adds a mock failure log for checkpoint validation and implements the DeepSeek4HyperHead class in src/maxtext/layers/mhc.py to collapse hyper-connection streams. Feedback suggests sharding the first dimension of the self.hc_fn parameter along the "activation_embed" axis to prevent potential out-of-memory (OOM) issues caused by replicating the entire weight matrix across all devices.
| in_axis=0, | ||
| out_axis=1, | ||
| ), | ||
| out_sharding=(None, None), |
There was a problem hiding this comment.
The parameter self.hc_fn is currently initialized with out_sharding=(None, None), which replicates the entire weight matrix across all devices. Since its first dimension is self.hc_mult * config.emb_dim (which can be quite large), replicating it can lead to significant memory overhead and potential OOMs during training. Sharding the first dimension along the "activation_embed" axis is more efficient and aligns with how other projection weights (like self.res_alpha, self.pre_alpha, and self.post_alpha) are sharded in this file.
| out_sharding=(None, None), | |
| out_sharding=("activation_embed", None), |
1 participant
Fix DeepSeek4HyperHead missing from maxtext.layers.mhc