Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / utils / layer.py: 99%
130 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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"""Layer descriptor used throughout SAPP-PPB: time, memory and recomputation metadata."""
16import json
17import os
18from enum import Enum
19from typing import Any, Dict, Optional
21import hyper_parallel.auto_parallel.sapp_ppb.utils.recompute as Recompute
22from hyper_parallel.auto_parallel.sapp_ppb.utils.computation_analyzer import ComputationAnalyzer
23from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger
26class Layer:
27 """
28 Mandatory parameter:
29 name_ (str): name of the layer
30 type_ (LayerType): type of the layer 'HEAD', 'BODY', 'TAIL'
31 nb_layer_ (int): number of layer to schedule
32 time_ (float): total time that a layer take
34 Optional (auto-compute) parameter:
35 forward_time_ (float): forward time for the layer (auto-derived from ``time_`` when not provided)
36 backward_time_rec_ (dict[Recompute.Type, float]): backward time (2/3 of time) per recomputation
37 recompute_considered_: dict[Recompute.Type, bool] set recomputations when considered
39 Optional memory parameter (for recompute):
40 memory_parameter_ (float): memory used by the layer (all kind)
41 memory_activation_rec_ (dict[Recompute.Type, float]): activation memory per recomputation
43 Not manage yet parameter (for multimodal):
44 model_name_ (str): name of the model the layer be part of (for multimodal)
45 """
47 type_enum = Enum("LayerType", ["UNKNOWN", "HEAD", "BODY", "TAIL"])
48 backward_default_ratio = 2 # of forward time
49 name_: str
50 model_name_: str
51 type_: type_enum
52 nb_layer_: int
53 time_: float
54 memory_parameter_: float
55 memory_activation_rec_: dict[Recompute.TYPE, float]
56 forward_time_: float
57 backward_time_rec_: dict[Recompute.TYPE, float]
58 backward_coef_rec_: dict[Recompute.TYPE, float]
59 recompute_considered_: dict[Recompute.TYPE, bool]
61 def __init__(
62 self,
63 model_name: str = "misc",
64 name: str = "misc",
65 ltype: type_enum = type_enum.UNKNOWN,
66 nb_layer: int = 0,
67 time: float = 0.0,
68 forward_time: Optional[float] = None,
69 backward_time_rec: Optional[Dict[Recompute.TYPE, float]] = None,
70 backward_coef_rec: Optional[Dict[Recompute.TYPE, float]] = None,
71 memory_parameter: float = 0.0,
72 memory_activation_rec: Optional[Dict[Recompute.TYPE, float]] = None,
73 ) -> None:
74 """Build a :class:`Layer` record.
76 Args:
77 model_name (str): Name of the owning model (``"misc"`` for ad-hoc entries).
78 name (str): Layer name.
79 ltype (Layer.type_enum): HEAD / BODY / TAIL / UNKNOWN classification.
80 nb_layer (int): Number of such layers present in the model.
81 time (float): Total time (forward + backward) for one layer.
82 forward_time (Optional[float], optional): Optional forward time; ``None`` means
83 derive from ``time`` via :meth:`compute_internal_time`. Default: ``None``.
84 backward_time_rec (Optional[Dict[Recompute.TYPE, float]], optional): Per-recomputation-type backward
85 times (``None`` -> zeros). Default: ``None``.
86 backward_coef_rec (Optional[Dict[Recompute.TYPE, float]], optional): Per-recomputation-type overhead
87 coefficients (``None`` -> zeros). Default: ``None``.
88 memory_parameter (float, optional): Parameter memory in MB. Default: ``0.0``.
89 memory_activation_rec (Optional[Dict[Recompute.TYPE, float]], optional): Per-recomputation-type
90 activation memory (``None`` -> zeros). Default: ``None``.
91 """
92 if backward_time_rec is None:
93 backward_time_rec = {r: 0 for r in Recompute.TYPE}
94 if backward_coef_rec is None:
95 backward_coef_rec = {r: 0 for r in Recompute.TYPE}
96 if memory_activation_rec is None:
97 memory_activation_rec = {r: 0.0 for r in Recompute.TYPE}
98 self.name_ = name
99 self.model_name_ = model_name
100 self.type_ = ltype
101 self.nb_layer_ = nb_layer
102 self.time_ = time
103 self.memory_activation_rec_ = memory_activation_rec
104 self.memory_parameter_ = memory_parameter
105 self.backward_time_rec_ = backward_time_rec
106 self.backward_coef_rec_ = backward_coef_rec
107 self.forward_time_ = forward_time
108 self.recompute_considered_ = self.find_recompute_considered()
109 self.compute_internal_time()
111 def __str__(self) -> str:
112 """Return a multi-line, human-readable description of the layer."""
113 result = "Layer Description:\n"
114 result += " name = " + self.name_ + "\n"
115 result += " model_name = " + str(self.model_name_) + "\n"
116 result += " type = " + self.type_.name + "\n"
117 result += " nb_layer = " + str(self.nb_layer_) + "\n"
118 result += " time = " + str(self.time_) + "\n"
119 result += " memory_parameter = " + str(self.memory_parameter_) + "\n"
120 for r in Recompute.TYPE:
121 if self.recompute_considered_[r]:
122 result += " " + Recompute.JSON_MEMORY_NAME[r] + " = "
123 result += str(self.memory_activation_rec_[r]) + "\n"
124 result += " forward_time = "
125 result += str(self.forward_time_) + "\n"
126 for r in Recompute.TYPE:
127 if self.recompute_considered_[r]:
128 result += " " + Recompute.JSON_TIME_NAME[r] + " = "
129 result += str(self.backward_time_rec_[r]) + "\n"
130 return result
132 def dump(self, dump_file: str) -> None:
133 """Dump the layer to ``dump_file`` as JSON (currently a placeholder)."""
134 logger.error("dump file (%s) Not implemented yet!!!", dump_file)
136 def to_json(self) -> None:
137 """Generate the JSON representation of the layer (currently a placeholder)."""
138 logger.error("Not implemented yet!!!")
140 def find_recompute_considered(self) -> Dict[Recompute.TYPE, bool]:
141 """Return which recomputation types have valid activation-memory data."""
142 recompute_considered = {rec: False for rec in Recompute.TYPE}
144 for rec in Recompute.TYPE:
145 if self.memory_activation_rec_[rec] is not None:
146 recompute_considered[rec] = True
148 return recompute_considered
150 def compute_internal_time(
151 self,
152 back_ratio: float = backward_default_ratio,
153 force_fb: bool = False,
154 ) -> None:
155 """Derive forward/backward times from ``time_`` if not already set."""
156 if force_fb or self.forward_time_ is None:
157 self.forward_time_ = self.time_
158 self.backward_time_ = back_ratio * self.time_
160 for rec in Recompute.TYPE:
161 if self.recompute_considered_[rec]:
162 if (
163 self.backward_time_rec_[rec] is None
164 or self.backward_time_rec_[rec] == 0
165 ):
166 if self.backward_coef_rec_[rec] is None:
167 self.backward_time_rec_[rec] = (
168 1 + Recompute.DEFAULT_COEF[rec]
169 ) * self.backward_time_
170 else:
171 self.backward_time_rec_[rec] = (
172 1 + self.backward_coef_rec_[rec]
173 ) * self.backward_time_
175 def update_internal_time_for_seqpp(
176 self,
177 back_ratio: float = backward_default_ratio,
178 force_fb: bool = False,
179 ) -> None:
180 """Adjust ``forward_time_``/``backward_time_`` for the sequence-pipeline mode."""
181 if force_fb or self.forward_time_ is None:
182 self.forward_time_ = (1 - back_ratio) * self.time_
183 self.backward_time_ = back_ratio * self.time_
185 for rec in Recompute.TYPE:
186 if self.recompute_considered_[rec]:
187 if self.backward_coef_rec_[rec] is None:
188 self.backward_time_rec_[rec] = (
189 1 + Recompute.DEFAULT_COEF[rec]
190 ) * self.backward_time_
191 else:
192 self.backward_time_rec_[rec] = (
193 1 + self.backward_coef_rec_[rec]
194 ) * self.backward_time_
196 def compute_timer(
197 self, timeline_folder: str = "./timeline", tmp_layer_info: Optional[dict] = None
198 ) -> None:
199 """Populate ``time_`` from profiling timelines stored in ``timeline_folder``."""
200 layer_time = ComputationAnalyzer(
201 timeline_folder,
202 self.model_name_,
203 num_of_micro_batch=0,
204 layer_list=tmp_layer_info,
205 )
206 self.time_ = layer_time.layer_with_cost_list.get(self.name_)
207 self.compute_internal_time(force_fb=True)
209 def compute_memory(self, memory_folder: str = "./memory") -> None:
210 """Compute the memory information from ``memory_folder`` dry-run logs (placeholder)."""
211 logger.error(
212 "compute_memory (%s) Not implemented yet!!!", memory_folder
213 )
216# Helper functions on layer list
219def generate_layers_list(layer_folder: str, model_name: str) -> list[Layer]:
220 """ "Parse layer_folder/model_name.json to generate a list of layer"""
221 layers = []
222 json_layer = os.path.join(layer_folder, model_name + ".json")
223 with open(json_layer, encoding="utf-8") as json_file:
224 layer_data_json = json.load(json_file)
225 if "layers_description" in layer_data_json:
226 for layer_data in layer_data_json["layers_description"]:
227 new_layer = Layer(
228 name=layer_data["name"],
229 ltype=Layer.type_enum[layer_data["type"]],
230 nb_layer=layer_data["nb_layer"],
231 time=layer_data["time"],
232 model_name=layer_data.get("model_name"),
233 forward_time=layer_data.get("forward_time"),
234 backward_time_rec={
235 r: layer_data.get(Recompute.JSON_TIME_NAME[r])
236 for r in Recompute.TYPE
237 },
238 backward_coef_rec={
239 r: layer_data.get(Recompute.JSON_COEF_NAME[r])
240 for r in Recompute.TYPE
241 },
242 memory_activation_rec={
243 r: layer_data.get(Recompute.JSON_MEMORY_NAME[r])
244 for r in Recompute.TYPE
245 },
246 memory_parameter=layer_data.get("memory_parameter"),
247 )
248 new_layer.compute_internal_time()
249 layers.append(new_layer)
250 else:
251 logger.error(
252 'ERROR: File "%s" doesn\'t have layers_description to parse.\n',
253 json_layer,
254 )
255 return layers
258def filter_layer_type(
259 layers: list[Layer], layer_type: Layer.type_enum
260) -> list[Layer]:
261 """Filters all layers of layer_type in layers."""
262 kept_layers = []
263 for layer in layers:
264 if layer.type_ == layer_type:
265 kept_layers.append(layer)
266 return kept_layers
269def aggregate(layers: list[Layer]) -> Layer:
270 """Aggregate all layers into one."""
272 def add_none(a: Optional[Any], b: Optional[Any]) -> Any:
273 """Add ``a`` and ``b``, returning whichever is not ``None`` when one is missing."""
274 if a is None:
275 return b
276 if b is None:
277 return a
278 return a + b
280 def add_rec_dict(a: Dict[Recompute.TYPE, Any],
281 b: Dict[Recompute.TYPE, Any]) -> Dict[Recompute.TYPE, Any]:
282 """Element-wise add two per-recomputation-type dictionaries."""
283 return {i: a[i] + b[i] for i in Recompute.TYPE}
285 aggregation = layers[0]
286 layers.pop(0)
287 for layer in layers:
288 aggregation.time_ += layer.time_
289 aggregation.backward_time_rec_ = add_rec_dict(
290 aggregation.backward_time_rec_, layer.backward_time_rec_
291 )
292 aggregation.memory_activation_rec_ = add_rec_dict(
293 aggregation.memory_activation_rec_, layer.memory_activation_rec_
294 )
295 aggregation.memory_parameter_ = add_none(
296 aggregation.memory_parameter_, layer.memory_parameter_
297 )
298 aggregation.nb_layer_ = add_none(
299 aggregation.nb_layer_, layer.nb_layer_
300 )
301 return aggregation