Skip to content

Commit bf22ce9

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request #4163 from AI-Hypercomputer:chengnuojin-2k-scale
PiperOrigin-RevId: 936265019
2 parents ae15d1c + 3c9f5c4 commit bf22ce9

7 files changed

Lines changed: 3481 additions & 0 deletions

File tree

‎docs/guides/optimization/custom_mesh_and_rule.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,10 @@ This rule utilizes the `data`, `stage`, `fsdp`, and `expert` axes. Its defining
5151

5252
Similar in philosophy to `ep-as-cp.yml`, this configuration explicitly includes the `context` axis in the mesh layout alongside `data`, `stage`, `fsdp`, and `expert`. While context sharding is mapped to the `context` axis globally, within MoE components, this `context` axis dynamically shifts to perform expert parallelism instead of FSDP. This custom rule supports using CP and EP together.
5353

54+
### `ep-as-dp.yml`
55+
56+
Different with the rule in `base.yml`, this rule configures expert physical axis to function as data parallelism rather than FSDP. This removes the constraint where FSDPxEP is limited by specific model dimensions, particularly for small tensors such as attention projections. Ultimately, this change benefits large-scale training.
57+
5458
### `pipeline-large-moe.yml`
5559

5660
Designed specifically to optimize pipeline parallelism for extremely large-scale MoE jobs (such as DeepSeek models). It defines the physical axes: `data`, `stage`, `fsdp`, `tensor`, `context`, and `expert`. To prevent dimension limit errors, it intentionally disables expert weight sharding on the (typically small) `q_lora` dimension. Furthermore, tensor and expert parallelism are strictly preserved to support advanced pipelining features like `pipeline_fsdp_ag_one` and `pipeline_fsdp_ag_per_repeat`.

‎src/maxtext/common/common_types.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,3 +153,4 @@ class CustomRule(enum.Enum):
153153
EP_AS_CP = "ep-as-cp" # Support EP only
154154
PIPELINE_LARGE_MOE = "pipeline-large-moe"
155155
FSDP_2D = "2d-fsdp"
156+
EP_AS_DP = "ep-as-dp"

‎src/maxtext/configs/custom_mesh_and_rule/ep-as-dp.yml‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# This rule uses data, FSDP, FSDP_transpose and expert. Expert axis acts as
16+
# data parallelism in components except core MoE part (between EP all2all).
17+
18+
mesh_axes: ['data', 'fsdp', 'fsdp_transpose', 'expert']
19+
data_sharding: [['data', 'fsdp', 'fsdp_transpose', 'expert']]
20+
logical_axis_rules: [
21+
# ==========================================
22+
# Vocabulary Embedding
23+
# ==========================================
24+
# Vocab Activations
25+
['activation_embed_and_logits_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
26+
['activation_embed_and_logits_batch_sequence', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
27+
['activation_vocab', []],
28+
# Vocab Weights
29+
['vocab', []],
30+
['embed_vocab', ['fsdp', 'fsdp_transpose']],
31+
# ==========================================
32+
# Attention
33+
# ==========================================
34+
# Attention Activations
35+
['activation_batch_attn', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
36+
['activation_heads', []],
37+
['activation_kv_heads', []],
38+
['activation_length_attn', ['context']],
39+
['activation_q_length', ['context']],
40+
['activation_kv_length', []],
41+
['activation_embed_attn', ['tensor', 'tensor_transpose']],
42+
['activation_kv', []],
43+
['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
44+
['activation_kv_head_dim', []],
45+
# Attention Weights
46+
['q_lora', ['fsdp']],
47+
["q_lora_up_proj", []],
48+
['kv_lora', ['fsdp']],
49+
["kv_lora_up_proj", []],
50+
# ==========================================
51+
# Mixture of Experts (MoE)
52+
# ==========================================
53+
# MoE Activations
54+
['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
55+
['activation_length_moe', []],
56+
['activation_norm_length_moe', []],
57+
['activation_embed_moe', []],
58+
['activation_mlp_moe', []],
59+
['activation_exp', ['expert']],
60+
# MoE Weights
61+
['exp', 'expert'],
62+
['mlp_moe', ['fsdp_transpose']],
63+
['embed_moe', ['fsdp']],
64+
# ==========================================
65+
# Standard MLP / Dense Layers / Model Structure
66+
# ==========================================
67+
# Dense Activations
68+
['activation_mlp', []],
69+
# Note activation batch and length also get used in vocab
70+
['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']],
71+
['activation_length', []],
72+
['activation_norm_length', []],
73+
['activation_embed', []],
74+
['activation_stage', []],
75+
# General Weights
76+
['mlp', ['fsdp_transpose']],
77+
['embed', ['fsdp', 'fsdp_transpose']],
78+
['embed', ['fsdp']],
79+
# ==========================================
80+
# Deprecated / Scheduled for Removal
81+
# ==========================================
82+
['exp_with_fsdp', 'fsdp'],
83+
]

‎tests/utils/sharding_dump.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,13 @@
6969
"2d-fsdp",
7070
("ici_fsdp_parallelism=-1", "ici_fsdp_transpose_parallelism=2"),
7171
),
72+
(
73+
"deepseek2-16b",
74+
"tpu7x-16",
75+
1,
76+
"ep-as-dp",
77+
("ici_fsdp_parallelism=-1", "ici_expert_parallelism=2", "use_ring_of_experts=true"),
78+
),
7279
("qwen3-0.6b", "tpu7x-16", 1, "", ()),
7380
("gpt-oss-20b", "tpu7x-16", 1, "", ()),
7481
("gpt-oss-20b", "tpu7x-16", 1, "", ("ici_fsdp_parallelism=-1", "ici_expert_parallelism=2")),

‎tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/rule_ep-as-dp_ici_fsdp_parallelism=-1_ici_expert_parallelism=2_use_ring_of_experts=true/input_shardings.json‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
{
2+
"Activation Sharding Dump": [
3+
{
4+
"deepseek/inputs: bfloat16[192,2048,2048]": {
5+
"logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
6+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
7+
}
8+
},
9+
{
10+
"deepseek/pre_attention_norm: bfloat16[192,2048,2048]": {
11+
"logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
12+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
13+
}
14+
},
15+
{
16+
"attention_mla/inputs_q: bfloat16[192,2048,2048]": {
17+
"logic_axes": "('activation_batch_attn', 'activation_length', 'activation_embed')",
18+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
19+
}
20+
},
21+
{
22+
"attention_mla/inputs_kv: bfloat16[192,2048,2048]": {
23+
"logic_axes": "('activation_batch_attn', 'activation_length', 'activation_embed')",
24+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
25+
}
26+
},
27+
{
28+
"attention_mla/q_nope: bfloat16[192,2048,16,128]": {
29+
"logic_axes": "('activation_kv_batch', 'activation_length', 'activation_kv_heads', 'activation_kv_head_dim')",
30+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
31+
}
32+
},
33+
{
34+
"attention_mla/q_pe: bfloat16[192,2048,16,64]": {
35+
"logic_axes": "('activation_kv_batch', 'activation_length', 'activation_kv_heads', 'activation_kv_head_dim')",
36+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
37+
}
38+
},
39+
{
40+
"attention_mla/query: bfloat16[192,2048,16,192]": {
41+
"logic_axes": "('activation_kv_batch', 'activation_length', 'activation_kv_heads', 'activation_kv_head_dim')",
42+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
43+
}
44+
},
45+
{
46+
"attention_mla/key_nope: bfloat16[192,2048,16,128]": {
47+
"logic_axes": "('activation_kv_batch', 'activation_length', 'activation_kv_heads', 'activation_kv_head_dim')",
48+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
49+
}
50+
},
51+
{
52+
"attention_mla/key_rope: bfloat16[192,2048,16,64]": {
53+
"logic_axes": "('activation_kv_batch', 'activation_length', 'activation_kv_heads', 'activation_kv_head_dim')",
54+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
55+
}
56+
},
57+
{
58+
"attention_mla/key: bfloat16[192,2048,16,192]": {
59+
"logic_axes": "('activation_kv_batch', 'activation_length', 'activation_kv_heads', 'activation_kv_head_dim')",
60+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
61+
}
62+
},
63+
{
64+
"attention_mla/value: bfloat16[192,2048,16,128]": {
65+
"logic_axes": "('activation_kv_batch', 'activation_length', 'activation_kv_heads', 'activation_kv_head_dim')",
66+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
67+
}
68+
},
69+
{
70+
"attention_op/arr: int8[1,4,4]": {
71+
"logic_axes": "Unknown",
72+
"PartitionSpec": "P(None, None)"
73+
}
74+
},
75+
{
76+
"attention_op/arr: int32[2048]": {
77+
"logic_axes": "Unknown",
78+
"PartitionSpec": "P(None,)"
79+
}
80+
},
81+
{
82+
"attention_op/query: bfloat16[192,16,2048,192]": {
83+
"logic_axes": "Unknown",
84+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
85+
}
86+
},
87+
{
88+
"attention_op/key: bfloat16[192,16,2048,192]": {
89+
"logic_axes": "Unknown",
90+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
91+
}
92+
},
93+
{
94+
"attention_op/value: bfloat16[192,16,2048,128]": {
95+
"logic_axes": "Unknown",
96+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
97+
}
98+
},
99+
{
100+
"attention_mla/out: bfloat16[192,2048,16,128]": {
101+
"logic_axes": "('activation_batch_attn', 'activation_length', 'activation_heads', 'activation_kv')",
102+
"PartitionSpec": "P(('fsdp', 'expert'), None, None, None)"
103+
}
104+
},
105+
{
106+
"deepseek/attention_result: bfloat16[192,2048,2048]": {
107+
"logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
108+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
109+
}
110+
},
111+
{
112+
"deepseek/post_attention_norm: bfloat16[192,2048,2048]": {
113+
"logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
114+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
115+
}
116+
},
117+
{
118+
"linears/x: bfloat16[192,2048,10944]": {
119+
"logic_axes": "('activation_batch', 'activation_length', 'activation_mlp')",
120+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
121+
}
122+
},
123+
{
124+
"deepseek/mlp: bfloat16[192,2048,2048]": {
125+
"logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
126+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
127+
}
128+
},
129+
{
130+
"deepseek/x: bfloat16[192,2048,2048]": {
131+
"logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
132+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
133+
}
134+
},
135+
{
136+
"moe/inputs: bfloat16[192,2048,2048]": {
137+
"logic_axes": "('activation_batch', 'activation_norm_length', None)",
138+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
139+
}
140+
},
141+
{
142+
"moe/gate_logits: bfloat16[192,2048,64]": {
143+
"logic_axes": "('activation_batch', 'activation_norm_length', None)",
144+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
145+
}
146+
},
147+
{
148+
"moe/w0_kernel: bfloat16[64,2048,1408]": {
149+
"logic_axes": "Unknown",
150+
"PartitionSpec": "P('expert', None, None)"
151+
}
152+
},
153+
{
154+
"moe/w1_kernel: bfloat16[64,2048,1408]": {
155+
"logic_axes": "Unknown",
156+
"PartitionSpec": "P('expert', None, None)"
157+
}
158+
},
159+
{
160+
"moe/wo_kernel: bfloat16[64,1408,2048]": {
161+
"logic_axes": "Unknown",
162+
"PartitionSpec": "P('expert', None, None)"
163+
}
164+
},
165+
{
166+
"linears/x: bfloat16[192,2048,2816]": {
167+
"logic_axes": "('activation_batch', 'activation_length', 'activation_mlp')",
168+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
169+
}
170+
},
171+
{
172+
"deepseek/mlp_lnx: bfloat16[192,2048,2048]": {
173+
"logic_axes": "('activation_batch', 'activation_norm_length', 'activation_embed')",
174+
"PartitionSpec": "P(('fsdp', 'expert'), None, None)"
175+
}
176+
}
177+
]
178+
}

0 commit comments

Comments
 (0)