Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / global_config.py: 95%
177 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 typing import Any, Dict, Optional
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, parent=None):
32 self.wrap = CWrap(config)
33 self.ccfg = self.wrap.ccfg
34 # The other submodules of a multimodal model, which run on the same
35 # stages as the one the search drives: each takes the candidate's
36 # degrees, with its own layers placed and its own recompute adapted.
37 self.siblings = self._pipeline_siblings(config, parent, mppb)
39 if dimensions is not None:
40 logger.debug("dimensions = %s", str(dimensions))
41 self.dimensions = dimensions
42 else:
43 logger.debug("dimensions = %s", str(Dim.ALL_DIMS))
44 self.dimensions = Dim.ALL_DIMS.copy()
45 logger.debug("self.dimensions = %s", str(self.dimensions))
46 logger.debug("layer_num_for_offset = %d", self.layer_num_for_offset())
47 logger.debug("total layer num = %d", self.total_layer_num())
48 self.balancing = BA.BalancingAdapter(
49 self.layer_num_for_offset(),
50 copy.deepcopy(self.ccfg.offset),
51 copy.deepcopy(self.ccfg.full_rec),
52 mppb,
53 )
55 def dim_val(self, dim, parallel_config):
56 """Get the value of a parallel dimension"""
57 if parallel_config.has_dim(dim):
58 return parallel_config.val(dim)
59 return dim.from_config(self.ccfg)
61 def accumulates_grads(self) -> bool:
62 """Whether the run accumulates gradients over micro-batches without pipeline parallelism."""
63 return bool(getattr(self.ccfg, "accumulates_grads", False))
65 def global_batch_size(self, parallel_config):
66 """Compute global batch size from hyperparameters.
68 Micro-batches multiply it under pipeline parallelism, and without it
69 in a run that accumulates gradients over them.
70 """
71 dp = self.dim_val(Dim.DP, parallel_config)
72 pp = self.dim_val(Dim.PP, parallel_config)
73 mb = self.dim_val(Dim.MBN, parallel_config)
74 bs = self.dim_val(Dim.MBS, parallel_config)
75 if pp > 1 or self.accumulates_grads():
76 logger.info("GBS = %dDP * %dMB * %dBS", dp, mb, bs)
77 return dp * mb * bs
78 logger.info("GBS = %dDP * %dBS", dp, bs)
79 return dp * bs
81 def layer_num_for_offset(self):
82 """Compute layer number including MTP when necessary for offset"""
83 layer_num = self.ccfg.n_lay
84 if self.ccfg.emb_out_in_offset:
85 layer_num += 2
86 if self.ccfg.is_mtp_in_offset:
87 layer_num += self.ccfg.n_mtp
88 return layer_num
90 def total_layer_num(self):
91 """Compute total layer number, always including MTP"""
92 layer_num = self.ccfg.n_lay + self.ccfg.n_mtp
93 return layer_num
95 def adapt_config_balancing(self, new_pp, new_vpp):
96 """Adapt the layer-to-stage assignment to different PP"""
97 logger.debug("new_pp=%d, new_vpp=%d", new_pp, new_vpp)
99 new_recompute_config = self.balancing.treat_recompute(new_pp, new_vpp)
100 logger.debug("adapted recompute config: %s", str(new_recompute_config))
101 new_offset = self.balancing.treat_offset(new_pp, new_vpp)
102 logger.debug("adapted offset: %s", str(new_offset))
103 ok = self.balancing.offset_checker(new_pp, new_vpp, new_offset)
104 if not ok:
105 logger.error("Offset {%s} NOT VALID", str(new_offset))
106 return new_offset, new_recompute_config
108 def adapt_config(self, pp, vpp):
109 """Adapt configuration to different parallel config"""
110 return self.adapt_config_balancing(pp, vpp)
112 def state_recompute(self, full: Optional[bool]) -> None:
113 """Give every candidate *full* as its full recompute, on every config of its pipeline.
115 Args:
116 full: The full recompute a recompute dimension's mode states, or
117 None for the recompute each candidate's balancing derives.
118 """
119 self.balancing.stated_recompute = full
120 for _, balancing in self.siblings:
121 if balancing is not None:
122 balancing.stated_recompute = full
124 def write(self, folder, parallel_config):
125 """Dump config into a yaml file"""
126 if folder:
127 file_name = parallel_config.unique_name()
128 self.ccfg.config.dump(file_name, folder)
130 def moe_valid(self, parallel_config):
131 """Whether a MoE model's EP fits its experts and the ranks they spread over (:meth:`max_ep`)."""
132 expert_num = self.ccfg.n_exp
133 if expert_num > 1:
134 ep = self.dim_val(Dim.EP, parallel_config)
135 dp = self.dim_val(Dim.DP, parallel_config)
136 mp = self.dim_val(Dim.TP, parallel_config)
137 cp = self.dim_val(Dim.CP, parallel_config)
138 logger.debug(
139 "moe valid ? EP %d <= E %d & EP %d <= %d, DP %d MP %d CP %d",
140 ep,
141 expert_num,
142 ep,
143 self.max_ep(dp, mp, cp),
144 dp,
145 mp,
146 cp,
147 )
148 return ep <= min(expert_num, self.max_ep(dp, mp, cp))
149 return True
151 def max_ep(self, dp: int, tp: int, cp: int = 1) -> int:
152 """Compute bound for dimension EP: the ranks a stage spreads its experts over.
154 They are its DP x TP ranks, and under HyperParallel its DP x CP x TP
155 ranks: it builds its expert mesh ``(edp_replicate, edp_shard, ep)``
156 over the whole ``(dp, cp, tp)`` device mesh
157 (``MeshContext._build_expert_parallel_mesh``), the domain its FSDP
158 shard spans as well (``shard_spans_cp``). Bounded by DP x TP, a
159 strategy under CP could not spread its experts over CP's ranks,
160 though the runtime does.
161 """
162 return dp * tp * (cp if getattr(self.ccfg, "shard_spans_cp", False) else 1)
164 def ep_constraints_valid(self, parallel_config):
165 """Check EP-specific divisibility constraints.
167 Runs only for MoE models (n_exp > 1). C1 ensures experts can be
168 evenly partitioned across EP ranks; C2 ensures the expert FFN hidden
169 dim can be evenly sharded by the expert TP degree. Both checks use
170 architecture constants from ``self.ccfg`` and the candidate values
171 from ``parallel_config``.
173 C3 (device count) is intentionally skipped here because the search
174 loop borrows EP from the dp*tp budget, so dp*tp*pp*cp already
175 equals total_devices by construction.
177 Args:
178 parallel_config: candidate ``Dim.Dimensions``.
180 Returns:
181 bool: True if all applicable EP constraints pass (or the model
182 is dense), False otherwise.
183 """
184 if self.ccfg.n_exp <= 1:
185 return True
186 # pylint: disable=C0415
187 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.validators.ep_constraints import EpConstraints
188 ep = self.dim_val(Dim.EP, parallel_config)
189 r1 = EpConstraints.check_ep_divisibility(self.ccfg.n_exp, ep)
190 if not r1:
191 logger.warning("EP constraint C1 failed: %s", r1.message)
192 return False
193 tp = self.dim_val(Dim.TP, parallel_config)
194 etp = max(getattr(self.ccfg, "etp", 0), 0)
195 t_exp = max(etp, 1) if etp > 1 else max(tp, 1)
196 hff_exp = max(getattr(self.ccfg, "hff_exp", 0), 0)
197 r2 = EpConstraints.check_expert_hidden_divisibility(hff_exp, t_exp)
198 if not r2:
199 logger.warning("EP constraint C2 failed: %s", r2.message)
200 return False
201 # C4: an expert shard the run states divides its expert data-parallel
202 # group, the stage's ranks over EP, as HyperParallel requires. A run
203 # that shards over the whole group takes each strategy's own.
204 shard = getattr(self.ccfg, "expert_shard", None) or 1
205 if ep > 1 and shard > 1 and not getattr(self.ccfg, "expert_shard_group", False):
206 group = self.dim_val(Dim.DP, parallel_config) * self.dim_val(Dim.CP, parallel_config) * tp // ep
207 if group % shard:
208 logger.warning("EP constraint C4 failed: an expert shard of %d in a group of %d", shard, group)
209 return False
210 return True
212 def make_parallel_config_args(self, **kwargs):
213 """Create a parallel config from parallel values"""
214 logger.debug("dimensions considered: %s", str(self.dimensions))
216 dims = []
217 # dims.append((Dim.DP, dp))
218 for dim in self.dimensions:
219 dims.append((dim, kwargs.get(dim.lname())))
221 # The micro-batch count is the search's own, the global batch over DP
222 # and MBS, whatever -l names. Kept from the yaml where -l left it out,
223 # every DP but the yaml's own made another batch, so the search
224 # refused every candidate or kept a single DP.
225 if Dim.MBN not in self.dimensions:
226 dims.append((Dim.MBN, kwargs.get(Dim.MBN.lname())))
227 self.dimensions.append(Dim.MBN)
228 return Dim.Dimensions(dims, all_dims=self.dimensions, accumulates=self.accumulates_grads())
230 def make_parallel_config(self, dtpc_p, mbsn, evos_p):
231 """Create a parallel config from parallel values"""
232 logger.debug("dimensions considered: %s", str(self.dimensions))
233 (dp, mp, pp, cp) = dtpc_p
234 (mbs, mbn) = mbsn
235 (ep, vpp, op, sp) = evos_p
236 return self.make_parallel_config_args(
237 dp=dp,
238 mp=mp,
239 pp=pp,
240 cp=cp,
241 mbs=mbs,
242 mb=mbn,
243 ep=ep,
244 vpp=vpp,
245 op=op,
246 sp=sp,
247 )
249 def set_parallel_config(self, parallel_config):
250 """Set a given parallel configuration in the config"""
251 kwargs = {}
252 ok = True
253 new_pp = self.dim_val(Dim.PP, parallel_config)
254 new_vp = self.dim_val(Dim.VPP, parallel_config)
255 new_offset, new_recompute = self.adapt_config(new_pp, new_vp)
256 kwargs["offset"] = new_offset
257 kwargs["full_rec"] = new_recompute
258 # kwargs["sel_rec"] = sel_rec
259 for dim, value in parallel_config.dims_val.items():
260 kwargs[dim.name.lower()] = value
262 self.ccfg.set_strategy(**kwargs)
263 self.siblings_take(kwargs)
264 if not self.ccfg.multimodal:
265 if not self.ccfg.hooks_dict:
266 logger.info(
267 "'hook_cls' not specified,"
268 "search in predefined arch_hooks"
269 )
270 check_and_apply_custom_hook(self.ccfg)
271 else:
272 logger.info("Apply hooks")
273 hook = list(self.ccfg.hooks_dict.values())[0]
274 hook(self.wrap)
276 return ok
278 @staticmethod
279 def _pipeline_siblings(config, parent, mppb):
280 """Every other config of *config*'s pipeline: a multimodal parent's other submodules, and the parent.
282 Each submodule comes with the balancing of its own layers, which
283 adapts its recompute to a candidate's pipeline; the parent has no
284 layers of its own and takes the degrees alone.
285 """
286 if parent is None or not parent.multimodal:
287 return []
288 siblings = []
289 for name in parent.mm_order:
290 sub = parent.mm_ccfgs[name]
291 if sub is not config:
292 siblings.append((sub, BA.BalancingAdapter(
293 sub.n_lay + sub.n_mtp,
294 copy.deepcopy(sub.offset),
295 copy.deepcopy(sub.full_rec),
296 mppb,
297 )))
298 return siblings + [(parent, None)]
300 def siblings_take(self, kwargs: Dict[str, Any]) -> None:
301 """Give every other config of a shared pipeline the degrees of *kwargs*.
303 A vision tower is trained on the same devices as the language model
304 the search drives, so it takes the candidate's degrees; its layers
305 stay on the first stages, where the parser places them, and its
306 recompute is adapted to the candidate's pipeline rather than the
307 language model's, which counts other layers. The parent, which
308 holds the strategy its submodules share and combines their
309 partitions, takes the degrees and no layer of its own.
310 """
311 degrees = {key: value for key, value in kwargs.items() if key not in ("offset", "full_rec")}
312 for sibling, balancing in self.siblings:
313 pipeline = (degrees.get("pp", sibling.p), degrees.get("vpp", sibling.vp))
314 stated = {"offset": BA.front_loaded_offset(balancing.layers if balancing else 0, *pipeline)}
315 if balancing is not None:
316 stated["full_rec"] = balancing.treat_recompute(*pipeline)
317 sibling.set_shared_strategy(**degrees, **stated)
319 def space(self, dim, divide, reverse=False):
320 """Generate the space for a given dimension"""
321 if dim in self.dimensions:
322 if dim.get_bound() is not None:
323 logger.debug(
324 "Space of bounded dim %s is %s",
325 str(dim),
326 str(
327 Hard.all_divisors(
328 divide, reverse=reverse, max_bound=dim.get_bound()
329 )
330 ),
331 )
332 return Hard.all_divisors(
333 divide, reverse=reverse, max_bound=dim.get_bound()
334 )
335 logger.debug(
336 "Space of dim %s is %s",
337 str(dim),
338 str(Hard.all_divisors(divide, reverse=reverse)),
339 )
340 return Hard.all_divisors(divide, reverse=reverse)
341 logger.debug(
342 "Space of original dim %s is [%s]",
343 str(dim),
344 str(dim.from_config(self.ccfg)),
345 )
346 return [dim.from_config(self.ccfg)]
348 def range_space(self, dim, bound):
349 """Generate the space for a given dimension"""
350 if dim in self.dimensions:
351 return range(1, bound + 1)
352 return [dim.from_config(self.ccfg)]
354 def bool_space(self, dim):
355 """Generate the space for a given boolean dimension"""
356 if dim in self.dimensions:
357 return [False, True]
358 return [dim.from_config(self.ccfg)]
360 def max_op(self, dp, tp, ep, cp=1):
361 """Compute bound for dimension OP.
363 OP is the runtime's ``dp_shard_size``, and all the runtime asks of it
364 is that the data-parallel group divide into a replicate axis and a
365 shard axis (``distributed/mesh.py``, ``MeshContext.build_meshs``), so
366 every divisor of DP is reachable. HyperParallel's group is DP times
367 CP, its FSDP domain, where a shard spans CP's ranks
368 (``shard_spans_cp``): bounded by DP alone, a strategy under CP could
369 not be sharded wider than DP, though the runtime shards it over the
370 whole domain.
372 Under Muon this used to narrow to a greatest common divisor over the
373 expert count and the attention widths. Nothing in the runtime asks
374 for that: Newton-Schulz runs on a matrix's last two dimensions and the
375 optimizer all-gathers them whenever the shard falls there, running
376 locally when it does not (``core/optimizer/muon.py``,
377 ``_classify_parameters_for_step``; ``core/optimizer/
378 sharding_category.py``, ``is_last2d_sharded``), and a width that does
379 not divide a parameter is only logged, never refused. The cap came in
380 with the original search import and cost the strategies engineers run:
381 at 64 devices with EP 16 it allowed no shard wider than 4, and at 16
382 devices with EP 16 none at all. What Muon does change, one momentum
383 per matrix in place of two moments, the parser carries as
384 ``optimizer_states``.
385 """
386 del tp, ep # the shard is bounded by the data-parallel group alone
387 return dp * cp if getattr(self.ccfg, "shard_spans_cp", False) else dp