Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / utils / config.py: 89%
334 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-description JSON generator and YAML-based dry-run configuration parser."""
16import json
17import os
18import random
19from dataclasses import asdict, dataclass
20from math import ceil
21from typing import Any, Dict, List, Optional, Tuple, Union
23import numpy as np
24import yaml
26import hyper_parallel.auto_parallel.sapp_ppb.utils.recompute as Recompute
27from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_solver import SappSolver
28from hyper_parallel.auto_parallel.sapp_ppb.utils.check_rules import check_yaml_depth_before_loading
29from hyper_parallel.auto_parallel.sapp_ppb.utils.computation_analyzer import ComputationAnalyzer
30from hyper_parallel.auto_parallel.sapp_ppb.utils.compute_memory import ComputeMemory, Stage
31from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import Layer
32from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger
34random.seed()
37@dataclass
38class LayersDescription:
39 """layers description"""
41 name: str
42 type: Layer.type_enum
43 model_name: str
44 time: int
45 nb_layer: int
46 memory_parameter: int
48 def __init__(
49 self, layer_type: Layer.type_enum, time: int, nb_layer: int, model_name: str
50 ) -> None:
51 """Store layer metadata using the enum ``layer_type.name`` for both ``type`` and ``name``."""
52 self.type = layer_type.name
53 self.time = time
54 self.name = layer_type.name
55 self.nb_layer = nb_layer
56 self.model_name = model_name
59@dataclass
60class ModelInfo:
61 """basic info of a model"""
63 name: str
64 stage_const_mem: int
65 layers_description: list[LayersDescription]
67 def __init__(self, model_name: str, head_time: int, body_time: int,
68 tail_time: int, nb_layer: int) -> None:
69 """Record HEAD/BODY/TAIL timing and initialise the stage-constant memory to 0."""
70 self.name = model_name
71 self.stage_const_mem = 0
72 self.to_json_ = None
73 self.layers_description = []
74 self.layers_description.append(
75 LayersDescription(Layer.type_enum.HEAD, head_time, 1, model_name)
76 )
77 self.layers_description.append(
78 LayersDescription(Layer.type_enum.BODY, body_time, nb_layer, model_name)
79 )
80 self.layers_description.append(
81 LayersDescription(Layer.type_enum.TAIL, tail_time, 1, model_name)
82 )
84 def get_layer_by_type(self, layer_type: Layer.type_enum) -> Optional[LayersDescription]:
85 """Return the layer description matching ``layer_type`` or ``None`` if absent."""
86 for layer in self.layers_description:
87 if layer.type == layer_type.name:
88 return layer
89 return None
91 def set_stage_const_mem(self, mem_const: int) -> None:
92 """Override the per-stage constant memory component."""
93 self.stage_const_mem = mem_const
95 def layer_memory_update(self, mem_act: Dict[Recompute.TYPE, int], mem_par: int,
96 mem_head: int, mem_tail: int) -> None:
97 """Apply the solver-derived memory values onto the internal JSON payload."""
98 self.get_layer_by_type(Layer.type_enum.HEAD).memory_parameter = mem_head
99 self.get_layer_by_type(Layer.type_enum.TAIL).memory_parameter = mem_tail
100 self.get_layer_by_type(Layer.type_enum.BODY).memory_parameter = mem_par
102 json_data = asdict(self)
104 for rec in Recompute.TYPE:
105 rec_mem = mem_act.get(rec)
106 if rec_mem is None:
107 continue
108 for layer in json_data["layers_description"]:
109 if layer["type"] == Layer.type_enum.BODY.name:
110 layer[Recompute.JSON_MEMORY_NAME[rec]] = rec_mem
111 self.to_json_ = json_data
113 def dump_json(self, file_name: str) -> None:
114 """Write the current ``to_json_`` payload to ``file_name`` as indented JSON."""
115 with open(file_name, "w", encoding="utf-8") as json_file:
116 json.dump(self.to_json_, json_file, indent=4)
119def _as_list(value: Union[str, List[str]]) -> List[str]:
120 """Wrap a bare string in a single-element list, leaving other values unchanged."""
121 return [value] if isinstance(value, str) else value
124def _accumulate_layer_times(cost_list: Dict[str, float], predefined_layers: Dict[str, int],
125 head_time: float, body_time: float, tail_time: float
126 ) -> Tuple[float, float, float]:
127 """Add each measured layer cost to the HEAD/TAIL/BODY total based on its marker."""
128 for layer, time in cost_list.items():
129 marker = predefined_layers.get(layer)
130 if marker == 0:
131 head_time += time
132 elif marker == -1:
133 tail_time += time
134 else:
135 body_time += time
136 return head_time, body_time, tail_time
139def time_parser(file_name: str, model_name: str) -> Tuple[float, float, float]:
140 """Parse the HEAD/BODY/TAIL timing values from a SAPP-PPB YAML config file."""
141 if file_name is None:
142 logger.error("input file cannot be none")
143 raise ValueError("input file cannot be none")
145 if not file_name.endswith(("yaml", "yml")):
146 logger.error("Only accept yaml as input format")
147 raise ValueError(f"Only accept yaml as input format. not {file_name}")
149 filepath = os.path.realpath(file_name)
150 with open(filepath, encoding="utf-8") as fp:
151 check_yaml_depth_before_loading(fp)
152 fp.seek(0)
153 cfg_dict = yaml.safe_load(fp)
155 head_time = 0
156 body_time = 0
157 tail_time = 0
159 if "time_config" in cfg_dict:
160 head_time = cfg_dict["time_config"].get("head")
161 body_time = cfg_dict["time_config"].get("body")
162 tail_time = cfg_dict["time_config"].get("tail")
163 if {"head", "body", "tail"}.issubset(cfg_dict["time_config"]):
164 return head_time, body_time, tail_time
166 if cfg_dict.get("profiling_config"):
167 head_layers = cfg_dict["profiling_config"].get("head_layers", ["LlamaEmbedding"])
168 body_layers = cfg_dict["profiling_config"].get("body_layers", ["LLamaDecodeLayer"])
169 tail_layers = cfg_dict["profiling_config"].get("tail_layers", ["lm_head-Linear", "LlamaRMSNorm"])
170 head_layers = _as_list(head_layers)
171 tail_layers = _as_list(tail_layers)
172 body_layers = _as_list(body_layers)
174 num_layer = cfg_dict["pipeline_config"]["num_layer"]
175 micro_batch_num = cfg_dict["profiling_config"]["micro_batch_num"]
176 timeline_folder_path = cfg_dict["profiling_config"]["folder_path"]
177 layer_list = {
178 "pre_defined_layer": {
179 **dict.fromkeys(head_layers, 0),
180 **dict.fromkeys(tail_layers, -1),
181 },
182 "auto_partition_layer": dict.fromkeys(body_layers, num_layer),
183 }
184 analyzer = ComputationAnalyzer(timeline_folder_path, model_name, micro_batch_num, layer_list)
185 cost_list = analyzer.layer_with_cost_list
186 logger.info(cost_list)
187 head_time, body_time, tail_time = _accumulate_layer_times(
188 cost_list, layer_list["pre_defined_layer"], head_time, body_time, tail_time
189 )
191 logger.info("head_time: %s, body_time: %s, tail_time: %s", head_time, body_time, tail_time)
193 return head_time, body_time, tail_time
196def process_offset(offset: Union[int, List[int], List[List[int]]],
197 pipeline_num: int) -> Tuple[Union[List[int], List[List[int]]], int]:
198 """Normalise the YAML ``offset`` field into a list (or list of lists) and return the rounds."""
199 rounds = 1
200 if isinstance(offset, int) and offset == 0:
201 offset = [0] * pipeline_num
202 # if offset is list of lists (usually when pp=4)
203 elif isinstance(offset, list) and any(isinstance(item, list) for item in offset):
204 tmp_offset = []
205 for item in offset:
206 if isinstance(item, int) and item == 0:
207 tmp_offset.append([0] * pipeline_num)
208 elif not (isinstance(item, list) and len(item) == pipeline_num):
209 raise ValueError(f"Unsupported input format offset: {item},",
210 "please check the length of your offset list and the pipeline number")
211 else:
212 tmp_offset.append(item)
213 offset = tmp_offset
214 rounds = len(offset)
215 elif not (isinstance(offset, list) and len(offset) == pipeline_num):
216 raise ValueError(f"Unsupported input format offset: {offset},",
217 "please check the length of your offset list and the pipeline number")
219 return offset, rounds
222def process_rec_config(
223 layer_per_stage: int, pipeline_num: int, offset: List[int],
224 rec_config: Optional[Union[bool, List[int]]]
225) -> Optional[List[List[int]]]:
226 """Normalise a per-recomputation-type YAML entry into a two-dimensional list."""
227 if rec_config is None or offset is None:
228 return None
229 if isinstance(rec_config, bool):
230 if rec_config:
231 rec_config = [layer_per_stage] * pipeline_num
232 rec_config = [recom + bias for recom, bias in (rec_config, offset)]
233 else:
234 rec_config = [0] * pipeline_num
235 rec_config = [rec_config]
236 elif isinstance(rec_config, list) and len(rec_config) == pipeline_num:
237 # in order to be compatible with internal_from_yaml, change list into double list
238 rec_config = [rec_config]
239 else:
240 raise ValueError(f"Unsupported input format recompute: {rec_config}, please check the length of list")
242 return rec_config
245def instantiate_stage(stage_id: int, pipeline_num: int, nb_layer: int,
246 layer_per_recompute: Dict[Recompute.TYPE, List[List[int]]],
247 memory: int) -> Stage:
248 """Build a :class:`Stage` snapshot from per-recomputation-type layer assignments."""
249 stage = Stage(
250 sid=stage_id,
251 nb_stage=pipeline_num,
252 nb_layer=nb_layer,
253 nb_layer_rec={
254 Recompute.TYPE.COMM: layer_per_recompute[Recompute.TYPE.COMM][0][stage_id],
255 Recompute.TYPE.FULL: layer_per_recompute[Recompute.TYPE.FULL][0][stage_id],
256 Recompute.TYPE.SLCT: layer_per_recompute[Recompute.TYPE.SLCT][0][stage_id],
257 Recompute.TYPE.BOTH: layer_per_recompute[Recompute.TYPE.BOTH][0][stage_id],
258 },
259 memory_usage=memory,
260 )
261 return stage
264def memory_parser(file_name: str) -> Tuple[int, List[Stage], int]:
265 """Parse a SAPP-PPB memory-usage YAML file into solver-ready stage observations."""
266 if file_name is None:
267 logger.error("input file cannot be none")
268 raise ValueError("input file cannot be none")
269 if not file_name.endswith("yaml") and not file_name.endswith("yml"):
270 logger.error("Only accept yaml as input format")
271 raise ValueError(f"Only accept yaml as input format. not {file_name}")
273 filepath = os.path.realpath(file_name)
274 with open(filepath, encoding="utf-8") as fp:
275 check_yaml_depth_before_loading(fp)
276 fp.seek(0)
277 cfg_dict = yaml.safe_load(fp)
279 # get pipeline config
280 pipeline_num = cfg_dict["pipeline_config"]["pipeline_num"]
281 num_layer = cfg_dict["pipeline_config"]["num_layer"]
282 offset = cfg_dict["pipeline_config"]["offset"]
284 offset, rounds = process_offset(offset, pipeline_num)
286 layer_per_stage = int(num_layer / pipeline_num)
288 # get recompute config
289 if rounds > 1:
290 layer_per_recompute = []
291 for i in range(rounds):
292 rec_config = {}
293 for rec in Recompute.YAML_NAME.values():
294 rec_list = cfg_dict["recompute_config"].get(rec)
295 if rec_list is None:
296 continue
297 rec_config[rec] = process_rec_config(layer_per_stage, pipeline_num, offset[i], rec_list[i])
298 rec_config[Recompute.OFFSET] = [offset[i]]
299 layer_per_recompute.append(
300 Recompute.internal_from_yaml(1, pipeline_num, rec_config, [[layer_per_stage] * pipeline_num])
301 )
302 else:
303 rec_config = {}
304 for rec in Recompute.YAML_NAME.values():
305 rec_list = cfg_dict["recompute_config"].get(rec)
306 rec_config[rec] = process_rec_config(layer_per_stage, pipeline_num, offset, rec_list)
307 rec_config[Recompute.OFFSET] = [offset]
308 layer_per_recompute = Recompute.internal_from_yaml(1, pipeline_num, rec_config,
309 [[layer_per_stage] * pipeline_num])
310 # get memory usage
311 stage_id = cfg_dict["memory_usage"]["body_memories"]["stage_id"]
312 mem_head_stage = cfg_dict["memory_usage"]["head_memory"]
313 mem_tail_stage = cfg_dict["memory_usage"]["tail_memory"]
314 body_memories = cfg_dict["memory_usage"]["body_memories"]["memories"]
315 stages_a = []
316 if rounds > 1:
317 for i in range(rounds):
318 for idx, sg_id in enumerate(stage_id[i]):
319 stages_a.append(
320 instantiate_stage(
321 sg_id, pipeline_num,
322 layer_per_stage + offset[i][sg_id],
323 layer_per_recompute[i], body_memories[i][idx],
324 )
325 )
326 stages_a.append(
327 instantiate_stage(
328 0, pipeline_num, layer_per_stage + offset[i][0],
329 layer_per_recompute[i], mem_head_stage,
330 )
331 )
332 stages_a.append(
333 instantiate_stage(
334 pipeline_num - 1, pipeline_num,
335 layer_per_stage + offset[i][pipeline_num - 1],
336 layer_per_recompute[i], mem_tail_stage,
337 )
338 )
339 else:
340 for idx, sg_id in enumerate(stage_id):
341 stages_a.append(
342 instantiate_stage(
343 sg_id, pipeline_num, layer_per_stage + offset[sg_id],
344 layer_per_recompute, body_memories[idx],
345 )
346 )
347 stages_a.append(
348 instantiate_stage(
349 0, pipeline_num, layer_per_stage + offset[0],
350 layer_per_recompute, mem_head_stage,
351 )
352 )
353 stages_a.append(
354 instantiate_stage(
355 pipeline_num - 1, pipeline_num, layer_per_stage + offset[pipeline_num - 1],
356 layer_per_recompute, mem_tail_stage,
357 )
358 )
360 return pipeline_num, stages_a, num_layer
363def initialize_layer_json(model_name: str, file_name: str) -> None:
364 """Derive a layer-description JSON from a dry-run YAML and dump it to ``./layers``."""
365 num_stage, stages_a, num_layer = memory_parser(file_name)
366 head_time, body_time, tail_time = time_parser(file_name, model_name)
367 comp_mem = ComputeMemory(number_of_stage=num_stage, stages_a=stages_a)
368 mi = ModelInfo(model_name, head_time, body_time, tail_time, num_layer)
370 mem_act = {}
371 for r in Recompute.TYPE:
372 if comp_mem.recompute_considered_[r]:
373 mem_act[r] = int(comp_mem.get_memory_activation(r))
374 logger.info(
375 "[INFO] %s = %f", Recompute.JSON_MEMORY_NAME[r],
376 int(comp_mem.get_memory_activation(r))
377 )
378 if comp_mem.get_memory_const() is not None:
379 mem_const = int(comp_mem.get_memory_const())
380 logger.info("[INFO] memory_const = %s", mem_const)
381 mi.set_stage_const_mem(mem_const)
382 mem_par = int(comp_mem.get_memory_parameter())
383 mem_tail = int(comp_mem.get_memory_tail())
384 mem_head = int(comp_mem.get_memory_head())
385 logger.info("[INFO] memory_parameter = %s", mem_par)
386 logger.info("[INFO] memory_tail = %s", mem_tail)
387 logger.info("[INFO] memory_head = %s", mem_head)
389 mi.layer_memory_update(mem_act, mem_par, mem_head, mem_tail)
390 mi.dump_json(os.path.join("./layers", model_name + ".json"))
393def get_stage_const_mem(layer_folder: str, model_name: str) -> int:
394 """Read ``stage_const_mem`` from ``<layer_folder>/<model_name>.json`` (0 if missing)."""
395 json_layer = os.path.join(layer_folder, model_name + '.json')
396 with open(json_layer, encoding="utf-8") as json_file:
397 json_data = json.load(json_file)
398 if "stage_const_mem" in json_data:
399 return json_data["stage_const_mem"]
400 return 0
403def _generate_offset_config(rounds, unknowns, target_sum, array_length):
404 """Generate legal random offset arrays"""
405 while True:
406 offset_config_list = []
407 flat = []
408 for _ in range(rounds):
409 offset_config = _generate_offset_array(target_sum, array_length)
410 offset_config_slice = offset_config[1:]
411 offset_config_slice = offset_config_slice[:-1]
412 flat.append(offset_config_slice)
413 offset_config_list.append(offset_config)
414 flat = np.array(flat)
415 flat = flat.flatten()[0 : unknowns + 1]
416 if not np.all(flat == flat[0]):
417 return offset_config_list
420def _generate_offset_array(target_sum, array_length):
421 """Generate a random offset array"""
422 if target_sum == array_length:
423 return [0] * array_length
425 if target_sum < array_length:
426 logger.error("number of layers must be larger than stage number")
427 return None
429 random_array = np.random.randint(1, 10, size=array_length)
430 total_sum = random_array.sum()
431 scaled_array = (random_array / total_sum) * target_sum
432 scaled_array = np.round(scaled_array).astype(int)
433 scaled_array = np.maximum(scaled_array, 1)
434 current_sum = scaled_array.sum()
435 diff = target_sum - current_sum
437 if diff >= 0:
438 for i in range(abs(diff)):
439 scaled_array[i] += 1
440 else:
441 count = 0
442 for i in range(len(scaled_array)):
443 # backwards iteration to avoid infinite loop
444 if scaled_array[-1 - i] > 1:
445 scaled_array[-1 - i] -= 1
446 count += 1
447 if count == abs(diff):
448 break
450 baseline = target_sum // array_length
451 offset = scaled_array - baseline
452 return offset.tolist()
455def _get_coef_matrix(pp, layer_per_stage, offset_config_list, rec_config_list, considered_rec):
456 """get coef matrix of equations"""
457 activation_nums = SappSolver.compute_activation_nums(pp, 1, 0)[0]
458 coef_matrix = []
459 rounds = len(offset_config_list)
460 for round_ in range(rounds):
461 for stage in range(pp):
462 if stage not in [0, pp - 1]:
463 coef_matrix.append(
464 [1, layer_per_stage + offset_config_list[round_][stage]]
465 + Recompute.to_list(
466 {
467 rec: rec_config_list[round_][rec][stage]
468 * activation_nums[stage]
469 for rec in considered_rec
470 }
471 )
472 )
473 if len(coef_matrix) == 2 + len(considered_rec):
474 return coef_matrix
475 return None
478def print_dryrun_config(offset_config_list: List[List[int]],
479 rec_config_list: List[Dict[Recompute.TYPE, List[int]]]) -> None:
480 """Log the dry-run YAML config that the user should execute to collect memory data."""
481 logger.output(
482 "Please dryrun following config, %s round(s) is needed",
483 len(offset_config_list),
484 )
486 for round_, offset_config in enumerate(offset_config_list):
487 yaml_config = {
488 Recompute.OFFSET: [],
489 Recompute.YAML_NAME[Recompute.TYPE.FULL]: [],
490 Recompute.YAML_NAME[Recompute.TYPE.SLCT]: [],
491 Recompute.YAML_NAME[Recompute.TYPE.COMM]: [],
492 }
493 yaml_config[Recompute.OFFSET] = offset_config
494 pp = len(offset_config)
495 slct = rec_config_list[round_].get(Recompute.TYPE.SLCT, [0] * pp)
496 comm = rec_config_list[round_].get(Recompute.TYPE.COMM, [0] * pp)
497 full = rec_config_list[round_].get(Recompute.TYPE.FULL, [0] * pp)
498 both = rec_config_list[round_].get(Recompute.TYPE.BOTH, [0] * pp)
500 yaml_config[Recompute.YAML_NAME[Recompute.TYPE.FULL]] = full
501 yaml_config[Recompute.YAML_NAME[Recompute.TYPE.SLCT]] = [
502 x + y + z for x, y, z in zip(slct, both, full)
503 ]
504 yaml_config[Recompute.YAML_NAME[Recompute.TYPE.COMM]] = [
505 x + y + z for x, y, z in zip(comm, both, full)
506 ]
507 yaml_results = f"for round {round_ + 1}, please dryrun config:"
508 for y, v in yaml_config.items():
509 yaml_results += f"\n\t{y}: {v}"
510 logger.output(yaml_results)
513def generate_solvable_config(
514 pp: int, num_layers: int,
515 considered_rec: List[Recompute.TYPE]
516) -> Optional[Tuple[List[List[int]], List[Dict[Recompute.TYPE, List[int]]]]]:
517 """Generate offset / recompute configs whose coefficient matrix is full-rank."""
518 if pp == 2:
519 logger.error("pp = 2 is not supported yet")
520 return None
522 considered_rec.append(Recompute.TYPE.NONE)
523 rounds = ceil((2 + len(considered_rec)) / (pp - 2))
524 layer_per_stage = num_layers // pp
525 is_solvable = False
526 offset_config_list = _generate_offset_config(
527 rounds, 2 + len(considered_rec), num_layers, pp
528 )
529 while not is_solvable:
530 rec_config_list = []
531 for round_ in range(rounds):
532 offset_config = offset_config_list[round_]
533 layer_per_recompute = {r: [0] * pp for r in considered_rec}
534 for rec in considered_rec:
535 stage_sum = [
536 sum(col) for col in zip(*layer_per_recompute.values())
537 ] # summation of each rec in each stage
538 layers_left = [
539 offset_config[i] + layer_per_stage - stage_sum[i] for i in range(pp)
540 ]
541 layer_per_recompute[rec] = [
542 random.randint(0, layers_left[i]) for i in range(pp)
543 ]
544 rec_config_list.append(layer_per_recompute)
546 coef_matrix = _get_coef_matrix(
547 pp, layer_per_stage, offset_config_list, rec_config_list, considered_rec
548 )
549 coef_rank = np.linalg.matrix_rank(coef_matrix)
550 if coef_rank == len(considered_rec) + 2:
551 is_solvable = True
553 return offset_config_list, rec_config_list
556def parse_training_config(yaml_path: str) -> Optional[Dict[str, Any]]:
557 """Extract the training-YAML fields needed to compute operator shapes for seqpp."""
558 try:
559 with open(yaml_path, "r", encoding="utf-8") as file:
560 config = yaml.safe_load(file)
562 # Extract the requested values
563 model_config = config["model"]["model_config"]
564 parallel_config = config["parallel_config"]
565 runner_config = config["runner_config"]
567 # Create a dictionary with the requested parameters
568 extracted_params = {
569 "num_heads": model_config["num_heads"],
570 "hidden_size": model_config["hidden_size"],
571 "head_dim": int(model_config["hidden_size"] / model_config["num_heads"]),
572 "seq_length": model_config["seq_length"],
573 "batch_size": runner_config["batch_size"],
574 "vocab_size": model_config["vocab_size"],
575 "data_parallel": parallel_config.get("data_parallel", 1),
576 "model_parallel": parallel_config.get("model_parallel", 1),
577 "context_parallel": parallel_config.get("context_parallel", 1),
578 }
579 logger.output("Extracted training parameters:")
580 for key, value in extracted_params.items():
581 logger.output("%s: %s", key, value)
583 return extracted_params
585 except (yaml.YAMLError, FileNotFoundError, PermissionError) as e:
586 logger.error("Error parsing Training YAML file: %s", e)
587 return None