Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / parallelize.py: 77%
325 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2024 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"""find parallelization"""
17import time
18import copy
19import multiprocessing as proc
20import json
21import os
22import logging
24from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2
25from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.estimate import estimate_performance
27from hyper_parallel.auto_parallel.sapp_nd.nd.global_config import GlobalConfig
28from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger
29import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
30import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
31import hyper_parallel.auto_parallel.sapp_nd.nd.debug as Debug
33# logger = proc.log_to_stderr()
34# logger.setLevel(proc.SUBDEBUG)
37class ParallelizeLayer:
38 """Parallelize one layer type"""
40 def __init__(
41 self,
42 evaluator,
43 machine,
44 global_batch_size=None,
45 dimensions=None,
46 **extra_config,
47 ):
49 self.enable_debug = logger.level < logging.CRITICAL
50 self.machine = machine
51 if "mppb" in extra_config:
52 manual_ppb = extra_config.pop("mppb")
53 else:
54 manual_ppb = False
56 self.mem_eval = evaluator
58 self.model_name = self.mem_eval._ccfg.model_name
59 logger.debug("model is %s", self.model_name)
61 if "mem_for_ppb" in extra_config:
62 reserve_mem = extra_config.pop("mem_for_ppb")
63 self.mem_eval._ccfg.device_capacity.decrease(reserve_mem)
65 if "max_mem" in extra_config:
66 max_mem = extra_config.pop("max_mem")
67 if max_mem is not None:
68 self.mem_eval._ccfg.device_capacity.set(max_mem)
70 logger.debug("before global config init")
72 if "sub_model" in extra_config:
73 sub_model = extra_config.pop("sub_model")
74 if sub_model is not None:
75 self.config = GlobalConfig(
76 self.mem_eval._ccfg.mm_ccfgs[sub_model],
77 dimensions,
78 mppb=manual_ppb,
79 )
80 else:
81 self.config = GlobalConfig(
82 self.mem_eval._ccfg, dimensions, mppb=manual_ppb
83 )
84 else:
85 self.config = GlobalConfig(
86 self.mem_eval._ccfg, dimensions, mppb=manual_ppb
87 )
89 self.mem_eval.set_passes(**extra_config)
91 self.machine.update_num_if_none(
92 self.config.ccfg.strategy_num_devices()
93 )
95 if global_batch_size:
96 self.global_batch_size = global_batch_size
97 else:
98 self.global_batch_size = self.config.ccfg.gbs
100 self.bound_space()
102 def bound_space(self):
103 """Set bounds for parallel dimensions"""
104 vpp = (
105 1
106 if Dim.VPP in self.config.dimensions
107 else Dim.VPP.from_config(self.config.ccfg)
108 )
109 pp_bound = min(
110 self.machine.pipeline_bound(),
111 self.config.total_layer_num() // vpp,
112 self.global_batch_size,
113 )
114 Dim.PP.set_bound(pp_bound)
115 logger.info(
116 "PP bound is %d, machine bound = %d, L = %d, VPP = %d, B = %d",
117 pp_bound,
118 self.machine.pipeline_bound(),
119 self.config.total_layer_num(),
120 vpp,
121 self.global_batch_size,
122 )
123 Dim.EP.set_bound(self.config.ccfg.n_exp)
124 # if (
125 # self.config.dimensions.count(Dim.EP) > 0
126 # and Dim.EP.from_config(self.config.ccfg) <= 1
127 # ):
128 # Dim.EP.set_bound(1)
129 # self.config.dimensions.remove(Dim.EP)
130 kv_heads = self.config.ccfg.n_kv
131 if kv_heads:
132 Dim.TP.set_bound(kv_heads)
133 logger.warning(
134 "Because of n_kv_heads, MP will be limited to %s",
135 str(kv_heads),
136 )
137 else:
138 # num_head % (TP * UP) == 0. Add UP later
139 Dim.TP.set_bound(
140 Hard.highest_power_of_2_divisor(self.config.ccfg.a)
141 )
143 def filtered_out(self, _):
144 """Manual conditions to remove config patterns"""
145 # if parallel_config.has_dim(Dim.EP):
146 # if self.config.dim_val(Dim.EP, parallel_config) < 8:
147 # return True
148 return False
150 def is_valid(self, parallel_config):
151 """Check configuration validity"""
152 if not parallel_config.is_valid():
153 logger.warning("configuration %s not valid", str(parallel_config))
154 return False
155 if not self.config.moe_valid(parallel_config):
156 logger.warning("expert parallel is higher than expert number")
157 return False
158 if self.filtered_out(parallel_config):
159 logger.warning("Config manually filtered out")
160 return False
161 gbs = self.config.global_batch_size(parallel_config)
162 if not gbs == self.global_batch_size:
163 logger.error(
164 "wrong global batch size: ccfg is %d, instead of %d",
165 gbs,
166 self.global_batch_size,
167 )
168 return False
169 return True
171 def memory_estim(self, debugger=None):
172 """Whether the config fits memory"""
173 logger.debug("estimate_peak")
174 verbose = logger.level < logging.INFO
175 self.mem_eval.set_config(self.config.ccfg) # = self.config.ccfg
176 # self.mem_eval = EvaluatorV2(self.config)
177 logger.debug("ccfg = %s", str(self.config.ccfg))
178 peak = self.mem_eval.estimate_peak(
179 verbose=verbose
180 ) # (logger.level>2))
181 logger.debug("peak memory = %d", peak)
182 if debugger and debugger.is_enabled():
183 debugger.info[Debug.MemParts.TOTAL] = peak
184 return peak
186 def generate_search_space(self, folder, threads_num):
187 """Return a search space computed with memory estimation"""
188 space = ({}, 0)
189 configs = []
190 results = {}
191 if threads_num:
192 with proc.Pool(processes=threads_num) as pool:
193 logger.debug("before loops")
194 results, size = self.device_loops(space, pool)
195 logger.debug("%d results", len(results))
196 for config, result in results.items():
197 logger.debug("result = %s", str(result))
198 logger.debug(
199 "before get: is ready ? %s", str(result.ready())
200 )
201 logger.debug(
202 "before get: is successful ? %s",
203 str(result.successful()),
204 )
205 # if result.successful():
206 peak_mem = result.get(1)
207 logger.debug(
208 "after get: is ready ? %s", str(result.ready())
209 )
210 logger.debug(
211 "after get: is successful ? %s",
212 str(result.successful()),
213 )
214 logger.debug("peak_mem = %s", str(peak_mem))
215 if self.mem_eval.mem_fit(peak_mem):
216 configs.append((config, peak_mem))
217 pool.close()
218 pool.join()
219 else:
220 results, size = self.device_loops(space, None)
221 for config, peak_mem in results.items():
222 if self.mem_eval.mem_fit(peak_mem):
223 configs.append((config, peak_mem))
224 if folder:
225 self.config.write(folder, config)
226 logger.output("%d valid configurations generated", size)
227 logger.output("%d configuration fitting memory to order", len(configs))
229 return configs
231 def device_loops(self, space, pool):
232 """Exploration loop nest level 0: parallel dimensions dividing devices"""
233 for tp in self.config.space(Dim.TP, self.machine.number):
234 for pp in self.config.space(Dim.PP, self.machine.number // tp):
235 for cp in self.config.space(
236 Dim.CP, self.machine.number // tp // pp
237 ):
238 logger.debug(
239 "dp = %d / %d / %d / %d",
240 self.machine.number,
241 tp,
242 cp,
243 pp,
244 )
245 dp = self.machine.number // tp // cp // pp
246 if dp < 1:
247 break
248 space = self.batch_loops(space, pool, (dp, tp, pp, cp))
249 return space
251 def batch_loops(self, space, pool, dtpc_p):
252 """Exploration loop nest level 1: dimensions dividing batch (except already processed DP)"""
253 dp, _, pp, _ = dtpc_p
254 # if pp > 1:
255 for mbs in self.config.space(
256 Dim.MBS, self.global_batch_size // pp // dp
257 ):
258 logger.debug("mbn= %d / %d / %d", self.global_batch_size, dp, mbs)
259 mbn = self.global_batch_size // dp // mbs
260 space = self.parallel_loops(space, pool, (dtpc_p, (mbs, mbn)))
261 # else:
262 # logger.debug("no pipeline so mbn = 1")
263 # mbs = self.global_batch_size // dp
264 # space = self.parallel_loops(space, pool, (dtpc_p, (mbs, 1)))
265 return space
267 def parallel_loops(self, space, pool, dims):
268 """Exploration loop nest level 2: dimensions dependent on others"""
269 dtpc_p, mbsn = dims
270 dp, tp, pp, _ = dtpc_p
271 for ep in self.config.space(Dim.EP, dp * tp):
272 for vpp in self.config.range_space(
273 Dim.VPP, min(4, pp, self.config.total_layer_num() // pp)
274 ):
275 for op in self.config.space(
276 Dim.OP, self.config.max_op(dp, tp, ep)
277 ):
278 for sp in self.config.bool_space(Dim.SP):
279 space = self.inside_loop_nest(
280 space,
281 pool,
282 (dtpc_p, mbsn, (ep, vpp, op, sp)),
283 )
284 return space
286 def inside_loop_nest(self, space, pool, dims):
287 """Exploration loop nest statements"""
288 dtpc_p, mbsn, evos_p = dims
289 configs, size = space
290 parallel_config = self.config.make_parallel_config(
291 dtpc_p, mbsn, evos_p
292 )
293 logger.info("test config %d : %s", size, str(parallel_config))
294 size += 1
296 if self.is_valid(parallel_config) and self.config.set_parallel_config(
297 parallel_config
298 ):
299 if pool is None:
300 if self.enable_debug:
301 mem_debugger = Debug.Debug(
302 parallel_config,
303 info_type=Debug.MemParts,
304 enable=self.enable_debug,
305 output_file="debug_mem.csv",
306 )
307 # try:
308 peak = self.memory_estim(mem_debugger)
309 mem_debugger.write()
310 else:
311 peak = self.memory_estim()
312 # except:
313 # logger.error()
314 # return (configs, size)
315 else:
316 # logger.debug("before evaluator copy")
317 # evaluator = copy.deepcopy(self.mem_eval)
318 logger.debug("before apply_async")
319 peak = pool.apply_async(
320 pool_estimate_memory,
321 args=(copy.deepcopy(self.config.ccfg),),
322 # args=(evaluator,),
323 # self.memory_estim,
324 )
325 logger.debug("after apply_async")
326 configs[parallel_config] = peak
328 return (configs, size)
330 def order_search_space(self, space, threads_num, cache_file):
331 """Sort the search space computed with performance estimation"""
332 if not space:
333 return ([], [])
334 multiproc = False
335 if threads_num and threads_num > 5 * len(space):
336 multiproc = True
337 scored_space = []
338 debug_parts = []
339 for config, mem in space:
340 self.config.set_parallel_config(config)
341 values = []
342 if multiproc:
343 with proc.Pool(processes=threads_num) as pool:
344 score = pool.apply_async(
345 pool_estimate_performance,
346 args=(
347 copy.deepcopy(self.config),
348 self.machine.device,
349 cache_file,
350 ),
351 )
352 else:
353 if self.enable_debug:
354 debugger = Debug.Debug(
355 config,
356 info_type=Debug.PerfParts,
357 enable=self.enable_debug,
358 )
359 score = estimate_performance(
360 self.config.ccfg,
361 debugger=debugger,
362 device_type=self.machine.device,
363 memory=mem,
364 cache_file=cache_file,
365 )
366 debugger.write()
367 debug_parts = list(debugger.info.keys())
368 values = list(debugger.info.values())
369 del values[-2:]
370 del debug_parts[-2:]
371 else:
372 score = estimate_performance(
373 self.config.ccfg,
374 device_type=self.machine.device,
375 memory=mem,
376 )
377 scored_space.append((config, mem, score, values))
379 logger.info("config %s has score %f", str(config), score)
381 if multiproc:
382 new_scored_space = []
383 pool.close()
384 pool.join()
385 for config, mem, score, values in scored_space:
386 new_scored_space.append((config, mem, score.get(), values))
387 else:
388 new_scored_space = scored_space
389 return (sorted(new_scored_space, key=lambda x: x[2]), debug_parts)
391 def order_space_test_comm_classified(self, space, order_by=2):
392 """Order the given space with performance estimation"""
393 scored_space = []
394 debug_parts = []
395 for config, real_time, real_comm_wait in space:
396 debugger = Debug.Debug(
397 config, info_type=Debug.PerfParts, enable=self.enable_debug
398 )
399 self.config.set_parallel_config(config)
400 peak_mem = self.memory_estim()
401 score = estimate_performance(
402 self.config.ccfg,
403 debugger=debugger,
404 device_type=self.machine.device,
405 stage_focused=0,
406 ) # , memory = mem)
407 debugger.write()
408 debug_parts = list(debugger.info.keys())
409 values = list(debugger.info.values())
410 del values[-2:]
411 scored_space.append(
412 (config, peak_mem, real_time, score, values, real_comm_wait)
413 )
415 logger.info("config %s has score %f", str(config), score)
416 del debug_parts[-2:]
417 return (sorted(scored_space, key=lambda x: x[order_by]), debug_parts)
419 def order_space_test(self, space, order_by=2):
420 """Order the given space with performance estimation"""
421 scored_space = []
422 debug_parts = []
423 for config, real_time in space:
424 debugger = Debug.Debug(
425 config, info_type=Debug.PerfParts, enable=self.enable_debug
426 )
427 logger.info("Test config %s", str(config))
428 self.config.set_parallel_config(config)
429 logger.debug(self.mem_eval.get_strategy())
430 peak_mem = self.memory_estim()
431 score = estimate_performance(
432 self.config.ccfg,
433 debugger=debugger,
434 device_type=self.machine.device,
435 ) # , memory = mem)
436 debugger.write()
437 debug_parts = list(debugger.info.keys())
438 values = list(debugger.info.values())
439 del values[-2:]
440 scored_space.append((config, peak_mem, real_time, score, values))
442 logger.info("config %s has score %f", str(config), score)
443 del debug_parts[-2:]
444 return (sorted(scored_space, key=lambda x: x[order_by]), debug_parts)
446 def plot_title(self):
447 """Generate plot title"""
448 return (
449 f"{self.model_name} on {self.machine.number}"
450 + f" {self.machine.device} with {self.global_batch_size} GBS"
451 )
453 def run_generation_to_ordering(
454 self, yaml_folder, threads_num=None, top_num=None, cache_file=None
455 ):
456 """Test some functions"""
457 start = time.time()
458 space = self.generate_search_space(yaml_folder, threads_num)
459 generation = time.time()
460 scored_space, dbg = self.order_search_space(
461 space, threads_num, cache_file=cache_file
462 )
463 ordering = time.time()
464 logger.output(
465 space_to_string(scored_space, max_num=top_num, debug_parts=dbg)
466 )
467 logger.output(
468 "Space generation took %.2fs and ordering took %.2fs",
469 generation - start,
470 ordering - generation,
471 )
472 is_not = " NOT" if not self.config.balancing.from_config else ""
473 logger.output(
474 "Offset & Recompute were%s computed from config info", is_not
475 )
476 logger.output(
477 "Device number is %d, global batch size is %d, dimensions are %s",
478 self.machine.number,
479 self.global_batch_size,
480 str(self.config.dimensions),
481 )
482 if self.enable_debug:
483 file_path = os.path.dirname(os.path.realpath(__file__))
484 output_path = os.path.join(file_path, "output")
485 if scored_space:
486 Debug.plot_nd(
487 scored_space,
488 output_path,
489 dbg,
490 title=self.plot_title(),
491 max_num=top_num,
492 )
493 return scored_space
495 def to_ppb(self, scored_space, k, cfg_name):
496 """Create an input file for pipeline balancing"""
497 parallel_config = scored_space[k][0]
498 self.config.set_parallel_config(parallel_config)
499 self.mem_eval.update_config(self.config)
500 m = cfg_name + "_nd_to_ppb_" + str(k)
501 s = self.config.dim_val(Dim.PP, parallel_config)
502 mb = self.config.dim_val(Dim.MBN, parallel_config)
503 i = self.config.dim_val(Dim.VPP, parallel_config)
504 mem = str(self.config.ccfg.device_capacity.to_mb)
505 filename = (
506 os.path.dirname(os.path.realpath(__file__))
507 + "/../pipeline_balance/layers/"
508 + m
509 + ".json"
510 )
511 with open(filename, "w+", encoding="utf-8") as fp:
512 json.dump(
513 self.mem_eval.estimate_layer_memory(
514 device_type=self.machine.device
515 ),
516 fp,
517 indent=4,
518 )
519 logger.output(
520 "To run pipeline balancing on configuration %s:"
521 "\npython run_pipeline_balance.py "
522 "-m %d -s %d -mb %d -i %d -mem %d",
523 parallel_config,
524 m,
525 s,
526 mb,
527 i,
528 mem,
529 )
530 logger.output("Warning: currently select_recompute_memory \
531 should be removed & layer time need to be added")
533 def test_from_csv(self, csv_f, output_path=None):
534 """Run estimation tests against a real run profiling in csv format"""
535 configs, row_num = Debug.get_real_data(csv_f)
536 configs_estimated, debug_parts = self.order_space_test(
537 configs, order_by=2
538 )
539 if output_path is not None:
540 Debug.plot_vs_real(
541 configs_estimated,
542 csv_f,
543 output_path,
544 debug_parts,
545 title=self.plot_title(),
546 )
547 correl, topk = Debug.correlation_topk(configs_estimated, csv_f)
548 return correl, topk, row_num
550 def test_from_csv_comm_classified(
551 self, csv_f, output_path=None, plot_idle=False
552 ):
553 """Run test to compare estimation with detailed profiling"""
554 configs = Debug.get_comm_classified_data(csv_f, plot_idle=plot_idle)
555 configs_estimated, debug_parts = self.order_space_test_comm_classified(
556 configs, order_by=2
557 )
559 if output_path is not None:
560 Debug.plot_vs_real_comm_classified(
561 configs_estimated,
562 csv_f,
563 output_path,
564 debug_parts,
565 title=self.plot_title(),
566 plot_idle=plot_idle,
567 )
569 return Debug.correlation_with_classified_comms(configs_estimated)
572class ParallelizeMultiModal(ParallelizeLayer):
573 """Parallelize a MultiModel"""
575 def __init__(
576 self,
577 evaluator,
578 machine,
579 global_batch_size=None,
580 dimensions=None,
581 **extra_config,
582 ):
584 super().__init__(
585 evaluator,
586 machine,
587 global_batch_size=global_batch_size,
588 dimensions=dimensions,
589 sub_model="deepseekv3",
590 **extra_config,
591 )
594class Parallelize: # pylint: disable=R0903
595 """Main class instantiated by one of the above two"""
597 def __init__(
598 self,
599 framework,
600 config,
601 machine,
602 **extra_config,
603 ):
604 logger.debug("before evaluator init")
605 if "model" in extra_config:
606 model_name = extra_config.pop("model")
607 mem_eval = EvaluatorV2(
608 config, framework=framework, hook_cls=model_name, machine=machine
609 )
610 else:
611 mem_eval = EvaluatorV2(config, framework=framework, machine=machine)
613 if "global_batch_size" in extra_config:
614 global_batch_size = extra_config.pop("global_batch_size")
615 else:
616 global_batch_size = None
618 if "dimensions" in extra_config:
619 dimensions = extra_config.pop("dimensions")
620 else:
621 dimensions = None
623 if mem_eval.ccfg.multimodal:
624 logger.debug("MultiModal is triggered")
625 self.instance = ParallelizeMultiModal(
626 mem_eval,
627 machine,
628 global_batch_size=global_batch_size,
629 dimensions=dimensions,
630 **extra_config,
631 )
632 else:
633 self.instance = ParallelizeLayer(
634 mem_eval,
635 machine,
636 global_batch_size=global_batch_size,
637 dimensions=dimensions,
638 sub_model=None,
639 **extra_config,
640 )
642 def __getattr__(self, name):
643 return self.instance.__getattribute__(name)
646def space_to_string(space, max_num=None, debug_parts=None):
647 """Space printer"""
648 i = 0
649 s = ""
650 if max_num is not None:
651 s += "Top " + str(max_num) + " configurations:\n"
652 else:
653 s += "\n"
654 if len(space) == 0:
655 return s
656 s += "\t"
657 for d in space[0][0].all_dims:
658 s += str(d) + " " * (6 - len(str(d)))
659 s += "Memory Performance score "
660 if debug_parts is not None:
661 for dbg_part in debug_parts:
662 s += "\t" + dbg_part.short_name()
663 s += "\n"
664 for config in space:
665 if max_num is not None and max_num == i:
666 break
667 s += "\t"
668 for v in config[0].values():
669 s += v + " " * (6 - len(v))
670 s += str(config[1]) + " MB " # + str(config[2])
671 s += f"{(config[2]):16.12e}"
672 for v in config[3]:
673 s += f"\t{(100*v/config[2]):.2f}%"
674 s += "\n"
675 i += 1
676 return s
679def pool_estimate_memory(config):
680 """Calls memory estimation for multiprocessing"""
681 logger.debug("estimate_peak")
682 # print("estimate_peak")
683 e = EvaluatorV2(config)
684 return e.estimate_peak()
687# def pool_estimate_memory(evaluator):
688# """Calls memory estimation for multiprocessing"""
689# logger.debug("estimate_peak")
690# return evaluator.estimate_peak()
693def pool_estimate_performance(config, device):
694 """Calls performance estimation for multiprocessing"""
695 return estimate_performance(config, device_type=device)