« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025 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"""Backward overhead module"""
17from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
20class _BackwardOverhead:
21 """Backward overhead class"""
23 def __init__(self, backbone, ccfg, ctx, inner_dyn_fun):
24 self.backbone = backbone
25 self._ccfg = ccfg
26 self._ctx = ctx
27 self._inner_dynamic_mem = inner_dyn_fun
29 def _fetch_node_and_switch_env(
30 self, stages, record_lay_types, stage_id, chunk_id, lay_id
31 ):
32 """Fetch one node and restore the evaluation context for it."""
33 # if negative index access, convert to positive index
34 if stage_id < 0:
35 stage_id = len(stages) + stage_id
36 if chunk_id < 0:
37 chunk_id = len(stages[stage_id]) + chunk_id
38 if lay_id < 0:
39 lay_id = len(stages[stage_id][chunk_id]) + lay_id
40 ccfg, ctx, hook = record_lay_types[(stage_id, chunk_id, lay_id)]
41 # print("here-",id(ccfg))
42 # print("here0",id(self._ccfg))
43 # print("here00",id(self.backbone._ccfg))
44 self._ccfg = ccfg
45 self._ctx = ctx
46 self._ctx.current_stage_id = stage_id
47 self._ctx.current_chunk_id = chunk_id
48 self._ctx.current_lay_id = lay_id
49 # print("here1",id(self._ccfg))
50 self.backbone.apply_hook(hook, ccfg=self._ccfg, ctx=self._ctx)
51 # print("here11",id(self.backbone._ccfg))
52 return stages[stage_id][chunk_id][lay_id]
54 def estimate(
55 self, stages: list, stage_id: int, record_lay_types: dict
56 ) -> float:
57 """estimate stage's end-of-warmup overhead"""
58 # self._ctx.enable_node_log = False
59 res = 0
60 if self._ccfg.pp_sched in ["1f1b", "seqpipe", "seqsmartvpp"]:
61 res = self.__stage_bwd_overhead_1f1b(
62 stages, stage_id, record_lay_types
63 )
64 if self._ccfg.pp_sched == "zero_bubble_v":
65 res = self.__stage_bwd_overhead_zbv(
66 stages, stage_id, record_lay_types
67 )
68 return res
70 def vpp_1f1b_steady_overhead(self, stage_id, dyn_mem_i):
71 """potential overhead due to imbalanced chunks"""
73 def dyn(chunk_id):
74 """chunk total mem"""
75 return sum(dyn_mem_i[chunk_id])
77 # less mem
78 micro_left = self._ccfg.m - self._ccfg.p
79 vpp = self._ccfg.vp
80 max_overhead = 0
81 if self._ctx.vpp_less_mem:
82 top_triangles_chunks = [(vpp - 1 - v, v) for v in range(vpp // 2)]
83 bot_triangles_chunks = [
84 (vpp - 2 - v, v) for v in range((vpp - 1) // 2)
85 ]
86 print("top_triangles_chunks", top_triangles_chunks)
87 print("bot_triangles_chunks", bot_triangles_chunks)
88 for c0, c1 in top_triangles_chunks:
89 overhead = (micro_left - stage_id) * abs(dyn(c0) - dyn(c1))
90 max_overhead = max(max_overhead, overhead)
91 for c0, c1 in bot_triangles_chunks:
92 overhead = stage_id * abs(dyn(c0) - dyn(c1))
93 max_overhead = max(max_overhead, overhead)
94 # bigmem
95 else:
96 top_triangles_chunks = [
97 (vpp - 1 - v, v + 1) for v in range((vpp - 1) // 2)
98 ]
99 diamond_chunks = [(vpp - 1 - v, v) for v in range(vpp // 2)]
100 bot_triangles_chunks = [
101 (vpp - 2 - v, v) for v in range((vpp - 1) // 2)
102 ]
103 print("top_triangles_chunks", top_triangles_chunks)
104 print("diamond_chunks", diamond_chunks)
105 print("bot_triangles_chunks", bot_triangles_chunks)
106 for c0, c1 in top_triangles_chunks:
107 overhead = (micro_left - stage_id) * abs(dyn(c0) - dyn(c1))
108 max_overhead = max(max_overhead, overhead)
109 for c0, c1 in diamond_chunks:
110 overhead = (micro_left - stage_id) * abs(dyn(c0) - dyn(c1))
111 max_overhead = max(max_overhead, overhead)
112 for c0, c1 in bot_triangles_chunks:
113 overhead = stage_id * abs(dyn(c0) - dyn(c1))
114 max_overhead = max(max_overhead, overhead)
115 return max_overhead
117 def __stage_bwd_overhead_1f1b(
118 self, stages: list, stage_id: int, record_lay_types: dict
119 ) -> float:
120 """1f1b end of warmup"""
121 res = 0
122 if stages[stage_id][self._ccfg.vp - 1]:
123 last_node = self._fetch_node_and_switch_env(
124 stages, record_lay_types, stage_id, -1, -1
125 )
126 # full rec -> not rec + grad
127 # not rec -> grad
128 if last_node == LayerType.FULL_REC_LAYER:
129 self._ctx.current_lay_id = f"rec_{self._ctx.current_lay_id}"
130 self._ctx.current_node = LayerType.NOT_REC_LAYER
131 res = sum(self._inner_dynamic_mem(default_micro_factor=1))
132 else:
133 self._ctx.current_node = last_node
134 self._ctx.current_lay_id = f"G_{self._ctx.current_lay_id}"
135 res = sum(self._inner_dynamic_mem(default_micro_factor=1))
136 if (
137 last_node == LayerType.OUTPUT_LAYER
138 and self._ccfg.n_mtp > 0
139 ):
140 last_mtp = self._fetch_node_and_switch_env(
141 stages, record_lay_types, stage_id, -1, -2
142 )
143 if last_mtp == LayerType.FULL_REC_LAYER:
144 self._ctx.current_lay_id = (
145 f"rec_{self._ctx.current_lay_id}"
146 )
147 self._ctx.current_node = LayerType.NOT_REC_LAYER
148 else:
149 self._ctx.current_lay_id = (
150 f"G_{self._ctx.current_lay_id}"
151 )
152 self._ctx.current_node = last_mtp
153 res += sum(self._inner_dynamic_mem(default_micro_factor=1))
154 return res
156 def __stage_bwd_overhead_zbv(
157 self, stages: list, stage_id: int, record_lay_types: dict
158 ) -> float:
159 """dualpipeV end of warmup"""
160 res = 0
161 # Overlapping first/last chunks FWD/BWD
162 # fwd first + bwd last
163 fwd_first, bwd_last = 0, 0
164 for lay_id, lay in enumerate(stages[stage_id][0]):
165 self._ctx.current_node = lay
166 result = self._fetch_node_and_switch_env(
167 stages, record_lay_types, stage_id, 0, lay_id
168 )
169 if result is None:
170 raise RuntimeError("_fetch_node_and_switch_env returned None.")
171 fwd_first += sum(self._inner_dynamic_mem(default_micro_factor=1))
172 for lay_id, lay in enumerate(stages[stage_id][1]):
173 self._ctx.current_node = LayerType.NOT_REC_LAYER
174 result = self._fetch_node_and_switch_env(
175 stages, record_lay_types, stage_id, 1, lay_id
176 )
177 if result is None:
178 raise RuntimeError("_fetch_node_and_switch_env returned None.")
179 if lay == LayerType.FULL_REC_LAYER:
180 bwd_last = max(
181 bwd_last,
182 sum(self._inner_dynamic_mem(default_micro_factor=1)),
183 )
184 # fwd last + bwd first
185 fwd_last, bwd_first = 0, 0
186 for lay_id, _ in enumerate(stages[stage_id][1]):
187 self._ctx.current_node = LayerType.NOT_REC_LAYER
188 result = self._fetch_node_and_switch_env(
189 stages, record_lay_types, stage_id, 1, lay_id
190 )
191 if result is None:
192 raise RuntimeError("_fetch_node_and_switch_env returned None.")
193 fwd_last = sum(self._inner_dynamic_mem(default_micro_factor=1))
194 for lay_id, lay in enumerate(stages[stage_id][0]):
195 self._ctx.current_node = LayerType.NOT_REC_LAYER
196 result = self._fetch_node_and_switch_env(
197 stages, record_lay_types, stage_id, 0, lay_id
198 )
199 if result is None:
200 raise RuntimeError("_fetch_node_and_switch_env returned None.")
201 if lay == LayerType.FULL_REC_LAYER:
202 bwd_first = max(
203 bwd_first,
204 sum(self._inner_dynamic_mem(default_micro_factor=1)),
205 )
207 # print(self.mb(overlap1),self.mb(overlap2))
208 res = max(fwd_first + bwd_last, fwd_last + bwd_first)
209 # res = 0
210 return res