Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / modules / mega_moe / backward / graph.py: 100%
63 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 05:03 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
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# http://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"""MoE-FFN backward compute graph."""
17from typing import Optional
19from hyper_multicore._core_spec import DEFAULT_CORE_SPEC
20from hyper_multicore.scheduler.config import TaskSplitValue
21from hyper_multicore.scheduler.graph import (
22 ComputeGraph, OperatorNode, TensorSpec, SplitSpec, OpType,
23)
24from hyper_multicore.tasks.alltoall import AllToAllFillConfig, AllToAllType
25from hyper_multicore.tasks.gmm import GmmFillConfig
26from hyper_multicore.tasks.swiglu import SwiGLUFillConfig
28# Tiling buffer positions in the C++ runtime descriptor array.
29_TILING_POS_ACT_GRAD = 20
30_TILING_POS_GATE_GRAD = 21
31_TILING_POS_W1_GRAD = 22
32_TILING_POS_W2_GRAD = 23
33_TILING_POS_SWIGLU_GRAD = 24
36def _build_bwd_tensor_specs(
37 tsv, hidden_size, intermediate_size, dtype_size, expert_weight_layout
38):
39 """Create all TensorSpec objects for the backward graph."""
40 sre = tsv.single_rank_expert_num
41 weight_is_eoi = expert_weight_layout == "EOI"
42 target = TensorSpec(
43 "target", [tsv.per_rank_seq, hidden_size],
44 param_position=0, dtype_size=dtype_size, tensor_type=1, is_dynamic=True)
45 target_offset = TensorSpec(
46 "target_offset", [tsv.all_expert_num],
47 param_position=1, dtype_size=8)
48 src = TensorSpec(
49 "src", [tsv.per_rank_seq, hidden_size],
50 param_position=2, dtype_size=dtype_size, tensor_type=1)
51 src_offset = TensorSpec(
52 "src_offset", [tsv.all_expert_num],
53 param_position=3, dtype_size=8)
54 size_d = TensorSpec(
55 "size_d", [tsv.all_expert_num],
56 param_position=4, dtype_size=4)
57 w2_grad_x1 = TensorSpec(
58 "w2_grad_x1", [tsv.per_rank_seq, intermediate_size],
59 param_position=5, dtype_size=dtype_size,
60 tensor_type=1, is_dynamic=True, transpose=not weight_is_eoi)
61 w2_grad_y = TensorSpec(
62 "w2_grad_y", ([sre, hidden_size, intermediate_size]
63 if weight_is_eoi else [sre, intermediate_size, hidden_size]),
64 param_position=6, dtype_size=dtype_size, tensor_type=1)
65 act_grad_weight = TensorSpec(
66 "act_grad_weight",
67 ([sre, hidden_size, intermediate_size]
68 if weight_is_eoi else [sre, intermediate_size, hidden_size]),
69 param_position=7, dtype_size=dtype_size, tensor_type=1,
70 transpose=not weight_is_eoi)
71 act_grad_y = TensorSpec(
72 "act_grad_y", [tsv.per_rank_seq, intermediate_size],
73 param_position=8, dtype_size=dtype_size, tensor_type=1, is_dynamic=True)
74 swiglu_dy = TensorSpec(
75 "swiglu_dy", [tsv.per_rank_seq, intermediate_size * 2],
76 param_position=9, dtype_size=dtype_size, tensor_type=1, is_dynamic=True)
77 swiglu_out = TensorSpec(
78 "swiglu_out", [tsv.per_rank_seq, intermediate_size * 2],
79 param_position=10, dtype_size=dtype_size, tensor_type=1, is_dynamic=True)
80 gate_grad_weight = TensorSpec(
81 "gate_grad_weight",
82 ([sre, intermediate_size * 2, hidden_size]
83 if weight_is_eoi else [sre, hidden_size, intermediate_size * 2]),
84 param_position=11, dtype_size=dtype_size, tensor_type=1,
85 transpose=not weight_is_eoi)
86 gate_grad_y = TensorSpec(
87 "gate_grad_y", [tsv.per_rank_seq, hidden_size],
88 param_position=12, dtype_size=dtype_size, tensor_type=1, is_dynamic=True)
89 combine_out = TensorSpec(
90 "combine_out", [tsv.per_rank_seq, hidden_size],
91 param_position=13, dtype_size=dtype_size, tensor_type=1)
92 target_offset_c = TensorSpec(
93 "target_offset_c", [tsv.all_expert_num],
94 param_position=14, dtype_size=8)
95 src_offset_c = TensorSpec(
96 "src_offset_c", [tsv.all_expert_num],
97 param_position=15, dtype_size=8)
98 size_c = TensorSpec(
99 "size_c", [tsv.all_expert_num],
100 param_position=16, dtype_size=4)
101 w1_grad_x1 = TensorSpec(
102 "w1_grad_x1", [tsv.per_rank_seq, hidden_size],
103 param_position=17, dtype_size=dtype_size,
104 tensor_type=1, is_dynamic=True, transpose=not weight_is_eoi)
105 w1_grad_y = TensorSpec(
106 "w1_grad_y", ([sre, intermediate_size * 2, hidden_size]
107 if weight_is_eoi else [sre, hidden_size, intermediate_size * 2]),
108 param_position=18, dtype_size=dtype_size, tensor_type=1)
109 glist = TensorSpec(
110 "glist", [sre],
111 param_position=19, dtype_size=8)
112 return (target, target_offset, src, src_offset, size_d,
113 w2_grad_x1, w2_grad_y, act_grad_weight, act_grad_y,
114 swiglu_dy, swiglu_out, gate_grad_weight, gate_grad_y,
115 combine_out, target_offset_c, src_offset_c, size_c,
116 w1_grad_x1, w1_grad_y, glist)
119def _build_bwd_ops_first(tsv, specs, *, dispatch_sv, act_grad_sv, w2_grad_sv,
120 swiglu_sv, num_cube_cores, swiglu_limit=None, expert_weight_layout="EIO"):
121 """Create dispatch, act_grad, w2_grad, swiglu_grad operator nodes."""
122 (target, target_offset, src, src_offset, size_d,
123 w2_grad_x1, w2_grad_y, act_grad_weight, act_grad_y,
124 swiglu_dy, swiglu_out, *_) = specs
125 glist = specs[-1]
126 dispatch = OperatorNode(
127 name="dispatch", op_type=OpType.ALLTOALL,
128 diagnostic_name="DispatchGrad",
129 inputs=[target_offset, src, src_offset, size_d], outputs=[target],
130 param_positions=[1, 2, 3, 4, 0], split_value=dispatch_sv,
131 split_spec=SplitSpec(
132 split_inputs=None, split_output_dims=[0],
133 task_num_fn=lambda tsv: tsv.all_expert_num * (tsv.per_expert_seq_to_other // dispatch_sv),
134 ),
135 tiling_position=-1,
136 fill_config=AllToAllFillConfig(moe_type=AllToAllType.DISPATCH, advance="vector",
137 event_group=tsv.all_expert_num),
138 )
139 act_grad = OperatorNode(
140 name="act_grad", op_type=OpType.GMM,
141 diagnostic_name="ActGrad",
142 inputs=[target, act_grad_weight, glist], outputs=[act_grad_y],
143 param_positions=[0, 7, 19, 8], split_value=act_grad_sv,
144 split_spec=SplitSpec(
145 split_inputs=[(0, 0)], split_output_dims=[0],
146 task_num_fn=lambda tsv, _ncc=num_cube_cores: _ncc * tsv.single_rank_expert_num,
147 ),
148 tiling_position=_TILING_POS_ACT_GRAD,
149 fill_config=GmmFillConfig(offset_inputs={0}, rank_in_event=True, global_trigger=False,
150 out_offset=True, advance="cube_only", num_cube_cores=num_cube_cores),
151 )
152 weight_is_eoi = expert_weight_layout == "EOI"
153 w2_inputs = [target, w2_grad_x1, glist] if weight_is_eoi else [w2_grad_x1, target, glist]
154 w2_grad = OperatorNode(
155 name="w2_grad", op_type=OpType.GMM,
156 diagnostic_name="W2Grad",
157 inputs=w2_inputs, outputs=[w2_grad_y],
158 param_positions=[tensor.param_position for tensor in w2_inputs] + [6], split_value=w2_grad_sv,
159 split_spec=SplitSpec(
160 split_inputs=[(0 if weight_is_eoi else 1, 0)], split_output_dims=[0],
161 task_num_fn=lambda tsv, _ncc=num_cube_cores: _ncc * tsv.single_rank_expert_num,
162 ),
163 tiling_position=_TILING_POS_W2_GRAD,
164 fill_config=GmmFillConfig(offset_inputs={0, 1}, transpose_inputs={0} if weight_is_eoi else set(),
165 rank_in_event=True, global_trigger=True,
166 out_offset=False, advance="cube_custom",
167 event_delta=tsv.single_rank_expert_num, num_cube_cores=num_cube_cores),
168 )
169 swiglu_grad = OperatorNode(
170 name="swiglu_grad", op_type=OpType.SWIGLU_GRAD,
171 diagnostic_name="SwiGLUGrad",
172 inputs=[act_grad_y, swiglu_dy], outputs=[swiglu_out],
173 param_positions=[8, 9, 10], split_value=swiglu_sv,
174 split_spec=SplitSpec(
175 split_inputs=[(0, 0)], split_output_dims=[0],
176 task_num_fn=lambda tsv: (tsv.per_expert_seq // swiglu_sv) * tsv.single_rank_expert_num,
177 ),
178 tiling_position=_TILING_POS_SWIGLU_GRAD,
179 fill_config=SwiGLUFillConfig(clamp_limit=swiglu_limit),
180 )
181 return dispatch, act_grad, w2_grad, swiglu_grad
184def _build_bwd_ops_second(specs, *, gate_grad_sv, w1_grad_sv, combine_sv, num_cube_cores,
185 expert_weight_layout="EIO"):
186 """Create gate_grad, combine, w1_grad operator nodes."""
187 (*_, swiglu_out, gate_grad_weight, gate_grad_y,
188 combine_out, target_offset_c, src_offset_c, size_c,
189 w1_grad_x1, w1_grad_y, glist) = specs
190 gate_grad = OperatorNode(
191 name="gate_grad", op_type=OpType.GMM,
192 diagnostic_name="GateGrad",
193 inputs=[swiglu_out, gate_grad_weight, glist], outputs=[gate_grad_y],
194 param_positions=[10, 11, 19, 12], split_value=gate_grad_sv,
195 split_spec=SplitSpec(
196 split_inputs=[(0, 0)], split_output_dims=[0],
197 task_num_fn=lambda tsv, _ncc=num_cube_cores: _ncc * tsv.single_rank_expert_num,
198 ),
199 tiling_position=_TILING_POS_GATE_GRAD,
200 fill_config=GmmFillConfig(offset_inputs={0}, rank_in_event=False, global_trigger=False,
201 out_offset=True, advance="cube", num_cube_cores=num_cube_cores),
202 )
203 combine = OperatorNode(
204 name="combine", op_type=OpType.ALLTOALL,
205 diagnostic_name="CombineGrad",
206 inputs=[target_offset_c, gate_grad_y, src_offset_c, size_c], outputs=[combine_out],
207 param_positions=[14, 12, 15, 16, 13], split_value=combine_sv,
208 split_spec=SplitSpec(
209 split_inputs=[(1, 0)], split_output_dims=[0],
210 task_num_fn=lambda tsv: tsv.all_expert_num * (tsv.per_expert_seq_to_other // combine_sv),
211 ),
212 tiling_position=-1,
213 fill_config=AllToAllFillConfig(moe_type=AllToAllType.COMBINE, advance="vector_only", event_group=1),
214 )
215 weight_is_eoi = expert_weight_layout == "EOI"
216 w1_inputs = [swiglu_out, w1_grad_x1, glist] if weight_is_eoi else [w1_grad_x1, swiglu_out, glist]
217 w1_grad = OperatorNode(
218 name="w1_grad", op_type=OpType.GMM,
219 diagnostic_name="W1Grad",
220 inputs=w1_inputs, outputs=[w1_grad_y],
221 param_positions=[tensor.param_position for tensor in w1_inputs] + [18], split_value=w1_grad_sv,
222 split_spec=SplitSpec(
223 split_inputs=[(0 if weight_is_eoi else 1, 0)], split_output_dims=[0],
224 task_num_fn=lambda tsv, _ncc=num_cube_cores: _ncc * tsv.single_rank_expert_num,
225 ),
226 tiling_position=_TILING_POS_W1_GRAD,
227 fill_config=GmmFillConfig(offset_inputs={0, 1}, transpose_inputs={0} if weight_is_eoi else set(),
228 rank_in_event=False, global_trigger=True,
229 out_offset=False, advance="cube_custom",
230 event_delta=1, num_cube_cores=num_cube_cores),
231 )
232 return gate_grad, combine, w1_grad
235def build_backward_graph(tsv: TaskSplitValue, *,
236 dispatch_sv: int = 128,
237 act_grad_sv: int = 4096,
238 w2_grad_sv: int = 4096,
239 swiglu_sv: int = 128,
240 gate_grad_sv: int = 4096,
241 w1_grad_sv: int = 4096,
242 combine_sv: int = 128,
243 hidden_size: int = 7168,
244 intermediate_size: int = 2048,
245 dtype_size: int = 2,
246 num_cube_cores: int = DEFAULT_CORE_SPEC.num_cube_cores,
247 swiglu_limit: Optional[float] = None,
248 expert_weight_layout: str = "EIO") -> ComputeGraph:
249 """Build the MoE-FFN backward DAG.
251 Execution order:
252 dispatch -> {act_grad, w2_grad} -> swiglu_grad -> gate_grad -> {combine, w1_grad}
253 act_grad and w2_grad run in parallel (both depend on dispatch).
254 combine and w1_grad run in parallel (both depend on gate_grad).
256 Param positions (C++ memory slots):
257 dispatch: [target_offset=1, src=2, src_offset=3, size=4 | target=0]
258 act_grad: [x=0, weight=7, glist=19 | y=8]
259 w2_grad EIO: [x1=5, x2=0, glist=19 | y=6]
260 w2_grad EOI: [x1=0, x2=5, glist=19 | y=6]
261 swiglu_grad: [x=8, dy=9 | out=10]
262 gate_grad: [x=10, weight=11, glist=19 | y=12]
263 combine: [target_offset=14, src=12, src_offset=15, size=16 | target=13]
264 w1_grad EIO: [x1=17, x2=10, glist=19 | y=18]
265 w1_grad EOI: [x1=10, x2=17, glist=19 | y=18]
267 EOI writes gradients directly in parameter layout. A transpose view itself
268 does not copy data; returning it can require a later layout copy during
269 gradient accumulation.
271 Args:
272 tsv: TaskSplitValue carrying TP/EP/seq partition metadata.
273 dispatch_sv / act_grad_sv / w2_grad_sv / swiglu_sv / gate_grad_sv /
274 w1_grad_sv / combine_sv: tile sizes per operator.
275 hidden_size: model hidden dimension.
276 intermediate_size: FFN intermediate dimension after SwiGLU halving.
277 dtype_size: bytes per activation element (2=bf16, 4=fp32).
278 num_cube_cores: number of AIC cube cores on the target device.
279 swiglu_limit: optional positive clamp limit for SwiGLU; ``None`` keeps
280 the original unclamped operator path.
281 expert_weight_layout: Logical expert-weight axis order, either ``"EIO"`` or ``"EOI"``.
283 Returns:
284 A fully-connected ComputeGraph ready for propagate_splits().
285 """
286 if expert_weight_layout not in ("EIO", "EOI"):
287 raise ValueError(f"expert_weight_layout must be 'EIO' or 'EOI', got {expert_weight_layout!r}.")
288 specs = _build_bwd_tensor_specs(
289 tsv, hidden_size, intermediate_size, dtype_size, expert_weight_layout
290 )
291 dispatch, act_grad, w2_grad, swiglu_grad = _build_bwd_ops_first(
292 tsv, specs,
293 dispatch_sv=dispatch_sv, act_grad_sv=act_grad_sv,
294 w2_grad_sv=w2_grad_sv, swiglu_sv=swiglu_sv, num_cube_cores=num_cube_cores,
295 swiglu_limit=swiglu_limit,
296 expert_weight_layout=expert_weight_layout,
297 )
298 gate_grad, combine, w1_grad = _build_bwd_ops_second(
299 specs,
300 gate_grad_sv=gate_grad_sv, w1_grad_sv=w1_grad_sv,
301 combine_sv=combine_sv, num_cube_cores=num_cube_cores,
302 expert_weight_layout=expert_weight_layout,
303 )
304 graph = ComputeGraph()
305 (graph.add_op(dispatch).add_op(act_grad).add_op(w2_grad)
306 .add_op(swiglu_grad).add_op(gate_grad).add_op(combine).add_op(w1_grad)
307 .add_edge(dispatch, act_grad)
308 .add_edge(dispatch, w2_grad)
309 .add_edge(act_grad, swiglu_grad)
310 .add_edge(swiglu_grad, gate_grad)
311 .add_edge(gate_grad, combine)
312 .add_edge(gate_grad, w1_grad))
313 return graph