Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / run_nd.py: 93%
84 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2024-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"""run parallelization"""
17import argparse
18import os
19import sys
21from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
22from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger, set_verbose_level
23import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as Par
24import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
25import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
28def _run_hyper_v2_search(cli_parser, cli_args):
29 """Run the HyperParallel V2 strategy search via ``config_adapter``.
31 This branch is activated when ``-f hyper_v2`` is combined with
32 ``-s/--search-config``. It reads the Search Config YAML, validates
33 it, runs the ND search engine through :func:`search_strategies`,
34 and writes the resolved strategy back into a copy of the original
35 ``train.yaml``.
37 Args:
38 cli_parser: The :class:`argparse.ArgumentParser` (used for ``error()``).
39 cli_args: The parsed CLI namespace. Requires ``yaml_config``,
40 ``search_config``, and optionally ``output_dir``.
42 Raises:
43 SystemExit: If validation fails (via ``parser.error``).
44 """
45 # pylint: disable=import-outside-toplevel
46 from hyper_parallel.auto_parallel.config_adapter import (
47 read_search_config,
48 validate,
49 search_strategies,
50 write_resolved_yaml,
51 )
53 if not os.path.isfile(cli_args.search_config):
54 cli_parser.error(f"search-config not found: {cli_args.search_config}")
55 if not os.path.isfile(cli_args.yaml_config):
56 cli_parser.error(f"yaml-config not found: {cli_args.yaml_config}")
58 set_verbose_level(cli_args.verbosity)
60 search_cfg = read_search_config(cli_args.search_config)
62 if cli_args.global_batch_size is not None:
63 search_cfg.constraint["global_batch_size"] = cli_args.global_batch_size
64 if cli_args.devices is not None:
65 cards_per_node = max(1, search_cfg.cluster_spec.get("cards_per_node", 8))
66 search_cfg.cluster_spec["num_nodes"] \
67 = max(1, cli_args.devices // cards_per_node)
69 errors = validate(search_cfg)
70 hard_errors = [e for e in errors if e.severity == "error"]
71 warnings = [e for e in errors if e.severity == "warning"]
72 for w in warnings:
73 logger.warning("%s: %s", w.field_path, w.message)
74 if hard_errors:
75 for e in hard_errors:
76 logger.error("%s: %s", e.field_path, e.message)
77 cli_parser.error(
78 f"Search config validation failed with {len(hard_errors)} error(s)."
79 )
81 result = search_strategies(search_cfg)
82 search_cfg.resolved_strategy = result
84 output_dir = cli_args.output_dir or "."
85 if not os.path.isdir(output_dir):
86 os.makedirs(output_dir, exist_ok=True)
87 resolve_path = os.path.join(output_dir, "resolved.yaml")
88 write_resolved_yaml(search_cfg, cli_args.yaml_config, resolve_path)
89 logger.output("Resolved strategy written to %s", resolve_path)
90 logger.output(
91 "Optimal strategy: dp=%(dp)s tp=%(tp)s pp=%(pp)s "
92 "cp=%(cp)s ep=%(ep)s mb_num=%(micro_batch_num)s "
93 "mem=%(memory_estimate_mb).0f MB score=%(score).2e",
94 result,
95 )
98if __name__ == "__main__":
99 parser = argparse.ArgumentParser(
100 prog="python run_nd.py",
101 description=("Provides a degree to *N* parallelism dimensions"),
102 epilog="",
103 )
105 parser.add_argument(
106 "-y",
107 "--yaml_config",
108 type=str,
109 required=True,
110 help="Path to yaml configuration file",
111 )
112 parser.add_argument(
113 "-f",
114 "--framework",
115 default="mindformers",
116 type=str,
117 required=False,
118 help="Framework to evaluate in "
119 "[mindformers, mindspeed, hyperparallel, hyper_v2, torchtitan]",
120 )
121 parser.add_argument(
122 "-d",
123 "--devices",
124 type=int,
125 default=None,
126 help="Number of devices. Takes yaml value if unspecified",
127 )
128 parser.add_argument(
129 "-b",
130 "--global_batch_size",
131 type=int,
132 default=None,
133 help="Global batch size. Takes yaml value if unspecified",
134 )
135 parser.add_argument(
136 "-m",
137 "--model",
138 type=str,
139 default=None,
140 help="Model Name to use. Takes yaml value if unspecified",
141 )
142 # parser.add_argument(
143 # "-g",
144 # "--generate_yaml_in",
145 # type=str,
146 # default=None,
147 # help="Generate all fitting yaml configurations in the given folder",
148 # )
149 # parser.add_argument(
150 # "-c",
151 # "--csv",
152 # type=str,
153 # default=None,
154 # help="Computes correlation coefficient from csv results file",
155 # )
156 parser.add_argument(
157 "-l",
158 "--dimensions",
159 nargs="*",
160 type=str,
161 default=None,
162 help="list of varying (output) dimensions",
163 )
164 # parser.add_argument(
165 # "-j",
166 # "--threads_num",
167 # type=int,
168 # default=None,
169 # help="Number of threads for the space generation",
170 # )
171 parser.add_argument(
172 "-v",
173 "--verbosity",
174 type=int,
175 default=2,
176 help="Level of verbosity in range [0,6], "
177 "0 being no output and 6 being debug level output. "
178 "Plot and debug csv are generated from 2",
179 )
180 # parser.add_argument(
181 # "-k",
182 # "--ppb_k",
183 # type=int,
184 # default=None,
185 # help="choose configuration number k for ppb",
186 # )
187 parser.add_argument(
188 "-A",
189 "--device_type",
190 default="A2",
191 help="choose device type between A2 or A3",
192 )
193 parser.add_argument(
194 "-swap_os",
195 "--swap_opt_state",
196 action=argparse.BooleanOptionalAction,
197 default=False,
198 help="Activate swap optimiezr state",
199 )
200 # parser.add_argument(
201 # "-lm",
202 # "--less_memory",
203 # action=argparse.BooleanOptionalAction,
204 # default=False,
205 # help="Activate less memory schedule",
206 # )
207 parser.add_argument(
208 "-mppb",
209 "-–manual_pipeline_balance",
210 action=argparse.BooleanOptionalAction,
211 default=False,
212 help="Takes offset and recompute from yaml",
213 )
214 parser.add_argument(
215 "-t",
216 "--top_config_number",
217 type=int,
218 default=None,
219 help="Number of top configs to print & plot",
220 )
221 parser.add_argument(
222 "-mem",
223 "--mem_for_ppb",
224 type=str,
225 default="0GB",
226 help="Memory to reserve for pipeline balancing. "
227 "Will be decreased from the memory budget allowed by ND (default 0GB)",
228 )
229 parser.add_argument(
230 "-c",
231 "--cache_file",
232 type=str,
233 default=None,
234 help="Cache file with ratios to recalibrate ND scores. "
235 "Will be defaulted to 'None'.",
236 )
238 parser.add_argument(
239 "-M",
240 "--max_mem",
241 type=str,
242 default=None,
243 help="Memory to reserve for pipeline balancing. "
244 "Will be decreased from the memory budget allowed by ND (default 0GB)",
245 )
246 parser.add_argument(
247 "--train-yaml",
248 type=str,
249 default=None,
250 help="Path to training configuration yaml file (for hyperparallel2)",
251 )
252 parser.add_argument(
253 "--accelerate-yaml",
254 type=str,
255 default=None,
256 help="Path to accelerate configuration yaml file (for hyperparallel2)",
257 )
258 parser.add_argument(
259 "-s",
260 "--search-config",
261 type=str,
262 default=None,
263 help="Path to Search Config YAML for fine-grained search-space control "
264 "(hyper_v2 only). Scalar=fixed, list=candidates, 'auto'=ND decides.",
265 )
266 parser.add_argument(
267 "-o",
268 "--output-dir",
269 type=str,
270 default=None,
271 help="Directory for output files when using --search-config "
272 "(default: current directory).",
273 )
275 args = parser.parse_args()
277 max_mem = (
278 Memory.from_string(args.max_mem.strip())
279 if args.max_mem is not None
280 else None
281 )
283 if args.cache_file is not None:
284 if not os.path.exists(args.cache_file):
285 logger.error(
286 f"cache file not found:"
287 f" {args.cache_file}"
288 "\nProceeding without cache file..."
289 )
290 args.cache_file = None
292 if args.framework == "hyper_v2" and args.search_config:
293 _run_hyper_v2_search(parser, args)
294 sys.exit(0)
296 set_verbose_level(args.verbosity)
297 dims = Dim.get_dims(args.dimensions)
298 YAML_FOLDER = None # args.generate_yaml_in
299 machine = Hard.Machine(args.devices, args.device_type)
301 if args.framework == "hyperparallel2":
302 if args.yaml_config is None or args.train_yaml is None or args.accelerate_yaml is None:
303 parser.error("-y (model yaml), --train-yaml, and --accelerate-yaml are required for hyperparallel2")
304 input_config = {
305 "model": args.yaml_config,
306 "train": args.train_yaml,
307 "accelerate": args.accelerate_yaml,
308 "machine": args.devices
309 }
310 elif args.framework == "torchtitan":
311 module, config = args.yaml_config.split(":")
312 input_config = {
313 "module": module,
314 "config": config,
315 "machine": machine,
316 }
317 else:
318 input_config = args.yaml_config
320 nd_runner = Par.Parallelize(
321 args.framework,
322 input_config,
323 machine,
324 global_batch_size=args.global_batch_size,
325 dimensions=dims,
326 swap_os=args.swap_opt_state,
327 mppb=args.mppb,
328 model=args.model,
329 # model="Telecom", # args.model ====ONLY FOR XINYU BRANCH====
330 max_mem=max_mem,
331 mem_for_ppb=Memory.from_string(args.mem_for_ppb.strip()),
332 # vpp_less_mem=args.less_memory,
333 )
335 if YAML_FOLDER and not os.path.exists(YAML_FOLDER):
336 os.makedirs(YAML_FOLDER)
338 space = nd_runner.run_generation_to_ordering(
339 YAML_FOLDER,
340 threads_num=None, # args.threads_num
341 top_num=args.top_config_number,
342 cache_file=args.cache_file,
343 )