Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / utils / recompute.py: 99%
150 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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"""Recomputation taxonomy and conversion helpers between internal dicts and the YAML schema."""
16from enum import IntEnum
17from typing import Any, Dict, List, Optional
19from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger
21TYPE = IntEnum("RecomputeType", ["NONE", "SLCT", "COMM", "BOTH", "FULL"], start=0)
22OFFSET = "offset"
24DEFAULT_COEF = {
25 TYPE.NONE: 0,
26 TYPE.SLCT: 0.04,
27 TYPE.COMM: 0.125,
28 TYPE.BOTH: 0.165,
29 TYPE.FULL: 0.5,
30}
32YAML_NAME = {
33 TYPE.NONE: "",
34 TYPE.COMM: "select_comm_recompute",
35 TYPE.SLCT: "select_recompute",
36 TYPE.BOTH: "both_comm_select",
37 TYPE.FULL: "recompute",
38}
40JSON_MEMORY_NAME = {
41 TYPE.NONE: "memory_activation",
42 TYPE.COMM: "memory_select_comm",
43 TYPE.BOTH: "memory_both_comm_select",
44 TYPE.SLCT: "memory_select_rec",
45 TYPE.FULL: "memory_recompute",
46}
48JSON_MEMORY_NAME_ALIGNED = {
49 TYPE.NONE: "memory_activation ",
50 TYPE.COMM: "memory_select_comm",
51 TYPE.BOTH: "memory_both_comm_select",
52 TYPE.SLCT: "memory_select_rec ",
53 TYPE.FULL: "memory_recompute ",
54}
57JSON_TIME_NAME = {
58 TYPE.NONE: "backward_time",
59 TYPE.COMM: "select_comm_time",
60 TYPE.BOTH: "both_comm_select_time",
61 TYPE.SLCT: "select_rec_time",
62 TYPE.FULL: "recompute_time ",
63}
65JSON_COEF_NAME = {
66 TYPE.NONE: "backward_coef",
67 TYPE.SLCT: "select_rec_coef",
68 TYPE.BOTH: "both_comm_select_coef",
69 TYPE.COMM: "select_comm_coef",
70 TYPE.FULL: "recompute_coef",
71}
74def sums(rec_dict: Dict[TYPE, int]) -> int:
75 """Return the sum of layer counts across every :class:`TYPE` key in ``rec_dict``."""
76 x = 0
77 for r in TYPE:
78 x += rec_dict[r]
79 return x
82def zero_if_none_var(v: Any, i: int, s: int) -> int:
83 """Read ``int(v[i][s].varValue)`` guarding against ``None`` at any step."""
84 if v is not None and v[i][s].varValue is not None:
85 return int(v[i][s].varValue)
86 return 0
89def zero_if_none(v: Any, i: int, s: int) -> int:
90 """Read ``int(v[i][s])`` guarding against ``None`` at any step."""
91 if v is not None and v[i][s] is not None:
92 return int(v[i][s])
93 return 0
96def yaml_from_internal(vpp: int, pp: int,
97 lp_variables: Dict[TYPE, Any],
98 nass: List[List[int]]) -> Dict[str, List[List[int]]]:
99 """Convert solver variables into the MindFormers YAML schema.
101 Args:
102 vpp: Number of virtual pipeline (VPP) chunks.
103 pp: Number of physical pipeline stages.
104 lp_variables: Solver variables keyed by :class:`TYPE`.
105 nass: Naive layer assignments per ``(vpp_chunk, stage)``.
107 Returns:
108 A mapping from YAML field name to a 2-D list of ``(vpp, pp)`` integers.
109 """
110 slct_is = 0
111 comm_is = 0
112 both_is = 0
113 full_is = 0
115 yaml_out: Dict[str, List[List[int]]] = {
116 OFFSET: [],
117 YAML_NAME[TYPE.FULL]: [],
118 YAML_NAME[TYPE.SLCT]: [],
119 YAML_NAME[TYPE.COMM]: [],
120 }
121 logger.debug("pp = %s, vpp = %s", pp, vpp)
122 for i in range(vpp):
123 for _, v in yaml_out.items():
124 v.append([])
125 for s in range(pp):
126 gass_i_s = 0
127 for r in TYPE:
128 gass_i_s += zero_if_none_var(lp_variables[r], i, s)
129 slct_is = zero_if_none_var(lp_variables[TYPE.SLCT], i, s)
130 comm_is = zero_if_none_var(lp_variables[TYPE.COMM], i, s)
131 both_is = zero_if_none_var(lp_variables[TYPE.BOTH], i, s)
132 full_is = zero_if_none_var(lp_variables[TYPE.FULL], i, s)
133 yaml_out[OFFSET][i].append(gass_i_s - nass[i][s])
134 yaml_out[YAML_NAME[TYPE.FULL]][i].append(full_is)
135 yaml_out[YAML_NAME[TYPE.SLCT]][i].append(slct_is + both_is + full_is)
136 yaml_out[YAML_NAME[TYPE.COMM]][i].append(comm_is + both_is + full_is)
138 logger.debug("yaml = %s", yaml_out)
139 return yaml_out
142def internal_from_yaml(vpp: int, pp: int,
143 yaml_in: Dict[str, Any],
144 nass: List[List[int]]) -> Dict[TYPE, List[List[int]]]:
145 """Convert a MindFormers YAML schema back into per-type layer counts.
147 Args:
148 vpp: Number of virtual pipeline chunks.
149 pp: Number of physical pipeline stages.
150 yaml_in: YAML mapping with ``offset`` and per-recomputation-type fields.
151 nass: Naive layer assignments per ``(vpp_chunk, stage)``.
153 Returns:
154 A mapping from :class:`TYPE` to a 2-D list of ``(vpp, pp)`` integers.
155 """
156 slct_is = 0
157 comm_is = 0
158 full_is = 0
159 layer_per_recompute: Dict[TYPE, List[List[int]]] = {r: [] for r in TYPE}
160 if yaml_in[OFFSET] == 0:
161 yaml_in[OFFSET] = [[0] * pp for _ in range(vpp)]
163 for rec in [TYPE.SLCT, TYPE.COMM, TYPE.FULL]:
164 if (
165 YAML_NAME[rec] not in yaml_in
166 or yaml_in[YAML_NAME[rec]] is False
167 or yaml_in[YAML_NAME[rec]] == 0
168 ):
169 yaml_in[YAML_NAME[rec]] = [[0] * pp for _ in range(vpp)]
170 if yaml_in[YAML_NAME[rec]] is True:
171 yaml_in[YAML_NAME[rec]] = [
172 [a + b for a, b in zip(list1, list2)]
173 for list1, list2 in zip(nass, yaml_in[OFFSET])
174 ]
176 for i in range(vpp):
177 for _, v in layer_per_recompute.items():
178 v.append([])
179 for s in range(pp):
180 slct_is = zero_if_none(yaml_in[YAML_NAME[TYPE.SLCT]], i, s)
181 comm_is = zero_if_none(yaml_in[YAML_NAME[TYPE.COMM]], i, s)
182 full_is = zero_if_none(yaml_in[YAML_NAME[TYPE.FULL]], i, s)
183 layer_per_recompute[TYPE.FULL][i].append(full_is)
184 layer_per_recompute[TYPE.BOTH][i].append(
185 max(min(slct_is - full_is, comm_is - full_is), 0)
186 )
187 layer_per_recompute[TYPE.SLCT][i].append(
188 max(slct_is - full_is - layer_per_recompute[TYPE.BOTH][i][s], 0)
189 )
190 layer_per_recompute[TYPE.COMM][i].append(
191 max(comm_is - full_is - layer_per_recompute[TYPE.BOTH][i][s], 0)
192 )
193 layer_per_recompute[TYPE.NONE][i].append(
194 (
195 yaml_in[OFFSET][i][s]
196 + nass[i][s]
197 - layer_per_recompute[TYPE.FULL][i][s]
198 - layer_per_recompute[TYPE.BOTH][i][s]
199 - layer_per_recompute[TYPE.SLCT][i][s]
200 - layer_per_recompute[TYPE.COMM][i][s]
201 )
202 )
204 logger.debug("layer_per_recompute = %s", layer_per_recompute)
205 return layer_per_recompute
208def to_list(rec_dict: Dict[TYPE, Any]) -> List[Any]:
209 """Return the values of ``rec_dict`` in :class:`TYPE` enum order."""
210 return list(rec_dict.values())
213def right_extend(ll: List[List[int]], n: int) -> List[List[int]]:
214 """Return ``ll`` extended by appending each of ``range(n)`` to every sub-list.
216 Args:
217 ll: List of partially built index vectors.
218 n: Number of values (``0..n-1``) to append.
220 Returns:
221 A new list where each input sub-list appears ``n`` times, each with one of the new values.
222 """
223 all_l: List[List[int]] = []
224 for i in range(n):
225 for sublist in ll:
226 all_l += [sublist + [i]]
227 return all_l
230def make_all_indexes_local(used_rec: Dict[TYPE, bool], num_of_interleave: int,
231 all_indexes: List[List[int]], r: TYPE) -> List[List[int]]:
232 """Recursive helper behind :func:`make_all_indexes`.
234 Args:
235 used_rec: Which recomputation types are currently considered.
236 num_of_interleave: Interleave (VPP) degree.
237 all_indexes: Accumulated partial assignments.
238 r: The current :class:`TYPE` being processed.
240 Returns:
241 The completed list of index vectors once the last :class:`TYPE` is reached.
242 """
243 if r >= len(TYPE) - 1:
244 if used_rec[r]:
245 all_indexes = right_extend(all_indexes, num_of_interleave)
246 return all_indexes
247 if used_rec[r]:
248 return make_all_indexes_local(
249 used_rec,
250 num_of_interleave,
251 right_extend(all_indexes, num_of_interleave),
252 TYPE(r + 1),
253 )
254 return make_all_indexes_local(used_rec, num_of_interleave, all_indexes, TYPE(r + 1))
257def make_all_indexes(used_rec: Dict[TYPE, bool], num_of_interleave: int) -> List[List[int]]:
258 """Enumerate all per-recomputation-type assignments across ``num_of_interleave`` chunks."""
259 return make_all_indexes_local(used_rec, num_of_interleave, [[]], TYPE.NONE)
262def recomputes_from_indexes(used_rec: Dict[TYPE, bool],
263 indexes: List[List[int]]) -> List[Dict[TYPE, Optional[int]]]:
264 """Decode index vectors produced by :func:`make_all_indexes` into per-type dictionaries."""
265 recomputes: List[Dict[TYPE, Optional[int]]] = []
266 for idx in indexes:
267 recompute: Dict[TYPE, Optional[int]] = {r: None for r in TYPE}
268 for r in TYPE:
269 if used_rec[r]:
270 recompute[r] = idx[0]
271 idx.pop(0)
272 recomputes.append(recompute)
273 return recomputes
276def average(rec_list: List[Dict[TYPE, Optional[float]]]) -> Dict[TYPE, Optional[float]]:
277 """Return the per-type mean of a list of per-type recomputation dicts.
279 Args:
280 rec_list: Mapping from :class:`TYPE` to a numeric value or ``None``.
282 Returns:
283 A new dict holding the arithmetic mean for each :class:`TYPE` (``None`` propagates).
284 """
285 num = len(rec_list)
286 if num == 0:
287 return rec_list
288 rec_1 = rec_list.pop(0)
289 for rec_i in rec_list:
290 for r in TYPE:
291 if rec_1[r] is not None and rec_i[r] is not None:
292 rec_1[r] = rec_1[r] + rec_i[r]
293 elif not (rec_1[r] is None and rec_i[r] is None):
294 logger.warning(
295 "WARNING: Recomputation %s is not taken into consideration by all body layers",
296 r.name,
297 )
298 for r in TYPE:
299 if rec_1[r] is not None:
300 rec_1[r] = rec_1[r] / num
301 return rec_1
304def assign_used(values: List[int], unused_rec: List[TYPE]) -> Dict[TYPE, Optional[int]]:
305 """Associate each value with its recomputation type, skipping ``unused_rec`` entries."""
306 assignment: Dict[TYPE, Optional[int]] = {r: None for r in TYPE}
307 value_idx = 0
308 for r in TYPE:
309 if r not in unused_rec:
310 assignment[r] = values[value_idx]
311 value_idx += 1
312 return assignment
315def get_used_list(recompute_considered: Dict[TYPE, bool]) -> List[TYPE]:
316 """Return recomputation types flagged as enabled in ``recompute_considered``."""
317 used_rec: List[TYPE] = []
318 for rec in TYPE:
319 if recompute_considered[rec]:
320 used_rec.append(rec)
321 return used_rec
324def get_unused_list(recompute_considered: Dict[TYPE, bool]) -> List[TYPE]:
325 """Return recomputation types flagged as disabled (or missing) in ``recompute_considered``."""
326 unused_rec: List[TYPE] = []
327 for rec in TYPE:
328 if rec not in recompute_considered or not recompute_considered[rec]:
329 unused_rec.append(rec)
330 return unused_rec
333def least_recomputed(recompute_considered: Dict[TYPE, bool]) -> TYPE:
334 """Return the lowest-index enabled recomputation :class:`TYPE`."""
335 rec = TYPE.NONE
336 for r in TYPE:
337 if recompute_considered[r]:
338 rec = r
339 break
340 return rec
343def most_recomputed(recompute_considered: Dict[TYPE, bool]) -> TYPE:
344 """Return the highest-index enabled recomputation :class:`TYPE`."""
345 rec = TYPE.FULL
346 for r in TYPE:
347 if recompute_considered[r]:
348 rec = r
349 return rec