Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / recompute / potential.py: 97%
101 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Sort, retain and export recompute plans at fixed parallel degrees.
17What a plan costs here is the cost model's own estimate of the step under
18it, the same estimate the search ranks strategies with. A measured run can
19disagree with it, so the sequence reads as what the cost model expects each
20further recomputed layer to cost, not as a measured or guaranteed figure.
21"""
23import csv
24import json
25import math
26from collections import Counter
27from pathlib import Path
28from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
30from hyper_parallel.auto_parallel.sapp_nd.nd.dimensions import Dimensions
31from hyper_parallel.auto_parallel.sapp_nd.recompute.candidate import MODES, RecomputeChoice, to_records, trainer_plan
33# The configurations a plot shows when the run asked for no particular number.
34DEFAULT_TOP = 20
36# The columns of the CSV a retained sequence is written to.
37_FIXED_COLUMNS = (
38 "changed_layer", "full_layers", "memory_mb", "score", "recompute_mode", "recompute_layers",
39 "layer_recompute", "stage_memory_mb", "stage_savings", "plan_scope",
40)
43def sort_recompute_results(results: Sequence[RecomputeChoice]) -> List[RecomputeChoice]:
44 """Return recompute plans fastest first by the cost model's estimate, peak memory breaking ties.
46 Args:
47 results: Scored plans to order, which are not modified.
49 Returns:
50 The plans in ascending score order.
52 Raises:
53 ValueError: A plan has no finite score.
54 """
55 if any(result.score is None or not math.isfinite(result.score) for result in results):
56 raise ValueError("recompute results must have finite performance scores before sorting")
57 return sorted(results, key=lambda result: (result.score, result.memory))
60def _layer_settings(choice: RecomputeChoice) -> Dict[int, Tuple[Any, str, Any]]:
61 """Expand a named plan per layer, independently of how its consecutive ranges are grouped."""
62 defined: Dict[int, Tuple[Any, str, Any]] = {}
63 for item in choice.ranges:
64 if item.first < 0 or item.count <= 0 or item.mode not in MODES:
65 raise ValueError("recompute potentials require valid ranges with a named mode per layer")
66 for index in range(item.first, item.first + item.count):
67 if index in defined:
68 raise ValueError("recompute layer ranges must not overlap")
69 defined[index] = (item.kind, item.mode, item.option)
70 return defined
73def _one_more_full(
74 defined: Mapping[int, Tuple[Any, str, Any]], proposed: Mapping[int, Tuple[Any, str, Any]]
75) -> bool:
76 """Whether *proposed* runs one more layer of one kind full and leaves every other count alone."""
77 if any(proposed[index][0] != defined[index][0] or proposed[index][2].link_bandwidth for index in defined):
78 return False
79 before = Counter((kind, mode) for kind, mode, _ in defined.values())
80 after = Counter((kind, mode) for kind, mode, _ in proposed.values())
81 removed, added = before - after, after - before
82 if sum(removed.values()) != 1 or sum(added.values()) != 1:
83 return False
84 old_kind, old_mode = next(iter(removed))
85 new_kind, new_mode = next(iter(added))
86 return old_kind == new_kind and old_mode != "full" and new_mode == "full"
89def keep_recompute_potential(
90 current: RecomputeChoice,
91 kept: List[RecomputeChoice],
92 results: Sequence[RecomputeChoice],
93) -> Optional[RecomputeChoice]:
94 """Append and return the fastest unused plan that recomputes one more layer of one kind.
96 Args:
97 current: The plan whose degrees, kinds and other mode counts are kept.
98 kept: The plans retained so far; a qualifying plan is appended in place.
99 results: The scored alternatives, which need not be sorted.
101 Returns:
102 The appended plan, or ``None`` when no unused qualifying plan is left.
104 Raises:
105 ValueError: A plan has no finite score, or no valid plan per layer.
106 """
107 defined = _layer_settings(current)
108 previous = [(choice.dimensions, _layer_settings(choice)) for choice in kept]
109 for result in sort_recompute_results(results):
110 if result.dimensions != current.dimensions:
111 continue
112 proposed = _layer_settings(result)
113 if proposed.keys() != defined.keys() or not _one_more_full(defined, proposed):
114 continue
115 if any(dimensions == result.dimensions and layers == proposed for dimensions, layers in previous):
116 continue
117 kept.append(result)
118 return result
119 return None
122def recompute_next_plot_groups(
123 scored_space: Sequence[tuple],
124 results: Mapping[Dimensions, Sequence[RecomputeChoice]],
125 top_num: Optional[int] = None,
126) -> List[Tuple[Dimensions, List[RecomputeChoice]]]:
127 """Select the plans below each next strategy's score, and the next best beyond it.
129 Args:
130 scored_space: The parallel configurations in the search's order.
131 results: The scored alternatives of each configuration, its own plan
132 among them.
133 top_num: How many configurations to show; :data:`DEFAULT_TOP` when
134 omitted.
136 Returns:
137 The groups in the search's order, each sorted by score. A group holds
138 every plan the cost model scores below the next strategy, and the
139 first plan at or above it. It keeps at least two plans where there
140 are two, so a group shows an alternative even when the scores tie.
141 The last group takes the mean size of the groups before it, rounded
142 half up, at least two and at most what it has; with one
143 configuration alone it takes up to two plans.
144 """
145 top = scored_space[:max(0, DEFAULT_TOP if top_num is None else top_num)]
146 if not top:
147 return []
148 groups = []
149 for entry, following in zip(top, top[1:]):
150 config = entry[0]
151 choices = sort_recompute_results(results.get(config, ()))
152 count = sum(choice.score < following[2] for choice in choices)
153 groups.append((config, choices[:max(2, count + 1)]))
154 count = math.floor(sum(len(choices) for _, choices in groups) / len(groups) + 0.5) if groups else 2
155 config = top[-1][0]
156 groups.append((config, sort_recompute_results(results.get(config, ()))[:max(2, count)]))
157 return groups
160def recompute_plot_data(
161 groups: Sequence[Tuple[Dimensions, Sequence[RecomputeChoice]]],
162) -> List[Tuple[Dimensions, List[dict]]]:
163 """Prepare the selected plans as chart records, without changing their configuration ranks.
165 Args:
166 groups: The selected plans, grouped by fixed parallel configuration.
168 Returns:
169 The groups with each plan's score, its parts, its peak memory, the
170 count of fully recomputed layers and the plan the trainer reads.
171 """
172 prepared = []
173 for config, choices in groups:
174 records = []
175 for choice in choices:
176 mode, layers = trainer_plan(choice) if choice.ranges or choice.mode is not None else (None, None)
177 records.append({
178 "score": choice.score,
179 "parts": choice.score_parts,
180 "memory": choice.memory,
181 "mode": mode,
182 "layers": layers,
183 "full_layers": _full_layers(choice),
184 })
185 prepared.append((config, records))
186 return prepared
189def _full_layers(choice: RecomputeChoice) -> Optional[int]:
190 """How many layers the plan fully recomputes; ``None`` for a plan that is not stated per layer."""
191 if not choice.ranges:
192 return None
193 return sum(item.count for item in choice.ranges if item.mode == "full")
196def _plan_scope(choice: RecomputeChoice) -> str:
197 """How wide a plan is stated: per layer, for the whole model, with offload, or not at all."""
198 if any(item.option.link_bandwidth for item in choice.ranges):
199 return "offload"
200 if choice.ranges:
201 return "per_layer"
202 return "whole_model" if choice.mode is not None else "unavailable"
205def _csv_row(
206 choice: RecomputeChoice,
207 rank: Any,
208 step: int,
209 previous: Mapping[int, Tuple[Any, str, Any]],
210) -> Dict[str, Any]:
211 """One row of the sequence: the plan, what it changed against the step before, and its estimate."""
212 defined = _layer_settings(choice)
213 changed = [index for index, setting in defined.items() if step and setting != previous[index]]
214 mode, layers = trainer_plan(choice) if choice.ranges or choice.mode is not None else ("", {})
215 row = dict(choice.dimensions)
216 row.update({
217 "TOP": rank,
218 "step": step,
219 "changed_layer": ",".join(map(str, changed)),
220 "full_layers": _full_layers(choice) if choice.ranges else "",
221 "memory_mb": choice.memory,
222 "score": choice.score,
223 "recompute_mode": mode,
224 "recompute_layers": json.dumps(layers),
225 "layer_recompute": json.dumps(to_records(choice)),
226 "stage_memory_mb": json.dumps(choice.stage_memory),
227 "stage_savings": json.dumps(choice.stage_savings),
228 "plan_scope": _plan_scope(choice),
229 })
230 return row
233def write_recompute_csv(potentials: Mapping[Dimensions, Sequence[RecomputeChoice]], path: str) -> None:
234 """Replace a CSV with each fixed strategy's retained recompute sequence.
236 The score of every row is the cost model's estimate of the step under
237 that plan, so a difference between two rows is what the cost model
238 expects the extra recomputed layer to cost.
240 Args:
241 potentials: The retained plans in sequence order, grouped by
242 parallel configuration.
243 path: The CSV to write, whose parent directory is created where it
244 is missing.
245 """
246 target = Path(path)
247 target.parent.mkdir(parents=True, exist_ok=True)
248 degrees = list(dict.fromkeys(
249 name for choices in potentials.values() for choice in choices for name, _ in choice.dimensions
250 ))
251 with target.open("w", encoding="utf-8", newline="") as handle:
252 writer = csv.DictWriter(handle, fieldnames=["TOP", "step"] + degrees + list(_FIXED_COLUMNS))
253 writer.writeheader()
254 for config, choices in potentials.items():
255 previous: Mapping[int, Tuple[Any, str, Any]] = {}
256 for step, choice in enumerate(choices):
257 writer.writerow(_csv_row(choice, config.rank, step, previous))
258 previous = _layer_settings(choice)