Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / dimensions.py: 95%
219 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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, Any, Optional, 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, of a run that *accumulates* gradients over micro-batches or not.
172 Attributes:
173 rank: One-based position in the full memory-fitting search ranking,
174 or None before the configuration has been ranked.
175 """
177 def __init__(self, config, all_dims=None, accumulates=False, recompute=None):
178 self.accumulates = accumulates
179 # The activation checkpoint mode of a search with a recompute
180 # dimension, which is no degree (recompute_dimension); None elsewhere.
181 self.recompute = recompute
182 self.rank: Optional[int] = None
183 if isinstance(config, list):
184 self.all_dims = [d for d, _ in config]
185 self.dims_val = dict(config)
186 # elif isinstance(config, dict):
187 # self.all_dims = ALL_DIMS
188 # self.dims_val = {d: d.from_config(config) for d in self.all_dims}
189 elif isinstance(config, bool):
190 self.all_dims = ALL_DIMS
191 self.dims_val = {d: d.default for d in self.all_dims}
192 else:
193 raise TypeError(
194 f"Dimensions cannot be constructed from type {type(config)}"
195 )
196 if all_dims:
197 self.all_dims = all_dims
199 def __str__(self):
200 mode = getattr(self, "recompute", None)
201 return f"{self.dims_val} recompute={mode}" if mode else str(self.dims_val)
203 def __repr__(self):
204 return str(self)
206 def keys(self):
207 """Return dimensions"""
208 return list(self.dims_val)
210 def global_batch_size(self):
211 """Compute the global batch size"""
212 gbs = self.dims_val[DP] * self.dims_val[MBS]
213 has_pp = self.has_dim(PP) and self.dims_val[PP] > 1
214 has_mbn = self.has_dim(MBN) and self.dims_val[MBN] > 1
215 if has_pp and has_mbn:
216 gbs *= self.dims_val[MBN]
217 return gbs
219 def values(self):
220 """Return dimension value"""
221 return [str(self.dims_val[d]) for d in self.dims_val]
223 def unique_name(self):
224 """Return all values as a unique string, the recompute mode last where there is one"""
225 mode = getattr(self, "recompute", None)
226 return "_".join(self.values() + ([mode] if mode else []))
228 def has_dim(self, d):
229 """Check that this dimension has a value in the parallel config"""
230 return d in self.dims_val
232 @staticmethod
233 def _check_mbn_pp(dims_val, all_dims, accumulates=False):
234 """Return True if MBN/PP combination is valid.
236 A pipeline needs a micro-batch per stage at least. Without one,
237 micro-batches are gradient accumulation steps, which only a run
238 that *accumulates* its gradients over them takes.
239 """
240 if MBN not in dims_val or PP not in all_dims:
241 return True
242 valid = dims_val[MBN] >= dims_val[PP]
243 valid = valid and (dims_val[PP] > 1 or dims_val[MBN] == 1 or accumulates)
244 if not valid:
245 logger.warning("PP and MBN were deemed not suitable")
246 return valid
248 @staticmethod
249 def _check_power_of_two(dim_obj, value):
250 """Return True if *value* is a power of 2 (for TP / OP)."""
251 if not value & (value - 1) == 0:
252 logger.warning("%s must be a power of 2", str(dim_obj))
253 return False
254 return True
256 def is_valid(self):
257 """Check if all dimensions values are valid"""
258 if not self._check_mbn_pp(self.dims_val, self.all_dims, getattr(self, "accumulates", False)):
259 return False
260 if TP in self.all_dims and not self._check_power_of_two(TP, self.dims_val[TP]):
261 return False
262 for d in self.dims_val:
263 if not d.is_valid(self.dims_val[d]):
264 logger.warning("Dimension %d is not valid", d)
265 return False
266 if SP in self.all_dims and CP in self.all_dims:
267 if self.dims_val[SP] and self.dims_val[CP] > 1:
268 logger.warning("SP & CP cannot coexist")
269 return False
270 if OP in self.all_dims and not self._check_power_of_two(OP, self.dims_val[OP]):
271 return False
272 return True
274 def val(self, dim):
275 """Get Dimension value"""
276 return self.dims_val[dim]
278 def set(self, dim, val):
279 """Get Dimension value"""
280 self.dims_val[dim] = val
282 def steal(self, factor, dim_from, dim_to):
283 """Assign a dimension factor to another dimension"""
284 self.dims_val[dim_from] = self.dims_val[dim_from] // factor
285 self.dims_val[dim_to] = self.dims_val[dim_to] * factor
288def get_dim(acronym):
289 """Return the dimension of the given string acronym"""
290 dname = str(acronym).upper()
291 for d in ALL_DIMS:
292 if d.name == dname:
293 return d
294 raise ValueError(f"Dimension {dname} does NOT exist")
297def get_dims(dims):
298 """Return all dimensions considered"""
299 if dims is None:
300 return ALL_DIMS
301 return [get_dim(acronym) for acronym in dims]
304def _cp_ok_result(**overrides):
305 """Build a passing CPValidationResult with sensible defaults."""
306 defaults = {
307 "is_valid": True, "error_message": None, "warning_message": None,
308 "seq_len_divisible": True, "topology_feasible": True,
309 "device_sufficient": True, "memory_within_limit": True,
310 "recommended_cp_max": None, "topology_penalty": None,
311 "unsupported_reason": None,
312 }
313 defaults.update(overrides)
314 return CPValidationResult(**defaults)
317def _cp_check_ulysses_heads(p):
318 """Check Ulysses KV-head divisibility. Returns error message or None.
320 Ulysses CP shards KV heads across tp×cp ranks, so num_kv_heads must be
321 divisible by tp×cp. When TP replicates KV (GQA with tp ≤ kv-head-groups)
322 the check is conservative but safe — it never lets a crashing config
323 through. num_kv_heads=0 falls back to attention_heads (MHA semantics,
324 matching compute_kv_dim).
325 """
326 if p.cp_algo != "ulysses_cp":
327 return None
328 kv_heads = p.num_kv_heads if p.num_kv_heads > 0 else p.attention_heads
329 if kv_heads <= 0:
330 return None
331 shards = p.tp_degree * p.cp_degree
332 if shards <= 0:
333 return None
334 if kv_heads % shards != 0:
335 return (
336 f"Ulysses CP requires num_kv_heads ({kv_heads}) divisible by "
337 f"tp_degree × cp_degree ({p.tp_degree} × {p.cp_degree} = {shards})."
338 )
339 return None
342def _cp_check_device_sufficiency(p):
343 """Check tp*cp*pp <= total_devices. Returns error message or None."""
344 if p.total_devices > 0 and p.tp_degree * p.cp_degree * p.pp_degree > p.total_devices:
345 return (
346 f"tp×cp×pp ({p.tp_degree}×{p.cp_degree}×{p.pp_degree} = "
347 f"{p.tp_degree * p.cp_degree * p.pp_degree}) exceeds "
348 f"total available devices ({p.total_devices})."
349 )
350 return None
353def _cp_collect_warnings(p):
354 """Collect warning messages and derived topology/recommendation info."""
355 warnings = []
356 topology_feasible = True
357 topology_penalty = None
358 recommended_cp_max = None
360 if p.tp_degree * p.cp_degree > p.device_per_node:
361 topology_feasible = False
362 topology_penalty = 1.0 - (p.bw_inter / p.bw_intra)
363 warnings.append(
364 f"CP will cross node boundary (tp={p.tp_degree} × cp={p.cp_degree} = "
365 f"{p.tp_degree * p.cp_degree} > {p.device_per_node} devices/node). "
366 f"Communication will use slower inter-node bandwidth."
367 )
369 if p.seq_len < 8192:
370 warnings.append(
371 f"CP not recommended for short sequences (seq_len={p.seq_len} < 8192). "
372 f"Communication overhead may outweigh memory benefits."
373 )
375 attn_upper = p.attention_type_str.upper()
376 if attn_upper == "MLA":
377 recommended_cp_max = 16
378 elif attn_upper == "GQA":
379 recommended_cp_max = 8
380 else:
381 recommended_cp_max = 4
383 if recommended_cp_max is not None and p.cp_degree > recommended_cp_max:
384 warnings.append(
385 f"cp_degree ({p.cp_degree}) exceeds recommended max ({recommended_cp_max}) "
386 f"for {attn_upper} attention."
387 )
389 return warnings, topology_feasible, topology_penalty, recommended_cp_max
392def _cp_check_memory(p, warnings, topology_feasible, topology_penalty, recommended_cp_max):
393 """Check CP memory capacity. Returns CPValidationResult or None."""
394 if not (p.cp_memory_per_layer > 0 and p.device_capacity > 0 and p.num_layers > 0):
395 return None
396 total_cp_memory = p.cp_memory_per_layer * p.num_layers
397 if total_cp_memory <= p.device_capacity:
398 return None
399 return _cp_ok_result(
400 is_valid=False,
401 error_message=(
402 f"CP memory per card ({total_cp_memory / 1e6:.1f} MB) exceeds "
403 f"device capacity ({p.device_capacity / 1e6:.1f} MB)."
404 ),
405 warning_message="; ".join(warnings) if warnings else None,
406 topology_feasible=topology_feasible,
407 topology_penalty=topology_penalty,
408 recommended_cp_max=recommended_cp_max,
409 memory_within_limit=False,
410 )
413_CP_POSITIONAL_FIELDS = (
414 "cp_degree", "tp_degree", "pp_degree", "device_per_node", "attention_type_str",
415 "bw_intra", "bw_inter", "total_devices", "cp_memory_per_layer", "device_capacity",
416 "num_layers", "cp_algo", "attention_heads", "sp_enabled", "num_kv_heads",
417)
420def _cp_params_from_args(seq_len, args, kwargs):
421 """Preserve the legacy positional field order when building CP parameters."""
422 if len(args) > len(_CP_POSITIONAL_FIELDS):
423 raise TypeError("validate_cp_constraints accepts at most 16 positional arguments")
424 options = dict(zip(_CP_POSITIONAL_FIELDS, args))
425 duplicates = options.keys() & kwargs.keys()
426 if duplicates:
427 raise TypeError(f"Multiple values for CP arguments: {sorted(duplicates)}")
428 unknown = kwargs.keys() - set(_CP_POSITIONAL_FIELDS)
429 if unknown:
430 raise TypeError(f"Unexpected CP arguments: {sorted(unknown)}")
431 if isinstance(seq_len, CPConstraintParams):
432 return seq_len
433 return CPConstraintParams(seq_len=seq_len, **{"cp_degree": 1, **options, **kwargs})
436def validate_cp_constraints(
437 seq_len: Union[CPConstraintParams, int],
438 *args: Any,
439 **kwargs: Any,
440) -> CPValidationResult:
441 """Validate CP constraints for a given parallel configuration.
443 Accepts either a CPConstraintParams dataclass or its fields as positional
444 and keyword arguments for backward compatibility.
446 Args:
447 seq_len: Sequence length, or a CPConstraintParams dataclass.
448 *args: Legacy positional fields, from cp_degree through num_kv_heads.
449 The final two fields retain the order sp_enabled, num_kv_heads.
450 **kwargs: The other CPConstraintParams fields, read only when seq_len
451 is an int. An omitted field takes the dataclass default, except
452 cp_degree, which defaults to 1.
454 cp_degree: CP degree.
455 tp_degree: TP degree.
456 pp_degree: PP degree.
457 device_per_node: Number of devices per node.
458 attention_type_str: Attention type string ("mha", "gqa", "mla").
459 bw_intra: Intra-node bandwidth in GB/s (default: 300.0 for Ascend A2).
460 bw_inter: Inter-node bandwidth in GB/s (default: 25.0).
461 total_devices: Total number of available devices (0 = skip check).
462 cp_memory_per_layer: CP memory per layer in bytes (0 = skip check).
463 device_capacity: Device memory capacity in bytes (0 = skip check).
464 num_layers: Number of transformer layers (0 = skip check).
465 cp_algo: CP algorithm ("colossalai_cp", "hybrid_cp", "ulysses_cp").
466 attention_heads: Number of attention heads (0 = skip Ulysses head check).
467 sp_enabled: Whether sequence parallelism is enabled (SP and CP are incompatible).
468 num_kv_heads: Number of KV heads for Ulysses divisibility check
469 (0 = fall back to attention_heads, matching compute_kv_dim).
471 Returns:
472 CPValidationResult with validation outcome.
473 """
474 p = _cp_params_from_args(seq_len, args, kwargs)
476 if p.cp_degree <= 1:
477 return _cp_ok_result()
479 if p.sp_enabled:
480 return _cp_ok_result(
481 is_valid=False,
482 error_message=(
483 f"Context Parallelism (cp={p.cp_degree}) and Sequence Parallelism "
484 f"are incompatible and cannot be used together."
485 ),
486 unsupported_reason="CP+SP incompatible",
487 )
489 if p.seq_len % (p.cp_degree * 2) != 0:
490 return _cp_ok_result(
491 is_valid=False, seq_len_divisible=False,
492 error_message=(
493 f"Sequence length {p.seq_len} must be divisible by "
494 f"cp_degree × 2 = {p.cp_degree * 2}. "
495 f"Current remainder: {p.seq_len % (p.cp_degree * 2)}"
496 ),
497 )
499 ulysses_err = _cp_check_ulysses_heads(p)
500 if ulysses_err:
501 return _cp_ok_result(
502 is_valid=False, error_message=ulysses_err,
503 unsupported_reason="Ulysses insufficient heads",
504 )
506 device_err = _cp_check_device_sufficiency(p)
507 if device_err:
508 return _cp_ok_result(
509 is_valid=False, error_message=device_err, device_sufficient=False,
510 )
512 warnings, topology_feasible, topology_penalty, recommended_cp_max = (
513 _cp_collect_warnings(p)
514 )
516 mem_result = _cp_check_memory(
517 p, warnings, topology_feasible, topology_penalty, recommended_cp_max)
518 if mem_result is not None:
519 return mem_result
521 return _cp_ok_result(
522 warning_message="; ".join(warnings) if warnings else None,
523 topology_feasible=topology_feasible,
524 topology_penalty=topology_penalty,
525 recommended_cp_max=recommended_cp_max,
526 )