Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / utils / computation_analyzer.py: 92%
216 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# ============================================================================
16"""Computation cost analyzer for pipeline balancing."""
18import json
19import os
20import re
21import sys
22from itertools import chain
23from typing import Dict, List, Optional
25from tqdm import tqdm
27from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger
29UNSTABLE_STEPS = 2
32class ComputationAnalyzer:
33 """Parser & Analyzer for profiling timelines"""
35 is_msprof_file = False
37 def __init__(self, timeline_folder_path: str, model_name: str,
38 num_of_micro_batch: int, layer_list: Optional[dict] = None) -> None:
39 """Load profiling timelines and pre-compute the per-layer cost maps.
41 Args:
42 timeline_folder_path: Directory containing timeline JSON files.
43 model_name: Model name used to look up layer metadata in ``cfgs/model_layers.json``.
44 num_of_micro_batch: Micro-batch count for the observed run (used for
45 normalizing totals).
46 layer_list: Optional pre-parsed layer metadata. When ``None`` it is loaded from disk.
47 """
48 self.timeline_folder_path = timeline_folder_path
49 self.model_name = model_name
50 self.num_of_micro_batch = num_of_micro_batch
51 self.counted_steps = 0
52 self.step_time = 0.0
53 self.select_step_number = UNSTABLE_STEPS + 1
54 if layer_list:
55 self.layer_list = layer_list
56 else:
57 self.layer_list = self._get_layer_list()
58 self.timeline_data = self._get_timeline_data()
59 logger.info("parsing layer objs")
60 self.auto_partition_layer_objects, self.pre_defined_layer_objects = (
61 self._parse_layer_objects())
62 logger.info("parsing auto partition layer name")
63 self.auto_partition_layer_name_list = self.parse_auto_partition_layer_name_list()
64 logger.info("parse layer with computation time list")
65 self.layer_with_computation_time_list = self.parse_layer_with_computation_time_list()
66 logger.info("transform layer with cost list")
67 self.layer_with_cost_list = self.transform_layer_with_cost_list()
69 def _get_layer_list(self):
70 """Return cfgs from model config file"""
72 model_config_file = os.path.join(os.getcwd(), "cfgs", "model_layers.json")
73 with open(model_config_file, encoding="utf-8") as json_file:
74 model_layers_data = json.load(json_file)
75 for layer_list in model_layers_data:
76 if self.model_name in layer_list["name"]:
77 return layer_list
78 logger.info("ERROR: Not found model in model config file")
79 return False
81 def _get_timeline_data(self):
82 """Return timeline objects from json file."""
83 logger.info("loading timeline data")
84 timeline_data = []
85 for file_name in [file for file in os.listdir(self.timeline_folder_path) if
86 file.endswith(".json")]:
87 if file_name.startswith("trace_view"):
88 self.is_msprof_file = True
89 elif file_name.startswith("msprof"):
90 self.is_msprof_file = True
91 else:
92 self.is_msprof_file = False
93 logger.error("ERROR: Not support timeline file type")
94 with open(os.path.join(self.timeline_folder_path, file_name), encoding="utf-8") as json_file:
95 timeline_data.append(json.load(json_file))
96 return timeline_data
98 def _parse_step_duration(self, timeline_data):
99 """Return timeline objects during a training step."""
101 op_name = ""
102 step_start = 0.0
103 step_end = 0.0
104 cpt = 0
105 for obj in timeline_data:
106 if "MatMul-op" in obj["name"]:
107 op_name = obj["name"]
108 break
109 for obj in timeline_data:
110 if obj["name"] == op_name:
111 cpt += 1
112 if cpt == self.select_step_number:
113 step_start = float(obj["ts"])
114 if cpt == (self.select_step_number + 1):
115 step_end = float(obj["ts"])
116 step_time = step_end - step_start
117 self.step_time = step_time
118 return (step_start, step_end)
120 def _load_json_data(self, file_path):
121 with open(file_path, encoding="utf-8") as json_file:
122 return json.load(json_file)
124 def _initialize_step_duration(self, timeline_data, step_start, step_end):
125 if step_start == 0 or step_end == 0:
126 step_start, step_end = self._parse_step_duration(timeline_data)
127 return step_start, step_end
129 def _add_layer_object(self, objects_list, condition, obj):
130 if condition and obj not in objects_list:
131 objects_list.append(obj)
133 def _is_counted(self, default_table: list, step_start, step_end, cell_object):
134 """Check if cell in under forward scope"""
135 if float(cell_object["ts"]) < step_start or float(cell_object["ts"]) + float(cell_object["dur"]) > step_end:
136 return False
138 is_counted = False
139 for duration in default_table:
140 start = float(cell_object["ts"])
141 end = float(cell_object["ts"]) + float(cell_object["dur"])
142 if start >= duration[0] and end <= duration[1]:
143 is_counted = True
144 break
145 return is_counted
147 def _forward_parser(self, timeline_data):
148 """Parse time range of forward operators"""
149 logger.info("parsing forward scope")
150 scope_pid = 3
151 default_durations = []
152 cell_durations = {}
153 step_range = []
154 op_name = ""
155 for obj in timeline_data:
156 if obj["name"] == "Scope Layer":
157 scope_pid = obj["pid"]
158 break
159 for obj in tqdm(timeline_data):
160 if op_name == "" and "MatMul-op" in obj["name"]:
161 op_name = obj["name"]
162 if obj["name"] == op_name:
163 step_range.append(float(obj["ts"]))
164 if obj["pid"] != scope_pid:
165 continue
166 if obj["name"] == "Default" and obj["tid"] == 0:
167 start = float(obj["ts"])
168 end = float(obj["ts"]) + float(obj["dur"])
169 default_durations.append((start, end))
170 continue
171 for layer_name in chain(self.layer_list["pre_defined_layer"], self.layer_list["auto_partition_layer"]):
172 if layer_name in obj["name"]:
173 layer_time = cell_durations.get(layer_name)
174 if layer_time is None:
175 cell_durations[layer_name] = []
176 cell_durations[layer_name].append(obj)
178 # step times of first 2 steps are not stable
179 # so we don't consider them when enough steps are given
180 steps = len(step_range)
181 logger.info("There are %s steps in given timeline data", steps)
182 if steps == 0:
183 raise ValueError("Failed to parse timeline")
184 if steps == 1:
185 select_step_start = 0.0
186 select_step_end = sys.float_info.max
187 else:
188 select_step_start = step_range[min(len(step_range) - UNSTABLE_STEPS, UNSTABLE_STEPS)]
189 select_step_end = step_range[-1]
190 logger.info("select_step_start: %f", select_step_start)
191 logger.info("select_step_end: %f", select_step_end)
192 self.counted_steps = max(len(step_range) - (UNSTABLE_STEPS + 1), 1)
193 logger.info("counted_steps: %f", self.counted_steps)
194 return default_durations, cell_durations, select_step_start, select_step_end
196 def _process_timelines(self, timeline_data, step_start, step_end, pre_defined_layer_objects,
197 auto_partition_layer_objects):
198 """_process_file"""
199 logger.info("processing timeline. %d objects in it.", len(timeline_data))
200 if not self.is_msprof_file:
201 step_start, step_end = self._initialize_step_duration(self.timeline_data, step_start, step_end)
202 default_durations, cell_durations, step_start, step_end = self._forward_parser(timeline_data)
203 for cell_name, cell_objs in cell_durations.items():
204 for obj in cell_objs:
205 if not self._is_counted(default_durations, step_start, step_end, obj):
206 continue
207 for layer_name in self.layer_list["pre_defined_layer"]:
208 if layer_name in cell_name:
209 self._add_layer_object(pre_defined_layer_objects, True, obj)
210 for layer_name in self.layer_list["auto_partition_layer"]:
211 if layer_name in cell_name:
212 self._add_layer_object(auto_partition_layer_objects, True, obj)
213 return step_start, step_end
215 def _parse_layer_objects(self):
216 auto_partition_layer_objects = []
217 pre_defined_layer_objects = []
218 step_start, step_end = 0, 0
219 for timeline in self.timeline_data:
220 step_start, step_end = self._process_timelines(timeline, step_start, step_end,
221 pre_defined_layer_objects,
222 auto_partition_layer_objects)
223 return auto_partition_layer_objects, pre_defined_layer_objects
225 def parse_auto_partition_layer_name_list(self) -> List[str]:
226 """example: [42-TransformerEncoderLayer,43-TransformerEncoderLayer]"""
227 auto_partition_layer_name_list = []
228 for auto_partition_name in self.layer_list["auto_partition_layer"]:
229 for obj in [item for timeline in self.timeline_data for item in timeline]:
230 object_name = obj["name"]
231 if auto_partition_name in object_name:
232 if self.is_msprof_file:
233 find_layer_name = re.findall(r"[0-9]*-" + auto_partition_name,
234 object_name)
235 layer_name = find_layer_name[0]
236 else:
237 layer_name = object_name
238 if layer_name not in auto_partition_layer_name_list:
239 auto_partition_layer_name_list.append(layer_name)
240 return auto_partition_layer_name_list
242 def parse_layer_with_computation_time_list(self) -> Dict[str, float]:
243 """
244 Map each layer_name with its duration time.
245 For example: [46-TransformerEncoderLayer':37.24729124999999, '47-TransformerEncoderLayer': 37.36572429687501]
246 """
247 layer_with_computation_time_list = {}
248 for pre_defined_layer_name in self.layer_list["pre_defined_layer"]:
249 layer_with_computation_time_list[pre_defined_layer_name] = 0
250 for auto_partition_layer_name in self.auto_partition_layer_name_list:
251 layer_with_computation_time_list[auto_partition_layer_name] = 0
253 for obj in self.pre_defined_layer_objects:
254 for pre_defined_layer_name in self.layer_list["pre_defined_layer"]:
255 if pre_defined_layer_name in obj["name"]:
256 layer_with_computation_time_list[pre_defined_layer_name] += (float(obj["dur"]) / 1000)
257 for obj in self.auto_partition_layer_objects:
258 for auto_partition_layer_name in self.auto_partition_layer_name_list:
259 # if auto_partition_layer_name in obj["name"]:
260 if re.search(rf"\b{re.escape(auto_partition_layer_name)}", obj["name"]):
261 layer_with_computation_time_list[auto_partition_layer_name] += (
262 float(obj["dur"]) / 1000)
263 return layer_with_computation_time_list
265 def transform_layer_with_cost_list(self) -> Dict[str, float]:
266 """calculating the average value of layer cost"""
267 total_cost_auto_partition_layer = {}
268 number_of_auto_partition_layer = {}
269 transform_layer_with_cost_list = {}
270 for pre_defined_layer_name in self.layer_list["pre_defined_layer"]:
271 transform_layer_with_cost_list[pre_defined_layer_name] = 0
273 for auto_partition_layer_name in self.layer_list["auto_partition_layer"]:
274 total_cost_auto_partition_layer[auto_partition_layer_name] = 0
275 number_of_auto_partition_layer[auto_partition_layer_name] = 0
276 transform_layer_with_cost_list[auto_partition_layer_name] = 0
278 # test
279 for layer_name, layer_time in self.layer_with_computation_time_list.items():
280 for pre_defined_layer_name in self.layer_list["pre_defined_layer"]:
281 if pre_defined_layer_name in layer_name:
282 var_tmp = layer_time / self.counted_steps / self.num_of_micro_batch
283 transform_layer_with_cost_list[pre_defined_layer_name] += var_tmp
284 for auto_partition_layer_name in self.layer_list["auto_partition_layer"]:
285 if auto_partition_layer_name in layer_name:
286 # assuming that the duration time of a layer can not exceed 10% of step time
287 # in order to
288 # avoid some specific long time layers that caused from the error caused by
289 # time_line.json
290 if not self.is_msprof_file and layer_time > self.step_time / 1000 / 10:
291 continue
292 total_cost_auto_partition_layer[auto_partition_layer_name] += \
293 layer_time
294 number_of_auto_partition_layer[auto_partition_layer_name] += 1
296 for auto_partition_layer_name in self.layer_list["auto_partition_layer"]:
297 transform_layer_with_cost_list[auto_partition_layer_name] = (
298 total_cost_auto_partition_layer[auto_partition_layer_name] /
299 number_of_auto_partition_layer[auto_partition_layer_name] /
300 self.counted_steps / self.num_of_micro_batch)
301 return transform_layer_with_cost_list
304if __name__ == "__main__":
305 path = "/your/path/here"
306 example_model_name = "LLaMA_prof"
307 comp1 = ComputationAnalyzer(path, example_model_name, 8)
308 logger.info(comp1.layer_with_computation_time_list)
309 logger.info(comp1.layer_with_cost_list)