Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / common / hardware.py: 97%
194 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"""hardware abstraction"""
17from __future__ import annotations
19from dataclasses import dataclass, replace
20from typing import Any, Mapping, Optional
22import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim
23from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger
26@dataclass(frozen=True)
27class HostLink:
28 """How fast a device moves activations to its host's memory and back, as offload prices it.
30 A3's bandwidth and copy cost below are measured; every other figure is an
31 assumption.
32 Measure a link with every die of a job copying at once, as
33 ``nd_golden/offload/bench_host_link.py`` does, since a job copies at its
34 slowest die, and state it here or with ``run_nd --host_link_gibps``.
36 Attributes:
37 gib_per_s: The sustained copy bandwidth between the device and
38 pinned host memory, in GiB/s, with one copy stream carrying both
39 directions, as hyper_offload and the trainer's input swap run it.
40 sustained_tflops: The device's sustained dense throughput at the
41 training precision, in TFLOP/s, which turns a copy's seconds into
42 the performance estimate's units where no calibration states
43 how long one unit takes.
44 overlap: The share of the forward time the copy stream may take:
45 1 where the copies meet nothing, less where the step's compute
46 or its collectives share what the copies need. 0.8 is an
47 assumption; ``run_nd --host_link_overlap`` states another.
48 ms_per_unit: The milliseconds one unit of the forward the
49 performance estimate prices stands for: a calibration round's
50 FORWARD ratio where it measured the forward on its own, else its
51 COMPUTE ratio. Stated, it turns a copy's seconds into the
52 estimate's units in place of *sustained_tflops*: the estimate
53 counts a step's matmul FLOPs, a fraction of its work, so a
54 device rate makes every forward look several times shorter than
55 it runs; and it may split the forward from the backward
56 otherwise than the device does, so that COMPUTE can make the
57 forward the copies hide in look longer.
58 copy_cost_ms_per_gib: What the copies cost the step, in milliseconds
59 a GiB moved either way: the time the compute and the collectives
60 beside them lose to the copies, which hide in the forward but
61 still take part of what those need. 0 prices offload as free.
62 score_ms_per_unit: The milliseconds one unit of the score stands
63 for, a calibration round's COMPUTE ratio, which turns that cost
64 into the score's units; unset, the cost converts as a copy's
65 seconds do.
66 """
68 gib_per_s: float
69 sustained_tflops: float
70 overlap: float = 0.8
71 ms_per_unit: Optional[float] = None
72 copy_cost_ms_per_gib: float = 0.0
73 score_ms_per_unit: Optional[float] = None
75 def __post_init__(self) -> None:
76 """Refuse a figure that cannot price a copy."""
77 if self.gib_per_s <= 0 or self.sustained_tflops <= 0 or not 0 < self.overlap <= 1:
78 raise ValueError(f"a host link needs positive figures and an overlap in (0, 1], not {self}")
79 if self.ms_per_unit is not None and self.ms_per_unit <= 0:
80 raise ValueError(f"a calibrated host link needs a positive ms_per_unit, not {self}")
81 if self.copy_cost_ms_per_gib < 0 or (self.score_ms_per_unit is not None and self.score_ms_per_unit <= 0):
82 raise ValueError(f"a host link needs a copy cost of 0 or more and a positive score_ms_per_unit, not {self}")
84 def seconds_per_byte(self) -> float:
85 """The seconds one byte takes over the link, one way."""
86 return 1.0 / (self.gib_per_s * 2 ** 30)
88 def cost_seconds_per_byte(self) -> float:
89 """The step time one byte moved costs, either way, in seconds."""
90 return self.copy_cost_ms_per_gib / 1000.0 / 2 ** 30
92 def flops_per_second(self) -> float:
93 """The device's sustained FLOP/s."""
94 return self.sustained_tflops * 10 ** 12
96 @classmethod
97 def of(cls, link: Optional["HostLink"], figures: Mapping[str, Any]) -> "HostLink":
98 """*link* with the *figures* stated replacing its own, those ``None`` left as they are.
100 Args:
101 link: A device's link, or ``None`` for a device that states none.
102 figures: Any of the fields, ``None`` for one not stated.
104 Raises:
105 ValueError: Without *link*, where *figures* leave out the
106 bandwidth or the throughput.
107 """
108 stated = {name: value for name, value in figures.items() if value is not None}
109 if link is not None:
110 return replace(link, **stated)
111 if "gib_per_s" not in stated or "sustained_tflops" not in stated:
112 raise ValueError("it states no host link; give both --host_link_gibps and --sustained_tflops for -ao")
113 return cls(**stated)
115 def units_per_second(self, bytes_p: float) -> float:
116 """The performance estimate's units a second, calibrated where *ms_per_unit* is stated.
118 Args:
119 bytes_p: The training precision's bytes, by which the estimate
120 multiplies its FLOPs.
121 """
122 if self.ms_per_unit is not None:
123 return 1000.0 / self.ms_per_unit
124 return self.flops_per_second() * bytes_p
126 def score_units_per_second(self, bytes_p: float) -> float:
127 """The score's units a second, for the copies' cost: *score_ms_per_unit*'s, else as a copy converts.
129 Args:
130 bytes_p: The training precision's bytes, by which the estimate
131 multiplies its FLOPs.
132 """
133 if self.score_ms_per_unit is not None:
134 return 1000.0 / self.score_ms_per_unit
135 return self.units_per_second(bytes_p)
138class Type:
139 """Machine type"""
141 name: str
142 levels: int # levels in hierarchy
143 level_bound_number: list[int] # devices per level
144 level_bandwidth: list[int] # bandwidth (GB/s) per level
145 host_link: Optional[HostLink] # the copy path to the host, for offload
147 def __init__(self, name, bounds, bandwidths, host_link=None):
148 self.name = name
149 self.level_bound_number = bounds
150 self.level_bandwidth = bandwidths
151 if len(bounds) != len(bandwidths):
152 raise ValueError("bounds and bandwidths must have the same length")
153 self.levels = len(bounds)
154 self.host_link = host_link
156 def __str__(self):
157 return self.name
159 def __repr__(self):
160 return str(self)
162 def devices_below_level(self, level):
163 """Number of devices below the given hierarchy level"""
164 devices = 1
165 for lvl in range(min(level, self.levels)):
166 devices *= self.level_bound_number[lvl]
167 return devices
169 def intra_node_num(self):
170 """Number of devices in a node"""
171 return self.devices_below_level(1)
173 def levels_used(self, device_number):
174 """Number of hierarchy level used"""
175 devices = 1
176 for lvl in range(self.levels):
177 if self.level_bound_number[lvl]:
178 devices *= self.level_bound_number[lvl]
179 if device_number <= devices:
180 return lvl
181 else:
182 return lvl
183 return self.levels
185 def level_assign(self, dp=1, tp=1, cp=1, pp=1, ep=1):
186 """device assignment of the different parallel dimensions"""
187 # EP borrows from DP, not counted in total devices; kept for
188 # topology tracking only — callers should NOT pass ep > 1
189 # unless they also account for EP-in-DP convention.
190 device_number = dp * tp * cp * pp * ep
191 logger.debug("DP = %d, TP = %d, EP = %d, CP = %d, PP = %d", dp, tp, ep, cp, pp)
192 # Order matters: each dim takes the devices the previous ones left.
193 degrees = {Dim.TP: tp, Dim.EP: ep, Dim.CP: cp, Dim.DP: dp, Dim.PP: pp}
194 assignment = {dim: [] for dim in degrees}
195 for level in range(self.levels):
196 bound = self.level_bound_number[level]
197 if bound:
198 remaining_devices = max(min(device_number, bound), 1)
199 device_number = device_number // bound
200 else:
201 remaining_devices = max(device_number, 1)
202 for dim, degree in degrees.items():
203 dim_level = min(degree, remaining_devices)
204 assignment[dim].append(dim_level)
205 degrees[dim] = degree // dim_level
206 remaining_devices = remaining_devices // dim_level
208 return assignment
211# A3's link is measured: 6.92 GiB/s is the slowest die of a two-node job copying
212# to the host on 7 October 2026, where one node alone measured 10.33 to 13.29
213# (nd_golden/offload/host_link_*_1007.txt). The copies to the host bind, in the
214# forward; those back run in the longer backward. Its cost is measured too:
215# inside the Demo 2 step on 8 October, four nodes, 4 GiB each way at layer 0's
216# points cost the step 22.5 ms, 2.8 ms a GiB moved, where that job's slowest die
217# copied 8.61 GiB/s alone (nd_golden/offload/probe_4node_1008.txt). The other links are
218# placeholders: 16 GiB/s is hyper_offload's own default before it profiles the
219# link. The throughputs are about half of each device's dense BF16 or FP16
220# peak, which a calibration replaces (HostLink.ms_per_unit).
221# Device_A2 = Machine(devices_per_node=8, inter_node_bw=10, intra_node_bw=50)
222Device_A2 = Type(
223 name="A2", bounds=[8, None], bandwidths=[50, 10], host_link=HostLink(gib_per_s=16.0, sustained_tflops=140.0)
224)
225Device_A3 = Type(
226 name="A3",
227 bounds=[16, 24, None],
228 bandwidths=[200, 25, 10],
229 host_link=HostLink(gib_per_s=6.92, sustained_tflops=160.0, copy_cost_ms_per_gib=2.8),
230)
231device_map = {
232 "A2": Device_A2,
233 "A3": Device_A3,
234 "V100": Type(
235 name="V100", bounds=[8, None], bandwidths=[50, 10], host_link=HostLink(gib_per_s=12.0, sustained_tflops=60.0)
236 ),
237}
240class Machine:
241 """Hardware description"""
243 number: int
244 device: Type
246 def __init__(self, number, device):
247 self.number = number
248 if isinstance(device, int):
249 if device == 2:
250 self.device = Device_A2
251 elif device == 3:
252 self.device = Device_A3
253 else:
254 raise ValueError(f"Ascend A{device} unknown")
255 elif isinstance(device, str):
256 if device not in device_map:
257 raise ValueError(
258 f"Device {device} is not supported. "
259 f"Supported devices: {list(device_map.keys())}"
260 )
261 self.device = device_map[device]
262 else:
263 self.device = device
265 def update_num_if_none(self, num):
266 """Assign number of device if not already precised"""
267 if self.number is None:
268 self.number = num
270 def pipeline_bound(self):
271 """Return pipeline bound from hardware topology because as pipeline may currently not cross hierarchy levels"""
272 max_bound = 1
273 devices = self.number
274 while devices > 1:
275 max_bound = max(
276 max_bound,
277 devices
278 // self.device.devices_below_level(
279 self.device.levels_used(devices)
280 ),
281 )
282 devices = devices // 2
283 # devices = self.devices_below_level(self.levels_used(device_number))
284 # return device_number // devices
285 return max_bound
288def prime_factors(n):
289 """Decompose n into a product of prime factors"""
290 divisor = 2
291 factors = []
292 while n > 1:
293 while n % divisor != 0:
294 divisor += 1
295 factors.append(divisor)
296 n = n // divisor
297 return factors
300def all_factors_combinations(factors):
301 """Computes all divisors from a prime factor list"""
302 def rec_factors(n, factors):
303 combinations = {n}
304 for u in set(factors):
305 remaining = factors.copy()
306 remaining.remove(u)
307 combinations = combinations.union(rec_factors(n * u, remaining))
308 return combinations
309 return rec_factors(1, factors)
312def all_divisors(n, reverse=False, min_bound=1, max_bound=float("inf")):
313 """Computes all divisors of an integer n"""
314 divisors = sorted(
315 all_factors_combinations(prime_factors(n)), reverse=reverse
316 )
317 div_in_bound = []
318 for d in divisors:
319 if min_bound <= d <= max_bound:
320 div_in_bound.append(d)
322 return div_in_bound
325def from_prime_factors(factors):
326 """Compute a number from its prime factor decomposition"""
327 number = 1
328 for f in factors:
329 number *= f
330 return number
333def split_node(n, device):
334 """Split decompositions into intra & inter devices"""
335 devices_per_node = device.intra_node_num()
336 nodes = prime_factors(max(1, n // devices_per_node))
337 intra = prime_factors(min(n, devices_per_node))
338 return [intra, nodes]
341def unique_factors(factors):
342 """Remove duplicates. Factors are sorted"""
343 offset = 0
344 for i, f in enumerate(factors[:-1]):
345 j = i - offset
346 if factors[j + 1] == f:
347 factors.pop(j)
348 offset += 1
349 return factors
352def highest_power_of_2_divisor(divisor_of):
353 """Compute the highest number that is both a divisor of 'divisor_of' and a power of 2"""
354 divisor = 1
355 factors = prime_factors(divisor_of)
356 for f in factors:
357 if f == 2:
358 divisor *= f
359 return divisor
362def get_cp_topology(tp_degree: int, cp_degree: int, device_per_node: int) -> tuple:
363 """Determine CP topology and effective bandwidth.
365 Args:
366 tp_degree: Tensor parallelism degree.
367 cp_degree: Context parallelism degree.
368 device_per_node: Number of devices per node.
370 Returns:
371 Tuple of (topology_type, effective_bandwidth, is_intra_node).
372 - topology_type: "intra-node" or "cross-node"
373 - effective_bandwidth: Bandwidth in GB/s
374 - is_intra_node: True if CP stays within node
375 """
376 total_devices_needed = tp_degree * cp_degree
378 if total_devices_needed <= device_per_node:
379 topology_type = "intra-node"
380 is_intra_node = True
381 effective_bandwidth = 300.0
382 else:
383 topology_type = "cross-node"
384 is_intra_node = False
385 effective_bandwidth = 25.0
387 return topology_type, effective_bandwidth, is_intra_node
390def get_cp_bandwidth(topology_type: str, device_type: str = "A2") -> float:
391 """Get effective bandwidth for CP communication based on topology.
393 Args:
394 topology_type: "intra-node" or "cross-node"
395 device_type: Device type string (e.g., "A2", "A3")
397 Returns:
398 Bandwidth in GB/s
399 """
400 device = device_map.get(device_type, Device_A2)
402 if topology_type == "intra-node":
403 return device.level_bandwidth[0] if device.level_bandwidth else 300.0
404 return device.level_bandwidth[1] if len(device.level_bandwidth) > 1 else 25.0
407def recommend_cp_max_by_attention(attention_type: str) -> int:
408 """Recommend maximum CP degree based on attention type.
410 Args:
411 attention_type: "mla", "gqa", or "mha"
413 Returns:
414 Recommended maximum CP degree
415 """
416 attention_type_upper = attention_type.upper()
417 if attention_type_upper == "MLA":
418 return 16
419 if attention_type_upper == "GQA":
420 return 8
421 return 4