Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / dimensions.py: 94%
208 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +0800
1# Copyright 2025-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"""parallel dimensions"""
16from __future__ import annotations
18import sys
19from typing import TYPE_CHECKING, Union
21from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger
22from hyper_parallel.auto_parallel.sapp_nd.nd.common.cp_types import (
23 CPValidationResult,
24 CPConstraintParams,
25)
27if TYPE_CHECKING:
28 from hyper_parallel.auto_parallel.sapp_nd.nd.common._cost_model_variables import _CostModVar
31class Dimension:
32 """Output dimension"""
34 def __init__(
35 self,
36 acronym,
37 cost_model_var_name,
38 from_str,
39 default=1,
40 ):
41 self.name = acronym
42 self.cost_model_var = cost_model_var_name
43 self.default = default
44 self.bound = None
45 self.from_str = from_str
47 def __str__(self):
48 return self.name
50 def __repr__(self):
51 return str(self)
53 def lname(self):
54 """lower case name"""
55 return self.name.lower()
57 def from_config(self, ccfg: _CostModVar):
58 """Get dimension value from cost model config"""
59 try:
60 value = ccfg.__dict__[self.cost_model_var]
61 except KeyError:
62 logger.error(
63 "variable %s does not exist in the cost model: %s",
64 self.cost_model_var,
65 str(ccfg.__dict__),
66 )
67 sys.exit(1)
68 return value
70 def reset_bound(self):
71 """Reset bound the dimension space"""
72 self.bound = None
74 def set_bound(self, bound):
75 """Bound the dimension space"""
76 if self.bound:
77 logger.debug(
78 "bound(%s) = min (%d, %d)", self.name, bound, self.bound
79 )
80 self.bound = min(bound, self.bound)
81 else:
82 logger.debug("bound(%s) = %d", self.name, bound)
83 self.bound = bound
85 def get_bound(self):
86 """Return dimension bound"""
87 return self.bound
89 def is_valid(self, value):
90 """Check dimension value validity"""
91 invalid = False
92 if isinstance(value, bool):
93 return not invalid
94 if isinstance(value, int):
95 invalid = self.bound and value > self.bound
96 invalid = invalid or value < 1
97 if invalid:
98 logger.warning(
99 "Dimension %s = %s is invalid", self.name, str(value)
100 )
101 return not invalid
104DP = Dimension(
105 "DP",
106 "d",
107 default=1,
108 from_str=int,
109)
110EP = Dimension(
111 "EP",
112 "ep",
113 default=1,
114 from_str=int,
115)
116TP = Dimension(
117 "MP",
118 "t",
119 default=1,
120 from_str=int,
121)
122CP = Dimension(
123 "CP",
124 "cp",
125 default=1,
126 from_str=int,
127)
128PP = Dimension(
129 "PP",
130 "p",
131 default=1,
132 from_str=int,
133)
134MBN = Dimension(
135 "MB",
136 "m",
137 default=1,
138 from_str=int,
139)
140MBS = Dimension(
141 "MBS",
142 "b",
143 default=1,
144 from_str=int,
145)
146SP = Dimension(
147 "SP",
148 "sp",
149 default=True,
150 from_str=bool,
151)
152OP = Dimension(
153 "OP",
154 "os_max_shard",
155 # "op_weight_shard",
156 default=1,
157 from_str=int,
158)
159VPP = Dimension(
160 "VPP",
161 "vp",
162 default=1,
163 from_str=int,
164)
166ALL_DIMS = [DP, EP, TP, CP, PP, VPP, MBN, MBS, SP, OP]
169class Dimensions:
170 """All output dimensions"""
172 def __init__(self, config, all_dims=None):
173 if isinstance(config, list):
174 self.all_dims = [d for d, _ in config]
175 self.dims_val = dict(config)
176 # elif isinstance(config, dict):
177 # self.all_dims = ALL_DIMS
178 # self.dims_val = {d: d.from_config(config) for d in self.all_dims}
179 elif isinstance(config, bool):
180 self.all_dims = ALL_DIMS
181 self.dims_val = {d: d.default for d in self.all_dims}
182 else:
183 raise TypeError(
184 f"Dimensions cannot be constructed from type {type(config)}"
185 )
186 if all_dims:
187 self.all_dims = all_dims
188 self._reset_all_dims()
190 def _reset_all_dims(self):
191 for d in self.all_dims:
192 d.reset_bound()
194 def __str__(self):
195 return str(self.dims_val)
197 def __repr__(self):
198 return str(self)
200 def keys(self):
201 """Return dimensions"""
202 return list(self.dims_val)
204 def global_batch_size(self):
205 """Compute the global batch size"""
206 gbs = self.dims_val[DP] * self.dims_val[MBS]
207 has_pp = self.has_dim(PP) and self.dims_val[PP] > 1
208 has_mbn = self.has_dim(MBN) and self.dims_val[MBN] > 1
209 if has_pp and has_mbn:
210 gbs *= self.dims_val[MBN]
211 return gbs
213 def values(self):
214 """Return dimension value"""
215 return [str(self.dims_val[d]) for d in self.dims_val]
217 def unique_name(self):
218 """Return all values as a unique string"""
219 return "_".join(self.values())
221 def has_dim(self, d):
222 """Check that this dimension has a value in the parallel config"""
223 return d in self.dims_val
225 @staticmethod
226 def _check_mbn_pp(dims_val, all_dims):
227 """Return True if MBN/PP combination is valid."""
228 if MBN not in dims_val or PP not in all_dims:
229 return True
230 valid = dims_val[MBN] >= dims_val[PP]
231 valid = valid and not (dims_val[PP] == 1 and dims_val[MBN] > 1)
232 if not valid:
233 logger.warning("PP and MBN were deemed not suitable")
234 return valid
236 @staticmethod
237 def _check_power_of_two(dim_obj, value):
238 """Return True if *value* is a power of 2 (for TP / OP)."""
239 if not value & (value - 1) == 0:
240 logger.warning("%s must be a power of 2", str(dim_obj))
241 return False
242 return True
244 def is_valid(self):
245 """Check if all dimensions values are valid"""
246 if not self._check_mbn_pp(self.dims_val, self.all_dims):
247 return False
248 if TP in self.all_dims and not self._check_power_of_two(TP, self.dims_val[TP]):
249 return False
250 for d in self.dims_val:
251 if not d.is_valid(self.dims_val[d]):
252 logger.warning("Dimension %d is not valid", d)
253 return False
254 if SP in self.all_dims and CP in self.all_dims:
255 if self.dims_val[SP] and self.dims_val[CP] > 1:
256 logger.warning("SP & CP cannot coexist")
257 return False
258 if OP in self.all_dims and not self._check_power_of_two(OP, self.dims_val[OP]):
259 return False
260 return True
262 def val(self, dim):
263 """Get Dimension value"""
264 return self.dims_val[dim]
266 def set(self, dim, val):
267 """Get Dimension value"""
268 self.dims_val[dim] = val
270 def steal(self, factor, dim_from, dim_to):
271 """Assign a dimension factor to another dimension"""
272 self.dims_val[dim_from] = self.dims_val[dim_from] // factor
273 self.dims_val[dim_to] = self.dims_val[dim_to] * factor
276def get_dim(acronym):
277 """Return the dimension of the given string acronym"""
278 dname = str(acronym).upper()
279 for d in ALL_DIMS:
280 if d.name == dname:
281 return d
282 raise ValueError(f"Dimension {dname} does NOT exist")
285def get_dims(dims):
286 """Return all dimensions considered"""
287 if dims is None:
288 return ALL_DIMS
289 return [get_dim(acronym) for acronym in dims]
292def _cp_params_from_args(
293 seq_len, cp_degree, tp_degree, pp_degree, device_per_node,
294 attention_type_str, bw_intra, bw_inter, total_devices,
295 cp_memory_per_layer, device_capacity, num_layers,
296 cp_algo, attention_heads, sp_enabled, num_kv_heads=0,
297):
298 """Resolve CPConstraintParams from flexible args."""
299 if isinstance(seq_len, CPConstraintParams):
300 return seq_len
301 return CPConstraintParams(
302 seq_len=seq_len, cp_degree=cp_degree, tp_degree=tp_degree,
303 pp_degree=pp_degree, device_per_node=device_per_node,
304 attention_type_str=attention_type_str, bw_intra=bw_intra,
305 bw_inter=bw_inter, total_devices=total_devices,
306 cp_memory_per_layer=cp_memory_per_layer,
307 device_capacity=device_capacity, num_layers=num_layers,
308 cp_algo=cp_algo, attention_heads=attention_heads,
309 num_kv_heads=num_kv_heads, sp_enabled=sp_enabled,
310 )
313def _cp_ok_result(**overrides):
314 """Build a passing CPValidationResult with sensible defaults."""
315 defaults = {
316 "is_valid": True, "error_message": None, "warning_message": None,
317 "seq_len_divisible": True, "topology_feasible": True,
318 "device_sufficient": True, "memory_within_limit": True,
319 "recommended_cp_max": None, "topology_penalty": None,
320 "unsupported_reason": None,
321 }
322 defaults.update(overrides)
323 return CPValidationResult(**defaults)
326def _cp_check_ulysses_heads(p):
327 """Check Ulysses KV-head divisibility. Returns error message or None.
329 Ulysses CP shards KV heads across tp×cp ranks, so num_kv_heads must be
330 divisible by tp×cp. When TP replicates KV (GQA with tp ≤ kv-head-groups)
331 the check is conservative but safe — it never lets a crashing config
332 through. num_kv_heads=0 falls back to attention_heads (MHA semantics,
333 matching compute_kv_dim).
334 """
335 if p.cp_algo != "ulysses_cp":
336 return None
337 kv_heads = p.num_kv_heads if p.num_kv_heads > 0 else p.attention_heads
338 if kv_heads <= 0:
339 return None
340 shards = p.tp_degree * p.cp_degree
341 if shards <= 0:
342 return None
343 if kv_heads % shards != 0:
344 return (
345 f"Ulysses CP requires num_kv_heads ({kv_heads}) divisible by "
346 f"tp_degree × cp_degree ({p.tp_degree} × {p.cp_degree} = {shards})."
347 )
348 return None
351def _cp_check_device_sufficiency(p):
352 """Check tp*cp*pp <= total_devices. Returns error message or None."""
353 if p.total_devices > 0 and p.tp_degree * p.cp_degree * p.pp_degree > p.total_devices:
354 return (
355 f"tp×cp×pp ({p.tp_degree}×{p.cp_degree}×{p.pp_degree} = "
356 f"{p.tp_degree * p.cp_degree * p.pp_degree}) exceeds "
357 f"total available devices ({p.total_devices})."
358 )
359 return None
362def _cp_collect_warnings(p):
363 """Collect warning messages and derived topology/recommendation info."""
364 warnings = []
365 topology_feasible = True
366 topology_penalty = None
367 recommended_cp_max = None
369 if p.tp_degree * p.cp_degree > p.device_per_node:
370 topology_feasible = False
371 topology_penalty = 1.0 - (p.bw_inter / p.bw_intra)
372 warnings.append(
373 f"CP will cross node boundary (tp={p.tp_degree} × cp={p.cp_degree} = "
374 f"{p.tp_degree * p.cp_degree} > {p.device_per_node} devices/node). "
375 f"Communication will use slower inter-node bandwidth."
376 )
378 if p.seq_len < 8192:
379 warnings.append(
380 f"CP not recommended for short sequences (seq_len={p.seq_len} < 8192). "
381 f"Communication overhead may outweigh memory benefits."
382 )
384 attn_upper = p.attention_type_str.upper()
385 if attn_upper == "MLA":
386 recommended_cp_max = 16
387 elif attn_upper == "GQA":
388 recommended_cp_max = 8
389 else:
390 recommended_cp_max = 4
392 if recommended_cp_max is not None and p.cp_degree > recommended_cp_max:
393 warnings.append(
394 f"cp_degree ({p.cp_degree}) exceeds recommended max ({recommended_cp_max}) "
395 f"for {attn_upper} attention."
396 )
398 return warnings, topology_feasible, topology_penalty, recommended_cp_max
401def _cp_check_memory(p, warnings, topology_feasible, topology_penalty, recommended_cp_max):
402 """Check CP memory capacity. Returns CPValidationResult or None."""
403 if not (p.cp_memory_per_layer > 0 and p.device_capacity > 0 and p.num_layers > 0):
404 return None
405 total_cp_memory = p.cp_memory_per_layer * p.num_layers
406 if total_cp_memory <= p.device_capacity:
407 return None
408 return _cp_ok_result(
409 is_valid=False,
410 error_message=(
411 f"CP memory per card ({total_cp_memory / 1e6:.1f} MB) exceeds "
412 f"device capacity ({p.device_capacity / 1e6:.1f} MB)."
413 ),
414 warning_message="; ".join(warnings) if warnings else None,
415 topology_feasible=topology_feasible,
416 topology_penalty=topology_penalty,
417 recommended_cp_max=recommended_cp_max,
418 memory_within_limit=False,
419 )
422def validate_cp_constraints(
423 seq_len: Union[CPConstraintParams, int],
424 cp_degree: int = 1,
425 tp_degree: int = 1,
426 pp_degree: int = 1,
427 device_per_node: int = 8,
428 attention_type_str: str = "mha",
429 bw_intra: float = 300.0,
430 bw_inter: float = 25.0,
431 total_devices: int = 0,
432 cp_memory_per_layer: float = 0.0,
433 device_capacity: float = 0.0,
434 num_layers: int = 0,
435 cp_algo: str = "colossalai_cp",
436 attention_heads: int = 0,
437 sp_enabled: bool = False,
438 num_kv_heads: int = 0,
439) -> CPValidationResult:
440 """Validate CP constraints for a given parallel configuration.
442 Accepts either a CPConstraintParams dataclass or individual keyword
443 arguments for backward compatibility.
445 Args:
446 seq_len: Sequence length, or a CPConstraintParams dataclass.
447 cp_degree: CP degree.
448 tp_degree: TP degree.
449 pp_degree: PP degree.
450 device_per_node: Number of devices per node.
451 attention_type_str: Attention type string ("mha", "gqa", "mla").
452 bw_intra: Intra-node bandwidth in GB/s (default: 300.0 for Ascend A2).
453 bw_inter: Inter-node bandwidth in GB/s (default: 25.0).
454 total_devices: Total number of available devices (0 = skip check).
455 cp_memory_per_layer: CP memory per layer in bytes (0 = skip check).
456 device_capacity: Device memory capacity in bytes (0 = skip check).
457 num_layers: Number of transformer layers (0 = skip check).
458 cp_algo: CP algorithm ("colossalai_cp", "hybrid_cp", "ulysses_cp").
459 attention_heads: Number of attention heads (0 = skip Ulysses head check).
460 sp_enabled: Whether sequence parallelism is enabled (SP and CP are incompatible).
461 num_kv_heads: Number of KV heads for Ulysses divisibility check
462 (0 = fall back to attention_heads, matching compute_kv_dim).
464 Returns:
465 CPValidationResult with validation outcome.
466 """
467 p = _cp_params_from_args(
468 seq_len, cp_degree, tp_degree, pp_degree, device_per_node,
469 attention_type_str, bw_intra, bw_inter, total_devices,
470 cp_memory_per_layer, device_capacity, num_layers,
471 cp_algo, attention_heads, sp_enabled, num_kv_heads,
472 )
474 if p.cp_degree <= 1:
475 return _cp_ok_result()
477 if p.sp_enabled:
478 return _cp_ok_result(
479 is_valid=False,
480 error_message=(
481 f"Context Parallelism (cp={p.cp_degree}) and Sequence Parallelism "
482 f"are incompatible and cannot be used together."
483 ),
484 unsupported_reason="CP+SP incompatible",
485 )
487 if p.seq_len % (p.cp_degree * 2) != 0:
488 return _cp_ok_result(
489 is_valid=False, seq_len_divisible=False,
490 error_message=(
491 f"Sequence length {p.seq_len} must be divisible by "
492 f"cp_degree × 2 = {p.cp_degree * 2}. "
493 f"Current remainder: {p.seq_len % (p.cp_degree * 2)}"
494 ),
495 )
497 ulysses_err = _cp_check_ulysses_heads(p)
498 if ulysses_err:
499 return _cp_ok_result(
500 is_valid=False, error_message=ulysses_err,
501 unsupported_reason="Ulysses insufficient heads",
502 )
504 device_err = _cp_check_device_sufficiency(p)
505 if device_err:
506 return _cp_ok_result(
507 is_valid=False, error_message=device_err, device_sufficient=False,
508 )
510 warnings, topology_feasible, topology_penalty, recommended_cp_max = (
511 _cp_collect_warnings(p)
512 )
514 mem_result = _cp_check_memory(
515 p, warnings, topology_feasible, topology_penalty, recommended_cp_max)
516 if mem_result is not None:
517 return mem_result
519 return _cp_ok_result(
520 warning_message="; ".join(warnings) if warnings else None,
521 topology_feasible=topology_feasible,
522 topology_penalty=topology_penalty,
523 recommended_cp_max=recommended_cp_max,
524 )