Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / debug / _comm_debug_mode.py: 61%
251 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""CommDebugMode — context manager that traces DTensor ops and collectives.
17Public API mirrors ``torch.distributed.tensor.debug.CommDebugMode``:
19 get_comm_counts()
20 get_total_counts()
21 get_parameter_info()
22 get_sharding_info()
23 generate_comm_debug_tracing_table(noise_level)
24 log_comm_debug_tracing_table_to_file(file_name, noise_level)
26Record retention is bounded on purpose. A traced step appends one node per
27DTensor op, so an unbounded tree would grow with the step count and pin host
28memory for the whole process lifetime. Use :meth:`CommDebugMode.clear` to
29release the collected records as soon as they have been consumed.
30"""
31# pylint: disable=C9006,C9007
32import json
33import logging
34import re
35from collections import defaultdict
36from typing import Any, Dict, List, Optional
38import torch
40from hyper_parallel.core.dtensor.debug._call_records import (
41 AnnotateCall,
42 CollectiveCall,
43 DebugCall,
44 OpCall,
45 TensorInfo,
46)
47from hyper_parallel.core.dtensor.debug._collective_tracer import CollectiveTracer
48from hyper_parallel.core.dtensor.debug._module_tracker import ModuleTracker
50logger = logging.getLogger(__name__)
51Tensor = torch.Tensor
53# Argument index of the process group for each traced collective method.
54# Derived from the ``_utils`` function signatures:
55# differentiable_all_gather_concat(data, group, concat_size, concat_dim, ...)
56# differentiable_all_to_all(input_data, output_shape, group)
57# differentiable_all_reduce(data, op, group)
58# differentiable_reduce_scatter(data, dev_num, axis, op, group)
59# differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group)
60# differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group)
61_COLLECTIVE_GROUP_ARG_INDEX: Dict[str, int] = {
62 "differentiable_all_gather_concat": 1,
63 "differentiable_all_to_all": 2,
64 "differentiable_all_reduce": 2,
65 "differentiable_reduce_scatter": 4,
66 "differentiable_all_to_all_single": 3,
67 "differentiable_all_to_all_single_async": 3,
68}
70# Hard ceiling on the number of live record nodes. Tracing an entire training
71# run is not a supported use case — the resulting table is unreadable and the
72# nodes cost host memory — so past this point records are counted and dropped
73# instead of retained. The counters (comm_counts, dropped count) stay exact.
74_MAX_RECORDS = 200000
76# Deepest call nesting that still gets a node. Deeper calls are counted but not
77# retained, which also stops recursion in the rendering helpers from growing
78# with the traced program's call depth.
79_MAX_DEPTH = 128
82class CommDebugMode:
83 """Context manager that records DTensor operator dispatches and collective
84 communication operations, producing a hierarchical call tree.
86 Usage::
88 with CommDebugMode() as mode:
89 output = model(input_dtensor)
90 print(mode.generate_comm_debug_tracing_table())
91 print(mode.get_comm_counts())
93 Long-running training:
94 Tracing keeps one record node per dispatched op, so a mode object that
95 spans a whole training loop retains every op of every step. Trace a
96 bounded slice instead, read the results, then release them::
98 mode = CommDebugMode()
99 for step in range(total_steps):
100 # Re-entering resets the previous window, so this alone stays
101 # bounded — but only if each window is read before the next.
102 if step % report_every == 0:
103 with mode:
104 loss = train_step()
105 print(f"step {step}: {mode.get_comm_counts()}")
106 mode.clear() # release this window's records now
107 else:
108 loss = train_step()
110 ``clear()`` is what actually releases the memory; without it the last
111 window stays resident until the next ``__enter__``. Entries returned by
112 ``get_parameter_info()`` alias the model's parameter storages, so they
113 pin the model in memory for as long as they are held — ``clear()``
114 drops those too. If a window is left traced too long, records are
115 dropped once the budget is exhausted; ``get_dropped_record_count()``
116 reports how many, and both the counters and a ``logger.warning`` stay
117 accurate, but the rendered table will be incomplete.
119 Args:
120 module: Optional ``nn.Module`` to track forward enter/exit events.
121 """
123 def __init__(self, module=None):
124 self._module = module
126 # ---- tracing state ----
127 self._call_stack: List[DebugCall] = []
128 self._root_records: List[DebugCall] = []
129 self._record_count = 0
130 self._dropped_records = 0
131 self._comm_counts: Dict[str, int] = defaultdict(int)
132 # ---- module-level info (populated when module is provided) ----
133 self._parameter_info: Dict[str, Dict[str, Any]] = {}
134 self._sharding_info: Dict[str, Any] = {}
136 # ---- internal handles ----
137 self._collective_tracer: Optional[CollectiveTracer] = None
138 self._module_tracker: Optional[ModuleTracker] = None
139 # Tokens are pushed in ``__enter__`` and popped LIFO in ``__exit__``, so
140 # out-of-order exits restore the observer that was actually active
141 # before this instance took over rather than a stale token.
142 self._observer_tokens: List = []
143 self._active = False
145 # ------------------------------------------------------------------
146 # Context manager protocol
147 # ------------------------------------------------------------------
149 def __enter__(self):
150 # pylint: disable=C0415
151 from hyper_parallel.core.shard._op_dispatch import _debug_mode_observer
153 # Release anything left over from a previous window before collecting
154 # into a fresh tree.
155 self._reset_records()
157 self._observer_tokens.append(_debug_mode_observer.set(self))
158 self._active = True
160 self._collective_tracer = CollectiveTracer(self._on_collective_call)
161 self._collective_tracer.install()
163 if self._module is not None:
164 self._module_tracker = ModuleTracker(self._module, self._on_module_event)
165 self._module_tracker.install()
166 self._collect_module_info()
168 return self
170 def __exit__(self, exc_type, exc_val, exc_tb):
171 # pylint: disable=C0415
172 from hyper_parallel.core.shard._op_dispatch import _debug_mode_observer
174 # Leaving the window stops collection. Anything that reads results
175 # (get_comm_counts, generate_comm_debug_tracing_table, ...) keeps
176 # working afterwards, so the records themselves are not dropped here —
177 # call clear() once they have been consumed.
178 self._active = False
180 if self._module_tracker is not None:
181 self._module_tracker.uninstall()
182 self._module_tracker = None
184 if self._collective_tracer is not None:
185 self._collective_tracer.uninstall()
186 self._collective_tracer = None
188 if self._observer_tokens:
189 _debug_mode_observer.reset(self._observer_tokens.pop())
191 # The stack is expected to be empty here. Unwinding it matters when a
192 # traced call raised: the matching exit callback never ran, so the
193 # leftovers would otherwise swallow every record of the next window.
194 self._call_stack.clear()
196 return False
198 def clear(self):
199 """Release all collected records and counters.
201 Safe to call at any point, including while the context is active.
202 The observer/hook installation is left untouched — only the data.
203 """
204 self._reset_records()
206 def _reset_records(self):
207 """Drop records, counters and module info."""
208 self._call_stack.clear()
209 self._root_records.clear()
210 self._record_count = 0
211 self._dropped_records = 0
212 self._comm_counts.clear()
213 self._parameter_info.clear()
214 self._sharding_info.clear()
216 def __repr__(self):
217 return f"CommDebugMode(get_total_counts()={self.get_total_counts()})"
219 # ------------------------------------------------------------------
220 # Observer callbacks (called from _op_dispatch.py — internal)
221 # ------------------------------------------------------------------
223 def _on_op_dispatch_enter(self, op_name: str, op_call, args, kwargs): # pylint: disable=W0613
224 """Called by OpDispatcher.dispatch() before the op executes."""
225 if not self._active:
226 return
228 record = self._make_record(
229 OpCall,
230 op_name=op_name,
231 input_infos=self._extract_tensor_infos(args),
232 )
233 if record is None:
234 return
236 self._attach(record)
238 def _on_op_dispatch_exit(self, op_name, result): # pylint: disable=W0613
239 """Called by OpDispatcher.dispatch() after the op executes."""
240 if not self._call_stack:
241 return
243 # Peek before popping: a nested scope that never saw an enter must not
244 # pop someone else's frame off the stack.
245 if not isinstance(self._call_stack[-1], OpCall):
246 return
248 record = self._call_stack.pop()
249 record.output_infos = self._extract_tensor_infos((result,))
251 # Keep old names as aliases for backward compatibility with tests.
252 on_op_dispatch_enter = _on_op_dispatch_enter
253 on_op_dispatch_exit = _on_op_dispatch_exit
255 # ------------------------------------------------------------------
256 # Record bookkeeping
257 # ------------------------------------------------------------------
259 def _make_record(self, record_cls, **fields) -> Optional[DebugCall]:
260 """Build a record node, or return None when it must not be retained.
262 Returns:
263 Optional[DebugCall]: The node, or None when the depth or node budget
264 is exhausted. Callers must not push anything onto ``_call_stack``
265 in that case.
266 """
267 depth = len(self._call_stack)
268 if depth >= _MAX_DEPTH or self._record_count >= _MAX_RECORDS:
269 self._dropped_records += 1
270 return None
272 self._record_count += 1
273 return record_cls(call_depth=depth, **fields)
275 def _attach(self, record: DebugCall):
276 """Link *record* under the current frame (or the roots) and open it."""
277 if self._call_stack:
278 self._call_stack[-1].children.append(record)
279 else:
280 self._root_records.append(record)
281 self._call_stack.append(record)
283 def get_dropped_record_count(self) -> int:
284 """Returns how many records were discarded because the budget ran out."""
285 return self._dropped_records
287 # ------------------------------------------------------------------
288 # Collective tracer callback
289 # ------------------------------------------------------------------
291 def _on_collective_call(self, method_name: str, args, kwargs, result): # pylint: disable=W0613
292 """Invoked by CollectiveTracer after a collective op completes."""
293 # The counter is authoritative and stays exact even once records are
294 # being dropped, so it is updated before anything else.
295 self._comm_counts[method_name] += 1
297 if not self._active:
298 return
300 input_shape = None
301 input_dtype = ""
302 if args and hasattr(args[0], "shape"):
303 input_shape = tuple(args[0].shape)
304 input_dtype = str(args[0].dtype)
306 output_shape = None
307 if result is not None and hasattr(result, "shape"):
308 output_shape = tuple(result.shape)
310 group_size = 0
311 group_str = None
312 group_idx = _COLLECTIVE_GROUP_ARG_INDEX.get(method_name)
313 if group_idx is not None and len(args) > group_idx:
314 group = args[group_idx]
315 if isinstance(group, str):
316 group_str = group
317 elif hasattr(group, "size"):
318 try:
319 group_size = group.size()
320 except Exception: # pylint: disable=W0703
321 pass
323 record = self._make_record(
324 CollectiveCall,
325 collective_type=method_name,
326 group_size=group_size,
327 group=group_str,
328 input_shape=input_shape,
329 output_shape=output_shape,
330 input_dtype=input_dtype,
331 )
332 if record is None:
333 return
335 # A collective is a leaf: it is linked but never pushed, so it cannot
336 # swallow the enclosing op's exit.
337 if self._call_stack:
338 self._call_stack[-1].children.append(record)
339 else:
340 self._root_records.append(record)
342 # ------------------------------------------------------------------
343 # Module tracker callback
344 # ------------------------------------------------------------------
346 def _on_module_event(self, module_fqn: str, event_type: str):
347 """Invoked by ModuleTracker on forward enter/exit."""
348 if not self._active:
349 return
351 if event_type == "exit":
352 # Only unwind a frame this tracer opened.
353 if self._call_stack and isinstance(self._call_stack[-1], AnnotateCall):
354 self._call_stack.pop()
355 return
357 record = self._make_record(
358 AnnotateCall,
359 module_fqn=module_fqn,
360 event_type=event_type,
361 )
362 if record is None:
363 return
365 self._attach(record)
367 # ------------------------------------------------------------------
368 # Module info collection
369 # ------------------------------------------------------------------
371 def _collect_module_info(self):
372 """Collect parameter and sharding info from the tracked module.
374 Note:
375 ``get_parameter_info()`` returns the parameter *storages* (via
376 ``param.data``), not copies, so the model stays pinned in memory
377 while those entries are held. Call :meth:`clear` to release them.
378 """
379 from hyper_parallel.core.dtensor.dtensor import ( # pylint: disable=C0415
380 DTensor, _distribute_module_named_modules, _distribute_module_named_parameters,
381 )
383 if self._module is None:
384 return
386 for fqn, mod in _distribute_module_named_modules(self._module):
387 name = fqn or "(root)"
388 params = {}
389 for param_name, param in _distribute_module_named_parameters(mod):
390 params[param_name] = param.data
391 if isinstance(param, DTensor):
392 key = f"{name}.{param_name}" if fqn else param_name
393 self._sharding_info[key] = param.placements
394 if params:
395 self._parameter_info[name] = params
397 # ------------------------------------------------------------------
398 # Tensor info extraction
399 # ------------------------------------------------------------------
401 def _extract_tensor_infos(self, args) -> List[TensorInfo]:
402 """Extract TensorInfo from args, handling DTensor and plain Tensor."""
403 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415
405 infos = []
406 for arg in args:
407 if isinstance(arg, DTensor):
408 placements = tuple(repr(p) for p in arg.placements) if hasattr(arg, "placements") else None
409 mesh_shape = None
410 if hasattr(arg, "device_mesh") and arg.device_mesh is not None:
411 mesh_shape = tuple(arg.device_mesh.shape) if hasattr(arg.device_mesh, "shape") else None
412 infos.append(TensorInfo(
413 shape=tuple(arg.shape),
414 dtype=str(arg.dtype),
415 is_dtensor=True,
416 placements=placements,
417 mesh_shape=mesh_shape,
418 ))
419 elif isinstance(arg, Tensor):
420 infos.append(TensorInfo(
421 shape=tuple(arg.shape),
422 dtype=str(arg.dtype),
423 ))
424 elif isinstance(arg, (tuple, list)):
425 infos.extend(self._extract_tensor_infos(arg))
426 return infos
428 # ------------------------------------------------------------------
429 # Public API (aligned with torch.distributed.tensor.debug.CommDebugMode)
430 # ------------------------------------------------------------------
432 def get_comm_counts(self) -> Dict[str, int]:
433 """Returns the communication counts as a dictionary.
435 Returns:
436 Dict[str, int]: Mapping from collective type name to invocation count.
437 """
438 return dict(self._comm_counts)
440 def get_total_counts(self) -> int:
441 """Returns the total number of collective calls recorded."""
442 return sum(self._comm_counts.values())
444 def get_parameter_info(self) -> Dict[str, Dict[str, Any]]:
445 """Returns parameter info collected from the tracked module.
447 Returns:
448 Dict mapping module FQN to a dict of ``{param_name: param_data}``.
449 The values are the live parameter storages, not copies. Only
450 available when a *module* was passed to the constructor.
451 """
452 return self._parameter_info
454 def get_sharding_info(self) -> Dict[str, Any]:
455 """Returns sharding info for DTensor parameters.
457 Returns:
458 Dict mapping ``module_fqn.param_name`` to its placements.
459 Only available when a *module* was passed to the constructor.
460 """
461 return self._sharding_info
463 def generate_comm_debug_tracing_table(self, noise_level: Optional[int] = None) -> str:
464 """Generate a formatted tracing table.
466 Args:
467 noise_level: 0 = collectives only, 1 = ops + collectives,
468 2 = full detail. Defaults to 1.
470 Returns:
471 str: Formatted multi-line table string.
472 """
473 if noise_level is None:
474 noise_level = 1
476 if self._dropped_records:
477 logger.warning(
478 "%d record(s) were dropped because the record budget was exhausted; "
479 "the table is incomplete. Only trace a bounded number of steps.",
480 self._dropped_records,
481 )
483 if noise_level >= 2 and self._module is None:
484 logger.warning(
485 "noise_level=2 shows module boundary annotations, but no module was passed "
486 "to CommDebugMode(). Pass CommDebugMode(module=model) to enable module tracking."
487 )
489 lines = []
490 for record in self._root_records:
491 self._collect_table_lines(record, lines, noise_level, indent=0)
493 if not lines:
494 return "(no operations recorded)"
496 header = f"{'Type':<20} {'Detail':<60}"
497 separator = "-" * 80
498 table_lines = [header, separator] + lines
499 return "\n".join(table_lines)
501 def log_comm_debug_tracing_table_to_file(
502 self, file_name: str = "comm_mode_log.txt", noise_level: Optional[int] = None
503 ) -> None:
504 """Write tracing table to a file (ANSI escape codes stripped).
506 Args:
507 file_name: Output file path.
508 noise_level: Verbosity level (see ``generate_comm_debug_tracing_table``).
509 """
510 ansi_escape = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
511 table = ansi_escape.sub("", self.generate_comm_debug_tracing_table(noise_level))
512 with open(file_name, "w", encoding="utf-8") as f:
513 f.write(table)
515 def generate_json_dump(
516 self, file_name: str = "comm_mode_log.json", noise_level: Optional[int] = None
517 ) -> None:
518 """Export tracing data as a JSON file.
520 Args:
521 file_name: Output file path.
522 noise_level: Verbosity level. Defaults to 1.
523 """
524 if noise_level is None:
525 noise_level = 1
527 def _record_to_dict(record: DebugCall) -> Optional[dict]:
528 entry: dict = {}
529 if isinstance(record, CollectiveCall):
530 entry["type"] = "collective"
531 entry["collective_type"] = record.collective_type
532 if record.group is not None:
533 entry["group"] = record.group
534 else:
535 entry["group_size"] = record.group_size
536 entry["input_shape"] = list(record.input_shape) if record.input_shape else None
537 entry["output_shape"] = list(record.output_shape) if record.output_shape else None
538 elif isinstance(record, OpCall):
539 if noise_level < 1:
540 return None
541 entry["type"] = "op"
542 entry["op_name"] = record.op_name
543 entry["inputs"] = [
544 {"shape": list(t.shape), "dtype": t.dtype, "is_dtensor": t.is_dtensor,
545 "placements": list(t.placements) if t.placements else None}
546 for t in record.input_infos
547 ]
548 entry["outputs"] = [
549 {"shape": list(t.shape), "dtype": t.dtype, "is_dtensor": t.is_dtensor}
550 for t in record.output_infos
551 ]
552 elif isinstance(record, AnnotateCall):
553 if noise_level < 2:
554 return None
555 entry["type"] = "module"
556 entry["module_fqn"] = record.module_fqn
557 entry["event_type"] = record.event_type
558 else:
559 return None
561 children = []
562 for child in record.children:
563 child_dict = _record_to_dict(child)
564 if child_dict is not None:
565 children.append(child_dict)
566 if children:
567 entry["children"] = children
569 return entry
571 data = {
572 "comm_counts": dict(self._comm_counts),
573 "total_counts": self.get_total_counts(),
574 "records": [],
575 }
577 if self._dropped_records:
578 data["dropped_records"] = self._dropped_records
580 if self._sharding_info:
581 data["sharding_info"] = {k: str(v) for k, v in self._sharding_info.items()}
583 for record in self._root_records:
584 entry = _record_to_dict(record)
585 if entry is not None:
586 data["records"].append(entry)
588 with open(file_name, "w", encoding="utf-8") as f:
589 json.dump(data, f, indent=2)
591 # Keep old name as alias for backward compatibility.
592 generate_tracing_table = generate_comm_debug_tracing_table
594 # ------------------------------------------------------------------
595 # Internal helpers
596 # ------------------------------------------------------------------
598 def _collect_table_lines(self, record: DebugCall, lines: List[str],
599 noise_level: int, indent: int):
600 """Recursively append one formatted table line per record in the subtree.
602 Collectives are always emitted; ``OpCall`` and ``AnnotateCall`` lines are
603 gated on *noise_level* so that lowering verbosity drops them while still
604 descending into their children.
606 Args:
607 record: Subtree root to render.
608 lines: Accumulator the formatted lines are appended to; passed through
609 the recursion rather than returned, so the caller can seed it.
610 noise_level: 0 = collectives only, 1 = ops + collectives,
611 2 = module annotations as well.
612 indent: Current nesting level; each level adds two leading spaces.
613 """
614 prefix = " " * indent
615 if isinstance(record, CollectiveCall):
616 lines.append(f"{prefix}{'Collective':<20} {record.render_self()}")
617 elif isinstance(record, OpCall) and noise_level >= 1:
618 lines.append(f"{prefix}{'Op':<20} {record.render_self()}")
619 elif isinstance(record, AnnotateCall) and noise_level >= 2:
620 lines.append(f"{prefix}{'Module':<20} {record.render_self()}")
622 for child in record.children:
623 self._collect_table_lines(child, lines, noise_level, indent + 1)