Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / trace_classify.py: 98%

1# Copyright 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"""Split profiled training steps into ND's real-profiling parts. 

16 

17``run_nd --real_csv`` (``ParallelizeLayer.compare_with_csv``) compares ND's 

18estimate with a measured run given as a CSV of per-part times (read by 

19``debug.get_comm_classified_data``). This module writes that CSV from 

20``torch.profiler`` Chrome traces, one per rank. 

21 

22Each ``ProfilerStep`` is split, on the thread that ran it, into: 

23 

24- ``<dim>_wait``: time blocked in a ``torch.distributed`` call or a ``Work.wait()``, 

25 typed by the nearest HyperParallel call site above it. FSDP gathers give 

26 ``op_wait``; FSDP reductions, gradient clipping and trainer syncs give 

27 ``dp_wait``; the tensor, sequence, context, expert and pipeline parallel 

28 modules give ``mp``, ``sp``, ``cp``, ``ep`` and ``pp`` waits. 

29- ``comp``: time inside top-level ATen ops that are not communication. 

30- ``idle``: the rest (Python, hooks, data loading). 

31 

32The call sites are Python frames, so traces need ``with_stack=True``. What is 

33measured is host-side blocking: exact for host-synchronous backends such as 

34gloo, while on accelerators a host wait is not the exposed device time. 

35 

36An Ascend run directory is read instead through ``step_trace_time.csv`` and 

37``communication.json``, which need neither Python frames nor host blocking; see 

38``split_ascend_output``. Pass the run directory in place of a trace file. 

39 

40Example: 

41 python -m hyper_parallel.auto_parallel.sapp_nd.nd.trace_classify rank*.pt.trace.json.gz \ 

42 --dims DP=4,PP=1,MB=2,MBS=1,OP=4 --csv real.csv 

43 

44 python -m hyper_parallel.auto_parallel.sapp_nd.nd.trace_classify profiling_dp64_ep16_op2 \ 

45 --dims DP=64,MP=1,PP=1,CP=1,EP=16,MB=1,OP=2 --csv real.csv 

46""" 

47 

48import argparse 

49import csv 

50import gzip 

51import json 

52import re 

53from collections import defaultdict 

54from dataclasses import dataclass, field 

55from pathlib import Path 

56from typing import Dict, FrozenSet, List, Optional, Sequence, Tuple 

57 

58import hyper_parallel.auto_parallel.sapp_nd.nd.debug as Debug 

59import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim 

60from hyper_parallel.auto_parallel.sapp_nd.nd.logger import logger 

61 

62TIME = "time" 

63COMP = str(Debug.RealParts.COMP) 

64DP_WAIT = str(Debug.RealParts.DP_WAIT) 

65MP_WAIT = str(Debug.RealParts.MP_WAIT) 

66EP_WAIT = str(Debug.RealParts.EP_WAIT) 

67CP_WAIT = str(Debug.RealParts.CP_WAIT) 

68PP_WAIT = str(Debug.RealParts.PP_WAIT) 

69# Columns real_in_parts folds into DP_WAIT and MP_WAIT. 

70OP_WAIT = "op_wait" 

71SP_WAIT = "sp_wait" 

72WAIT_COLUMNS = (DP_WAIT, MP_WAIT, EP_WAIT, CP_WAIT, PP_WAIT, OP_WAIT, SP_WAIT) 

73UNCLASSIFIED = "unclassified" 

74# "unsplit" carries device compute that the source cannot attribute to a pass. 

75PHASES = ("fw", "bw", "rec", "opt", "other", "unsplit") 

76 

77ASCEND_OUTPUT = "ASCEND_PROFILER_OUTPUT" 

78_STEP_TRACE_CSV = "step_trace_time.csv" 

79_COMMUNICATION_JSON = "communication.json" 

80_STEP_TRACE_COLUMNS = { 

81 "device": "Device_id", "step": "Step", "comp": "Computing", "free": "Free", "stage": "Stage", 

82 "exposed": "Communication(Not Overlapped)", "comm": "Communication", "bubble": "Bubble", 

83} 

84# hcom_allGather__800_0_1@5862276110395350800 -> kind allGather, group 5862276110395350800. 

85_HCOM_NAME = re.compile(r"^hcom_(?P<kind>[A-Za-z]+)_+\d[\d_]*@(?P<group>\S+)$") 

86_COLLECTIVE_KINDS = { 

87 "allgather": "all_gather", "allgatherv": "all_gather", "reducescatter": "reduce_scatter", 

88 "reducescatterv": "reduce_scatter", "allreduce": "all_reduce", "reduce": "reduce", 

89 "alltoall": "all_to_all", "alltoallv": "all_to_all", "alltoallvc": "all_to_all", 

90 "broadcast": "broadcast", "send": "p2p", "receive": "p2p", "batchsendrecv": "p2p", 

91} 

92# HCCS, PCIE and SIO are reported again inside SDMA, so only these two are summed. 

93_LINK_TOTALS = ("RDMA", "SDMA") 

94 

95_WAIT_PREFIX = "<built-in method wait of " 

96_C10D_API = re.compile(r"torch/distributed/distributed_c10d\.py\(\d+\): \w+$") 

97_C10D_OP = re.compile(r"^(c10d::|_c10d_functional::|c10d_functional::)") 

98_PHASE_PATTERNS = ( 

99 ("bw", re.compile(r"autograd::engine::evaluate_function|torch/autograd/__init__\.py\(\d+\): backward" 

100 r"|torch/_tensor\.py\(\d+\): backward")), 

101 ("rec", re.compile(r"torch/utils/checkpoint\.py|activation_checkpoint")), 

102 ("opt", re.compile(r"Optimizer\.step|torch/optim/|components/optim|clip_grad|zero_grad")), 

103) 

104# (wait column, call-site label, frame pattern): the innermost matching frame decides. 

105# A None column is FSDP, split by collective kind in _fsdp_site. 

106_RULES = ( 

107 (SP_WAIT, "sequence parallel", re.compile(r"sequence_parallel")), 

108 (MP_WAIT, "tensor parallel", re.compile(r"tensor_parallel|tp_collective_lowering")), 

109 (CP_WAIT, "context parallel", re.compile(r"context_parallel")), 

110 (EP_WAIT, "expert parallel", re.compile(r"expert_parallel")), 

111 (PP_WAIT, "pipeline p2p", re.compile(r"pipeline_parallel")), 

112 (DP_WAIT, "grad-norm all-reduce", re.compile(r"clip_grad")), 

113 (None, "fsdp", re.compile(r"core/fully_shard/")), 

114 (DP_WAIT, "trainer all-reduce", re.compile(r"hyper_parallel/trainer/")), 

115) 

116_KIND_HINTS = ( 

117 ("all_gather", re.compile(r"all_gather|allgather|unshard")), 

118 ("reduce_scatter", re.compile(r"reduce_scatter|post_backward|finalize_per_param_reductions")), 

119 ("all_reduce", re.compile(r"all_reduce|allreduce")), 

120 ("all_to_all", re.compile(r"all_to_all|alltoall")), 

121 ("p2p", re.compile(r"isend|irecv|send|recv|p2p")), 

122) 

123_DTYPE_BYTES = { 

124 "float": 4, "double": 8, "c10::BFloat16": 2, "c10::Half": 2, "long int": 8, 

125 "int": 4, "short int": 2, "bool": 1, "unsigned char": 1, "signed char": 1, 

126} 

127 

128 

129@dataclass 

130class SiteStats: 

131 """Communication attributed to one call site within one step. 

132 

133 Attributes: 

134 ms: Blocked time, in milliseconds. 

135 calls: ``torch.distributed`` calls issued from the site. 

136 waits: ``Work.wait()`` calls made outside those calls. 

137 payload_bytes: Logical tensor bytes of the collectives, largest argument each. 

138 transit_ms: Device time moving data; Ascend only, 0 from a Chrome trace. 

139 wait_ms: Device time waiting for peers; Ascend only, 0 from a Chrome trace. 

140 """ 

141 

142 ms: float = 0.0 

143 calls: int = 0 

144 waits: int = 0 

145 payload_bytes: int = 0 

146 transit_ms: float = 0.0 

147 wait_ms: float = 0.0 

148 

149 

150@dataclass 

151class StepSplit: 

152 """One rank's ``ProfilerStep`` split into ND parts, in milliseconds. 

153 

154 Attributes: 

155 rank: Rank read from the trace file name, -1 when absent. 

156 step: Profiler step name. 

157 time: Step duration. 

158 comp: Compute time per phase of ``PHASES``. 

159 waits: Blocked time per column of ``WAIT_COLUMNS``, plus ``UNCLASSIFIED``. 

160 sites: Per call-site detail, keyed by ``(column, label)``. 

161 """ 

162 

163 rank: int 

164 step: str 

165 time: float 

166 comp: Dict[str, float] 

167 waits: Dict[str, float] 

168 sites: Dict[Tuple[str, str], SiteStats] = field(default_factory=dict) 

169 

170 @property 

171 def idle(self) -> float: 

172 """Step time that is neither compute nor communication.""" 

173 return self.time - sum(self.comp.values()) - sum(self.waits.values()) 

174 

175 

176@dataclass(frozen=True) 

177class _Frame: 

178 """An open event and what its children inherit.""" 

179 

180 end: float 

181 name: str 

182 comm: Optional[Tuple[str, str]] 

183 under_op: bool 

184 phases: FrozenSet[str] 

185 

186 

187_ROOT = _Frame(float("inf"), "", None, False, frozenset()) 

188 

189 

190def load_trace_events(path: str) -> List[dict]: 

191 """Return the complete (``"X"``) events of a Chrome trace. 

192 

193 Args: 

194 path: ``.json`` or gzipped ``.json.gz`` trace written by ``torch.profiler``. 

195 

196 Returns: 

197 The trace events whose phase is ``"X"``. 

198 """ 

199 opener = gzip.open if str(path).endswith(".gz") else open 

200 with opener(path, "rt", encoding="utf-8") as handle: 

201 return [event for event in json.load(handle)["traceEvents"] if event.get("ph") == "X"] 

202 

203 

204def classify_call(stack_names: Sequence[str]) -> Tuple[str, str]: 

205 """Return the wait column and call-site label of a communication call. 

206 

207 Args: 

208 stack_names: Event names from the outermost open frame down to the call itself. 

209 

210 Returns: 

211 ``(column, label)`` decided by the innermost frame matching a rule, or 

212 ``(UNCLASSIFIED, UNCLASSIFIED)`` when no frame matches. 

213 """ 

214 for name in reversed(stack_names): 

215 for column, label, pattern in _RULES: 

216 if pattern.search(name): 

217 return _fsdp_site(stack_names) if column is None else (column, label) 

218 return UNCLASSIFIED, UNCLASSIFIED 

219 

220 

221def _collective_kind(stack_names: Sequence[str]) -> str: 

222 """Collective kind named by the innermost frame that names one.""" 

223 for name in reversed(stack_names): 

224 for kind, pattern in _KIND_HINTS: 

225 if pattern.search(name): 

226 return kind 

227 return "unknown" 

228 

229 

230def _fsdp_site(stack_names: Sequence[str]) -> Tuple[str, str]: 

231 """FSDP parameter gathers are ND's optimizer-parallel waits, reductions its DP waits.""" 

232 kind = _collective_kind(stack_names) 

233 if kind == "all_gather": 

234 return OP_WAIT, "fsdp param all-gather" 

235 return DP_WAIT, f"fsdp grad {kind.replace('_', '-')}" 

236 

237 

238def _is_comm_root(name: str, cat: Optional[str]) -> bool: 

239 """Whether an event blocks on communication: a c10d API frame, a wait, or a c10d op.""" 

240 if cat == "python_function": 

241 return name.startswith(_WAIT_PREFIX) or _C10D_API.search(name) is not None 

242 return cat == "cpu_op" and _C10D_OP.match(name) is not None 

243 

244 

245def _shape_numel(dims: list) -> int: 

246 """Element count of a shape, or the sum over a tensor list's shapes.""" 

247 if dims and all(isinstance(dim, list) for dim in dims): 

248 return sum(_shape_numel(dim) for dim in dims) 

249 if not dims or not all(isinstance(dim, int) for dim in dims): 

250 return 0 

251 numel = 1 

252 for dim in dims: 

253 numel *= dim 

254 return numel 

255 

256 

257def _payload_bytes(event: dict) -> int: 

258 """Bytes of the largest tensor argument of a c10d op.""" 

259 args = event.get("args", {}) 

260 sizes = [_shape_numel(dims) * _DTYPE_BYTES.get(dtype, 4) 

261 for dims, dtype in zip(args.get("Input Dims", []), args.get("Input type", [])) if dims] 

262 return max(sizes, default=0) 

263 

264 

265def _merge(intervals: Sequence[Tuple[float, float]]) -> List[List[float]]: 

266 """Sorted, disjoint union of ``[start, end)`` intervals.""" 

267 merged: List[List[float]] = [] 

268 for start, stop in sorted(intervals): 

269 if merged and start <= merged[-1][1]: 

270 merged[-1][1] = max(merged[-1][1], stop) 

271 else: 

272 merged.append([start, stop]) 

273 return merged 

274 

275 

276def _length_outside(intervals: Sequence[Tuple[float, float]], holes: Sequence[Tuple[float, float]]) -> float: 

277 """Length of the union of ``intervals`` once the union of ``holes`` is removed.""" 

278 kept, cut = _merge(intervals), _merge(holes) 

279 overlap, first = 0.0, 0 

280 for start, stop in kept: 

281 while first < len(cut) and cut[first][1] <= start: 

282 first += 1 

283 index = first 

284 while index < len(cut) and cut[index][0] < stop: 

285 overlap += min(stop, cut[index][1]) - max(start, cut[index][0]) 

286 index += 1 

287 return sum(stop - start for start, stop in kept) - overlap 

288 

289 

290class _StepSplitter: 

291 """Walks one thread's events in time order, keeping the open call stack.""" 

292 

293 def __init__(self, stop: float) -> None: 

294 """Start an empty walk of a step ending at ``stop`` (trace microseconds).""" 

295 self.stop = stop 

296 self.stack: List[_Frame] = [] 

297 self.comp: Dict[str, List[Tuple[float, float]]] = defaultdict(list) 

298 self.waits: Dict[Tuple[str, str], List[Tuple[float, float]]] = defaultdict(list) 

299 self.sites: Dict[Tuple[str, str], SiteStats] = defaultdict(SiteStats) 

300 

301 def visit(self, event: dict) -> None: 

302 """Account one event, given in ``(ts, -dur)`` order.""" 

303 while self.stack and self.stack[-1].end <= event["ts"]: 

304 self.stack.pop() 

305 parent = self.stack[-1] if self.stack else _ROOT 

306 name, cat = event["name"], event.get("cat") 

307 phases = set(parent.phases) | {phase for phase, pattern in _PHASE_PATTERNS if pattern.search(name)} 

308 if name.startswith("nn.Module: "): 

309 phases.add("fw") 

310 interval = (event["ts"], min(event["ts"] + event["dur"], self.stop)) 

311 comm = parent.comm 

312 if comm is None and _is_comm_root(name, cat): 

313 comm = classify_call([frame.name for frame in self.stack] + [name]) 

314 self.waits[comm].append(interval) 

315 if name.startswith(_WAIT_PREFIX): 

316 self.sites[comm].waits += 1 

317 else: 

318 self.sites[comm].calls += 1 

319 elif comm is None and cat == "cpu_op" and not parent.under_op: 

320 self.comp[self._phase(phases)].append(interval) 

321 if comm is not None and cat == "cpu_op" and _C10D_OP.match(name): 

322 self.sites[comm].payload_bytes += _payload_bytes(event) 

323 self.stack.append(_Frame(event["ts"] + event["dur"], name, comm, parent.under_op or cat == "cpu_op", 

324 frozenset(phases))) 

325 

326 @staticmethod 

327 def _phase(phases: set) -> str: 

328 """Compute phase of a top-level op from the phases of its frames.""" 

329 if "bw" in phases: 

330 return "rec" if "rec" in phases else "bw" 

331 return next((phase for phase in ("opt", "fw") if phase in phases), "other") 

332 

333 def split(self, rank: int, step: dict) -> StepSplit: 

334 """Close the step; a hook blocking inside an autograd node counts as wait only.""" 

335 holes = [interval for intervals in self.waits.values() for interval in intervals] 

336 comp = {phase: _length_outside(self.comp[phase], holes) / 1e3 for phase in PHASES} 

337 waits = dict.fromkeys(WAIT_COLUMNS + (UNCLASSIFIED,), 0.0) 

338 for key, intervals in self.waits.items(): 

339 self.sites[key].ms = sum(stop - start for start, stop in _merge(intervals)) / 1e3 

340 waits[key[0]] += self.sites[key].ms 

341 return StepSplit(rank, step["name"], step["dur"] / 1e3, comp, waits, dict(self.sites)) 

342 

343 

344def split_trace(path: str) -> List[StepSplit]: 

345 """Split every ``ProfilerStep`` of one rank's trace. 

346 

347 Args: 

348 path: Trace file; its rank is read from a ``rank<N>`` part of the file name. 

349 

350 Returns: 

351 One ``StepSplit`` per profiler step, in time order. 

352 """ 

353 events = load_trace_events(path) 

354 match = re.search(r"rank(\d+)", Path(path).name) 

355 rank = int(match.group(1)) if match else -1 

356 splits = [] 

357 for step in sorted((e for e in events if e["name"].startswith("ProfilerStep#")), key=lambda e: e["ts"]): 

358 start, stop = step["ts"], step["ts"] + step["dur"] 

359 splitter = _StepSplitter(stop) 

360 tid = step.get("tid") 

361 thread = (e for e in events if e.get("tid") == tid and e is not step and start <= e["ts"] < stop) 

362 for event in sorted(thread, key=lambda e: (e["ts"], -e["dur"])): 

363 splitter.visit(event) 

364 splits.append(splitter.split(rank, step)) 

365 return splits 

366 

367 

368def ascend_output_dir(path: str) -> Path: 

369 """Locate the ``ASCEND_PROFILER_OUTPUT`` directory of a profiling run. 

370 

371 Args: 

372 path: That directory, the ``*_ascend_pt`` directory above it, or the 

373 run directory above that. 

374 

375 Returns: 

376 The directory holding ``step_trace_time.csv``. 

377 

378 Raises: 

379 ValueError: If no such directory is found under ``path``, or if the run 

380 directory holds more than one. Re-profiling into a directory that 

381 already holds a run leaves both, and silently reading either one 

382 makes a stale profile look like a fresh one. 

383 """ 

384 base = Path(path) 

385 for candidate in (base, base / ASCEND_OUTPUT): 

386 if (candidate / _STEP_TRACE_CSV).is_file(): 

387 return candidate 

388 nested = [c for c in sorted(base.glob(f"*/{ASCEND_OUTPUT}")) 

389 if (c / _STEP_TRACE_CSV).is_file()] 

390 if len(nested) > 1: 

391 listed = "".join(f"\n {c.parent.name}" for c in nested) 

392 raise ValueError( 

393 f"{path} holds {len(nested)} profiling runs; pass the one to read:{listed}" 

394 ) 

395 if nested: 

396 return nested[0] 

397 raise ValueError(f"no {ASCEND_OUTPUT}/{_STEP_TRACE_CSV} under {path}") 

398 

399 

400def is_ascend_output(path: str) -> bool: 

401 """Whether ``path`` points at an Ascend profiling run rather than a Chrome trace.""" 

402 try: 

403 ascend_output_dir(path) 

404 except (ValueError, OSError): 

405 return False 

406 return True 

407 

408 

409def collective_kind(name: str) -> Tuple[str, str]: 

410 """Split an HCCL operator name into its collective kind and communication group. 

411 

412 Args: 

413 name: Operator name such as ``hcom_allGather__800_0_1@5862276110395350800``. 

414 

415 Returns: 

416 ``(kind, group)``; ``("unknown", "")`` when the name does not parse. 

417 """ 

418 match = _HCOM_NAME.match(name) 

419 if match is None: 

420 return "unknown", "" 

421 kind = match.group("kind").lower() 

422 return _COLLECTIVE_KINDS.get(kind, kind), match.group("group") 

423 

424 

425def ascend_site(kind: str, section: str, dims: Dict[str, str]) -> Tuple[str, str]: 

426 """Return the wait column and label of an HCCL operator. 

427 

428 Ascend traces carry no Python frames, so the axis comes from the collective 

429 kind: expert parallelism dispatches with all-to-all, FSDP gathers and reduces 

430 parameters, and the remaining all-reduces are data-parallel. Tensor and 

431 context parallelism also gather and reduce-scatter, so when either is active 

432 those operators need the rank sets of ``communication_matrix.json`` to be 

433 told apart and stay unclassified here. 

434 

435 Context parallelism sends point to point, as the Qwen3.5 Gated DeltaNet 

436 wrapper does, and exchanges all-to-all under Ulysses. Point to point is 

437 counted as context parallelism when there is a single pipeline stage, and 

438 all-to-all when there is no expert parallelism: the degrees leave it no 

439 other owner. Otherwise they keep the pipeline and expert columns. 

440 

441 Args: 

442 kind: Collective kind from ``collective_kind``. 

443 section: ``"collective"`` or ``"p2p"``, the section of ``communication.json``. 

444 dims: ND dimensions of the run, from ``parse_dims``. 

445 

446 Returns: 

447 ``(column, label)``. 

448 """ 

449 context = int(dims.get("CP", 1)) > 1 

450 if section == "p2p" or kind == "p2p": 

451 if context and int(dims.get("PP", 1)) <= 1: 

452 return CP_WAIT, "context p2p" 

453 return PP_WAIT, "pipeline p2p" 

454 if kind == "all_to_all": 

455 if context and int(dims.get("EP", 1)) <= 1: 

456 return CP_WAIT, "context all-to-all" 

457 return EP_WAIT, "expert all-to-all" 

458 if kind in ("all_gather", "reduce_scatter"): 

459 if int(dims.get("MP", 1)) > 1 or int(dims.get("CP", 1)) > 1: 

460 return UNCLASSIFIED, f"{kind} (MP or CP active)" 

461 return OP_WAIT, f"fsdp {kind.replace('_', '-')}" 

462 if kind == "all_reduce": 

463 return DP_WAIT, "all-reduce" 

464 return UNCLASSIFIED, kind 

465 

466 

467def _transit_bytes(bandwidth_info: dict) -> int: 

468 """Bytes moved by one collective, over the link types that do not double count.""" 

469 total = 0.0 

470 for link in _LINK_TOTALS: 

471 total += float(bandwidth_info.get(link, {}).get("Transit Size(MB)", 0.0)) 

472 return int(total * 2 ** 20) 

473 

474 

475def _step_collectives(step_comm: dict, dims: Dict[str, str]) -> Tuple[Dict[Tuple[str, str], SiteStats], float, set]: 

476 """Aggregate one step's HCCL operators per wait column, with their total elapse. 

477 

478 Every operator name starts with ``hcom_``. Anything else is one of the 

479 summary entries CANN writes beside them, such as ``Total Op Info``, which 

480 repeats the whole step and would double its communication; those are skipped 

481 and returned so the caller can say what it ignored. 

482 """ 

483 sites: Dict[Tuple[str, str], SiteStats] = defaultdict(SiteStats) 

484 total = 0.0 

485 skipped = set() 

486 for section, operators in step_comm.items(): 

487 if not isinstance(operators, dict): 

488 continue 

489 for name, entry in operators.items(): 

490 if not name.startswith("hcom_"): 

491 skipped.add(name) 

492 continue 

493 info = entry.get("Communication Time Info", {}) 

494 elapse = float(info.get("Elapse Time(ms)", 0.0)) 

495 stats = sites[ascend_site(collective_kind(name)[0], section, dims)] 

496 stats.ms += elapse 

497 stats.calls += 1 

498 stats.wait_ms += float(info.get("Wait Time(ms)", 0.0)) 

499 stats.transit_ms += float(info.get("Transit Time(ms)", 0.0)) 

500 stats.payload_bytes += _transit_bytes(entry.get("Communication Bandwidth Info", {})) 

501 total += elapse 

502 return sites, total, skipped 

503 

504 

505def _read_step_trace(path: Path) -> list: 

506 """Read ``step_trace_time.csv`` into one dict of floats per step. 

507 

508 Raises: 

509 ValueError: If a column the split needs is missing. 

510 """ 

511 rows = [] 

512 with open(path, newline="", encoding="utf-8") as handle: 

513 for row in csv.DictReader(handle): 

514 missing = [column for column in _STEP_TRACE_COLUMNS.values() if column not in row] 

515 if missing: 

516 raise ValueError(f"{path}: missing column(s) {', '.join(missing)}") 

517 rows.append({key: float(row[column]) for key, column in _STEP_TRACE_COLUMNS.items()}) 

518 return rows 

519 

520 

521def split_ascend_output(path: str, dims: Dict[str, str]) -> List[StepSplit]: 

522 """Split every step of an Ascend profiling run into ND parts. 

523 

524 ``step_trace_time.csv`` gives the closed top-level split of each step, in 

525 microseconds: computing, communication that was not overlapped with it, and 

526 free device time. ``communication.json`` then apportions the exposed 

527 communication over the parallelism axes in proportion to each axis's share of 

528 total HCCL elapse time, which assumes overlap is spread evenly across axes. 

529 Device compute is not attributed to a pass, so it lands in the ``unsplit`` 

530 phase. The pipeline bubble is added to ``pp_wait`` on top of any p2p time; 

531 with ``PP=1`` both are zero, and that path is untested against real data. 

532 

533 Args: 

534 path: The run directory, its ``*_ascend_pt`` directory, or the 

535 ``ASCEND_PROFILER_OUTPUT`` directory itself. 

536 dims: ND dimensions of the run, from ``parse_dims``. 

537 

538 Returns: 

539 One ``StepSplit`` per profiled step, in file order. 

540 """ 

541 base = ascend_output_dir(path) 

542 communication = {} 

543 comm_path = base / _COMMUNICATION_JSON 

544 if comm_path.is_file(): 

545 raw = json.loads(comm_path.read_text(encoding="utf-8")) 

546 communication = {int(re.sub(r"\D", "", key) or 0): value for key, value in raw.items()} 

547 splits = [] 

548 ignored = set() 

549 for row in _read_step_trace(base / _STEP_TRACE_CSV): 

550 step_number = int(row["step"]) 

551 sites, elapse, skipped = _step_collectives(communication.get(step_number, {}), dims) 

552 ignored |= skipped 

553 exposed, bubble = row["exposed"] / 1e3, row["bubble"] / 1e3 

554 # communication.json times every collective; the CSV says how much stayed exposed. 

555 # Concurrent operators overlap, so their elapse sums past the step: rescale to what 

556 # stayed exposed, and keep the wait and transit shares on that same basis. 

557 scale = exposed / elapse if elapse else 0.0 

558 waits = dict.fromkeys(WAIT_COLUMNS + (UNCLASSIFIED,), 0.0) 

559 for key, stats in sites.items(): 

560 stats.ms *= scale 

561 stats.wait_ms *= scale 

562 stats.transit_ms *= scale 

563 waits[key[0]] += stats.ms 

564 waits[PP_WAIT] += bubble 

565 comp = dict.fromkeys(PHASES, 0.0) 

566 comp["unsplit"] = row["comp"] / 1e3 

567 splits.append(StepSplit(int(row["device"]), f"step{step_number}", row["stage"] / 1e3, 

568 comp, waits, dict(sites))) 

569 if ignored: 

570 logger.info("%s: ignored %d summary entries: %s", base, len(ignored), ", ".join(sorted(ignored))) 

571 return splits 

572 

573 

574def mean_parts(splits: Sequence[StepSplit]) -> Dict[str, float]: 

575 """Mean over steps and ranks of every part, in milliseconds. 

576 

577 Args: 

578 splits: Non-empty sequence of step splits. 

579 

580 Returns: 

581 ``time``, ``comp``, ``comp_<phase>``, the wait columns, ``unclassified`` and ``idle``. 

582 

583 Raises: 

584 ValueError: If ``splits`` is empty. 

585 """ 

586 if not splits: 

587 raise ValueError("no ProfilerStep events to average") 

588 count = len(splits) 

589 mean = {TIME: sum(s.time for s in splits) / count, "idle": sum(s.idle for s in splits) / count} 

590 for phase in PHASES: 

591 mean[f"comp_{phase}"] = sum(s.comp[phase] for s in splits) / count 

592 mean[COMP] = sum(mean[f"comp_{phase}"] for phase in PHASES) 

593 for column in WAIT_COLUMNS + (UNCLASSIFIED,): 

594 mean[column] = sum(s.waits[column] for s in splits) / count 

595 return mean 

596 

597 

598def perf_parts(mean: Dict[str, float]) -> Dict[str, float]: 

599 """The measured step in the columns of ND's estimate (``debug.csv``). 

600 

601 Real data cannot split the pipeline bubble from P2P time, so all of ``pp_wait`` 

602 goes to ``PP_COMM``. FSDP's waits, ``op_wait``, go to ``DP_COMM`` and the 

603 all-reduces, ``dp_wait``, to ``DP_REDUCE``, as an Ascend profile splits them; 

604 a Chrome trace counts FSDP's reduce-scatters among the all-reduces. The 

605 columns after ``TOTAL`` hold what ND does not model, plus the device compute 

606 an Ascend run cannot attribute to a pass. 

607 

608 Args: 

609 mean: Output of ``mean_parts``. 

610 

611 Returns: 

612 Milliseconds per ``PerfParts`` name up to ``TOTAL``, then the unmodelled rest. 

613 """ 

614 parts = Debug.PerfParts 

615 return { 

616 str(parts.FW_COMPUTE): mean["comp_fw"], 

617 str(parts.BW_COMPUTE): mean["comp_bw"], 

618 str(parts.RECOMPUTE): mean["comp_rec"], 

619 str(parts.DP_COMM): mean[OP_WAIT], 

620 str(parts.DP_REDUCE): mean[DP_WAIT], 

621 str(parts.MP_COMM): mean[MP_WAIT] + mean[SP_WAIT], 

622 str(parts.EP_COMM): mean[EP_WAIT], 

623 str(parts.CP_COMM): mean[CP_WAIT], 

624 str(parts.PP_COMM): mean[PP_WAIT], 

625 str(parts.BUBBLE): 0.0, 

626 str(parts.TOTAL): mean[TIME], 

627 "OPTIMIZER": mean["comp_opt"], 

628 "OTHER_COMPUTE": mean["comp_other"], 

629 "UNSPLIT_COMPUTE": mean["comp_unsplit"], 

630 "UNCLASSIFIED_COMM": mean[UNCLASSIFIED], 

631 "IDLE": mean["idle"], 

632 } 

633 

634 

635def parse_dims(text: str) -> Dict[str, str]: 

636 """Parse ``DP=4,MB=2`` into ND dimension columns. 

637 

638 Args: 

639 text: Comma-separated ``NAME=VALUE`` pairs using ``dimensions`` acronyms. 

640 

641 Returns: 

642 Column name to value, in the given order. 

643 

644 Raises: 

645 ValueError: If a pair is malformed, a name is unknown or a value does not parse. 

646 ``SP`` is refused: ND reads it with ``bool()``, so any written value is True. 

647 """ 

648 dims = {} 

649 for item in (part.strip() for part in text.split(",")): 

650 if not item: 

651 continue 

652 name, sep, value = item.partition("=") 

653 if not sep: 

654 raise ValueError(f"dimension {item!r} must be written NAME=VALUE") 

655 dim = Dim.get_dim(name.strip()) 

656 if dim is Dim.SP: 

657 raise ValueError("SP cannot be written: ND parses it with bool(), so any value reads as True") 

658 dims[dim.name] = str(dim.from_str(value.strip())) 

659 return dims 

660 

661 

662def write_real_csv(splits: Sequence[StepSplit], dims: Dict[str, str], path: str) -> None: 

663 """Write the mean step in the format ``debug.get_comm_classified_data`` reads. 

664 

665 Args: 

666 splits: Step splits to average. 

667 dims: ND dimension columns of the run, from ``parse_dims``. 

668 path: Output CSV path. 

669 """ 

670 mean = mean_parts(splits) 

671 columns = [TIME, COMP, *WAIT_COLUMNS] 

672 with open(path, "w", newline="", encoding="utf-8") as handle: 

673 writer = csv.writer(handle) 

674 writer.writerow([*dims, *columns]) 

675 writer.writerow([*dims.values(), *(f"{mean[column]:.3f}" for column in columns)]) 

676 

677 

678def write_perf_parts_csv(splits: Sequence[StepSplit], dims: Dict[str, str], path: str) -> None: 

679 """Write the mean step in the columns of ND's ``debug.csv``, see ``perf_parts``. 

680 

681 Args: 

682 splits: Step splits to average. 

683 dims: ND dimension columns of the run, from ``parse_dims``. 

684 path: Output CSV path. 

685 """ 

686 parts = perf_parts(mean_parts(splits)) 

687 with open(path, "w", newline="", encoding="utf-8") as handle: 

688 writer = csv.writer(handle) 

689 writer.writerow([*dims, *parts]) 

690 writer.writerow([*dims.values(), *(f"{value:.3f}" for value in parts.values())]) 

691 

692 

693def write_detail_csv(splits: Sequence[StepSplit], path: str) -> None: 

694 """Write every rank and step with its per call-site split. 

695 

696 Args: 

697 splits: Step splits to write. 

698 path: Output CSV path. 

699 """ 

700 sites = sorted({site for split in splits for site in split.sites}) 

701 site_columns = [f"{column}:{label} {unit}" for column, label in sites for unit in ("ms", "calls", "waits", "MiB")] 

702 with open(path, "w", newline="", encoding="utf-8") as handle: 

703 writer = csv.writer(handle) 

704 writer.writerow(["rank", "step", TIME, *(f"comp_{p}" for p in PHASES), *WAIT_COLUMNS, UNCLASSIFIED, "idle", 

705 *site_columns]) 

706 for split in splits: 

707 row = [split.rank, split.step, f"{split.time:.3f}", *(f"{split.comp[p]:.3f}" for p in PHASES), 

708 *(f"{split.waits[c]:.3f}" for c in WAIT_COLUMNS + (UNCLASSIFIED,)), f"{split.idle:.3f}"] 

709 for site in sites: 

710 stats = split.sites.get(site, SiteStats()) 

711 row += [f"{stats.ms:.3f}", stats.calls, stats.waits, f"{stats.payload_bytes / 2**20:.1f}"] 

712 writer.writerow(row) 

713 

714 

715def format_report(splits: Sequence[StepSplit]) -> str: 

716 """Human-readable summary: ND real parts, ND estimate columns, then call sites. 

717 

718 Args: 

719 splits: Step splits to summarise. 

720 

721 Returns: 

722 The report, one line per part or call site, in milliseconds per step. 

723 """ 

724 mean = mean_parts(splits) 

725 total = mean[TIME] 

726 folded = [ 

727 (COMP, mean[COMP], ""), 

728 (DP_WAIT, mean[DP_WAIT] + mean[OP_WAIT], f"dp {mean[DP_WAIT]:.1f} + op {mean[OP_WAIT]:.1f}"), 

729 (MP_WAIT, mean[MP_WAIT] + mean[SP_WAIT], f"mp {mean[MP_WAIT]:.1f} + sp {mean[SP_WAIT]:.1f}"),