Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / debug / _comm_debug_mode.py: 60%
215 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-04 05:18 +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)
25"""
26import json
27import logging
28import re
29from collections import defaultdict
30from typing import Any, Dict, List, Optional
32from hyper_parallel.core.dtensor.debug._call_records import (
33 AnnotateCall,
34 CollectiveCall,
35 DebugCall,
36 OpCall,
37 TensorInfo,
38)
39from hyper_parallel.core.dtensor.debug._collective_tracer import CollectiveTracer
40from hyper_parallel.core.dtensor.debug._module_tracker import ModuleTracker
41from hyper_parallel.platform import get_platform
43logger = logging.getLogger(__name__)
44platform = get_platform()
45Tensor = platform.Tensor
47# Argument index of the process group for each traced collective method.
48# Derived from the platform method signatures:
49# differentiable_all_gather_concat(data, group, concat_size, concat_dim, ...)
50# differentiable_all_to_all(input_data, output_shape, group)
51# differentiable_all_reduce(data, op, group)
52# differentiable_reduce_scatter(data, dev_num, axis, op, group)
53# differentiable_all_to_all_single(input_tensor, input_splits, output_splits, group)
54# differentiable_all_to_all_single_async(input_tensor, input_splits, output_splits, group)
55_COLLECTIVE_GROUP_ARG_INDEX: Dict[str, int] = {
56 "differentiable_all_gather_concat": 1,
57 "differentiable_all_to_all": 2,
58 "differentiable_all_reduce": 2,
59 "differentiable_reduce_scatter": 4,
60 "differentiable_all_to_all_single": 3,
61 "differentiable_all_to_all_single_async": 3,
62}
65class CommDebugMode:
66 """Context manager that records DTensor operator dispatches and collective
67 communication operations, producing a hierarchical call tree.
69 Usage::
71 with CommDebugMode() as mode:
72 output = model(input_dtensor)
73 print(mode.generate_comm_debug_tracing_table())
74 print(mode.get_comm_counts())
76 Args:
77 module: Optional ``nn.Module`` to track forward enter/exit events.
78 """
80 def __init__(self, module=None):
81 self._module = module
83 # ---- tracing state ----
84 self._call_stack: List[DebugCall] = []
85 self._root_records: List[DebugCall] = []
86 self._comm_counts: Dict[str, int] = defaultdict(int)
87 # ---- module-level info (populated when module is provided) ----
88 self._parameter_info: Dict[str, Dict[str, Any]] = {}
89 self._sharding_info: Dict[str, Any] = {}
91 # ---- internal handles ----
92 self._collective_tracer: Optional[CollectiveTracer] = None
93 self._module_tracker: Optional[ModuleTracker] = None
94 self._observer_token = None
96 # ------------------------------------------------------------------
97 # Context manager protocol
98 # ------------------------------------------------------------------
100 def __enter__(self):
101 # pylint: disable=C0415
102 from hyper_parallel.core.shard._op_dispatch import _debug_mode_observer
104 self._comm_counts.clear()
105 self._root_records.clear()
106 self._call_stack.clear()
107 self._parameter_info.clear()
108 self._sharding_info.clear()
110 self._observer_token = _debug_mode_observer.set(self)
112 self._collective_tracer = CollectiveTracer(self._on_collective_call)
113 self._collective_tracer.install()
115 if self._module is not None:
116 self._module_tracker = ModuleTracker(self._module, self._on_module_event)
117 self._module_tracker.install()
118 self._collect_module_info()
120 return self
122 def __exit__(self, exc_type, exc_val, exc_tb):
123 # pylint: disable=C0415
124 from hyper_parallel.core.shard._op_dispatch import _debug_mode_observer
126 if self._module_tracker is not None:
127 self._module_tracker.uninstall()
128 self._module_tracker = None
130 if self._collective_tracer is not None:
131 self._collective_tracer.uninstall()
132 self._collective_tracer = None
134 if self._observer_token is not None:
135 _debug_mode_observer.reset(self._observer_token)
136 self._observer_token = None
138 def __repr__(self):
139 return f"CommDebugMode(get_total_counts()={self.get_total_counts()})"
141 # ------------------------------------------------------------------
142 # Observer callbacks (called from _op_dispatch.py — internal)
143 # ------------------------------------------------------------------
145 def _on_op_dispatch_enter(self, op_name: str, op_call, args, kwargs): # pylint: disable=W0613
146 """Called by OpDispatcher.dispatch() before the op executes."""
147 depth = len(self._call_stack)
148 record = OpCall(
149 call_depth=depth,
150 op_name=op_name,
151 input_infos=self._extract_tensor_infos(args),
152 )
154 if self._call_stack:
155 self._call_stack[-1].children.append(record)
156 else:
157 self._root_records.append(record)
159 self._call_stack.append(record)
161 def _on_op_dispatch_exit(self, op_name, result): # pylint: disable=W0613
162 """Called by OpDispatcher.dispatch() after the op executes."""
163 if not self._call_stack:
164 return
166 record = self._call_stack.pop()
167 if isinstance(record, OpCall):
168 record.output_infos = self._extract_tensor_infos((result,))
170 # Keep old names as aliases for backward compatibility with tests.
171 on_op_dispatch_enter = _on_op_dispatch_enter
172 on_op_dispatch_exit = _on_op_dispatch_exit
174 # ------------------------------------------------------------------
175 # Collective tracer callback
176 # ------------------------------------------------------------------
178 def _on_collective_call(self, method_name: str, args, kwargs, result): # pylint: disable=W0613
179 """Invoked by CollectiveTracer after a collective op completes."""
180 depth = len(self._call_stack)
182 input_shape = None
183 input_dtype = ""
184 if args and hasattr(args[0], "shape"):
185 input_shape = tuple(args[0].shape)
186 input_dtype = str(args[0].dtype)
188 output_shape = None
189 if result is not None and hasattr(result, "shape"):
190 output_shape = tuple(result.shape)
192 group_size = 0
193 group_str = None
194 group_idx = _COLLECTIVE_GROUP_ARG_INDEX.get(method_name)
195 if group_idx is not None and len(args) > group_idx:
196 group = args[group_idx]
197 if isinstance(group, str):
198 group_str = group
199 elif hasattr(group, "size"):
200 try:
201 group_size = group.size()
202 except Exception: # pylint: disable=W0703
203 pass
205 record = CollectiveCall(
206 call_depth=depth,
207 collective_type=method_name,
208 group_size=group_size,
209 group=group_str,
210 input_shape=input_shape,
211 output_shape=output_shape,
212 input_dtype=input_dtype,
213 )
215 if self._call_stack:
216 self._call_stack[-1].children.append(record)
217 else:
218 self._root_records.append(record)
220 self._comm_counts[method_name] += 1
222 # ------------------------------------------------------------------
223 # Module tracker callback
224 # ------------------------------------------------------------------
226 def _on_module_event(self, module_fqn: str, event_type: str):
227 """Invoked by ModuleTracker on forward enter/exit."""
228 depth = len(self._call_stack)
229 record = AnnotateCall(
230 call_depth=depth,
231 module_fqn=module_fqn,
232 event_type=event_type,
233 )
235 if event_type == "enter":
236 if self._call_stack:
237 self._call_stack[-1].children.append(record)
238 else:
239 self._root_records.append(record)
240 self._call_stack.append(record)
241 else: # "exit"
242 if self._call_stack:
243 self._call_stack.pop()
245 # ------------------------------------------------------------------
246 # Module info collection
247 # ------------------------------------------------------------------
249 def _collect_module_info(self):
250 """Collect parameter and sharding info from the tracked module."""
251 from hyper_parallel.core.dtensor.dtensor import ( # pylint: disable=C0415
252 DTensor, _distribute_module_named_modules, _distribute_module_named_parameters,
253 )
255 if self._module is None:
256 return
258 for fqn, mod in _distribute_module_named_modules(self._module):
259 name = fqn or "(root)"
260 params = {}
261 for param_name, param in _distribute_module_named_parameters(mod):
262 params[param_name] = param.data
263 if isinstance(param, DTensor):
264 key = f"{name}.{param_name}" if fqn else param_name
265 self._sharding_info[key] = param.placements
266 if params:
267 self._parameter_info[name] = params
269 # ------------------------------------------------------------------
270 # Tensor info extraction
271 # ------------------------------------------------------------------
273 def _extract_tensor_infos(self, args) -> List[TensorInfo]:
274 """Extract TensorInfo from args, handling DTensor and plain Tensor."""
275 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415
277 infos = []
278 for arg in args:
279 if isinstance(arg, DTensor):
280 placements = tuple(repr(p) for p in arg.placements) if hasattr(arg, "placements") else None
281 mesh_shape = None
282 if hasattr(arg, "device_mesh") and arg.device_mesh is not None:
283 mesh_shape = tuple(arg.device_mesh.shape) if hasattr(arg.device_mesh, "shape") else None
284 infos.append(TensorInfo(
285 shape=tuple(arg.shape),
286 dtype=str(arg.dtype),
287 is_dtensor=True,
288 placements=placements,
289 mesh_shape=mesh_shape,
290 ))
291 elif isinstance(arg, Tensor):
292 infos.append(TensorInfo(
293 shape=tuple(arg.shape),
294 dtype=str(arg.dtype),
295 ))
296 elif isinstance(arg, (tuple, list)):
297 infos.extend(self._extract_tensor_infos(arg))
298 return infos
300 # ------------------------------------------------------------------
301 # Public API (aligned with torch.distributed.tensor.debug.CommDebugMode)
302 # ------------------------------------------------------------------
304 def get_comm_counts(self) -> Dict[str, int]:
305 """Returns the communication counts as a dictionary.
307 Returns:
308 Dict[str, int]: Mapping from collective type name to invocation count.
309 """
310 return dict(self._comm_counts)
312 def get_total_counts(self) -> int:
313 """Returns the total number of collective calls recorded."""
314 return sum(self._comm_counts.values())
316 def get_parameter_info(self) -> Dict[str, Dict[str, Any]]:
317 """Returns parameter info collected from the tracked module.
319 Returns:
320 Dict mapping module FQN to a dict of ``{param_name: param_data}``.
321 Only available when a *module* was passed to the constructor.
322 """
323 return self._parameter_info
325 def get_sharding_info(self) -> Dict[str, Any]:
326 """Returns sharding info for DTensor parameters.
328 Returns:
329 Dict mapping ``module_fqn.param_name`` to its placements.
330 Only available when a *module* was passed to the constructor.
331 """
332 return self._sharding_info
334 def generate_comm_debug_tracing_table(self, noise_level: Optional[int] = None) -> str:
335 """Generate a formatted tracing table.
337 Args:
338 noise_level: 0 = collectives only, 1 = ops + collectives,
339 2 = full detail. Defaults to 1.
341 Returns:
342 str: Formatted multi-line table string.
343 """
344 if noise_level is None:
345 noise_level = 1
347 if noise_level >= 2 and self._module is None:
348 logger.warning(
349 "noise_level=2 shows module boundary annotations, but no module was passed "
350 "to CommDebugMode(). Pass CommDebugMode(module=model) to enable module tracking."
351 )
353 lines = []
354 for record in self._root_records:
355 self._collect_table_lines(record, lines, noise_level, indent=0)
357 if not lines:
358 return "(no operations recorded)"
360 header = f"{'Type':<20} {'Detail':<60}"
361 separator = "-" * 80
362 table_lines = [header, separator] + lines
363 return "\n".join(table_lines)
365 def log_comm_debug_tracing_table_to_file(
366 self, file_name: str = "comm_mode_log.txt", noise_level: Optional[int] = None
367 ) -> None:
368 """Write tracing table to a file (ANSI escape codes stripped).
370 Args:
371 file_name: Output file path.
372 noise_level: Verbosity level (see ``generate_comm_debug_tracing_table``).
373 """
374 ansi_escape = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
375 table = ansi_escape.sub("", self.generate_comm_debug_tracing_table(noise_level))
376 with open(file_name, "w", encoding="utf-8") as f:
377 f.write(table)
379 def generate_json_dump(
380 self, file_name: str = "comm_mode_log.json", noise_level: Optional[int] = None
381 ) -> None:
382 """Export tracing data as a JSON file.
384 Args:
385 file_name: Output file path.
386 noise_level: Verbosity level. Defaults to 1.
387 """
388 if noise_level is None:
389 noise_level = 1
391 def _record_to_dict(record: DebugCall) -> Optional[dict]:
392 entry: dict = {}
393 if isinstance(record, CollectiveCall):
394 entry["type"] = "collective"
395 entry["collective_type"] = record.collective_type
396 if record.group is not None:
397 entry["group"] = record.group
398 else:
399 entry["group_size"] = record.group_size
400 entry["input_shape"] = list(record.input_shape) if record.input_shape else None
401 entry["output_shape"] = list(record.output_shape) if record.output_shape else None
402 elif isinstance(record, OpCall):
403 if noise_level < 1:
404 return None
405 entry["type"] = "op"
406 entry["op_name"] = record.op_name
407 entry["inputs"] = [
408 {"shape": list(t.shape), "dtype": t.dtype, "is_dtensor": t.is_dtensor,
409 "placements": list(t.placements) if t.placements else None}
410 for t in record.input_infos
411 ]
412 entry["outputs"] = [
413 {"shape": list(t.shape), "dtype": t.dtype, "is_dtensor": t.is_dtensor}
414 for t in record.output_infos
415 ]
416 elif isinstance(record, AnnotateCall):
417 if noise_level < 2:
418 return None
419 entry["type"] = "module"
420 entry["module_fqn"] = record.module_fqn
421 entry["event_type"] = record.event_type
422 else:
423 return None
425 children = []
426 for child in record.children:
427 child_dict = _record_to_dict(child)
428 if child_dict is not None:
429 children.append(child_dict)
430 if children:
431 entry["children"] = children
433 return entry
435 data = {
436 "comm_counts": dict(self._comm_counts),
437 "total_counts": self.get_total_counts(),
438 "records": [],
439 }
441 if self._sharding_info:
442 data["sharding_info"] = {k: str(v) for k, v in self._sharding_info.items()}
444 for record in self._root_records:
445 entry = _record_to_dict(record)
446 if entry is not None:
447 data["records"].append(entry)
449 with open(file_name, "w", encoding="utf-8") as f:
450 json.dump(data, f, indent=2)
452 # Keep old name as alias for backward compatibility.
453 generate_tracing_table = generate_comm_debug_tracing_table
455 # ------------------------------------------------------------------
456 # Internal helpers
457 # ------------------------------------------------------------------
459 def _collect_table_lines(self, record: DebugCall, lines: List[str],
460 noise_level: int, indent: int):
461 """Recursively append formatted lines for *record* and its children."""
462 prefix = " " * indent
463 if isinstance(record, CollectiveCall):
464 lines.append(f"{prefix}{'Collective':<20} {record._render_self()}")
465 elif isinstance(record, OpCall) and noise_level >= 1:
466 lines.append(f"{prefix}{'Op':<20} {record._render_self()}")
467 elif isinstance(record, AnnotateCall) and noise_level >= 2:
468 lines.append(f"{prefix}{'Module':<20} {record._render_self()}")
470 for child in record.children:
471 self._collect_table_lines(child, lines, noise_level, indent + 1)