Onboard explicit sharding to deepseek [split version]

[NuojCheng]

Description

This PR introduces explicit sharding support for the DeepSeek model in MaxText. This allows for more granular control over how tensors are distributed across devices, which may lead to better performance and scalability.

This PR is a simplified version based on #2579 .

Key Changes:

• Enabled DeepSeek for Explicit Sharding: The deepseek decoder is now supported in ShardMode.EXPLICIT.
• Introduced shard_mode: A shard_mode parameter has been added to various layers, including attention, embeddings, and MoE, to control sharding behavior. This allows for more flexible and explicit sharding configurations.
• Refactored Sharding Logic: The existing sharding logic has been (minorly) refactored to use the new sharding utilities and NamedSharding objects, making the sharding implementation more explicit and maintainable.
• Updated Tests: The test suite has been updated to include tests for explicit sharding, ensuring the correctness and robustness of the implementation.

Detailed Changes by File:

• src/MaxText/common_types.py: Added Q_LORA_UP_PROJ and KV_LORA_UP_PROJ to support LoRA projections with explicit sharding.
• src/MaxText/configs/types.py: Enabled the deepseek decoder for explicit sharding.
• src/MaxText/layers/attention_mla.py: Updated to use out_sharding and _maybe_shard_with_logical for more explicit control over sharding in MLA.
• src/MaxText/layers/attentions.py: Added the shard_mode parameter to RotaryEmbedding and its variants. Refactored input and projection sharding to be more explicit.
• src/MaxText/layers/deepseek.py: Integrated explicit sharding into the DeepSeekDecoderLayer by using _maybe_shard_with_logical and passing out_sharding and intermediate_sharding to sub-layers.
• src/MaxText/layers/embeddings.py: Added the shard_mode parameter to all embedding classes to allow for explicit sharding configuration.
• src/MaxText/layers/moe.py: Added shard_mode to GateLogit and updated the MoeBlock to handle explicit sharding for weights and activations.
• src/MaxText/sharding.py: Introduced new utility functions for explicit sharding.
• tests/*: Updated various tests to include test cases for explicit sharding and to pass the new shard_mode and mesh parameters where necessary.

Tests

Model_name: deepseek3-test
Topology: v5p-8
JAX==0.8.1
TL;DR:

• Auto Sharding: Maintains performance parity with the main branch (deviations ≤ 0.3%).
• Explicit Sharding: While this introduces a slight expected overhead in TP case, it provides significant improvements for TP_transpose sharding.

Performance Impact Table

Configuration	Main Peak Memory (Mb)	Main Step Time (ms)	Auto Peak Memory (Mb)	Auto Peak Mem Change (%)	Auto Step Time (ms)	Auto Step Time Change (%)	Explicit Peak Memory (Mb)	Explicit Peak Mem Change (%)	Explicit Step Time (ms)	Explicit Step Time Change (%)
FSDP4	68748	2238	68748	0%	2238	0%	68961	0.31%	2284	2.01%
FSDP2+TP2	69447	2009	69448	0%	2014	0.25%	69617	0.24%	2378	18.37%
FSDP2+TPT2	93835	10742	93836	0%	10743	0.01%	69373	-26.07%	2714	-74.68%

Full FSDP

smoke_train model_name=deepseek3-test per_device_batch_size=1

• 0ea9b55: xprof
• shard_mode=auto: xprof
• shard_mode=explicit: xprof

FSDP + TP

smoke_train model_name=deepseek3-test per_device_batch_size=1 ici_tensor_parallelism=2

• 0ea9b55: xprof
• shard_mode=auto: xprof
• shard_mode=explicit: xprof

FSDP + TP_transpose

smoke_train model_name=deepseek3-test per_device_batch_size=1 ici_tensor_transpose_parallelism=2

• 0ea9b55: xprof
• shard_mode=auto: xprof
• shard_mode=explicit: xprof

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.

[RissyRan]

Thanks for simplifying it! Are you planning to merge this one instead?

If you could directly apply a few minor moments from old PR (just published). We should be good to go.

[NuojCheng]

Reply to @RissyRan:

Are you planning to merge this one instead?

Yes let's merge this one instead. The sharding refactoring parts are removed here.

If you could directly apply a few minor moments from old PR (just published). We should be good to go.

Thank you for the comments! The concerns from the previous PR are all related to the sharding refactoring, which are not included in this PR.

[github-actions]

🤖 Hi @NuojCheng, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This PR introduces explicit sharding for the DeepSeek model, which is a significant and valuable feature for performance and scalability. The changes are extensive and well-structured. I've identified a couple of minor logical issues and potential improvements in the implementation.

🔍 General Feedback

• The introduction of shard_mode and the refactoring to use _maybe_shard_with_logical is a clean way to handle explicit sharding.
• The tests have been updated comprehensively, which is great to see.
• Pay close attention to copy-paste errors, especially when dealing with similar variables like inputs_q and inputs_kv.

Overall, this is a solid contribution. Once the minor issues are addressed, this PR will be in great shape.

[github-actions]

🟡 The `mla_query_projection` function has a logic issue. The `return query` statement is inside the `if/else` block, making the following constraints unreachable. It seems like the return statement was misplaced during refactoring.

I've moved the return statement to the end of the function and removed the now-unnecessary final _maybe_shard_with_logical on query as it's already sharded correctly within the conditional blocks.

[NuojCheng]

I don't think the return query in line 544 is in an if-else statement. It is a little bit misleading.

[github-actions]

🟡 There appears to be a copy-paste error here. `inputs_kv` is being sharded using `inputs_q`'s sharding information. It should use its own.

[NuojCheng]

Good catch! Updated.

[github-actions]

🟢 Using `jnp.einsum` for complex multiplication might be less clear and potentially less performant than direct element-wise multiplication with `*`. The previous implementation used `*`, which is more idiomatic for this operation.

[NuojCheng]

Good idea. I will do jnp.broadcast_to and jnp.multiply to replace the jnp.einsum.