Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / run_nd.py: 96%
47 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"""run parallelization"""
17import argparse
18import os
20from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory
21from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger, set_verbose_level
22import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as Par
23import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
24import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
26if __name__ == "__main__":
27 parser = argparse.ArgumentParser(
28 prog="python run_nd.py",
29 description=("Provides a degree to *N* parallelism dimensions"),
30 epilog="",
31 )
33 parser.add_argument(
34 "-y",
35 "--yaml_config",
36 type=str,
37 required=True,
38 help="Path to yaml configuration file",
39 )
40 parser.add_argument(
41 "-f",
42 "--framework",
43 default="mindformers",
44 type=str,
45 required=False,
46 help="Framework to evaluate in "
47 "[mindformers, mindspeed, hyperparallel, torchtitan]",
48 )
49 parser.add_argument(
50 "-d",
51 "--devices",
52 type=int,
53 default=None,
54 help="Number of devices. Takes yaml value if unspecified",
55 )
56 parser.add_argument(
57 "-b",
58 "--global_batch_size",
59 type=int,
60 default=None,
61 help="Global batch size. Takes yaml value if unspecified",
62 )
63 parser.add_argument(
64 "-m",
65 "--model",
66 type=str,
67 default=None,
68 help="Model Name to use. Takes yaml value if unspecified",
69 )
70 # parser.add_argument(
71 # "-g",
72 # "--generate_yaml_in",
73 # type=str,
74 # default=None,
75 # help="Generate all fitting yaml configurations in the given folder",
76 # )
77 # parser.add_argument(
78 # "-c",
79 # "--csv",
80 # type=str,
81 # default=None,
82 # help="Computes correlation coefficient from csv results file",
83 # )
84 parser.add_argument(
85 "-l",
86 "--dimensions",
87 nargs="*",
88 type=str,
89 default=None,
90 help="list of varying (output) dimensions",
91 )
92 # parser.add_argument(
93 # "-j",
94 # "--threads_num",
95 # type=int,
96 # default=None,
97 # help="Number of threads for the space generation",
98 # )
99 parser.add_argument(
100 "-v",
101 "--verbosity",
102 type=int,
103 default=2,
104 help="Level of verbosity in range [0,6], "
105 "0 being no output and 6 being debug level output. "
106 "Plot and debug csv are generated from 2",
107 )
108 # parser.add_argument(
109 # "-k",
110 # "--ppb_k",
111 # type=int,
112 # default=None,
113 # help="choose configuration number k for ppb",
114 # )
115 parser.add_argument(
116 "-A",
117 "--device_type",
118 default="A2",
119 help="choose device type between A2 or A3",
120 )
121 parser.add_argument(
122 "-swap_os",
123 "--swap_opt_state",
124 action=argparse.BooleanOptionalAction,
125 default=False,
126 help="Activate swap optimiezr state",
127 )
128 # parser.add_argument(
129 # "-lm",
130 # "--less_memory",
131 # action=argparse.BooleanOptionalAction,
132 # default=False,
133 # help="Activate less memory schedule",
134 # )
135 parser.add_argument(
136 "-mppb",
137 "-–manual_pipeline_balance",
138 action=argparse.BooleanOptionalAction,
139 default=False,
140 help="Takes offset and recompute from yaml",
141 )
142 parser.add_argument(
143 "-t",
144 "--top_config_number",
145 type=int,
146 default=None,
147 help="Number of top configs to print & plot",
148 )
149 parser.add_argument(
150 "-mem",
151 "--mem_for_ppb",
152 type=str,
153 default="0GB",
154 help="Memory to reserve for pipeline balancing. "
155 "Will be decreased from the memory budget allowed by ND (default 0GB)",
156 )
157 parser.add_argument(
158 "-c",
159 "--cache_file",
160 type=str,
161 default=None,
162 help="Cache file with ratios to recalibrate ND scores. "
163 "Will be defaulted to 'None'.",
164 )
166 parser.add_argument(
167 "-M",
168 "--max_mem",
169 type=str,
170 default=None,
171 help="Memory to reserve for pipeline balancing. "
172 "Will be decreased from the memory budget allowed by ND (default 0GB)",
173 )
174 parser.add_argument(
175 "--train-yaml",
176 type=str,
177 default=None,
178 help="Path to training configuration yaml file (for hyperparallel2)",
179 )
180 parser.add_argument(
181 "--accelerate-yaml",
182 type=str,
183 default=None,
184 help="Path to accelerate configuration yaml file (for hyperparallel2)",
185 )
187 args = parser.parse_args()
189 max_mem = (
190 Memory.from_string(args.max_mem.strip())
191 if args.max_mem is not None
192 else None
193 )
195 if args.cache_file is not None:
196 if not os.path.exists(args.cache_file):
197 logger.error(
198 f"cache file not found:"
199 f" {args.cache_file}"
200 "\nProceeding without cache file..."
201 )
202 args.cache_file = None
204 set_verbose_level(args.verbosity)
205 dims = Dim.get_dims(args.dimensions)
206 YAML_FOLDER = None # args.generate_yaml_in
207 machine = Hard.Machine(args.devices, args.device_type)
209 if args.framework == "hyperparallel2":
210 if args.yaml_config is None or args.train_yaml is None or args.accelerate_yaml is None:
211 parser.error("-y (model yaml), --train-yaml, and --accelerate-yaml are required for hyperparallel2")
212 input_config = {
213 "model": args.yaml_config,
214 "train": args.train_yaml,
215 "accelerate": args.accelerate_yaml,
216 "machine": args.devices
217 }
218 elif args.framework == "torchtitan":
219 module, config = args.yaml_config.split(":")
220 input_config = {
221 "module": module,
222 "config": config,
223 "machine": machine,
224 }
225 else:
226 input_config = args.yaml_config
228 nd_runner = Par.Parallelize(
229 args.framework,
230 input_config,
231 machine,
232 global_batch_size=args.global_batch_size,
233 dimensions=dims,
234 swap_os=args.swap_opt_state,
235 mppb=args.mppb,
236 model=args.model,
237 # model="Telecom", # args.model ====ONLY FOR XINYU BRANCH====
238 max_mem=max_mem,
239 mem_for_ppb=Memory.from_string(args.mem_for_ppb.strip()),
240 # vpp_less_mem=args.less_memory,
241 )
243 if YAML_FOLDER and not os.path.exists(YAML_FOLDER):
244 os.makedirs(YAML_FOLDER)
246 space = nd_runner.run_generation_to_ordering(
247 YAML_FOLDER,
248 threads_num=None, # args.threads_num
249 top_num=args.top_config_number,
250 cache_file=args.cache_file,
251 )