« 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"""Utility submodule"""
16from __future__ import annotations
17from typing import TYPE_CHECKING, Any, Mapping
18import operator
19import ast
20import functools
21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
23if TYPE_CHECKING:
24 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
25 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
26 from typing import Union
28OPS_MAP = {
29 ast.Add: operator.add, # x + y
30 ast.Sub: operator.sub, # x - y
31 ast.Mult: operator.mul, # x * y
32 ast.Div: operator.truediv, # x / y
33 ast.FloorDiv: operator.floordiv, # x // y
34 ast.Mod: operator.mod, # x % y
35 ast.Pow: operator.pow, # x ** y
36 ast.USub: operator.neg, # -x
37 "min": min,
38 "max": max,
39 "abs": abs,
40 "round": round,
41}
44@functools.lru_cache(maxsize=256)
45def _parsed_expr(expr: str) -> ast.AST:
46 """The expression a communication formula sums, parsed once rather than once a layer and candidate.
48 The evaluation only reads the tree, so every call can share it.
49 """
50 return ast.parse(expr).body[0].value
53class EvalUtils:
54 """Utility methods class, PP Microbatch factor formulas"""
56 @staticmethod
57 def reduced_grads(ccfg: CostModelConfig, parts: tuple) -> tuple:
58 """What the gradients FSDP reduce-scatters weigh, whole and sharded.
60 Each part is its parameter count, the ranks the backward computes
61 its gradient over and the ranks the gradient is kept over. A part
62 no other rank shards is not reduce-scattered: the gradient the
63 backward computes is the one it keeps.
64 """
65 reduced = [(size, computed, kept) for size, computed, kept in parts if kept > computed]
66 whole = sum(size / computed for size, computed, _ in reduced)
67 sharded = sum(size / kept for size, _, kept in reduced)
68 return whole * ccfg.bytes_grad, sharded * ccfg.bytes_grad
70 @staticmethod
71 def census_bytes(ctx: Context, tokens: float, kept: float, held: float) -> float:
72 """What a census states for *tokens*: what a layer keeps, *kept* bytes per token, or its working set.
74 As its backward's working set, the most that backward holds, *held*
75 per token; less what the layer keeps, and never below it, where the
76 stage's dynamic memory counts that already, as for a layer that
77 does not recompute.
78 """
79 if not ctx.working_set:
80 return tokens * kept
81 return tokens * (max(0.0, held - kept) if ctx.working_on_saved else held)
83 @staticmethod
84 def mb(x: Union[float, dict, tuple]) -> int:
85 """Convert Byte to MB"""
86 if isinstance(x, dict):
87 return dict(
88 (
89 (k, int(sum(v) / 1024 / 1024))
90 if isinstance(v, tuple)
91 else (k, int(v / 1024 / 1024))
92 )
93 for k, v in x.items()
94 )
95 if isinstance(x, tuple):
96 return int(sum(x) / 1024 / 1024)
97 return int(x / 1024 / 1024)
99 # Layer Blocks
100 @staticmethod
101 def rec_coeff(rec_layer: bool, rec_op: bool) -> bool:
102 """Masking coefficient for select recompute"""
103 return int(not rec_layer) | rec_op
105 @staticmethod
106 def switches(ccfg: CostModelConfig, ctx: Context) -> Any:
107 """The switches of the layer *ctx* evaluates: the context's where it states them, else its config's."""
108 stated = getattr(ctx, "switches", None)
109 return stated if isinstance(stated, Mapping) else ccfg.rec_op
111 @staticmethod
112 def switch(ccfg: CostModelConfig, ctx: Context, name: str) -> Any:
113 """One recompute switch of the layer *ctx* evaluates, 1 to keep the op's activation and 0 to recompute it."""
114 return EvalUtils.state(EvalUtils.switches(ccfg, ctx), name)
116 @staticmethod
117 def state(stated: Any, name: str) -> Any:
118 """Switch *name* in *stated*, a mapping or a config's ``rec_op``: a switch it leaves out keeps its op.
120 Read from ``vars``, since a ``Config`` answers 0, recompute, for any
121 attribute it lacks.
122 """
123 values = stated if isinstance(stated, Mapping) else vars(stated) if stated is not None else {}
124 return values.get(name, 1)
126 @classmethod
127 def eval_expr_insight(cls, **kwargs):
128 """compute and categorize math expression"""
129 return cls.__eval_ast_mem(_parsed_expr(kwargs.get("expr")), 0, **kwargs)[0]
131 @staticmethod
132 def __log_ast_mem(ctx, cat, mem):
133 """Save AST memory evaluation result in the current context."""
134 ctx.save2log(cat, mem)
135 ctx.accu_mem_type[cat] += mem
137 @classmethod
138 def __eval_ast_name(cls, n: ast.Name, depth: int, wait: bool, **kwargs):
139 """Evaluate a named memory term."""
140 ctx = kwargs.get("ctx")
141 if n.id not in kwargs.get("mem_val") or n.id not in kwargs.get(
142 "mem_cat"
143 ):
144 raise AttributeError(
145 f"Unrecognized variable '{n.id}' "
146 f"from expr: '{kwargs['expr']}' "
147 f"(recognized: {list(kwargs.get('mem_val').keys())})"
148 )
149 mem = kwargs.get("mem_val")[n.id]
150 cat = kwargs.get("mem_cat")[n.id]
151 if depth <= 1 and not wait:
152 cls.__log_ast_mem(ctx, cat, mem)
153 return mem, cat
155 @classmethod
156 def __eval_ast_unary(cls, n: ast.UnaryOp, depth: int, wait: bool, **kwargs):
157 """Evaluate a unary expression."""
158 ctx = kwargs.get("ctx")
159 mem, cat = cls.__eval_ast_mem(
160 n.operand,
161 depth + int(not isinstance(n.operand, ast.BinOp)),
162 **kwargs,
163 )
164 mem = OPS_MAP[type(n.op)](mem)
165 if depth == 0 and not wait:
166 cls.__log_ast_mem(ctx, cat, mem)
167 return mem, cat
169 @classmethod
170 def __eval_ast_binop(cls, n: ast.BinOp, depth: int, **kwargs):
171 """Evaluate a binary expression."""
172 ctx = kwargs.get("ctx")
173 l_is_con = isinstance(n.left, ast.Constant) or not isinstance(
174 n.op, ast.Add
175 )
176 r_is_con = isinstance(n.right, ast.Constant) or not isinstance(
177 n.op, ast.Add
178 )
179 l_eval = cls.__eval_ast_mem(
180 n.left,
181 depth + int(not isinstance(n.left, ast.BinOp)),
182 wait=r_is_con,
183 **kwargs,
184 )
185 r_eval = cls.__eval_ast_mem(
186 n.right,
187 depth + int(not isinstance(n.right, ast.BinOp)),
188 wait=l_is_con,
189 **kwargs,
190 )
191 mem = OPS_MAP[type(n.op)](l_eval[0], r_eval[0])
192 cat = l_eval[1] if l_eval[1] else r_eval[1]
193 if depth == 0 and (r_is_con or l_is_con):
194 cls.__log_ast_mem(ctx, cat, mem)
195 return mem, cat
197 @classmethod
198 def __eval_ast_call(
199 cls, n: ast.Call, depth: int, wait: bool, **kwargs
200 ):
201 """Evaluate a supported function call."""
202 ctx = kwargs.get("ctx")
203 a_res = [[], []]
204 for x in n.args:
205 a_eval = cls.__eval_ast_mem(x, depth + 1, wait=True, **kwargs)
206 a_res[0] += [a_eval[0]]
207 a_res[1] += [a_eval[1]]
208 mem = OPS_MAP[n.func.id](*a_res[0])
209 cat = next(
210 (c for c in a_res[1] if a_res[0][a_res[1].index(c)] == mem),
211 a_res[1][0],
212 )
213 if depth <= 1 and not wait:
214 cls.__log_ast_mem(ctx, cat, mem)
215 return mem, cat
217 @classmethod
218 def __eval_ast_mem(
219 cls, n: ast.AST, depth: int, wait: bool = False, **kwargs
220 ):
221 """compute and categorize from AST"""
222 if isinstance(n, ast.Name):
223 return cls.__eval_ast_name(n, depth, wait, **kwargs)
224 if isinstance(n, ast.Constant):
225 return n.value, None
226 if isinstance(n, ast.UnaryOp):
227 return cls.__eval_ast_unary(n, depth, wait, **kwargs)
228 if isinstance(n, ast.BinOp):
229 return cls.__eval_ast_binop(n, depth, **kwargs)
230 if isinstance(n, ast.Call):
231 return cls.__eval_ast_call(n, depth, wait, **kwargs)
232 return 0, None
234 # PP MICRO FACTOR
236 @staticmethod
237 def pp_1f1b_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
238 """1F1B Warm-up microbatches count"""
239 stage_id, chunk_id = ctx.current_stage_id, ctx.current_chunk_id
240 # Warm_up micros num compute
241 micro_factor = 1
242 extra = 0
243 base_micro = min(ccfg.p, ccfg.m)
244 if ccfg.vp == 1:
245 micro_factor = base_micro - stage_id
246 else: # VPP
247 if 0 < chunk_id < ccfg.vp - 1: # Middle chunk
248 micro_factor = base_micro
249 else: # First/Last chunk
250 if not ctx.vpp_less_mem: # Big memory
251 # Balance micros between last chunk and next first chunk
252 extra = base_micro - stage_id - 1
253 last_chunk_micros = base_micro - stage_id
254 if last_chunk_micros < base_micro:
255 last_chunk_micros += min(1, extra)
256 extra = max(0, extra - 1)
257 if chunk_id == 0:
258 if stage_id == 0:
259 extra -= 1
260 micro_factor = base_micro + extra
261 else:
262 micro_factor = last_chunk_micros
263 else: # Less memory
264 if chunk_id == 0:
265 micro_factor = base_micro
266 else:
267 micro_factor = base_micro - stage_id
268 return micro_factor
270 @staticmethod
271 def pp_seq1f1b_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
272 """Seq1F1B Warm-up microbatches count, in sequence chunks.
274 Each chunk holds ``n_s_split``'th of a micro-batch's sequence: the
275 count goes to ``ctx.seq_chunks``, and the backbone evaluates the node
276 at that length.
277 """
278 stage_id, chunk_id = ctx.current_stage_id, ctx.current_chunk_id
279 # Warm_up micros num compute
280 micro_factor = 1
281 ctx.seq_chunks = max(1, ccfg.n_s_split)
282 base_micro = min(ccfg.p, ccfg.m)
283 if ccfg.vp == 1:
284 micro_factor = base_micro - stage_id + ccfg.n_s_split - 1
285 else: # VPP
286 if 0 < chunk_id < ccfg.vp - 1: # Middle chunk
287 micro_factor = base_micro
288 else: # First/Last chunk
289 if not ctx.vpp_less_mem: # Big memory
290 # Balance micros between last chunk and next first chunk
291 extra = base_micro - stage_id - 1
292 last_chunk_micros = base_micro - stage_id
293 last_chunk_micros += ccfg.n_s_split - 1
294 if last_chunk_micros > base_micro:
295 last_chunk_micros -= 1
296 extra += 1
297 elif last_chunk_micros < base_micro:
298 last_chunk_micros += min(1, extra)
299 extra = max(0, extra - 1)
300 if chunk_id == 0:
301 micro_factor = base_micro + extra
302 else:
303 micro_factor = last_chunk_micros
304 else: # Less memory
305 if chunk_id == 0:
306 micro_factor = base_micro
307 else:
308 micro_factor = base_micro - stage_id
309 micro_factor += ccfg.n_s_split - 1
310 return micro_factor
312 @staticmethod
313 def pp_dualpipe_v_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
314 """DualPipeV/ZeroBubbleV Warm-up microbatches count"""
315 stage_id, chunk_id = ctx.current_stage_id, ctx.current_chunk_id
316 # First half layer from stage 0->PP then second half from stage PP->0
317 if ccfg.vp > 2:
318 logger.warning("DualPipeV with VPP>2 not handled")
319 return 0
320 if chunk_id == 0:
321 return min(ccfg.p, ccfg.m) * 2 - 1 - stage_id
322 return stage_id
324 @staticmethod
325 def pp_gpipe_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
326 """GPipe warm-up microbatches count."""
327 _ = ctx
328 return ccfg.p # Minimum value