« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 typing import Tuple
19from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType
20from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.evaluators.body import EvalBody
23class _BackwardOverhead:
24 """Backward overhead class"""
26 def __init__(self, backbone, ccfg, ctx, inner_dyn_fun):
27 self.backbone = backbone
28 self._ccfg = ccfg
29 self._ctx = ctx
30 self._inner_dynamic_mem = inner_dyn_fun
32 def _fetch_node_and_switch_env(
33 self, stages, record_lay_types, stage_id, chunk_id, lay_id
34 ):
35 """Fetch one node and restore the evaluation context for it."""
36 # if negative index access, convert to positive index
37 if stage_id < 0:
38 stage_id = len(stages) + stage_id
39 if chunk_id < 0:
40 chunk_id = len(stages[stage_id]) + chunk_id
41 if lay_id < 0:
42 lay_id = len(stages[stage_id][chunk_id]) + lay_id
43 ccfg, ctx, kind, switches = record_lay_types[(stage_id, chunk_id, lay_id)]
44 # print("here-",id(ccfg))
45 # print("here0",id(self._ccfg))
46 # print("here00",id(self.backbone._ccfg))
47 self._ccfg = ccfg
48 self._ctx = ctx
49 self._ctx.switches = switches
50 self._ctx.current_stage_id = stage_id
51 self._ctx.current_chunk_id = chunk_id
52 self._ctx.current_lay_id = lay_id
53 # print("here1",id(self._ccfg))
54 self.backbone.apply_kind(kind, ccfg=self._ccfg, ctx=self._ctx)
55 # print("here11",id(self.backbone._ccfg))
56 return stages[stage_id][chunk_id][lay_id]
58 def _working_set(self, gathered: int = 2, on_saved: bool = False) -> float:
59 """The dynamic memory of the node set up, as the working set of its backward.
61 Under FSDP that reshards, a layer's backward holds *gathered* layers'
62 parameters: its own and the next layer's, prefetched. *on_saved*
63 says the stage's dynamic memory counts what the node keeps for this
64 backward already, as it does where the node does not recompute.
65 """
66 self._ctx.working_set = gathered
67 self._ctx.working_on_saved = on_saved
68 try:
69 return sum(self._inner_dynamic_mem(default_micro_factor=1))
70 finally:
71 self._ctx.working_set = 0
72 self._ctx.working_on_saved = False
74 def _last_backward_lay_ids(self, stages: list, stage_id: int, record_lay_types: dict) -> list:
75 """The stage's regular layers from the one whose backward a micro-batch's ends on, in stage order.
77 A stage runs its layers' backward in reverse, so the last is its
78 first layer. Where a vision tower shares the stage with the
79 language model the search drives, the tower's layers come first and
80 run last, and their working set is the smaller of the two; but the
81 moment the language model's own first layer runs its backward, every
82 activation of the tower is still held, and that is the moment the
83 stage has to fit. A stage of one module, or one holding
84 no layer of the main module, keeps its own first layers.
85 """
86 regular = [lay_id for lay_id, node in enumerate(stages[stage_id][0])
87 if self.backbone.is_regular_layer(node)]
88 main = getattr(self.backbone, "main_module", None)
89 if main is None:
90 return regular
91 of_main = [lay_id for lay_id in regular
92 if record_lay_types[(stage_id, 0, lay_id)][0].model_name == main]
93 return of_main or regular
95 def first_layer_working_set(self, stages: list, stage_id: int, record_lay_types: dict) -> Tuple[float, float]:
96 """The working set of the backward a micro-batch's ends on, the stage's first layer's, unlogged.
98 It holds its own gathered parameters, with no layer left to prefetch.
99 FSDP that reduces a layer's gradients while the next layer's backward
100 runs holds its whole gradients, and the second layer's, whose
101 reduction is still in flight. Also returns that layer's
102 reduce-scatter output, which does not exist yet. Where submodules
103 share the stage, it is the main module's first layer
104 (:meth:`_last_backward_lay_ids`).
105 """
106 lay_ids = self._last_backward_lay_ids(stages, stage_id, record_lay_types)
107 working, pending = 0, 0
108 for rank, lay_id in enumerate(lay_ids[:2]):
109 node = self._fetch_node_and_switch_env(stages, record_lay_types, stage_id, 0, lay_id)
110 ctx = self._ctx
111 saved = (ctx.current_node, ctx.current_lay_id, ctx.enable_node_log, dict(ctx.accu_mem_type))
112 ctx.current_node = LayerType.NOT_REC_LAYER if node == LayerType.FULL_REC_LAYER else node
113 ctx.enable_node_log = False
114 try:
115 if rank == 0:
116 working += self._working_set(gathered=1)
117 if self._ccfg.overlaps_grad_reduce:
118 whole, sharded = EvalBody.reduced_grad_layer(self._ccfg, ctx)
119 working += whole
120 pending = pending if rank else sharded
121 finally:
122 ctx.current_node, ctx.current_lay_id, ctx.enable_node_log, ctx.accu_mem_type = saved
123 return working, pending
125 def estimate(
126 self, stages: list, stage_id: int, record_lay_types: dict
127 ) -> float:
128 """estimate stage's end-of-warmup overhead"""
129 # self._ctx.enable_node_log = False
130 res = 0
131 if self._ccfg.pp_sched in ["1f1b", "seqpipe", "seqsmartvpp"]:
132 res = self.__stage_bwd_overhead_1f1b(
133 stages, stage_id, record_lay_types
134 )
135 if self._ccfg.pp_sched == "zero_bubble_v":
136 res = self.__stage_bwd_overhead_zbv(
137 stages, stage_id, record_lay_types
138 )
139 return res
141 def vpp_1f1b_steady_overhead(self, stage_id, dyn_mem_i):
142 """potential overhead due to imbalanced chunks"""
144 def dyn(chunk_id):
145 """chunk total mem"""
146 return sum(dyn_mem_i[chunk_id])
148 # less mem
149 micro_left = self._ccfg.m - self._ccfg.p
150 vpp = self._ccfg.vp
151 max_overhead = 0
152 if self._ctx.vpp_less_mem:
153 top_triangles_chunks = [(vpp - 1 - v, v) for v in range(vpp // 2)]
154 bot_triangles_chunks = [
155 (vpp - 2 - v, v) for v in range((vpp - 1) // 2)
156 ]
157 print("top_triangles_chunks", top_triangles_chunks)
158 print("bot_triangles_chunks", bot_triangles_chunks)
159 for c0, c1 in top_triangles_chunks:
160 overhead = (micro_left - stage_id) * abs(dyn(c0) - dyn(c1))
161 max_overhead = max(max_overhead, overhead)
162 for c0, c1 in bot_triangles_chunks:
163 overhead = stage_id * abs(dyn(c0) - dyn(c1))
164 max_overhead = max(max_overhead, overhead)
165 # bigmem
166 else:
167 top_triangles_chunks = [
168 (vpp - 1 - v, v + 1) for v in range((vpp - 1) // 2)
169 ]
170 diamond_chunks = [(vpp - 1 - v, v) for v in range(vpp // 2)]
171 bot_triangles_chunks = [
172 (vpp - 2 - v, v) for v in range((vpp - 1) // 2)
173 ]
174 print("top_triangles_chunks", top_triangles_chunks)
175 print("diamond_chunks", diamond_chunks)
176 print("bot_triangles_chunks", bot_triangles_chunks)
177 for c0, c1 in top_triangles_chunks:
178 overhead = (micro_left - stage_id) * abs(dyn(c0) - dyn(c1))
179 max_overhead = max(max_overhead, overhead)
180 for c0, c1 in diamond_chunks:
181 overhead = (micro_left - stage_id) * abs(dyn(c0) - dyn(c1))
182 max_overhead = max(max_overhead, overhead)
183 for c0, c1 in bot_triangles_chunks:
184 overhead = stage_id * abs(dyn(c0) - dyn(c1))
185 max_overhead = max(max_overhead, overhead)
186 return max_overhead
188 def __stage_bwd_overhead_1f1b(
189 self, stages: list, stage_id: int, record_lay_types: dict
190 ) -> float:
191 """1f1b end of warmup"""
192 res = 0
193 if stages[stage_id][self._ccfg.vp - 1]:
194 last_node = self._fetch_node_and_switch_env(
195 stages, record_lay_types, stage_id, -1, -1
196 )
197 # full rec -> not rec + grad
198 # not rec -> grad
199 if last_node == LayerType.FULL_REC_LAYER:
200 self._ctx.current_lay_id = f"rec_{self._ctx.current_lay_id}"
201 self._ctx.current_node = LayerType.NOT_REC_LAYER
202 res = self._working_set()
203 else:
204 self._ctx.current_node = last_node
205 self._ctx.current_lay_id = f"G_{self._ctx.current_lay_id}"
206 res = self._working_set(on_saved=True)
207 if (
208 last_node == LayerType.OUTPUT_LAYER
209 and self._ccfg.n_mtp > 0
210 ):
211 last_mtp = self._fetch_node_and_switch_env(
212 stages, record_lay_types, stage_id, -1, -2
213 )
214 if last_mtp == LayerType.FULL_REC_LAYER:
215 self._ctx.current_lay_id = (
216 f"rec_{self._ctx.current_lay_id}"
217 )
218 self._ctx.current_node = LayerType.NOT_REC_LAYER
219 else:
220 self._ctx.current_lay_id = (
221 f"G_{self._ctx.current_lay_id}"
222 )
223 self._ctx.current_node = last_mtp
224 res += self._working_set(on_saved=last_mtp != LayerType.FULL_REC_LAYER)
225 return res
227 def __stage_bwd_overhead_zbv(
228 self, stages: list, stage_id: int, record_lay_types: dict
229 ) -> float:
230 """dualpipeV end of warmup"""
231 res = 0
232 # Overlapping first/last chunks FWD/BWD
233 # fwd first + bwd last
234 fwd_first, bwd_last = 0, 0
235 for lay_id, lay in enumerate(stages[stage_id][0]):
236 self._ctx.current_node = lay
237 result = self._fetch_node_and_switch_env(
238 stages, record_lay_types, stage_id, 0, lay_id
239 )
240 if result is None:
241 raise RuntimeError("_fetch_node_and_switch_env returned None.")
242 fwd_first += sum(self._inner_dynamic_mem(default_micro_factor=1))
243 for lay_id, lay in enumerate(stages[stage_id][1]):
244 self._ctx.current_node = LayerType.NOT_REC_LAYER
245 result = self._fetch_node_and_switch_env(
246 stages, record_lay_types, stage_id, 1, lay_id
247 )
248 if result is None:
249 raise RuntimeError("_fetch_node_and_switch_env returned None.")
250 if lay == LayerType.FULL_REC_LAYER:
251 bwd_last = max(bwd_last, self._working_set())
252 # fwd last + bwd first
253 fwd_last, bwd_first = 0, 0
254 for lay_id, _ in enumerate(stages[stage_id][1]):
255 self._ctx.current_node = LayerType.NOT_REC_LAYER
256 result = self._fetch_node_and_switch_env(
257 stages, record_lay_types, stage_id, 1, lay_id
258 )
259 if result is None:
260 raise RuntimeError("_fetch_node_and_switch_env returned None.")
261 fwd_last = sum(self._inner_dynamic_mem(default_micro_factor=1))
262 for lay_id, lay in enumerate(stages[stage_id][0]):
263 self._ctx.current_node = LayerType.NOT_REC_LAYER
264 result = self._fetch_node_and_switch_env(
265 stages, record_lay_types, stage_id, 0, lay_id
266 )
267 if result is None:
268 raise RuntimeError("_fetch_node_and_switch_env returned None.")
269 if lay == LayerType.FULL_REC_LAYER:
270 bwd_first = max(bwd_first, self._working_set())
272 # print(self.mb(overlap1),self.mb(overlap2))
273 res = max(fwd_first + bwd_last, fwd_last + bwd_first)
274 # res = 0
275 return res