« 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"""Utility submodule"""
16from __future__ import annotations
17from typing import TYPE_CHECKING
18import operator
19import ast
20from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.logger import logger
22if TYPE_CHECKING:
23 from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import CostModelConfig
24 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation._context import Context
25 from typing import Union
27OPS_MAP = {
28 ast.Add: operator.add, # x + y
29 ast.Sub: operator.sub, # x - y
30 ast.Mult: operator.mul, # x * y
31 ast.Div: operator.truediv, # x / y
32 ast.FloorDiv: operator.floordiv, # x // y
33 ast.Mod: operator.mod, # x % y
34 ast.Pow: operator.pow, # x ** y
35 ast.USub: operator.neg, # -x
36 "min": min,
37 "max": max,
38 "abs": abs,
39 "round": round,
40}
43class EvalUtils:
44 """Utility methods class, PP Microbatch factor formulas"""
46 @staticmethod
47 def mb(x: Union[float, dict, tuple]) -> int:
48 """Convert Byte to MB"""
49 if isinstance(x, dict):
50 return dict(
51 (
52 (k, int(sum(v) / 1024 / 1024))
53 if isinstance(v, tuple)
54 else (k, int(v / 1024 / 1024))
55 )
56 for k, v in x.items()
57 )
58 if isinstance(x, tuple):
59 return int(sum(x) / 1024 / 1024)
60 return int(x / 1024 / 1024)
62 # Layer Blocks
63 @staticmethod
64 def rec_coeff(rec_layer: bool, rec_op: bool) -> bool:
65 """Masking coefficient for select recompute"""
66 return int(not rec_layer) | rec_op
68 @classmethod
69 def eval_expr_insight(cls, **kwargs):
70 """compute and categorize math expression"""
71 nodes = ast.parse(kwargs.get("expr"))
72 # print(ast.dump(nodes, indent=1))
73 # print(kwargs["expr"])
74 return cls.__eval_ast_mem(nodes.body[0].value, 0, **kwargs)[0]
76 @staticmethod
77 def __log_ast_mem(ctx, cat, mem):
78 """Save AST memory evaluation result in the current context."""
79 ctx.save2log(cat, mem)
80 ctx.accu_mem_type[cat] += mem
82 @classmethod
83 def __eval_ast_name(cls, n: ast.Name, depth: int, wait: bool, **kwargs):
84 """Evaluate a named memory term."""
85 ctx = kwargs.get("ctx")
86 if n.id not in kwargs.get("mem_val") or n.id not in kwargs.get(
87 "mem_cat"
88 ):
89 raise AttributeError(
90 f"Unrecognized variable '{n.id}' "
91 f"from expr: '{kwargs['expr']}' "
92 f"(recognized: {list(kwargs.get('mem_val').keys())})"
93 )
94 mem = kwargs.get("mem_val")[n.id]
95 cat = kwargs.get("mem_cat")[n.id]
96 if depth <= 1 and not wait:
97 cls.__log_ast_mem(ctx, cat, mem)
98 return mem, cat
100 @classmethod
101 def __eval_ast_unary(cls, n: ast.UnaryOp, depth: int, wait: bool, **kwargs):
102 """Evaluate a unary expression."""
103 ctx = kwargs.get("ctx")
104 mem, cat = cls.__eval_ast_mem(
105 n.operand,
106 depth + int(not isinstance(n.operand, ast.BinOp)),
107 **kwargs,
108 )
109 mem = OPS_MAP[type(n.op)](mem)
110 if depth == 0 and not wait:
111 cls.__log_ast_mem(ctx, cat, mem)
112 return mem, cat
114 @classmethod
115 def __eval_ast_binop(cls, n: ast.BinOp, depth: int, **kwargs):
116 """Evaluate a binary expression."""
117 ctx = kwargs.get("ctx")
118 l_is_con = isinstance(n.left, ast.Constant) or not isinstance(
119 n.op, ast.Add
120 )
121 r_is_con = isinstance(n.right, ast.Constant) or not isinstance(
122 n.op, ast.Add
123 )
124 l_eval = cls.__eval_ast_mem(
125 n.left,
126 depth + int(not isinstance(n.left, ast.BinOp)),
127 wait=r_is_con,
128 **kwargs,
129 )
130 r_eval = cls.__eval_ast_mem(
131 n.right,
132 depth + int(not isinstance(n.right, ast.BinOp)),
133 wait=l_is_con,
134 **kwargs,
135 )
136 mem = OPS_MAP[type(n.op)](l_eval[0], r_eval[0])
137 cat = l_eval[1] if l_eval[1] else r_eval[1]
138 if depth == 0 and (r_is_con or l_is_con):
139 cls.__log_ast_mem(ctx, cat, mem)
140 return mem, cat
142 @classmethod
143 def __eval_ast_call(
144 cls, n: ast.Call, depth: int, wait: bool, **kwargs
145 ):
146 """Evaluate a supported function call."""
147 ctx = kwargs.get("ctx")
148 a_res = [[], []]
149 for x in n.args:
150 a_eval = cls.__eval_ast_mem(x, depth + 1, wait=True, **kwargs)
151 a_res[0] += [a_eval[0]]
152 a_res[1] += [a_eval[1]]
153 mem = OPS_MAP[n.func.id](*a_res[0])
154 cat = next(
155 (c for c in a_res[1] if a_res[0][a_res[1].index(c)] == mem),
156 a_res[1][0],
157 )
158 if depth <= 1 and not wait:
159 cls.__log_ast_mem(ctx, cat, mem)
160 return mem, cat
162 @classmethod
163 def __eval_ast_mem(
164 cls, n: ast.AST, depth: int, wait: bool = False, **kwargs
165 ):
166 """compute and categorize from AST"""
167 if isinstance(n, ast.Name):
168 return cls.__eval_ast_name(n, depth, wait, **kwargs)
169 if isinstance(n, ast.Constant):
170 return n.value, None
171 if isinstance(n, ast.UnaryOp):
172 return cls.__eval_ast_unary(n, depth, wait, **kwargs)
173 if isinstance(n, ast.BinOp):
174 return cls.__eval_ast_binop(n, depth, **kwargs)
175 if isinstance(n, ast.Call):
176 return cls.__eval_ast_call(n, depth, wait, **kwargs)
177 return 0, None
179 # PP MICRO FACTOR
181 @staticmethod
182 def pp_1f1b_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
183 """1F1B Warm-up microbatches count"""
184 stage_id, chunk_id = ctx.current_stage_id, ctx.current_chunk_id
185 # Warm_up micros num compute
186 micro_factor = 1
187 extra = 0
188 base_micro = min(ccfg.p, ccfg.m)
189 if ccfg.vp == 1:
190 micro_factor = base_micro - stage_id
191 else: # VPP
192 if 0 < chunk_id < ccfg.vp - 1: # Middle chunk
193 micro_factor = base_micro
194 else: # First/Last chunk
195 if not ctx.vpp_less_mem: # Big memory
196 # Balance micros between last chunk and next first chunk
197 extra = base_micro - stage_id - 1
198 last_chunk_micros = base_micro - stage_id
199 if last_chunk_micros < base_micro:
200 last_chunk_micros += min(1, extra)
201 extra = max(0, extra - 1)
202 if chunk_id == 0:
203 if stage_id == 0:
204 extra -= 1
205 micro_factor = base_micro + extra
206 else:
207 micro_factor = last_chunk_micros
208 else: # Less memory
209 if chunk_id == 0:
210 micro_factor = base_micro
211 else:
212 micro_factor = base_micro - stage_id
213 return micro_factor
215 @staticmethod
216 def pp_seq1f1b_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
217 """Seq1F1B Warm-up microbatches count"""
218 stage_id, chunk_id = ctx.current_stage_id, ctx.current_chunk_id
219 # Warm_up micros num compute
220 micro_factor = 1
221 ccfg.s /= ccfg.n_s_split # Splitting seq length
222 base_micro = min(ccfg.p, ccfg.m)
223 if ccfg.vp == 1:
224 micro_factor = base_micro - stage_id + ccfg.n_s_split - 1
225 else: # VPP
226 if 0 < chunk_id < ccfg.vp - 1: # Middle chunk
227 micro_factor = base_micro
228 else: # First/Last chunk
229 if not ctx.vpp_less_mem: # Big memory
230 # Balance micros between last chunk and next first chunk
231 extra = base_micro - stage_id - 1
232 last_chunk_micros = base_micro - stage_id
233 last_chunk_micros += ccfg.n_s_split - 1
234 if last_chunk_micros > base_micro:
235 last_chunk_micros -= 1
236 extra += 1
237 elif last_chunk_micros < base_micro:
238 last_chunk_micros += min(1, extra)
239 extra = max(0, extra - 1)
240 if chunk_id == 0:
241 micro_factor = base_micro + extra
242 else:
243 micro_factor = last_chunk_micros
244 else: # Less memory
245 if chunk_id == 0:
246 micro_factor = base_micro
247 else:
248 micro_factor = base_micro - stage_id
249 micro_factor += ccfg.n_s_split - 1
250 return micro_factor
252 @staticmethod
253 def pp_dualpipe_v_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
254 """DualPipeV/ZeroBubbleV Warm-up microbatches count"""
255 stage_id, chunk_id = ctx.current_stage_id, ctx.current_chunk_id
256 # First half layer from stage 0->PP then second half from stage PP->0
257 if ccfg.vp > 2:
258 logger.warning("DualPipeV with VPP>2 not handled")
259 return 0
260 if chunk_id == 0:
261 return min(ccfg.p, ccfg.m) * 2 - 1 - stage_id
262 return stage_id
264 @staticmethod
265 def pp_gpipe_micro_factor(ccfg: CostModelConfig, ctx: Context) -> int:
266 """GPipe warm-up microbatches count."""
267 _ = ctx
268 return ccfg.p # Minimum value