Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / parallelize.py: 93%
764 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""find parallelization"""
17from collections import Counter
18from dataclasses import replace
19import time
20import copy
21import json
22import os
23import logging
24from typing import Any, Dict, Optional, Tuple
26from hyper_parallel.auto_parallel._exec_spec import ExecSpec
27from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.estimate_v2 import EvaluatorV2
28from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.estimate import LayerTimes, estimate_performance
29from hyper_parallel.auto_parallel.sapp_nd.recompute.candidate import (
30 MODES,
31 RecomputeAnalysis,
32 RecomputeChoice,
33 choose_recompute,
34 describe,
35 mode_ranges,
36 trainer_plan,
37 trainer_runs_layer_plans,
38 whole_modes,
39)
40from hyper_parallel.auto_parallel.sapp_nd.recompute.potential import (
41 DEFAULT_TOP,
42 keep_recompute_potential,
43 recompute_next_plot_groups,
44 recompute_plot_data,
45 sort_recompute_results,
46 write_recompute_csv,
47)
49from hyper_parallel.auto_parallel.sapp_nd.nd.global_config import GlobalConfig
50from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger
51import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
52import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
53import hyper_parallel.auto_parallel.sapp_nd.nd.debug as Debug
54from hyper_parallel.auto_parallel.sapp_nd.nd.dimensions import validate_cp_constraints
55from hyper_parallel.auto_parallel.sapp_nd.nd.common.apply_exec import apply_exec
56from hyper_parallel.auto_parallel.sapp_nd.nd.common.cost_model_preprocess import (
57 arm_strategy_guard,
58 detect_attention_type,
59 unset_reads_report,
60)
61from hyper_parallel.auto_parallel.sapp_nd.nd.recompute_dimension import (
62 HYPER_FRAMEWORKS,
63 RECOMPUTE_MODES,
64 parsed_recompute,
65 restore_recompute,
66 state_recompute_mode,
67)
69# logger = proc.log_to_stderr()
70# logger.setLevel(proc.SUBDEBUG)
72# The fewest candidates the workers of run_nd -j price; a search with fewer
73# prices them here. Spawned workers that import torch and parse the run are
74# ready about 5 s after a warm search starts generating on a laptop, and four
75# of them save about 3 ms a candidate across generation and ordering.
76POOL_FROM = 2000
79def state_machine_links(ccfg: Any, device: Any) -> None:
80 """Give a config, and each submodule's, its machine's ranks a node and link bandwidths.
82 The CP check, its topology penalty and the CP cost read them off the
83 config, where nothing assigned them, so they held the config's defaults,
84 8 ranks a node, 400 and 25 GB/s, whatever machine was chosen. The node
85 is the device's first level and its links the first two levels'
86 bandwidths, the second level's falling back to the first's.
88 Args:
89 ccfg: The config the search prices; a multimodal one's submodules
90 take the same machine.
91 device: The machine's ``Hard.Type``; nothing is stated without one.
92 """
93 levels = list(getattr(device, "level_bandwidth", None) or [])
94 if device is None or not levels:
95 return
96 for cfg in [ccfg, *(getattr(ccfg, "mm_ccfgs", None) or {}).values()]:
97 cfg.device_per_node = device.intra_node_num()
98 cfg.bw_intra = levels[0]
99 cfg.bw_inter = levels[1] if len(levels) > 1 else levels[0]
102class ParallelizeLayer:
103 """Parallelize one layer type"""
105 # No recompute dimension, no mode stated and no recompute chosen per
106 # candidate, also for an instance a test builds without __init__.
107 recompute_modes = None
108 recompute_dimension = None
109 _stated_mode = None
110 auto_recompute = False
111 # No workers: the candidates are priced in this process (run_nd -j starts them).
112 workers = None
114 def __init__(
115 self,
116 evaluator: Any,
117 machine: Any,
118 global_batch_size: Any = None,
119 dimensions: Any = None,
120 **extra_config: Any,
121 ) -> None:
122 """Build the search driver for one evaluator and machine."""
123 self.enable_debug = logger.level < logging.CRITICAL
124 self.machine = machine
125 if "mppb" in extra_config:
126 manual_ppb = extra_config.pop("mppb")
127 else:
128 manual_ppb = False
129 self._take_recompute_options(extra_config, manual_ppb)
130 # A recompute dimension: the activation checkpoint modes every
131 # candidate is priced under, each pair ranked on its own; None prices
132 # each candidate with the recompute it derives, as without one.
133 self.recompute_dimension = extra_config.pop("recompute_dimension", None)
134 if self.recompute_dimension is not None:
135 unknown = sorted(set(self.recompute_dimension) - set(RECOMPUTE_MODES))
136 if unknown or not self.recompute_dimension:
137 raise ValueError(f"recompute_dimension {list(self.recompute_dimension)}: expected some of "
138 f"{', '.join(RECOMPUTE_MODES)}")
139 if manual_ppb:
140 raise ValueError("recompute_dimension prices every candidate under each mode, and mppb takes "
141 "the recompute from the config: give one of them")
142 if self.auto_recompute:
143 raise ValueError("recompute_dimension ranks every candidate under each mode, and "
144 "auto_recompute chooses one for each: give one of them")
146 self.mem_eval = evaluator
147 state_machine_links(self.mem_eval._ccfg, getattr(machine, "device", None))
148 # The options chosen for each configuration the ordering scored.
149 self.recompute_choices = {}
150 # With the potentials report asked for: the plans retained for each
151 # configuration, and every alternative the analysis scored for it.
152 self.recompute_potentials = {}
153 self.recompute_results = {}
155 self.model_name = self.mem_eval._ccfg.model_name
156 logger.debug("model is %s", self.model_name)
158 # The cap replaces the device's capacity, and the reserve comes out
159 # of whichever capacity holds.
160 if "max_mem" in extra_config:
161 max_mem = extra_config.pop("max_mem")
162 if max_mem is not None:
163 self.mem_eval._ccfg.device_capacity.set(max_mem)
165 if "mem_for_ppb" in extra_config:
166 reserve_mem = extra_config.pop("mem_for_ppb")
167 self.mem_eval._ccfg.device_capacity.decrease(reserve_mem)
169 logger.debug("before global config init")
171 if "sub_model" in extra_config:
172 sub_model = extra_config.pop("sub_model")
173 if sub_model is not None:
174 self.config = GlobalConfig(
175 self.mem_eval._ccfg.mm_ccfgs[sub_model],
176 dimensions,
177 mppb=manual_ppb,
178 parent=self.mem_eval._ccfg,
179 )
180 else:
181 self.config = GlobalConfig(
182 self.mem_eval._ccfg, dimensions, mppb=manual_ppb
183 )
184 else:
185 self.config = GlobalConfig(
186 self.mem_eval._ccfg, dimensions, mppb=manual_ppb
187 )
189 self.mem_eval.set_passes(**extra_config)
190 # What a mode of the recompute dimension replaces, for when no mode is
191 # stated, and the mode stated now: None until one is.
192 self._parsed_recompute = parsed_recompute(self.mem_eval.ccfg)
193 self._stated_mode = None
195 self.machine.update_num_if_none(
196 self.config.ccfg.strategy_num_devices()
197 )
199 if global_batch_size:
200 self.global_batch_size = global_batch_size
201 else:
202 self.global_batch_size = self.config.ccfg.gbs
204 self.bound_space()
205 # From here on the configs this search owns take a strategy only
206 # through set_strategy; an estimator's copy of one starts unarmed.
207 arm_strategy_guard(self.mem_eval.ccfg)
209 def _take_recompute_options(self, extra_config: Dict[str, Any], manual_ppb: bool) -> None:
210 """Take the recompute and offload options out of *extra_config*, refusing those that do not go together."""
211 auto_recompute = extra_config.pop("auto_recompute", False)
212 if auto_recompute and manual_ppb:
213 raise ValueError(
214 "auto_recompute chooses every layer's recompute, so it cannot also take it from the config (mppb)"
215 )
216 # Choose every layer's recompute option for each candidate, rather
217 # than score it fully recomputed. A multimodal model, whose
218 # candidates are priced on every submodule, gets one mode for
219 # every layer, priced whole, and no option per layer: the options are
220 # built for one submodule's layers, whose budgets would leave the
221 # others out.
222 self.auto_recompute = bool(auto_recompute)
223 # With auto_recompute, the modes a runtime that runs every layer one
224 # way offers, of recompute.candidate.MODES; without them, each layer
225 # gets its own option.
226 self.recompute_modes = extra_config.pop("recompute_modes", None)
227 unknown = sorted(set(self.recompute_modes or ()) - set(MODES))
228 if unknown:
229 raise ValueError(f"unknown recompute modes {unknown}; expected some of {', '.join(MODES)}")
230 # The switches such a runtime's selective mode sets, where they are
231 # not the config's; see recompute.candidate.choose_recompute.
232 self.recompute_selective = extra_config.pop("recompute_selective", None)
233 # With recompute_modes, each layer gets the fastest of them that fits
234 # rather than one for every layer: a runtime that runs a mode per
235 # layer, as HyperParallel's trainer does with activation_checkpoint.layers.
236 self.recompute_mode_per_layer = bool(extra_config.pop("recompute_mode_per_layer", False))
237 if self.recompute_mode_per_layer and not (auto_recompute and self.recompute_modes):
238 raise ValueError("recompute_mode_per_layer gives each layer one of the recompute_modes auto_recompute "
239 "chooses among; give both")
240 # With auto_offload, a choice per layer may offload each stage's first
241 # layers over the host link: host_link, else the device's own.
242 auto_offload = extra_config.pop("auto_offload", False)
243 host_link = extra_config.pop("host_link", None)
244 self.offload_link = None
245 if auto_offload:
246 if not auto_recompute:
247 raise ValueError("auto_offload offloads in the choice per layer that auto_recompute makes")
248 self.offload_link = host_link or self.machine.device.host_link
249 if self.offload_link is None:
250 raise ValueError(f"device {self.machine.device} states no host link to offload over; give host_link")
252 def bound_space(self) -> None:
253 """Set bounds for parallel dimensions"""
254 # Bounds live on the module-level dimensions: start each search from
255 # none, or it keeps the tightest bound any earlier search set.
256 for dim in Dim.ALL_DIMS:
257 dim.reset_bound()
258 vpp = (
259 1
260 if Dim.VPP in self.config.dimensions
261 else Dim.VPP.from_config(self.config.ccfg)
262 )
263 pp_bound = min(
264 self.machine.pipeline_bound(),
265 self.config.total_layer_num() // vpp,
266 self.global_batch_size,
267 )
268 Dim.PP.set_bound(pp_bound)
269 logger.info(
270 "PP bound is %d, machine bound = %d, L = %d, VPP = %d, B = %d",
271 pp_bound,
272 self.machine.pipeline_bound(),
273 self.config.total_layer_num(),
274 vpp,
275 self.global_batch_size,
276 )
277 Dim.EP.set_bound(self.config.ccfg.n_exp)
278 # if (
279 # self.config.dimensions.count(Dim.EP) > 0
280 # and Dim.EP.from_config(self.config.ccfg) <= 1
281 # ):
282 # Dim.EP.set_bound(1)
283 # self.config.dimensions.remove(Dim.EP)
284 kv_heads = self.config.ccfg.n_kv
285 if kv_heads:
286 Dim.TP.set_bound(kv_heads)
287 logger.warning(
288 "Because of n_kv_heads, MP will be limited to %s",
289 str(kv_heads),
290 )
291 else:
292 # num_head % (TP * UP) == 0. Add UP later
293 Dim.TP.set_bound(
294 Hard.highest_power_of_2_divisor(self.config.ccfg.a)
295 )
296 self._report_bounds(vpp)
298 def _report_bounds(self, vpp: int) -> None:
299 """Name the bound each searched dimension takes, and why.
301 A degree over its bound is never generated, so no check refuses it
302 and no estimate drops it: this line is the only trace it leaves.
303 """
304 heads = "KV heads" if self.config.ccfg.n_kv else "attention heads' largest power of two"
305 bounds = (
306 (Dim.PP, f"the machine's {self.machine.pipeline_bound()}, {self.config.total_layer_num()} layers "
307 f"over VPP {vpp}, a batch of {self.global_batch_size}"),
308 (Dim.EP, "the model's experts"),
309 (Dim.TP, f"the model's {heads}"),
310 )
311 named = [f"{dim} up to {dim.get_bound()} ({why})" for dim, why in bounds if dim in self.config.dimensions]
312 if named:
313 logger.output("The search bounds %s", "; ".join(named))
315 @staticmethod
316 def filtered_out(_: Any) -> bool:
317 """Manual conditions to remove config patterns"""
318 # if parallel_config.has_dim(Dim.EP):
319 # if self.config.dim_val(Dim.EP, parallel_config) < 8:
320 # return True
321 return False
323 def _refuse(self, reason: str) -> bool:
324 """Count a candidate refused for *reason*, for the search's summary, and return False."""
325 self.__dict__.setdefault("refused", Counter())[reason] += 1
326 return False
328 def is_valid(self, parallel_config: Any) -> bool:
329 """Check configuration validity"""
330 if not parallel_config.is_valid():
331 logger.warning("configuration %s not valid", str(parallel_config))
332 return self._refuse("a degree out of bounds")
333 if not self.config.moe_valid(parallel_config):
334 logger.warning("expert parallel is higher than expert number")
335 return self._refuse("EP over the experts or the ranks they spread over")
336 if hasattr(self.config, 'ep_constraints_valid') and not self.config.ep_constraints_valid(parallel_config):
337 logger.warning("EP divisibility constraints not satisfied")
338 return self._refuse("EP divisibility")
339 if self.filtered_out(parallel_config):
340 logger.warning("Config manually filtered out")
341 return self._refuse("filtered out")
343 if hasattr(parallel_config, 'dims_val') and Dim.CP in parallel_config.dims_val:
344 cp_degree = parallel_config.dims_val[Dim.CP]
345 if cp_degree > 1:
346 seq_len = self.config.ccfg.s
347 tp_degree = parallel_config.dims_val.get(Dim.TP, 1)
348 pp_degree = parallel_config.dims_val.get(Dim.PP, 1)
349 device_per_node = self.machine.device.intra_node_num()
350 total_devices = self.machine.number
352 attention_type = detect_attention_type(self.config.ccfg).name.lower()
354 bw_intra = self.config.ccfg.bw_intra
355 bw_inter = self.config.ccfg.bw_inter
357 sp_enabled = bool(parallel_config.dims_val.get(Dim.SP, False))
359 cp_result = validate_cp_constraints(
360 seq_len=seq_len,
361 cp_degree=cp_degree,
362 tp_degree=tp_degree,
363 pp_degree=pp_degree,
364 device_per_node=device_per_node,
365 attention_type_str=attention_type,
366 bw_intra=bw_intra,
367 bw_inter=bw_inter,
368 total_devices=total_devices,
369 sp_enabled=sp_enabled,
370 cp_algo=getattr(self.config.ccfg, 'cp_algo', 'colossalai_cp'),
371 attention_heads=self.config.ccfg.a,
372 num_kv_heads=getattr(self.config.ccfg, 'n_kv', 0),
373 )
375 if not cp_result.is_valid:
376 logger.warning("CP constraints violated: %s", cp_result.error_message)
377 return self._refuse("CP constraints")
379 if cp_result.warning_message:
380 logger.info("CP warning: %s", cp_result.warning_message)
382 gbs = self.config.global_batch_size(parallel_config)
383 if not gbs == self.global_batch_size:
384 logger.error(
385 "wrong global batch size: ccfg is %d, instead of %d",
386 gbs,
387 self.global_batch_size,
388 )
389 return self._refuse("global batch")
390 return True
392 def priced(self) -> Any:
393 """The config a candidate is priced on: the whole model, every submodule of it.
395 The search drives one submodule's strategy (:class:`GlobalConfig`
396 gives the others the same degrees), and a candidate is priced on the
397 model the run trains: a vision tower's parameters, activations and
398 compute count toward the stage that holds them.
399 """
400 return self.mem_eval.ccfg
402 def _priced_on_submodules(self) -> bool:
403 """Whether a candidate is priced on several submodules, the one the search drives among them."""
404 return bool(getattr(self.priced(), "multimodal", False))
406 def memory_estim(self, debugger: Any = None) -> Any:
407 """Whether the config fits memory"""
408 logger.debug("estimate_peak")
409 verbose = logger.level < logging.INFO
410 self.mem_eval.set_config(self.priced())
411 # self.mem_eval = EvaluatorV2(self.config)
412 logger.debug("ccfg = %s", self.config.ccfg)
413 peak = self.mem_eval.estimate_peak(
414 verbose=verbose
415 ) # (logger.level>2))
416 logger.debug("peak memory = %d", peak)
417 if debugger and debugger.is_enabled():
418 debugger.info[Debug.MemParts.TOTAL] = peak
419 return peak
421 def _drop(self, config: Any, peak: float) -> None:
422 """Say that the memory budget drops *config*, and by how much its peak misses it."""
423 budget = self.mem_eval.get_max_device_memory()
424 logger.output(
425 "%s dropped: its peak of %.0f MB is over the %.0f MB budget by %.0f MB",
426 config,
427 peak,
428 budget,
429 peak - budget,
430 )
432 def _summarise_search(self, size: int, priced: int, fitting: int) -> None:
433 """The search's two counts: what was tested, refused and priced, and what fits."""
434 refused = self.__dict__.get("refused", Counter())
435 reasons = ", ".join(f"{reason} {count}" for reason, count in refused.most_common())
436 logger.output(
437 "%d configurations tested: %d refused%s, %d priced for memory",
438 size,
439 sum(refused.values()),
440 f" ({reasons})" if reasons else "",
441 priced,
442 )
443 logger.output(
444 "%d configuration fitting memory to order, %d dropped over the memory budget",
445 fitting,
446 priced - fitting,
447 )
449 def generate_search_space(self, folder: Any, threads_num: Any) -> Any:
450 """Return a search space computed with memory estimation.
452 A candidate the memory budget drops is named with its peak, and the
453 summary counts the candidates the checks refused apart from those the
454 budget dropped, each check by name. With *threads_num* workers that
455 run_nd started, they price the candidates (:meth:`_pooled`).
456 """
457 configs = []
458 self.refused = Counter()
459 if self._pooled(threads_num):
460 results, size = self._memory_pooled()
461 else:
462 with Debug.holding_open():
463 results, size = self.device_loops(({}, 0), None)
464 for config, peak_mem in results.items():
465 if self.mem_eval.mem_fit(peak_mem):
466 configs.append((config, peak_mem))
467 if folder:
468 self.config.write(folder, config)
469 else:
470 self._drop(config, peak_mem)
471 self._summarise_search(size, len(results), len(configs))
473 return configs
475 def _pooled(self, threads_num: Any) -> bool:
476 """Whether the workers run_nd started price this search's candidates, rather than this process.
478 A search that chooses a recompute per layer (-ar), has a recompute
479 dimension or prices a multimodal model carries more than the command
480 line a worker rebuilds its runner from, so it prices here.
481 """
482 return (bool(threads_num) and threads_num > 1 and self.workers is not None
483 and not self.auto_recompute and self.recompute_dimension is None
484 and not self._priced_on_submodules())
486 def _memory_pooled(self) -> Tuple[dict, int]:
487 """device_loops, with each candidate's memory priced in a worker.
489 The loop nest runs here and refuses and sets each candidate as it does
490 in one process, collecting the rest; the workers price them, and the
491 debug rows are written here in the loop's order. Fewer than
492 :data:`POOL_FROM` candidates are priced here instead, and the workers
493 stop unused.
494 """
495 collected = []
496 _, size = self.device_loops(({}, 0), collected)
497 if len(collected) >= POOL_FROM:
498 priced = self.workers.memory([(index, dims) for index, (_, dims) in enumerate(collected)])
499 else:
500 self.close_pool(abort=True)
501 priced = [self._memory_here(config) for config, _ in collected]
502 results = {}
503 with Debug.holding_open():
504 for (config, dims), (peak, info) in zip(collected, priced):
505 config.pool_dims = dims
506 if info is not None:
507 debugger = Debug.Debug(config, info_type=Debug.MemParts, enable=True, output_file="debug_mem.csv")
508 debugger.info = info
509 debugger.write()
510 results[config] = peak
511 return results, size
513 def _memory_here(self, config: Any) -> Tuple[Any, Optional[dict]]:
514 """A collected candidate's peak and debug parts, priced in this process as a worker prices it."""
515 self.config.set_parallel_config(config)
516 if self.enable_debug:
517 debugger = Debug.Debug(config, info_type=Debug.MemParts, enable=True, output_file="debug_mem.csv")
518 return self.memory_estim(debugger), dict(debugger.info)
519 return self.memory_estim(), None
521 def close_pool(self, abort: bool = False) -> None:
522 """Stop the workers run_nd started for this search, at once where the search failed."""
523 workers, self.workers = self.workers, None
524 if workers is not None:
525 workers.close(abort)
527 def device_loops(self, space: Any, collect: Optional[list]) -> Tuple[dict, int]:
528 """Exploration loop nest level 0: parallel dimensions dividing devices"""
529 for tp in self.config.space(Dim.TP, self.machine.number):
530 for pp in self.config.space(Dim.PP, self.machine.number // tp):
531 for cp in self.config.space(
532 Dim.CP, self.machine.number // tp // pp
533 ):
534 logger.debug(
535 "dp = %d / %d / %d / %d",
536 self.machine.number,
537 tp,
538 cp,
539 pp,
540 )
541 dp = self.machine.number // tp // cp // pp
542 if dp < 1:
543 break
544 space = self.batch_loops(space, collect, (dp, tp, pp, cp))
545 return space
547 def batch_reachable(self) -> bool:
548 """Whether any candidate the loops build makes the global batch; one line says so where none does.
550 Each candidate's batch is its DP x MB x MBS, MB the batch over DP and
551 MBS, so a batch that is no whole number of micro-batches of any DP x
552 MBS the loops reach is refused in every candidate, each with an error.
553 The search stops before building one, with this line instead.
554 """
555 number = self.machine.number
556 reached = set()
557 for tp in self.config.space(Dim.TP, number):
558 for pp in self.config.space(Dim.PP, number // tp):
559 for cp in self.config.space(Dim.CP, number // tp // pp):
560 dp = number // tp // cp // pp
561 if dp < 1:
562 break
563 for mbs in self.config.space(Dim.MBS, self.global_batch_size // pp // dp):
564 mbn = self.global_batch_size // dp // mbs
565 candidate = self.config.make_parallel_config((dp, tp, pp, cp), (mbs, mbn), (1, 1, 1, False))
566 if self.config.global_batch_size(candidate) == self.global_batch_size:
567 return True
568 reached.add(self.config.dim_val(Dim.DP, candidate) * mbs)
569 logger.error(
570 "No candidate makes a global batch of %d (the yaml states %d): the candidates the search builds "
571 "have DP x MBS of %s, and %d is no whole number of micro-batches of any. State -b as a multiple "
572 "of one of them, or let DP shrink by naming PP, CP or MP in -l.",
573 self.global_batch_size,
574 self.config.ccfg.gbs,
575 ", ".join(str(batch) for batch in sorted(reached)),
576 self.global_batch_size,
577 )
578 return False
580 def batch_loops(self, space: Any, collect: Optional[list], dtpc_p: Any) -> Tuple[dict, int]:
581 """Exploration loop nest level 1: dimensions dividing batch (except already processed DP)"""
582 dp, _, pp, _ = dtpc_p
583 # if pp > 1:
584 for mbs in self.config.space(
585 Dim.MBS, self.global_batch_size // pp // dp
586 ):
587 logger.debug("mbn= %d / %d / %d", self.global_batch_size, dp, mbs)
588 mbn = self.global_batch_size // dp // mbs
589 space = self.parallel_loops(space, collect, (dtpc_p, (mbs, mbn)))
590 # else:
591 # logger.debug("no pipeline so mbn = 1")
592 # mbs = self.global_batch_size // dp
593 # space = self.parallel_loops(space, pool, (dtpc_p, (mbs, 1)))
594 return space
596 def parallel_loops(self, space: Any, collect: Optional[list], dims: Any) -> Tuple[dict, int]:
597 """Exploration loop nest level 2: dimensions dependent on others"""
598 dtpc_p, mbsn = dims
599 dp, tp, pp, cp = dtpc_p
600 for ep in self.config.space(Dim.EP, self.config.max_ep(dp, tp, cp)):
601 for vpp in self.config.range_space(
602 Dim.VPP, min(4, pp, self.config.total_layer_num() // pp)
603 ):
604 for op in self.config.space(
605 Dim.OP, self.config.max_op(dp, tp, ep, cp)
606 ):
607 for sp in self.config.bool_space(Dim.SP):
608 space = self.inside_loop_nest(
609 space,
610 collect,
611 (dtpc_p, mbsn, (ep, vpp, op, sp)),
612 )
613 return space
615 def inside_loop_nest(self, space: Any, collect: Optional[list], dims: Any) -> Tuple[dict, int]:
616 """Exploration loop nest statements: price a candidate that passes the checks, or collect it.
618 With *collect*, a list, the candidate and its loop degrees are
619 appended for a worker to price (:meth:`_memory_pooled`), after the
620 same checks and the same setting of the config as here.
621 """
622 dtpc_p, mbsn, evos_p = dims
623 configs, size = space
624 parallel_config = self.config.make_parallel_config(
625 dtpc_p, mbsn, evos_p
626 )
627 logger.info("test config %d : %s", size, parallel_config)
628 size += 1
630 if self.is_valid(parallel_config) and self.config.set_parallel_config(
631 parallel_config
632 ):
633 if collect is not None:
634 collect.append((parallel_config, dims))
635 return (configs, size)
636 if self.enable_debug:
637 mem_debugger = Debug.Debug(
638 parallel_config,
639 info_type=Debug.MemParts,
640 enable=self.enable_debug,
641 output_file="debug_mem.csv",
642 )
643 peak = self.memory_estim(mem_debugger)
644 mem_debugger.write()
645 else:
646 peak = self.memory_estim()
647 configs[parallel_config] = peak
649 return (configs, size)
651 def order_search_space(self, space: Any, threads_num: Any, cache_file: Any) -> Any:
652 """Sort the search space computed with performance estimation.
654 With *threads_num* workers that run_nd started, they price the
655 candidates (:meth:`_pooled`).
656 """
657 if not space:
658 return ([], [])
659 self.recompute_choices = {}
660 if self._pooled(threads_num):
661 return self._order_pooled(space, cache_file)
662 scored_space = []
663 debug_parts = []
664 with Debug.holding_open():
665 for config, peak in space:
666 self.config.set_parallel_config(config)
667 values = []
668 mem, savings = peak, None
669 choice = self.choose_recompute(config)
670 priced = self.priced()
671 if choice is not None:
672 mem, savings = int(round(choice.memory)), choice.stage_savings
673 if choice.score is not None:
674 # The whole model, with every submodule running the mode.
675 priced, savings = self.priced_with_mode(choice.mode), None
676 if self.enable_debug:
677 debugger = Debug.Debug(
678 config,
679 info_type=Debug.PerfParts,
680 enable=self.enable_debug,
681 )
682 score = estimate_performance(
683 priced,
684 debugger=debugger,
685 device_type=self.machine.device,
686 memory=mem,
687 cache_file=cache_file,
688 stage_savings=savings,
689 )
690 debugger.write()
691 debug_parts = list(debugger.info.keys())
692 values = list(debugger.info.values())
693 del values[-2:]
694 del debug_parts[-2:]
695 else:
696 score = estimate_performance(
697 priced,
698 device_type=self.machine.device,
699 memory=mem,
700 cache_file=cache_file,
701 stage_savings=savings,
702 )
703 scored_space.append((config, mem, score, values))
704 logger.info("config %s has score %f", config, score)
705 return (sorted(scored_space, key=lambda x: x[2]), debug_parts)
707 def _order_pooled(self, space: Any, cache_file: Any) -> Any:
708 """order_search_space, with each candidate's score priced in a worker and written here in order.
710 A pooled search chooses no recompute (:meth:`_pooled`), so each
711 candidate is priced at the peak its generation found.
712 """
713 priced = self.workers.performance(
714 [(index, config.pool_dims, peak) for index, (config, peak) in enumerate(space)], cache_file)
715 scored_space = []
716 debug_parts = []
717 with Debug.holding_open():
718 for (config, peak), (score, info) in zip(space, priced):
719 values = []
720 if info is not None:
721 debugger = Debug.Debug(config, info_type=Debug.PerfParts, enable=True)
722 debugger.info = info
723 debugger.write()
724 debug_parts = list(info.keys())
725 values = list(info.values())
726 del values[-2:]
727 del debug_parts[-2:]
728 scored_space.append((config, peak, score, values))
729 logger.info("config %s has score %f", config, score)
730 return (sorted(scored_space, key=lambda x: x[2]), debug_parts)
732 def choose_recompute(self, parallel_config: Any) -> Optional[RecomputeChoice]:
733 """Every layer's recompute option for the configuration just set, with auto_recompute.
735 The options are the fastest that fit the device, and are kept in
736 ``recompute_choices`` under *parallel_config*.
738 Args:
739 parallel_config: The configuration the config was just set to.
741 Returns:
742 The choice; ``None`` without auto_recompute, or when there is
743 none to make and the configuration keeps its own recompute.
744 """
745 if not self.auto_recompute:
746 return None
747 if self._priced_on_submodules():
748 choice = self._one_mode_whole()
749 else:
750 self.mem_eval.set_config(self.config.ccfg)
751 choice = choose_recompute(self.mem_eval, self.machine.device, modes=self.recompute_modes,
752 link=self.offload_link, selective=self.recompute_selective,
753 per_layer=self.recompute_mode_per_layer)
754 if choice is not None:
755 self.recompute_choices[parallel_config] = choice
756 return choice
758 def recompute_potential_sequence(
759 self,
760 parallel_config: Any,
761 base_score: Optional[float] = None,
762 base_memory: Optional[float] = None,
763 cache_file: Any = None,
764 base_parts: Optional[Tuple[float, ...]] = None,
765 ) -> list:
766 """What the cost model expects each further recomputed layer to cost, at fixed degrees.
768 The configuration's own plan comes first, then the fastest plan that
769 fully recomputes one more layer, and so on while one is left. The
770 degrees and the memory budget do not move. Every figure is the cost
771 model's estimate of the step, the same estimate the search ranks
772 with; a measured run can disagree with it.
774 Args:
775 parallel_config: One configuration that fits memory, whose
776 degrees are kept.
777 base_score: Its score, already computed, reused where given.
778 base_memory: Its peak memory, used where it has no plan.
779 cache_file: Optional cache of ratios for the estimate.
780 base_parts: Its score components, already computed, reused for
781 the plot.
783 Returns:
784 The retained plans in order, the configuration's own first. A
785 plan for the whole model, or one that offloads, is retained
786 without extensions.
788 Raises:
789 ValueError: The search did not choose among named modes per layer.
790 """
791 if not self.auto_recompute or not self.recompute_mode_per_layer:
792 raise ValueError("recompute potentials need auto_recompute and a mode per layer among recompute_modes")
793 self.config.set_parallel_config(parallel_config)
794 choice = self._potential_baseline(parallel_config, base_memory)
795 dimensions = tuple((str(dim), self.config.dim_val(dim, parallel_config)) for dim in Dim.ALL_DIMS)
796 score, parts = self._baseline_estimate(parallel_config, choice, base_score, base_parts, cache_file)
797 kept = [replace(choice, score=score, dimensions=dimensions, score_parts=tuple(parts))]
798 every = list(kept)
799 if self._extends(choice):
800 self._extend_potentials(kept, every, dimensions, cache_file)
801 self.recompute_potentials[parallel_config] = kept
802 self.recompute_results[parallel_config] = sort_recompute_results(every)
803 return kept
805 def _potential_baseline(self, parallel_config: Any, base_memory: Optional[float]) -> RecomputeChoice:
806 """The configuration's own plan; a plan-less choice of its peak memory where the optimizer made none."""
807 choice = self.recompute_choices.get(parallel_config)
808 if choice is None:
809 choice = self.choose_recompute(parallel_config)
810 if choice is None:
811 memory = self.memory_estim() if base_memory is None else base_memory
812 choice = RecomputeChoice((), (float(memory),), ())
813 return choice
815 def _baseline_estimate(
816 self,
817 parallel_config: Any,
818 choice: RecomputeChoice,
819 base_score: Optional[float],
820 base_parts: Optional[Tuple[float, ...]],
821 cache_file: Any,
822 ) -> Tuple[float, Tuple[float, ...]]:
823 """The configuration's own score and its parts, priced here only where the search did not keep them."""
824 if base_score is None:
825 base_score = choice.score
826 if base_score is not None and base_parts is not None:
827 return base_score, tuple(base_parts)
828 debugger = Debug.Debug(parallel_config, info_type=Debug.PerfParts, enable=True)
829 priced = self.priced_with_mode(choice.mode) if choice.mode is not None else self.priced()
830 estimated = estimate_performance(
831 priced, device_type=self.machine.device, memory=int(round(choice.memory)),
832 stage_savings=choice.stage_savings or None, cache_file=cache_file, debugger=debugger,
833 )
834 parts = tuple(debugger.info[part] for part in Debug.PerfParts
835 if part not in {Debug.PerfParts.TOTAL, Debug.PerfParts.MEMORY})
836 return (estimated if base_score is None else base_score), (parts if base_parts is None else tuple(base_parts))
838 def _extends(self, choice: RecomputeChoice) -> bool:
839 """Whether a plan can have a layer turned up: stated per layer, without offload, with full to turn up to."""
840 return bool(
841 choice.mode is None and choice.ranges and "full" in (self.recompute_modes or ())
842 and not any(item.option.link_bandwidth for item in choice.ranges)
843 )
845 def _extend_potentials(self, kept: list, every: list, dimensions: tuple, cache_file: Any) -> None:
846 """Retain one further fully recomputed layer at a time, while the cost model offers a plan that fits."""
847 self.mem_eval.set_config(self.config.ccfg)
848 analysis = RecomputeAnalysis(
849 self.mem_eval, self.machine.device, self.recompute_modes,
850 selective=self.recompute_selective, dimensions=dimensions,
851 )
852 while True:
853 results = [self._score_potential(analysis, result, cache_file)
854 for result in analysis.alternatives(kept[-1])]
855 every.extend(results)
856 if keep_recompute_potential(kept[-1], kept, results) is None:
857 return
859 def _score_potential(self, analysis: RecomputeAnalysis, result: RecomputeChoice, cache_file: Any) -> \
860 RecomputeChoice:
861 """The cost model's estimate of the step under one alternative, with the parts a plot draws."""
862 debugger = Debug.Debug(None, Debug.PerfParts, enable=True)
863 score = estimate_performance(
864 analysis.config, device_type=self.machine.device, memory=int(round(result.memory)),
865 stage_savings=result.stage_savings, cache_file=cache_file, debugger=debugger,
866 )
867 parts = tuple(debugger.info[part] for part in Debug.PerfParts
868 if part not in {Debug.PerfParts.TOTAL, Debug.PerfParts.MEMORY})
869 return replace(result, score=score, score_parts=parts)
871 def _one_mode_whole(self) -> Optional[RecomputeChoice]:
872 """The fastest mode that fits, for every layer, each mode priced on the whole model.
874 A model priced on several submodules runs its runtime's mode on every
875 submodule's layers, a vision tower's as well as its language
876 model's, and a choice per layer, whose budgets are built for one
877 submodule's layers, would leave the others out. So each mode is
878 stated on every submodule's config, and the whole model priced under
879 it: its stages' memory by the memory model, its score by the
880 performance estimate. Without a runtime's modes, every one of
881 :data:`MODES` is weighed, selective with each submodule's switches.
883 Returns:
884 The fastest mode that fits, with the stage memory and the score
885 the whole model has under it; ``None`` when none fits.
886 """
887 whole = self.priced()
888 modes = whole_modes([whole.mm_ccfgs[name] for name in whole.mm_order], self.recompute_modes or MODES,
889 self.recompute_selective)
890 best = None
891 try:
892 for mode in modes:
893 # Each mode on a copy of the whole model, which the evaluator holds between them.
894 self.mem_eval.set_config(whole)
895 priced = self.priced_with_mode(mode)
896 self.mem_eval.set_config(priced)
897 stage_memory = tuple(insight["Static"] + insight["Dynamic"]
898 for insight in self.mem_eval.estimate_peak_insight())
899 if not self.mem_eval.mem_fit(max(stage_memory)):
900 continue
901 score = estimate_performance(priced, device_type=self.machine.device,
902 memory=int(round(max(stage_memory))))
903 if best is None or score < best.score:
904 best = RecomputeChoice((), stage_memory, (), mode=mode, score=score)
905 finally:
906 self.mem_eval.set_config(whole)
907 return best
909 def priced_with_mode(self, mode: str) -> Any:
910 """A copy of the config a candidate is priced on, every submodule's layers running *mode*.
912 Args:
913 mode: One of the runtime's modes; its selective mode sets the
914 runtime's own switches, where it states them, else each
915 submodule's.
916 """
917 priced = copy.deepcopy(self.priced())
918 for name in priced.mm_order:
919 config = priced.mm_ccfgs[name]
920 rec_op = getattr(config, "rec_op", None)
921 switches = self.recompute_selective or (vars(rec_op) if rec_op is not None else {})
922 apply_exec(config, ExecSpec(recompute=mode_ranges(mode, switches)))
923 return priced
925 def recompute_per_layer(self, parallel_config: Any) -> Tuple[Optional[RecomputeChoice], Optional[float]]:
926 """Each layer's own fastest option for one configuration, and the score it gives.
928 What the configuration gains when every layer may run its own way,
929 for a search that chose one mode for all of them.
931 Args:
932 parallel_config: The configuration.
934 Returns:
935 ``(choice, score)``, or ``(None, None)`` when there is no choice
936 to make, a multimodal model's among them.
937 """
938 if self._priced_on_submodules():
939 return None, None
940 self.config.set_parallel_config(parallel_config)
941 self.mem_eval.set_config(self.config.ccfg)
942 choice = choose_recompute(self.mem_eval, self.machine.device, link=self.offload_link)
943 if choice is None:
944 return None, None
945 score = estimate_performance(
946 self.config.ccfg,
947 device_type=self.machine.device,
948 memory=int(round(choice.memory)),
949 stage_savings=choice.stage_savings,
950 )
951 return choice, score
953 def order_space_test_comm_classified(self, space: Any, order_by: Any = 2) -> Any:
954 """Order the given space with performance estimation.
956 A measured configuration the cost model cannot represent, such as
957 expert parallelism wider than its stage, is left out and named, so
958 that it does not end the comparison of every other one.
959 """
960 scored_space = []
961 debug_parts = []
962 for config, real_time, real_comm_wait in space:
963 # The comparison needs the per-part split at any verbosity; debug.csv stays opt-in.
964 debugger = Debug.Debug(config, info_type=Debug.PerfParts, enable=True)
965 try:
966 self.set_recompute_mode(self._measured_mode(config))
967 self.config.set_parallel_config(config)
968 peak_mem = self.memory_estim()
969 score = estimate_performance(
970 self.priced(),
971 debugger=debugger,
972 device_type=self.machine.device,
973 stage_focused=0,
974 ) # , memory = mem)
975 except (TypeError, ValueError, ZeroDivisionError) as exc:
976 logger.output("ND cannot cost %s, left out of the comparison: %s", config, exc)
977 continue
978 if self.enable_debug:
979 debugger.write()
980 debug_parts = list(debugger.info.keys())
981 values = list(debugger.info.values())
982 del values[-2:]
983 scored_space.append(
984 (config, peak_mem, real_time, score, values, real_comm_wait)
985 )
987 logger.info("config %s has score %f", str(config), score)
988 self.set_recompute_mode(None)
989 del debug_parts[-2:]
990 return (sorted(scored_space, key=lambda x: x[order_by]), debug_parts)
992 def order_space_test(self, space: Any, order_by: Any = 2) -> Any:
993 """Order the given space with performance estimation"""
994 scored_space = []
995 debug_parts = []
996 for config, real_time in space:
997 debugger = Debug.Debug(
998 config, info_type=Debug.PerfParts, enable=self.enable_debug
999 )
1000 logger.info("Test config %s", str(config))
1001 self.set_recompute_mode(self._measured_mode(config))
1002 self.config.set_parallel_config(config)
1003 logger.debug(self.mem_eval.get_strategy())
1004 peak_mem = self.memory_estim()
1005 score = estimate_performance(
1006 self.priced(),
1007 debugger=debugger,
1008 device_type=self.machine.device,
1009 ) # , memory = mem)
1010 debugger.write()
1011 debug_parts = list(debugger.info.keys())
1012 values = list(debugger.info.values())
1013 del values[-2:]
1014 scored_space.append((config, peak_mem, real_time, score, values))
1016 logger.info("config %s has score %f", str(config), score)
1017 self.set_recompute_mode(None)
1018 del debug_parts[-2:]
1019 return (sorted(scored_space, key=lambda x: x[order_by]), debug_parts)
1021 def plot_title(self) -> str:
1022 """Generate plot title"""
1023 return (
1024 f"{self.model_name} on {self.machine.number}"
1025 + f" {self.machine.device} with {self.global_batch_size} GBS"
1026 )
1028 def set_recompute_mode(self, mode: Optional[str]) -> None:
1029 """Price what follows under one activation checkpoint mode, or None for the config's own recompute.
1031 Nothing changes when the mode is the one already stated, so a search
1032 without a recompute dimension never touches the config's recompute.
1034 Args:
1035 mode: One of the recompute dimension's modes, or None.
1036 """
1037 if mode == self._stated_mode:
1038 return
1039 self._stated_mode = mode
1040 if mode is None:
1041 restore_recompute(self._parsed_recompute)
1042 self.config.state_recompute(None)
1043 return
1044 state_recompute_mode(self.mem_eval.ccfg, mode)
1045 self.config.state_recompute(mode == "full")
1047 def _measured_mode(self, config: Any) -> Optional[str]:
1048 """The mode a measured configuration ran: its own, else the dimension's only one, else None."""
1049 mode = getattr(config, "recompute", None)
1050 if mode is None and self.recompute_dimension is not None and len(self.recompute_dimension) == 1:
1051 mode = self.recompute_dimension[0]
1052 return mode
1054 def _search_and_order(self, yaml_folder: Any, threads_num: Any, cache_file: Any) -> Tuple[list, list, float, float]:
1055 """Generate and order the space once, or once per mode of a recompute dimension.
1057 Each mode's pass prices and checks every candidate under that mode and
1058 tags it with the mode; the passes are then ranked together, so a mode
1059 that does not fit a candidate only removes that pair.
1061 Returns:
1062 The ordered entries, the parts of their scores, and the seconds
1063 generation and ordering took.
1064 """
1065 scored_space, dbg, generation, ordering = [], [], 0.0, 0.0
1066 if not self.batch_reachable():
1067 return scored_space, dbg, generation, ordering
1068 for mode in self.recompute_dimension or (None,):
1069 if mode is not None:
1070 logger.output("Search with recompute %s", mode)
1071 self.set_recompute_mode(mode)
1072 start = time.time()
1073 space = self.generate_search_space(yaml_folder, threads_num)
1074 generated = time.time()
1075 scored, parts = self.order_search_space(space, threads_num, cache_file=cache_file)
1076 generation += generated - start
1077 ordering += time.time() - generated
1078 dbg = parts or dbg
1079 for entry in scored:
1080 entry[0].recompute = mode
1081 scored_space += scored
1082 if self.recompute_dimension is not None:
1083 self.set_recompute_mode(None)
1084 scored_space.sort(key=lambda entry: entry[2])
1085 return scored_space, dbg, generation, ordering
1087 def _ranked_search(self, yaml_folder: Any, threads_num: Any, cache_file: Any,
1088 ranking_csv: Optional[str]) -> Tuple[list, list, float, float]:
1089 """Search and order the space, stop the search's workers, rank its entries and write the ranking.
1091 Returns:
1092 What :meth:`_search_and_order` returns, each entry's
1093 configuration given its rank.
1094 """
1095 try:
1096 scored_space, dbg, generation, ordering = self._search_and_order(
1097 yaml_folder, threads_num, cache_file
1098 )
1099 except BaseException:
1100 self.close_pool(abort=True)
1101 raise
1102 self.close_pool()
1103 for rank, entry in enumerate(scored_space, start=1):
1104 entry[0].rank = rank
1105 if ranking_csv:
1106 Debug.write_ranking_csv(scored_space, ranking_csv)
1107 logger.output(
1108 "ND's order of %d configuration(s) written to %s",
1109 len(scored_space),
1110 ranking_csv,
1111 )
1112 return scored_space, dbg, generation, ordering
1114 def _log_recompute_source(self, priced: Any) -> None:
1115 """Log where the offsets and the recompute came from, and a stated mode the candidates were not priced in."""
1116 if self.recompute_dimension is not None:
1117 logger.output(
1118 "Offset was NOT computed from config info; recompute was searched over %s",
1119 ", ".join(self.recompute_dimension),
1120 )
1121 return
1122 is_not = " NOT" if not self.config.balancing.from_config else ""
1123 logger.output(
1124 "Offset & Recompute were%s computed from config info", is_not
1125 )
1126 stated_mode = getattr(priced, "stated_ac_mode", None)
1127 # -ar chooses each candidate's recompute itself, and says how below.
1128 if (not self.config.balancing.from_config and not self.auto_recompute
1129 and stated_mode not in (None, "full")):
1130 # The balancing gives every candidate full recompute without
1131 # -mppb, whatever the run's mode, which cost an hour twice
1132 # because nothing said so.
1133 logger.output(
1134 "The run's activation checkpoint mode is %s, and every candidate above is priced fully "
1135 "recomputed: pass --recompute %s, or -mppb, to price the run's own",
1136 stated_mode,
1137 stated_mode,
1138 )
1140 def run_generation_to_ordering(
1141 self,
1142 yaml_folder: Any,
1143 threads_num: Any = None,
1144 top_num: Any = None,
1145 cache_file: Any = None,
1146 ranking_csv: Optional[str] = None,
1147 exhaustive: int = 0,
1148 force_exhaustive: int = 0,
1149 fforce_exhaustive: int = 0,
1150 dimensions: Optional[list] = None,
1151 recompute_potentials: bool = False,
1152 ) -> Any:
1153 """Search, order and print the configurations that fit memory.
1155 ``ranking_csv``, when given, receives every one of them in ND's order.
1156 It is written before anything is plotted, so a plot that fails cannot
1157 take the ranking with it. ``exhaustive`` sets the minimum number of
1158 plotted results with each requested dimension above degree one.
1159 When ``force_exhaustive >= exhaustive``, append that many qualifying
1160 results per dimension. Otherwise, append at least ``force_exhaustive``
1161 and count those additions toward the ``exhaustive`` minimum.
1162 ``fforce_exhaustive`` takes precedence and reserves distinct additions
1163 for each requested dimension. The workers run_nd started for the
1164 search stop when it ends.
1166 ``recompute_potentials`` adds the recompute potentials report, which
1167 the search does none of the work for when it is left out: for every
1168 configuration that fits, the cost model's estimate of what each
1169 further recomputed layer would cost, written to a CSV and to two
1170 comparison PDFs.
1171 """
1172 scored_space, dbg, generation, ordering = self._ranked_search(
1173 yaml_folder, threads_num, cache_file, ranking_csv
1174 )
1175 has_exhaustive = (
1176 exhaustive > 0 or force_exhaustive > 0 or fforce_exhaustive > 0
1177 )
1178 plot_space, exhaustive_additions = _plot_space(
1179 scored_space,
1180 dimensions if dimensions is not None else self.config.dimensions,
1181 top_num,
1182 (exhaustive, force_exhaustive, fforce_exhaustive),
1183 )
1185 logger.output(
1186 space_to_string(scored_space, max_num=top_num, debug_parts=dbg)
1187 )
1188 _log_exhaustive_additions(scored_space, exhaustive_additions, top_num, dbg)
1189 priced = self.priced()
1190 _log_priced_caveats(priced)
1191 logger.output(
1192 "Space generation took %.2fs and ordering took %.2fs",
1193 generation,
1194 ordering,
1195 )
1196 self._log_recompute_source(priced)
1197 if self.auto_recompute:
1198 self._log_recompute(scored_space)
1199 logger.output(
1200 "Device number is %d, global batch size is %d, dimensions are %s",
1201 self.machine.number,
1202 self.global_batch_size,
1203 str(self.config.dimensions),
1204 )
1205 if self.enable_debug:
1206 output_path = Debug.output_dir()
1207 if recompute_potentials:
1208 self._report_recompute_potentials(scored_space, dbg, top_num, output_path, cache_file)
1209 if plot_space:
1210 Debug.plot_nd(
1211 plot_space,
1212 output_path,
1213 dbg,
1214 title=self.plot_title(),
1215 max_num=(
1216 len(plot_space) if has_exhaustive else top_num
1217 ),
1218 include_all=has_exhaustive,
1219 top_result_count=(
1220 len(plot_space) - len(exhaustive_additions)
1221 if exhaustive_additions else None
1222 ),
1223 recompute_plans=self._plotted_plans(plot_space),
1224 )
1225 return scored_space
1227 def _report_recompute_potentials(
1228 self, scored_space: list, debug_parts: list, top_num: Optional[int], folder: str, cache_file: Any
1229 ) -> None:
1230 """Write the recompute potentials report: one sequence per configuration, as a CSV and two PDFs.
1232 Args:
1233 scored_space: Every configuration that fits, in the search's order.
1234 debug_parts: The performance components each estimate splits into.
1235 top_num: How many configurations the plots show.
1236 folder: Where the CSV and the PDFs go.
1237 cache_file: Optional cache of ratios for the estimate.
1238 """
1239 if not self.recompute_mode_per_layer:
1240 logger.warning("the recompute potentials report extends a plan per layer, and this search chose none: "
1241 "it needs -ar/--auto_recompute with the modes each layer chooses among")
1242 return
1243 self.recompute_potentials = {}
1244 self.recompute_results = {}
1245 for config, memory, score, parts in scored_space:
1246 self.recompute_potential_sequence(
1247 config, base_score=score, base_memory=memory, cache_file=cache_file, base_parts=tuple(parts),
1248 )
1249 path = os.path.join(folder, "recompute_potentials.csv")
1250 write_recompute_csv(self.recompute_potentials, path)
1251 logger.output(
1252 "%d retained recompute plans for %d of %d configurations written to %s; every figure is the cost "
1253 "model's estimate of the step, not a measured one",
1254 sum(len(choices) for choices in self.recompute_potentials.values()),
1255 len(self.recompute_potentials), len(scored_space), path,
1256 )
1257 self._plot_recompute_results(scored_space, debug_parts, top_num, folder)
1259 def _plot_recompute_results(
1260 self, scored_space: list, debug_parts: list, top_num: Optional[int], folder: str
1261 ) -> None:
1262 """Plot the retained sequence, and the plans the cost model scores below the next strategy."""
1263 top = scored_space[:max(0, DEFAULT_TOP if top_num is None else top_num)]
1264 retained = [(entry[0], self.recompute_potentials[entry[0]]) for entry in top]
1265 following = recompute_next_plot_groups(scored_space, self.recompute_results, top_num)
1266 for groups, filename, description in (
1267 (retained, "result_recompute_inc",
1268 "Retained recompute sequence at the same memory budget, as the cost model estimates it"),
1269 (following, "result_recompute_full_next",
1270 "Recompute plans the cost model scores below the next strategy, plus the next best"),
1271 ):
1272 if groups:
1273 Debug.plot_recompute(
1274 recompute_plot_data(groups), folder, debug_parts, filename,
1275 title=f"{self.plot_title()}\n{description}",
1276 )
1277 logger.output("Recompute plot written to %s", os.path.join(folder, filename + ".pdf"))
1279 def _plotted_plans(self, plot_space: list) -> Optional[dict]:
1280 """The plan per layer of each plotted configuration, as the trainer states it.
1282 Returns:
1283 ``{configuration: (mode, layers)}`` where the search chose among
1284 named modes per layer; ``None`` where it chose no such plan, so
1285 the plot leaves the plan row out.
1286 """
1287 if not getattr(self, "recompute_mode_per_layer", False):
1288 return None
1289 return {
1290 entry[0]: trainer_plan(self.recompute_choices[entry[0]])
1291 for entry in plot_space
1292 if entry[0] in self.recompute_choices
1293 }
1295 def _log_recompute(self, scored_space: Any) -> None:
1296 """Log how the recompute was chosen, and the best configuration's, as the trainer states it where it can.
1298 A plan per layer the installed trainer cannot run (it declares no
1299 ``activation_checkpoint.layers``) is logged as advice, beside the mode
1300 the trainer runs every layer instead.
1301 """
1302 if self.recompute_modes is None:
1303 how = "per layer"
1304 else:
1305 how = ("per layer among " if self.recompute_mode_per_layer else "among ") + ", ".join(self.recompute_modes)
1306 logger.output("Recompute was chosen %s for %d of %d configurations", how, len(self.recompute_choices),
1307 len(scored_space))
1308 best = self.recompute_choices.get(scored_space[0][0]) if scored_space else None
1309 if best is None:
1310 return
1311 logger.output("Recompute of the best configuration:\n%s", describe(best))
1312 if self.recompute_mode_per_layer and best.mode is None:
1313 mode, layers = trainer_plan(best)
1314 if not layers or trainer_runs_layer_plans():
1315 logger.output("As the trainer runs it: activation_checkpoint mode %s, layers %s", mode, layers)
1316 else:
1317 logger.output("As the trainer runs it: activation_checkpoint mode %s for every layer. The installed "
1318 "trainer's activation_checkpoint has no layers field, so it cannot run the plan per "
1319 "layer, which is advice: layers %s", mode, layers)
1321 def to_ppb(self, scored_space: Any, k: Any, cfg_name: Any, folder: Optional[str] = None) -> str:
1322 """Write the pipeline balancer's layer description of the k-th configuration.
1324 Every layer carries its forward time and the backward time of each of
1325 its recompute options, priced by the estimate the search scores with,
1326 in units of the forward time of a plain layer of the first body.
1328 Args:
1329 scored_space: The ordered search space.
1330 k: Rank of the configuration to describe.
1331 cfg_name: Prefix of the model name the balancer is given.
1332 folder: Where to write; ND's output directory by default.
1334 Returns:
1335 The path of the file written.
1336 """
1337 parallel_config = scored_space[k][0]
1338 self.set_recompute_mode(getattr(parallel_config, "recompute", None))
1339 self.config.set_parallel_config(parallel_config)
1340 self.mem_eval.set_config(self.config.ccfg)
1341 name = f"{cfg_name}_nd_to_ppb_{k}"
1342 folder = os.path.abspath(folder or Debug.output_dir())
1343 os.makedirs(folder, exist_ok=True)
1344 path = os.path.join(folder, name + ".json")
1345 description = self.mem_eval.estimate_layer_memory(
1346 device_type=self.machine.device,
1347 layer_times=LayerTimes(self.machine.device),
1348 )
1349 with open(path, "w", encoding="utf-8") as fp:
1350 json.dump(description, fp, indent=4)
1351 logger.output(
1352 "To run pipeline balancing on configuration %s:"
1353 "\npython run_pipeline_balance.py -lf %s -m %s -s %s -mb %s -i %s -mem %s",
1354 parallel_config,
1355 folder,
1356 name,
1357 self.config.dim_val(Dim.PP, parallel_config),
1358 self.config.dim_val(Dim.MBN, parallel_config),
1359 self.config.dim_val(Dim.VPP, parallel_config),
1360 int(self.config.ccfg.device_capacity.to_mb().size),
1361 )
1362 return path
1364 def test_from_csv(self, csv_f, output_path=None):
1365 """Run estimation tests against a real run profiling in csv format"""
1366 configs, row_num = Debug.get_real_data(csv_f)
1367 configs_estimated, debug_parts = self.order_space_test(
1368 configs, order_by=2
1369 )
1370 if output_path is not None:
1371 Debug.plot_vs_real(
1372 configs_estimated,
1373 csv_f,
1374 output_path,
1375 debug_parts,
1376 title=self.plot_title(),
1377 )
1378 correl, topk = Debug.correlation_topk(configs_estimated, csv_f)
1379 return correl, topk, row_num
1381 def test_from_csv_comm_classified(
1382 self, csv_f, output_path=None, plot_idle=False
1383 ):
1384 """Run test to compare estimation with detailed profiling"""
1385 return self.compare_with_csv(csv_f, output_path=output_path, plot_idle=plot_idle)[1]
1387 def compare_with_csv(
1388 self, csv_f: str, output_path: Optional[str] = None, plot_idle: bool = False
1389 ) -> Tuple[list, Any]:
1390 """Estimate every measured configuration of a classified profiling CSV.
1392 With ``output_path``, three files named after the CSV go there: the
1393 real-versus-estimate plot, ``<stem>.pdf``; with ``plot_idle``, the same
1394 plot without idle, ``<stem>_no_idle.pdf``, ordered by the measured step
1395 less idle; and ND's estimate of each configuration, memory included,
1396 ``<stem>_estimates.csv``. Idle is half the step or more on a short
1397 one and differs from run to run, while ND estimates none of it.
1399 Args:
1400 csv_f: CSV read by ``Debug.get_comm_classified_data``, e.g. written by
1401 ``nd.trace_classify``.
1402 output_path: Directory for ND's plots and estimates; none are written when None.
1403 plot_idle: Whether the measured bars include the idle remainder.
1405 Returns:
1406 ``(configs_estimated, metrics)``: one ``(config, peak_mem, real_time, score,
1407 values, real_parts)`` entry per measured configuration, ordered by measured
1408 time, and ``Debug.correlation_with_classified_comms`` over them.
1409 """
1410 configs = Debug.get_comm_classified_data(csv_f, plot_idle=plot_idle)
1411 configs_estimated, debug_parts = self.order_space_test_comm_classified(
1412 configs, order_by=2
1413 )
1414 if not configs_estimated:
1415 raise ValueError(f"ND cannot cost any configuration of {csv_f}")
1417 if output_path is not None:
1418 title = self.plot_title()
1419 Debug.plot_vs_real_comm_classified(
1420 configs_estimated,
1421 csv_f,
1422 output_path,
1423 debug_parts,
1424 title=title,
1425 plot_idle=plot_idle,
1426 )
1427 if plot_idle:
1428 Debug.plot_vs_real_comm_classified(
1429 sorted(configs_estimated, key=Debug.busy_time),
1430 csv_f,
1431 output_path,
1432 debug_parts,
1433 title=title,
1434 plot_idle=False,
1435 suffix="_no_idle",
1436 )
1437 stem = os.path.splitext(os.path.basename(csv_f))[0]
1438 estimates = os.path.join(output_path, f"{stem}_estimates.csv")
1439 Debug.write_estimates_csv(configs_estimated, estimates)
1440 logger.output("ND's estimate of every configuration written to %s", estimates)
1442 return configs_estimated, Debug.correlation_with_classified_comms(configs_estimated)
1445class ParallelizeMultiModal(ParallelizeLayer):
1446 """Parallelize a MultiModel"""
1448 def __init__(
1449 self,
1450 evaluator: Any,
1451 machine: Any,
1452 global_batch_size: Any = None,
1453 dimensions: Any = None,
1454 **extra_config: Any,
1455 ) -> None:
1456 """Drive the search from the submodule the parser marked as main."""
1457 super().__init__(
1458 evaluator,
1459 machine,
1460 global_batch_size=global_batch_size,
1461 dimensions=dimensions,
1462 sub_model=getattr(
1463 evaluator.ccfg, "mm_main", None
1464 ) or evaluator.ccfg.mm_order[-1],
1465 **extra_config,
1466 )
1469class Parallelize: # pylint: disable=R0903
1470 """Main class instantiated by one of the above two"""
1472 def __init__(self, framework: Any, config: Any, machine: Any, **extra_config: Any) -> None:
1473 """Dispatch to the unimodal or multimodal search driver.
1475 ``workers``, from run_nd -j, are the processes started to price the
1476 search's candidates (:class:`~hyper_parallel.auto_parallel.sapp_nd.nd._pool.Workers`).
1477 """
1478 if extra_config.get("recompute_dimension") is not None and framework not in HYPER_FRAMEWORKS:
1479 raise ValueError(
1480 f"recompute_dimension states HyperParallel's activation checkpoint modes, and the {framework} "
1481 f"parser states its own recompute: use one of {', '.join(HYPER_FRAMEWORKS)}"
1482 )
1483 workers = extra_config.pop("workers", None)
1484 logger.debug("before evaluator init")
1485 if "model" in extra_config:
1486 model_name = extra_config.pop("model")
1487 mem_eval = EvaluatorV2(
1488 config, framework=framework, hook_cls=model_name, machine=machine
1489 )
1490 else:
1491 mem_eval = EvaluatorV2(config, framework=framework, machine=machine)
1493 if "global_batch_size" in extra_config:
1494 global_batch_size = extra_config.pop("global_batch_size")
1495 else:
1496 global_batch_size = None
1498 if "dimensions" in extra_config:
1499 dimensions = extra_config.pop("dimensions")
1500 else:
1501 dimensions = None
1503 if mem_eval.ccfg.multimodal:
1504 logger.debug("MultiModal is triggered")
1505 self.instance = ParallelizeMultiModal(
1506 mem_eval,
1507 machine,
1508 global_batch_size=global_batch_size,
1509 dimensions=dimensions,
1510 **extra_config,
1511 )
1512 else:
1513 self.instance = ParallelizeLayer(
1514 mem_eval,
1515 machine,
1516 global_batch_size=global_batch_size,
1517 dimensions=dimensions,
1518 sub_model=None,
1519 **extra_config,
1520 )
1521 self.instance.workers = workers
1523 def __getattr__(self, name):
1524 return self.instance.__getattribute__(name)
1527def _plot_space(scored_space: list, dimensions: list, top_num: Optional[int],
1528 exhaustive_options: Tuple[int, int, int]) -> tuple:
1529 """The configurations to plot, and those the exhaustive options add to the top ones.
1531 Args:
1532 scored_space: The ordered search space.
1533 dimensions: The dimensions an exhaustive option covers.
1534 top_num: How many of the best configurations are plotted anyway.
1535 exhaustive_options: ``(exhaustive, force_exhaustive,
1536 fforce_exhaustive)``, as ``run_generation_to_ordering`` takes
1537 them; the last takes precedence.
1539 Returns:
1540 ``(plot_space, additions)``: the whole space and no addition when no
1541 option is set.
1542 """
1543 exhaustive, force_exhaustive, fforce_exhaustive = exhaustive_options
1544 if fforce_exhaustive > 0:
1545 return _fforce_exhaustive_plot_space(scored_space, dimensions, top_num, fforce_exhaustive)
1546 if exhaustive > 0 or force_exhaustive > 0:
1547 return _exhaustive_plot_space(scored_space, dimensions, top_num, exhaustive, force_exhaustive)
1548 return scored_space, []
1551def _log_exhaustive_additions(scored_space: list, additions: list, top_num: Optional[int], dbg: Any) -> None:
1552 """Log the configurations the exhaustive options add that the top ones printed leave out, and their count."""
1553 if not additions:
1554 return
1555 if top_num is not None:
1556 printed_ids = {
1557 config[0].unique_name()
1558 for config in scored_space[:max(0, top_num)]
1559 }
1560 additional_to_print = [
1561 config
1562 for config in additions
1563 if config[0].unique_name() not in printed_ids
1564 ]
1565 if additional_to_print:
1566 logger.output(
1567 space_to_string(
1568 additional_to_print,
1569 max_num=len(additional_to_print),
1570 debug_parts=dbg,
1571 heading="Additional exhaustive configurations",
1572 )
1573 )
1574 logger.output(
1575 "Exhaustive plot additions: %d unique configuration(s)",
1576 len(additions),
1577 )
1580def _log_priced_caveats(priced: Any) -> None:
1581 """Log what qualifies every number of the table: the length it prices, MTP layers unpriced, fields never set."""
1582 if not getattr(priced, "seq_len_stated", True):
1583 # A warning at parse time scrolls past a search's lines; this one
1584 # sits under the table it qualifies.
1585 logger.output(
1586 "Every number above is priced at %d tokens, the model's context limit: the config states no "
1587 "training sequence length (dataset.data_transform.max_seq_len or dataset.data_config.seq_length)",
1588 priced.s,
1589 )
1590 if getattr(priced, "mtp_unpriced", 0):
1591 # A stated MTP depth looked priced and was not.
1592 logger.output(
1593 "The model states %d MTP layer(s), which the AutoModels trainer does not build: no number above "
1594 "prices them",
1595 priced.mtp_unpriced,
1596 )
1597 for line in unset_reads_report():
1598 logger.output(line)
1601def space_to_string(
1602 space: Any,
1603 max_num: Any = None,
1604 debug_parts: Any = None,
1605 heading: Optional[str] = None,
1606) -> str:
1607 """Format ranked configurations for CLI output.
1609 Args:
1610 space: Configurations with memory, performance score and score components.
1611 TOP uses the rank stored on each configuration, falling back to its
1612 position in the supplied list when it has not been ranked.
1613 max_num: Maximum number of configurations to print.
1614 debug_parts: Performance components to name in the header.
1615 heading: Optional heading replacing the default top-results heading.
1616 """
1617 s = ""
1618 if heading is not None:
1619 s += heading + ":\n"
1620 elif max_num is not None:
1621 s += "Top " + str(max_num) + " configurations:\n"
1622 else:
1623 s += "\n"
1624 if len(space) == 0:
1625 return s
1626 # A search with a recompute dimension names each entry's mode after its degrees.
1627 moded = any(getattr(entry[0], "recompute", None) for entry in space)
1628 s += "\tTOP "
1629 for d in space[0][0].all_dims:
1630 s += str(d) + " " * (6 - len(str(d)))
1631 if moded:
1632 s += "Recompute "
1633 s += "Memory Performance score "
1634 if debug_parts is not None:
1635 for dbg_part in debug_parts:
1636 s += "\t" + dbg_part.short_name()
1637 s += "\n"
1638 for i, config in enumerate(space):
1639 if max_num is not None and max_num == i:
1640 break
1641 s += f"\t{config[0].rank or i + 1:<6}"
1642 for v in config[0].values():
1643 s += v + " " * (6 - len(v))
1644 if moded:
1645 s += f"{config[0].recompute:<11s}"
1646 s += str(config[1]) + " MB " # + str(config[2])
1647 s += f"{(config[2]):16.12e}"
1648 for v in config[3]:
1649 s += f"\t{(100*v/config[2]):.2f}%"
1650 s += "\n"
1651 return s
1654def _exhaustive_plot_space(
1655 scored_space: list,
1656 dimensions: list,
1657 top_num: Optional[int],
1658 exhaustive: int,
1659 force_exhaustive: int,
1660) -> tuple:
1661 """Add ranked configs to the normal plot to meet per-dimension counts."""
1662 normal_plot = Debug.top_plot_configs(scored_space, top_num)
1663 normal_ids = {entry[0].unique_name() for entry in normal_plot}
1664 selected_ids = set(normal_ids)
1666 for dimension in dimensions:
1667 present = sum(
1668 1
1669 for entry in normal_plot
1670 if entry[0].has_dim(dimension) and entry[0].val(dimension) > 1
1671 )
1672 if force_exhaustive >= exhaustive:
1673 needed = force_exhaustive
1674 else:
1675 needed = max(force_exhaustive, exhaustive - present)
1676 added = 0
1677 if needed == 0:
1678 continue
1679 for entry in scored_space:
1680 config_id = entry[0].unique_name()
1681 if config_id in normal_ids:
1682 continue
1683 if entry[0].has_dim(dimension) and entry[0].val(dimension) > 1:
1684 selected_ids.add(config_id)
1685 added += 1
1686 if added >= needed:
1687 break
1689 plot_space = [
1690 entry for entry in scored_space if entry[0].unique_name() in selected_ids
1691 ]
1692 additions = [
1693 entry for entry in plot_space if entry[0].unique_name() not in normal_ids
1694 ]
1695 return plot_space, additions
1698def _fforce_exhaustive_plot_space(
1699 scored_space: list,
1700 dimensions: list,
1701 top_num: Optional[int],
1702 fforce_exhaustive: int,
1703) -> tuple:
1704 """Add distinct ranked configs for each dimension to the normal plot."""
1705 normal_plot = Debug.top_plot_configs(scored_space, top_num)
1706 selected_ids = {entry[0].unique_name() for entry in normal_plot}
1707 added_by_dimension = {}
1709 for dimension in dimensions:
1710 added = 0
1711 for entry in scored_space:
1712 config_id = entry[0].unique_name()
1713 if config_id in selected_ids:
1714 continue
1715 if entry[0].has_dim(dimension) and entry[0].val(dimension) > 1:
1716 selected_ids.add(config_id)
1717 added += 1
1718 if added >= fforce_exhaustive:
1719 break
1720 added_by_dimension[dimension] = added
1722 for dimension, added in added_by_dimension.items():
1723 if added < fforce_exhaustive:
1724 logger.warning(
1725 "Only %d of %d unique exhaustive result(s) found for dimension %s",
1726 added,
1727 fforce_exhaustive,
1728 dimension,
1729 )
1731 plot_space = [
1732 entry for entry in scored_space if entry[0].unique_name() in selected_ids
1733 ]
1734 normal_ids = {entry[0].unique_name() for entry in normal_plot}
1735 additions = [
1736 entry for entry in plot_space if entry[0].unique_name() not in normal_ids
1737 ]
1738 return plot_space, additions