DTensor Debug 能力设计文档

1. 背景

1.1 问题描述

分布式训练中，DTensor 在执行算子时会自动触发 redistribute（重分布）操作，而 redistribute 内部会调用集合通信（all-reduce、all-gather、reduce-scatter、all-to-all 等）。这带来两个核心调试难题：

1. 通信黑盒问题：用户写的是高层 DTensor 算子（如 linear、matmul），框架内部自动产生集合通信，用户无从得知：

   • 一次前向/反向传播中触发了哪些集合通信？
   • 每次通信的数据量、组大小是多少？
   • 通信是由哪个 nn.Module、哪个 DTensor 算子引发的？

2. 操作追踪问题：框架调度链路（OpDispatcher → tensor_redistribution → platform 集合操作）跨越多层，缺少统一的"操作日志"，定位问题时需要手动打印大量中间状态。

1.2 需求目标

本次工作新增两项调试能力，统一通过 CommDebugMode 上下文管理器对外暴露：

• Comm Debug（集合通信调试）：统计和追踪所有集合通信操作，按类型计数，并以层次化调用树的形式展示通信发生的上下文（触发它的 DTensor 算子）。
• Debug 日志（算子分发日志）：在 OpDispatcher.dispatch() 入口/出口处挂钩，记录每个 DTensor 算子的名称、输入张量的 shape/placement/mesh 等元信息，形成结构化的调用记录。

这两项能力共同构建了"DTensor 算子 → redistribute → 集合通信"这一完整调用链路的可观测性。

---

2. 当前具体实现方案

2.1 整体架构

CommDebugMode（上下文管理器）
├── _debug_mode_observer（ContextVar）── 挂载到 OpDispatcher
│     ├── _on_op_dispatch_enter()     ── DTensor 算子分发入口回调
│     └── _on_op_dispatch_exit()      ── DTensor 算子分发出口回调
├── CollectiveTracer                  ── platform 层 monkey-patch
│     └── _on_collective_call()       ── 集合操作完成回调
├── ModuleTracker（可选）             ── nn.Module 前向 hook
│     └── _on_module_event()          ── 模块进入/退出回调

记录结果以树形结构保存在 _root_records: List[DebugCall] 中，集合通信节点自动嵌套在触发它的 DTensor 算子节点之下。

2.2 工作流程

┌─────────────────────────────────────────────────────────────────┐
│                   with CommDebugMode() as mode:                  │
│                         __enter__()                              │
│  ┌─────────────────────────────────────────────────────────┐    │
│  │ 1. _debug_mode_observer.set(self)   ← ContextVar 挂载   │    │
│  │ 2. CollectiveTracer.install()       ← monkey-patch 平台  │    │
│  │ 3. ModuleTracker.install()          ← 若传入 module      │    │
│  └─────────────────────────────────────────────────────────┘    │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                       model(x)  执行                             │
│                                                                  │
│  用户调用 DTensor 算子（如 linear）                               │
│         │                                                        │
│         ▼                                                        │
│  OpDispatcher.dispatch()                                         │
│    │  检查 _debug_mode_observer                                  │
│    ├─→ _on_op_dispatch_enter()                                   │
│    │     创建 OpCall，压入 _call_stack                            │
│    │                                                             │
│    │  算子执行（触发 tensor_redistribution）                      │
│    │         │                                                   │
│    │         ▼                                                   │
│    │  platform.differentiable_all_reduce()  ← 已被 monkey-patch │
│    │    │  原始通信执行                                           │
│    │    └─→ _on_collective_call()                                │
│    │          创建 CollectiveCall                                 │
│    │          挂到 _call_stack 栈顶（OpCall）的 children 下       │
│    │          _comm_counts[op] += 1                              │
│    │                                                             │
│    └─→ _on_op_dispatch_exit()                                   │
│          填充 output_infos，弹出 _call_stack                     │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                         __exit__()                               │
│  ┌─────────────────────────────────────────────────────────┐    │
│  │ 1. ModuleTracker.uninstall()        ← 移除模块 hook      │    │
│  │ 2. CollectiveTracer.uninstall()     ← 恢复平台原始方法   │    │
│  │ 3. _debug_mode_observer.reset()     ← ContextVar 复位    │    │
│  └─────────────────────────────────────────────────────────┘    │
└─────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────┐
│                         查询结果                                  │
│                                                                  │
│  mode.get_comm_counts()           → _comm_counts 字典            │
│  mode.get_total_counts()          → sum(_comm_counts)            │
│  mode.generate_comm_debug_tracing_table(noise_level)             │
│    遍历 _root_records 树，按 noise_level 过滤输出                 │
└─────────────────────────────────────────────────────────────────┘

执行期间 _call_stack 动态变化，最终结果沉淀到 _root_records：

执行中（_call_stack）                 执行后（_root_records）

[OpCall(linear)]          →           OpCall(linear)
[OpCall(linear),                        └── CollectiveCall(all_reduce)
 CollectiveCall]
[]  ← 算子退出后弹空

2.3 新增/修改文件清单

文件	状态	说明
hyper_parallel/core/dtensor/debug/__init__.py	新增	导出 CommDebugMode
hyper_parallel/core/dtensor/debug/_comm_debug_mode.py	新增	主上下文管理器实现
hyper_parallel/core/dtensor/debug/_call_records.py	新增	结构化调用记录类层次
hyper_parallel/core/dtensor/debug/_collective_tracer.py	新增	platform monkey-patch 拦截器
hyper_parallel/core/dtensor/debug/_module_tracker.py	新增	nn.Module 前向 hook 追踪器
hyper_parallel/core/shard/_op_dispatch.py	修改	新增 _debug_mode_observer ContextVar 及入口/出口回调
hyper_parallel/core/dtensor/__init__.py	修改	导出 CommDebugMode
hyper_parallel/core/dtensor/tensor_redistribution.py	修改	适配新文件头（不影响运行逻辑）

2.4 结构化调用记录（_call_records.py）

定义了一套类层次结构，将所有被追踪的调用统一抽象为 DebugCall：

DebugCall（基类）
├── call_depth: int          # 嵌套深度
├── timestamp: float         # 创建时间戳
├── children: List[DebugCall] # 子记录（树形结构）
└── render(indent) -> str    # 层次化文本渲染

OpCall(DebugCall)            # DTensor 算子调用
├── op_name: str             # 算子名，如 "linear"、"mm"、"gelu"
├── input_infos: List[TensorInfo]
└── output_infos: List[TensorInfo]

CollectiveCall(DebugCall)    # 集合通信操作
├── collective_type: str     # 如 "differentiable_all_reduce"
├── group_size: int
├── input_shape / output_shape
└── input_dtype: str

RedistributeCall(DebugCall)  # redistribute 操作（预留）
├── src_placements / dst_placements
└── tensor_shape

AnnotateCall(DebugCall)      # 模块边界标注
├── module_fqn: str          # 模块全限定名
└── event_type: str          # "enter" 或 "exit"

TensorInfo 是张量元信息快照：

@dataclass
class TensorInfo:
    shape: Tuple[int, ...]
    dtype: str
    is_dtensor: bool = False
    placements: Optional[Tuple] = None   # DTensor 专属
    mesh_shape: Optional[Tuple] = None   # DTensor 专属

2.5 集合操作拦截（_collective_tracer.py）

由于 hyper-parallel 的集合操作直接调用 platform 函数，无法通过算子注册表拦截，采用可恢复 monkey-patch方案：

拦截目标（platform 类的 staticmethod）：

differentiable_all_gather_concat
differentiable_all_to_all
differentiable_all_reduce
differentiable_reduce_scatter
differentiable_all_to_all_single
differentiable_all_to_all_single_async

关键实现细节：

• _patch_lock（threading.Lock）保证多线程安全
• 通过 cls.__dict__[name] 保存原始 descriptor（staticmethod 包装器），而非绑定函数，保证 uninstall() 时可精确还原
• 使用 type.__setattr__ 而非 setattr 还原，确保 descriptor 协议不被破坏
• wrapper 中捕获所有回调异常（except Exception: pass），保证生产逻辑不受调试代码影响

# 安装示例（简化）
def _make_wrapper(orig, cb, mname):
    def wrapper(*args, **kwargs):
        result = orig(*args, **kwargs)
        try:
            cb(mname, args, kwargs, result)
        except Exception:
            pass
        return result
    return wrapper
setattr(cls, name, staticmethod(wrapper))

# 还原
type.__setattr__(cls, name, original_descriptor)

2.6 OpDispatcher 钩子（_op_dispatch.py）

在 OpDispatcher.dispatch() 中增加了一个 ContextVar 类型的观察者槽，CommDebugMode 激活时注册自身，退出时清除，对非 debug 场景零开销：

# _op_dispatch.py 新增（约 5 行）
_debug_mode_observer: ContextVar = ContextVar('_debug_mode_observer', default=None)

# dispatch() 入口
observer = _debug_mode_observer.get()
if observer is not None:
    observer._on_op_dispatch_enter(op_name, op_call, args, kwargs)

# dispatch() 出口
if observer is not None:
    observer._on_op_dispatch_exit(op_name, result)

CommDebugMode.__enter__ 时：

self._observer_token = _debug_mode_observer.set(self)

CommDebugMode.__exit__ 时：

_debug_mode_observer.reset(self._observer_token)

利用 ContextVar 的 token 机制，天然支持嵌套 CommDebugMode（内层退出后外层自动恢复），无需手动管理嵌套栈。

2.7 模块追踪（_module_tracker.py）

ModuleTracker 在用户传入 module 参数时启用，通过 register_forward_pre_hook 和 register_forward_hook 对所有子模块打桩：

# 为每个子模块注册进入/退出 hook
handle_pre = mod.register_forward_pre_hook(_make_pre_hook(fqn))   # 触发 "enter"
handle_post = mod.register_forward_hook(_make_post_hook(fqn))      # 触发 "exit"

AnnotateCall("enter") 压入 _call_stack，AnnotateCall("exit") 将其弹出，从而使其内部的 DTensor 算子调用和集合通信调用自然成为该模块节点的子节点。

2.8 调用树构建机制

_call_stack: List[DebugCall] 充当当前嵌套上下文栈：

• _on_op_dispatch_enter：创建 OpCall，若栈非空则挂到栈顶 children，否则挂到 root；压栈
• _on_op_dispatch_exit：弹栈，填充 output_infos
• _on_collective_call：创建 CollectiveCall，挂到栈顶 children（或 root）；不压栈（集合通信是叶子节点）
• _on_module_event("enter")：创建 AnnotateCall，压栈
• _on_module_event("exit")：弹栈

由于 platform monkey-patch 的回调在集合通信函数返回后同步触发，此时 OpDispatcher 的 dispatch 尚未退出，_call_stack 中仍有该 DTensor 算子节点，集合通信因此自然成为该算子的子节点，构成正确的因果关系树。

---

3. 对外接口

3.1 导入路径

# 推荐导入
from hyper_parallel.core.dtensor.debug import CommDebugMode
from hyper_parallel.core.dtensor import CommDebugMode  # 也可通过 dtensor 顶层导入

3.2 CommDebugMode 构造参数

CommDebugMode(module=None)

参数	类型	说明
module	nn.Module（可选）	指定后启用模块追踪（ModuleTracker）

3.3 上下文管理器协议

with CommDebugMode(module=model) as mode:
    output = model(input)
# 退出后可查询结果

__enter__：

1. 清空所有追踪状态
2. 设置 _debug_mode_observer ContextVar（挂载到 OpDispatcher）
3. 安装 CollectiveTracer（monkey-patch platform 集合方法）
4. 若指定 module，安装 ModuleTracker 并收集初始参数/分片信息

__exit__：

1. 卸载 ModuleTracker
2. 卸载 CollectiveTracer（恢复原始 platform 方法）
3. 重置 _debug_mode_observer ContextVar

3.4 查询接口

接口	返回类型	说明
get_comm_counts()	Dict[str, int]	各类集合操作的调用次数，如 {"differentiable_all_reduce": 2}
get_total_counts()	int	集合操作总次数
get_parameter_info()	Dict[str, Dict[str, Any]]	模块参数信息（需传入 module）
get_sharding_info()	Dict[str, Any]	DTensor 参数的 placement 信息（需传入 module）

3.5 输出接口

generate_comm_debug_tracing_table(noise_level=None) -> str

生成格式化的多行追踪表，noise_level 控制输出详细程度（数值越大信息越多）：

noise_level	显示内容	适用场景
0	仅集合通信（CollectiveCall）	只关心"发生了哪些通信"
1	集合通信 + DTensor 算子（OpCall）（默认）	还想看"是哪个算子触发了通信"
2	以上 + 模块边界标注（AnnotateCall）	还想看"通信属于哪个 nn.Module"（需传入 module）

"noise"来自 PyTorch 的命名习惯，意为"信息噪声"——级别越高输出越嘈杂，用户根据调试需要自行权衡。

以下以 ColwiseParallel(fc1) + RowwiseParallel(fc2) MLP 前向为例展示各级别输出差异：

noise_level=0（只看通信）：

Type                 Detail
--------------------------------------------------------------------------------
  Collective         Collective(differentiable_all_reduce) group_size=2 input_shape=(4, 16) output_shape=(4, 16)

noise_level=1（看算子 + 通信）：

Type                 Detail
--------------------------------------------------------------------------------
Op                   Op(linear) inputs=[DTensor[4, 16], DTensor[32, 16]] outputs=[DTensor[4, 32]]
Op                   Op(gelu) inputs=[DTensor[4, 32]] outputs=[DTensor[4, 32]]
Op                   Op(linear) inputs=[DTensor[4, 32], DTensor[16, 32]] outputs=[DTensor[4, 16]]
  Collective         Collective(differentiable_all_reduce) group_size=2 input_shape=(4, 16) output_shape=(4, 16)

noise_level=2（看模块 + 算子 + 通信，需传入 module=model）：

Type                 Detail
--------------------------------------------------------------------------------
Module               Module((root)) [enter]
  Module             Module(fc1) [enter]
    Op               Op(linear) inputs=[DTensor[4, 16], DTensor[32, 16]] outputs=[DTensor[4, 32]]
  Op                 Op(gelu) inputs=[DTensor[4, 32]] outputs=[DTensor[4, 32]]
  Module             Module(fc2) [enter]
    Op               Op(linear) inputs=[DTensor[4, 32], DTensor[16, 32]] outputs=[DTensor[4, 16]]
      Collective     Collective(differentiable_all_reduce) group_size=2 input_shape=(4, 16) output_shape=(4, 16)

别名：generate_tracing_table(noise_level=None) 与此接口等价。

log_comm_debug_tracing_table_to_file(file_name="comm_mode_log.txt", noise_level=None) -> None

将追踪表写入文件，自动剥离 ANSI 转义码。

generate_json_dump(file_name="comm_mode_log.json", noise_level=None) -> None

导出 JSON 格式的完整追踪数据，结构如下：

{
  "comm_counts": {"differentiable_all_reduce": 1},
  "total_counts": 1,
  "sharding_info": {"fc2.weight": "Shard(dim=1)"},
  "records": [
    {
      "type": "op",
      "op_name": "linear",
      "inputs": [{"shape": [4, 32], "dtype": "torch.float32", "is_dtensor": true, "placements": ["Shard(-1)"]}],
      "outputs": [...],
      "children": [
        {"type": "collective", "collective_type": "differentiable_all_reduce", ...}
      ]
    }
  ]
}

3.6 完整使用示例

from hyper_parallel.core.dtensor.debug import CommDebugMode

# 基本用法：统计集合通信
with CommDebugMode() as mode:
    output = model(input_dtensor)

print(mode.get_comm_counts())      # {'differentiable_all_reduce': 2}
print(mode.get_total_counts())     # 2

# 查看详细追踪树（noise_level=1：算子 + 通信）
print(mode.generate_comm_debug_tracing_table(noise_level=1))

# 带模块追踪（可定位通信属于哪个 nn.Module）
with CommDebugMode(module=model) as mode:
    output = model(input_dtensor)

# 查看模块边界（noise_level=2）
print(mode.generate_comm_debug_tracing_table(noise_level=2))
print(mode.get_sharding_info())

# 导出 JSON
mode.generate_json_dump("debug_log.json")

# 写入文件
mode.log_comm_debug_tracing_table_to_file("debug_log.txt")

---

4. 使用约束

1. 嵌套使用时只有内层有效

当多个 CommDebugMode 嵌套使用时，_debug_mode_observer ContextVar 始终指向最内层实例，CollectiveTracer 的回调也只绑定到最内层，外层在内层存活期间不会收到任何算子或通信事件：

with CommDebugMode() as outer:
    with CommDebugMode() as inner:
        model(x)   # 所有事件只记录到 inner，outer 在此期间为空

2. 只追踪经过 OpDispatcher 的算子，只统计经由 platform 的集合通信

调用树中仅包含通过 OpDispatcher.dispatch() 分发的 DTensor 算子；_comm_counts 仅统计经由 platform 层集合方法（differentiable_all_reduce 等）发起的通信。直接通过 torch.distributed 或其他途径发起的通信不会被捕获。

3. 不支持图模式

torch.compile 或 MindSpore 图模式下，算子调度路径绕过了 OpDispatcher，CommDebugMode 无法追踪算子分发，集合通信也可能不经过 platform 层，导致追踪结果为空或不完整。

4. 支持 Torch / MindSpore 双平台

CommDebugMode 的核心追踪路径（OpDispatcher hook + CollectiveTracer monkey-patch）在两个平台上均可正常工作。MindSpore 平台缺少 differentiable_all_to_all_single 方法，CollectiveTracer 会安全跳过，不影响其他集合通信的追踪。

5. noise_level=2 需要传入 module 才能生效

noise_level=2 会在追踪表中显示模块边界标注（Module(fc1) [enter] 等），但这依赖 ModuleTracker 在各子模块上注册的 forward hook。若未向 CommDebugMode 传入 module，ModuleTracker 不会被安装，noise_level=2 的输出与 noise_level=1 完全相同，并会打印一条警告日志提示用户。

# 正确用法
with CommDebugMode(module=model) as mode:
    model(x)
print(mode.generate_comm_debug_tracing_table(noise_level=2))

---

5. 性能影响分析

5.1 修改范围

本次新增代码均位于 debug/ 子包中，属于纯新增模块，不影响原有流程。唯一对原有代码的修改是在 _op_dispatch.py 的 dispatch() 方法中增加了 3 行。

5.2 热路径改动

dispatch() 是每个 DTensor 算子必经的热路径，本次在其中新增：

# 入口处
observer = _debug_mode_observer.get()
if observer is not None:
    observer._on_op_dispatch_enter(op_name, op_call, args, kwargs)

# finally 块中
if observer is not None:
    observer._on_op_dispatch_exit(op_name, result)

5.3 非 debug 模式下的开销（生产训练）

操作	耗时量级	说明
ContextVar.get()	~50–200 ns	CPython C 层实现，类似 dict 查找
if observer is not None	可忽略	None 判断

每次 dispatch 增加约 2 次 ContextVar.get() + 2 次 None 判断，合计约 100–400 ns。相比实际 layout 推断、集合通信操作的微秒～毫秒级开销，影响可忽略不计。

CollectiveTracer 的 monkey-patch 仅在 CommDebugMode.__enter__ 时安装，正常训练时从未调用，对集合通信方法零开销。

5.4 debug 模式下的开销

当用户主动使用 CommDebugMode 时，每个算子会额外创建 OpCall 对象、执行 list append 和张量 shape/dtype 读取，每次集合通信会额外触发一次 Python 函数调用。这属于调试工具的预期代价，不适用于生产训练。

5.5 设计结论

ContextVar 是 Python 标准库为"可选激活"场景提供的专用工具，是当前能做到的最低开销方案。非 debug 模式下本次改动对原有训练流程无可感知的性能损耗。

---

6. 测试设计

6.1 测试文件列表

文件	类型	说明
tests/ut/core/dtensor/test_comm_debug_mode.py	单元测试	不依赖分布式环境，覆盖各组件行为
tests/torch/tensor_parallel/test_comm_debug_mode_mlp.py	集成测试入口	pytest 入口，通过 torchrun 启动 2-rank 进程
tests/torch/tensor_parallel/_test_comm_debug_mode_mlp.py	集成测试 Worker	在真实分布式环境中运行的测试逻辑

6.2 单元测试（test_comm_debug_mode.py）

运行环境：单进程，Torch 平台（HYPER_PARALLEL_PLATFORM=torch），无需分布式后端。

测试分组

组	测试类	覆盖要点
1	TestCallRecords	OpCall、CollectiveCall、DebugCall 的 _render_self() 和 render() 层次化输出
2	TestCollectiveTracer	monkey-patch 安装/卸载后方法精确还原（identity 检查）、回调触发、回调异常不传播
3	TestCommDebugModeContextManager	ContextVar 生命周期、嵌套 CommDebugMode、平台方法还原、空追踪输出
4	TestObserverHooks	on_op_dispatch_enter/exit 创建 OpCall 并维护树结构、集合通信计数、集合通信嵌套在 Op 下
5	TestDebugOutput	generate_comm_debug_tracing_table 包含层次内容、noise_level=0 过滤 Op、noise_level=1 包含 Op、__repr__
6	TestTensorInfoExtraction	普通张量 / 嵌套 tuple / 非张量参数的 _extract_tensor_infos 行为
7	TestZeroOverhead	非 debug 模式下 _debug_mode_observer 默认为 None

关键测试用例示例

# 验证集合通信嵌套在 Op 节点之下
def test_collective_nested_under_op(self):
    mode = CommDebugMode()
    with mode:
        mode.on_op_dispatch_enter("redistribute", ...)
        mode._on_collective_call("differentiable_all_gather_concat", ...)
        mode.on_op_dispatch_exit("redistribute", ...)
    # 期望：root_records[0].children[0] 是 CollectiveCall
    assert isinstance(mode._root_records[0].children[0], CollectiveCall)

# 验证 noise_level=0 只显示集合通信
def test_tracing_table_noise_level_0(self):
    table = mode.generate_tracing_table(noise_level=0)
    assert "differentiable_all_reduce" in table
    assert "linear" not in table

6.3 集成测试（MLP 分布式测试）

运行环境：torchrun 启动，2 个 rank，gloo CPU 后端，ColwiseParallel + RowwiseParallel 张量并行模型。

测试模型：

class TwoLayerMLP(nn.Module):
    def __init__(self, hidden, intermediate):
        self.fc1 = nn.Linear(hidden, intermediate)  # ColwiseParallel
        self.fc2 = nn.Linear(intermediate, hidden)  # RowwiseParallel

    def forward(self, x):
        return self.fc2(F.gelu(self.fc1(x)))

张量并行策略下的通信分析：

• fc1（ColwiseParallel）：weight Shard(0)，input Replicate，output Shard(-1)，无通信
• fc2（RowwiseParallel）：weight Shard(1)，input Shard(-1)，output 从 Partial 到 Replicate，触发 1 次 all-reduce

集成测试用例

测试函数	验证内容	期望结果
test_comm_debug_mode_captures_collectives	集合通信计数正确	get_total_counts() >= 1，包含 all-reduce 或 reduce-scatter
test_comm_debug_mode_debug_string	追踪表非空且包含内容	generate_comm_debug_tracing_table() 不为 "(no operations recorded)"
test_comm_debug_mode_tracing_table	noise_level 过滤效果	noise_level=1 的行数 >= noise_level=0 的行数
test_comm_debug_mode_with_module_tracker	模块边界标注	noise_level=2 追踪表包含 "Module" 字样
test_comm_debug_mode_restores_platform	平台方法还原	退出后 cls.__dict__[name] is orig（identity 检查）
test_comm_debug_mode_multiple_forwards	多次前向累积计数	3 次前向的 get_total_counts() == 1 次前向的 3 倍

6.4 测试覆盖矩阵

能力	单元测试	集成测试
调用记录类渲染	✅	—
平台 monkey-patch 安装/卸载	✅	✅
OpDispatcher ContextVar 生命周期	✅	✅
嵌套 CommDebugMode	✅	—
集合通信计数	✅	✅
调用树层次结构（通信嵌套于算子下）	✅	✅
模块追踪	—	✅
noise_level 过滤	✅	✅
多次前向累积	—	✅
回调异常隔离	✅	—

---