Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / global_config.py: 90%
135 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"""One configuration interface for parallelization"""
17import copy
18from math import gcd
20from hyper_parallel.auto_parallel.sapp_nd.nd.common.arch_hooks import CWrap, check_and_apply_custom_hook
21from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger
22import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
23import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard
24import hyper_parallel.auto_parallel.sapp_nd.nd.balancing_adapter as BA
27class GlobalConfig:
28 """Union of cost model & parallel config"""
30 def __init__(self, config, dimensions=None, mppb=False):
32 self.wrap = CWrap(config)
33 self.ccfg = self.wrap.ccfg
35 if dimensions is not None:
36 logger.debug("dimensions = %s", str(dimensions))
37 self.dimensions = dimensions
38 else:
39 logger.debug("dimensions = %s", str(Dim.ALL_DIMS))
40 self.dimensions = Dim.ALL_DIMS.copy()
41 logger.debug("self.dimensions = %s", str(self.dimensions))
42 logger.debug("layer_num_for_offset = %d", self.layer_num_for_offset())
43 logger.debug("total layer num = %d", self.total_layer_num())
44 self.balancing = BA.BalancingAdapter(
45 self.layer_num_for_offset(),
46 copy.deepcopy(self.ccfg.offset),
47 copy.deepcopy(self.ccfg.full_rec),
48 mppb,
49 )
51 def dim_val(self, dim, parallel_config):
52 """Get the value of a parallel dimension"""
53 if parallel_config.has_dim(dim):
54 return parallel_config.val(dim)
55 return dim.from_config(self.ccfg)
57 def global_batch_size(self, parallel_config):
58 """Compute global batch size from hyperparameters"""
59 dp = self.dim_val(Dim.DP, parallel_config)
60 pp = self.dim_val(Dim.PP, parallel_config)
61 mb = self.dim_val(Dim.MBN, parallel_config)
62 bs = self.dim_val(Dim.MBS, parallel_config)
63 if pp > 1:
64 logger.info("GBS = %dDP * %dMB * %dBS", dp, mb, bs)
65 return dp * mb * bs
66 logger.info("GBS = %dDP * %dBS", dp, bs)
67 return dp * bs
69 def layer_num_for_offset(self):
70 """Compute layer number including MTP when necessary for offset"""
71 layer_num = self.ccfg.n_lay
72 if self.ccfg.emb_out_in_offset:
73 layer_num += 2
74 if self.ccfg.is_mtp_in_offset:
75 layer_num += self.ccfg.n_mtp
76 return layer_num
78 def total_layer_num(self):
79 """Compute total layer number, always including MTP"""
80 layer_num = self.ccfg.n_lay + self.ccfg.n_mtp
81 return layer_num
83 def adapt_config_balancing(self, new_pp, new_vpp):
84 """Adapt the layer-to-stage assignment to different PP"""
85 logger.debug("new_pp=%d, new_vpp=%d", new_pp, new_vpp)
87 new_recompute_config = self.balancing.treat_recompute(new_pp, new_vpp)
88 logger.debug("adapted recompute config: %s", str(new_recompute_config))
89 new_offset = self.balancing.treat_offset(new_pp, new_vpp)
90 logger.debug("adapted offset: %s", str(new_offset))
91 ok = self.balancing.offset_checker(new_pp, new_vpp, new_offset)
92 if not ok:
93 logger.error("Offset {%s} NOT VALID", str(new_offset))
94 return new_offset, new_recompute_config
96 def adapt_config(self, pp, vpp):
97 """Adapt configuration to different parallel config"""
98 return self.adapt_config_balancing(pp, vpp)
100 def write(self, folder, parallel_config):
101 """Dump config into a yaml file"""
102 if folder:
103 file_name = parallel_config.unique_name()
104 self.ccfg.config.dump(file_name, folder)
106 def moe_valid(self, parallel_config):
107 """Check whether the model is MoE"""
108 expert_num = self.ccfg.n_exp
109 if expert_num > 1:
110 ep = self.dim_val(Dim.EP, parallel_config)
111 dp = self.dim_val(Dim.DP, parallel_config)
112 mp = self.dim_val(Dim.TP, parallel_config)
113 logger.debug(
114 "moe valid ? EP %d <= E %d & EP %d <= DP %d * MP %d",
115 ep,
116 expert_num,
117 ep,
118 dp,
119 mp,
120 )
121 return ep <= min(expert_num, dp * mp)
122 return True
124 def make_parallel_config_args(self, **kwargs):
125 """Create a parallel config from parallel values"""
126 logger.debug("dimensions considered: %s", str(self.dimensions))
128 dims = []
129 # dims.append((Dim.DP, dp))
130 for dim in self.dimensions:
131 dims.append((dim, kwargs.get(dim.lname())))
133 has_mbn_not_in = Dim.MBN not in self.dimensions
134 has_pp_in = Dim.PP in self.dimensions
135 has_dp_or_mbs_in = Dim.DP in self.dimensions or Dim.MBS in self.dimensions
136 if has_mbn_not_in and has_pp_in and has_dp_or_mbs_in:
137 dims.append((Dim.MBN, kwargs.get(Dim.MBN.lname())))
138 self.dimensions.append(Dim.MBN)
139 return Dim.Dimensions(dims, all_dims=self.dimensions)
141 def make_parallel_config(self, dtpc_p, mbsn, evos_p):
142 """Create a parallel config from parallel values"""
143 logger.debug("dimensions considered: %s", str(self.dimensions))
144 (dp, mp, pp, cp) = dtpc_p
145 (mbs, mbn) = mbsn
146 (ep, vpp, op, sp) = evos_p
147 return self.make_parallel_config_args(
148 dp=dp,
149 mp=mp,
150 pp=pp,
151 cp=cp,
152 mbs=mbs,
153 mb=mbn,
154 ep=ep,
155 vpp=vpp,
156 op=op,
157 sp=sp,
158 )
160 def set_parallel_config(self, parallel_config):
161 """Set a given parallel configuration in the config"""
162 kwargs = {}
163 ok = True
164 new_pp = self.dim_val(Dim.PP, parallel_config)
165 new_vp = self.dim_val(Dim.VPP, parallel_config)
166 new_offset, new_recompute = self.adapt_config(new_pp, new_vp)
167 kwargs["offset"] = new_offset
168 kwargs["full_rec"] = new_recompute
169 # kwargs["sel_rec"] = sel_rec
170 for dim, value in parallel_config.dims_val.items():
171 kwargs[dim.name.lower()] = value
173 self.ccfg.set_strategy(**kwargs)
174 if not self.ccfg.multimodal:
175 if not self.ccfg.hooks_dict:
176 logger.info(
177 "'hook_cls' not specified,"
178 "search in predefined arch_hooks"
179 )
180 check_and_apply_custom_hook(self.ccfg)
181 else:
182 logger.info("Apply hooks")
183 hook = list(self.ccfg.hooks_dict.values())[0]
184 hook(self.wrap)
186 return ok
188 def space(self, dim, divide, reverse=False):
189 """Generate the space for a given dimension"""
190 if dim in self.dimensions:
191 if dim.get_bound() is not None:
192 logger.debug(
193 "Space of bounded dim %s is %s",
194 str(dim),
195 str(
196 Hard.all_divisors(
197 divide, reverse=reverse, max_bound=dim.get_bound()
198 )
199 ),
200 )
201 return Hard.all_divisors(
202 divide, reverse=reverse, max_bound=dim.get_bound()
203 )
204 logger.debug(
205 "Space of dim %s is %s",
206 str(dim),
207 str(Hard.all_divisors(divide, reverse=reverse)),
208 )
209 return Hard.all_divisors(divide, reverse=reverse)
210 logger.debug(
211 "Space of original dim %s is [%s]",
212 str(dim),
213 str(dim.from_config(self.ccfg)),
214 )
215 return [dim.from_config(self.ccfg)]
217 def range_space(self, dim, bound):
218 """Generate the space for a given dimension"""
219 if dim in self.dimensions:
220 return range(1, bound + 1)
221 return [dim.from_config(self.ccfg)]
223 def bool_space(self, dim):
224 """Generate the space for a given boolean dimension"""
225 if dim in self.dimensions:
226 return [False, True]
227 return [dim.from_config(self.ccfg)]
229 def max_op(self, dp, tp, ep):
230 """Compute bound for dimension OP"""
231 if (
232 isinstance(self.ccfg.optimizer, str)
233 and "muon" not in self.ccfg.optimizer.lower()
234 ):
235 return dp
236 if self.ccfg.n_exp and self.ccfg.n_exp > 1:
237 exp_gcd = gcd(dp * tp // max(tp, ep), self.ccfg.n_exp)
238 else:
239 exp_gcd = dp
241 dc_kv_valid = self.ccfg.dc_kv and self.ccfg.dc_kv > 1
242 dhr_valid = self.ccfg.dhr and self.ccfg.dhr > 1
243 if dc_kv_valid and dhr_valid:
244 att_gcd = gcd(self.ccfg.h, self.ccfg.dc_kv + self.ccfg.dhr)
245 else:
246 att_gcd = self.ccfg.h
247 return gcd(exp_gcd, att_gcd)