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-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"""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 (1/3 of time)
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: float = 0.0,
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 (float, optional): Optional forward time, otherwise derived from ``time``. Default: ``0.0``.
83 backward_time_rec (Optional[Dict[Recompute.TYPE, float]], optional): Per-recomputation-type backward
84 times (``None`` -> zeros). Default: ``None``.
85 backward_coef_rec (Optional[Dict[Recompute.TYPE, float]], optional): Per-recomputation-type overhead
86 coefficients (``None`` -> zeros). Default: ``None``.
87 memory_parameter (float, optional): Parameter memory in MB. Default: ``0.0``.
88 memory_activation_rec (Optional[Dict[Recompute.TYPE, float]], optional): Per-recomputation-type
89 activation memory (``None`` -> zeros). Default: ``None``.
90 """
91 if backward_time_rec is None:
92 backward_time_rec = {r: 0 for r in Recompute.TYPE}
93 if backward_coef_rec is None:
94 backward_coef_rec = {r: 0 for r in Recompute.TYPE}
95 if memory_activation_rec is None:
96 memory_activation_rec = {r: 0.0 for r in Recompute.TYPE}
97 self.name_ = name
98 self.model_name_ = model_name
99 self.type_ = ltype
100 self.nb_layer_ = nb_layer
101 self.time_ = time
102 self.memory_activation_rec_ = memory_activation_rec
103 self.memory_parameter_ = memory_parameter
104 self.backward_time_rec_ = backward_time_rec
105 self.backward_coef_rec_ = backward_coef_rec
106 self.forward_time_ = forward_time
107 self.recompute_considered_ = self.find_recompute_considered()
108 self.compute_internal_time()
110 def __str__(self) -> str:
111 """Return a multi-line, human-readable description of the layer."""
112 result = "Layer Description:\n"
113 result += " name = " + self.name_ + "\n"
114 result += " model_name = " + str(self.model_name_) + "\n"
115 result += " type = " + self.type_.name + "\n"
116 result += " nb_layer = " + str(self.nb_layer_) + "\n"
117 result += " time = " + str(self.time_) + "\n"
118 result += " memory_parameter = " + str(self.memory_parameter_) + "\n"
119 for r in Recompute.TYPE:
120 if self.recompute_considered_[r]:
121 result += " " + Recompute.JSON_MEMORY_NAME[r] + " = "
122 result += str(self.memory_activation_rec_[r]) + "\n"
123 result += " forward_time = "
124 result += str(self.forward_time_) + "\n"
125 for r in Recompute.TYPE:
126 if self.recompute_considered_[r]:
127 result += " " + Recompute.JSON_TIME_NAME[r] + " = "
128 result += str(self.backward_time_rec_[r]) + "\n"
129 return result
131 def dump(self, dump_file: str) -> None:
132 """Dump the layer to ``dump_file`` as JSON (currently a placeholder)."""
133 logger.error("dump file (%s) Not implemented yet!!!", dump_file)
135 def to_json(self) -> None:
136 """Generate the JSON representation of the layer (currently a placeholder)."""
137 logger.error("Not implemented yet!!!")
139 def find_recompute_considered(self) -> Dict[Recompute.TYPE, bool]:
140 """Return which recomputation types have valid activation-memory data."""
141 recompute_considered = {rec: False for rec in Recompute.TYPE}
143 for rec in Recompute.TYPE:
144 if self.memory_activation_rec_[rec] is not None:
145 recompute_considered[rec] = True
147 return recompute_considered
149 def compute_internal_time(
150 self,
151 back_ratio: float = backward_default_ratio,
152 force_fb: bool = False,
153 ) -> None:
154 """Derive forward/backward times from ``time_`` if not already set."""
155 if force_fb or self.forward_time_ is None:
156 self.forward_time_ = self.time_
157 self.backward_time_ = back_ratio * self.time_
159 for rec in Recompute.TYPE:
160 if self.recompute_considered_[rec]:
161 if (
162 self.backward_time_rec_[rec] is None
163 or self.backward_time_rec_[rec] == 0
164 ):
165 if self.backward_coef_rec_[rec] is None:
166 self.backward_time_rec_[rec] = (
167 1 + Recompute.DEFAULT_COEF[rec]
168 ) * self.backward_time_
169 else:
170 self.backward_time_rec_[rec] = (
171 1 + self.backward_coef_rec_[rec]
172 ) * self.backward_time_
174 def update_internal_time_for_seqpp(
175 self,
176 back_ratio: float = backward_default_ratio,
177 force_fb: bool = False,
178 ) -> None:
179 """Adjust ``forward_time_``/``backward_time_`` for the sequence-pipeline mode."""
180 if force_fb or self.forward_time_ is None:
181 self.forward_time_ = (1 - back_ratio) * self.time_
182 self.backward_time_ = back_ratio * self.time_
184 for rec in Recompute.TYPE:
185 if self.recompute_considered_[rec]:
186 if self.backward_coef_rec_[rec] is None:
187 self.backward_time_rec_[rec] = (
188 1 + Recompute.DEFAULT_COEF[rec]
189 ) * self.backward_time_
190 else:
191 self.backward_time_rec_[rec] = (
192 1 + self.backward_coef_rec_[rec]
193 ) * self.backward_time_
195 def compute_timer(
196 self, timeline_folder: str = "./timeline", tmp_layer_info: Optional[dict] = None
197 ) -> None:
198 """Populate ``time_`` from profiling timelines stored in ``timeline_folder``."""
199 layer_time = ComputationAnalyzer(
200 timeline_folder,
201 self.model_name_,
202 num_of_micro_batch=0,
203 layer_list=tmp_layer_info,
204 )
205 self.time_ = layer_time.layer_with_cost_list.get(self.name_)
206 self.compute_internal_time(force_fb=True)
208 def compute_memory(self, memory_folder: str = "./memory") -> None:
209 """Compute the memory information from ``memory_folder`` dry-run logs (placeholder)."""
210 logger.error(
211 "compute_memory (%s) Not implemented yet!!!", memory_folder
212 )
215# Helper functions on layer list
218def generate_layers_list(layer_folder: str, model_name: str) -> list[Layer]:
219 """ "Parse layer_folder/model_name.json to generate a list of layer"""
220 layers = []
221 json_layer = os.path.join(layer_folder, model_name + ".json")
222 with open(json_layer, encoding="utf-8") as json_file:
223 layer_data_json = json.load(json_file)
224 if "layers_description" in layer_data_json:
225 for layer_data in layer_data_json["layers_description"]:
226 new_layer = Layer(
227 name=layer_data["name"],
228 ltype=Layer.type_enum[layer_data["type"]],
229 nb_layer=layer_data["nb_layer"],
230 time=layer_data["time"],
231 model_name=layer_data.get("model_name"),
232 forward_time=layer_data.get("forward_time"),
233 backward_time_rec={
234 r: layer_data.get(Recompute.JSON_TIME_NAME[r])
235 for r in Recompute.TYPE
236 },
237 backward_coef_rec={
238 r: layer_data.get(Recompute.JSON_COEF_NAME[r])
239 for r in Recompute.TYPE
240 },
241 memory_activation_rec={
242 r: layer_data.get(Recompute.JSON_MEMORY_NAME[r])
243 for r in Recompute.TYPE
244 },
245 memory_parameter=layer_data.get("memory_parameter"),
246 )
247 new_layer.compute_internal_time()
248 layers.append(new_layer)
249 else:
250 logger.error(
251 'ERROR: File "%s" doesn\'t have layers_description to parse.\n',
252 json_layer,
253 )
254 return layers
257def filter_layer_type(
258 layers: list[Layer], layer_type: Layer.type_enum
259) -> list[Layer]:
260 """Filters all layers of layer_type in layers."""
261 kept_layers = []
262 for layer in layers:
263 if layer.type_ == layer_type:
264 kept_layers.append(layer)
265 return kept_layers
268def aggregate(layers: list[Layer]) -> Layer:
269 """Aggregate all layers into one."""
271 def add_none(a: Optional[Any], b: Optional[Any]) -> Any:
272 """Add ``a`` and ``b``, returning whichever is not ``None`` when one is missing."""
273 if a is None:
274 return b
275 if b is None:
276 return a
277 return a + b
279 def add_rec_dict(a: Dict[Recompute.TYPE, Any],
280 b: Dict[Recompute.TYPE, Any]) -> Dict[Recompute.TYPE, Any]:
281 """Element-wise add two per-recomputation-type dictionaries."""
282 return {i: a[i] + b[i] for i in Recompute.TYPE}
284 aggregation = layers[0]
285 layers.pop(0)
286 for layer in layers:
287 aggregation.time_ += layer.time_
288 aggregation.backward_time_rec_ = add_rec_dict(
289 aggregation.backward_time_rec_, layer.backward_time_rec_
290 )
291 aggregation.memory_activation_rec_ = add_rec_dict(
292 aggregation.memory_activation_rec_, layer.memory_activation_rec_
293 )
294 aggregation.memory_parameter_ = add_none(
295 aggregation.memory_parameter_, layer.memory_parameter_
296 )
297 aggregation.nb_layer_ = add_none(
298 aggregation.nb_layer_, layer.nb_layer_
299 )
300 return aggregation