diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..ca2b667 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,68 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Repository scope + +This repo open-sources one piece of ByteDance's internal veScale: a DTensor library with a new **`RaggedShard`** placement (asymmetric sharding over flattened tensor storage). The active code is `vescale/` (~5.7k lines of pure Python). `legacy/` is the old, much larger veScale library (FSDP, optimizers, examples) moved here for reference only — do not develop against or reformat it. + +## Commands + +Install (Python >= 3.11; pins `torch==2.7.1` cu121 and `numpy<2.0.0`): + +```bash +pip3 install -r requirements.txt && pip3 install -e . +``` + +Run the full test suite (installs, then runs every `test_*.py` under `test/`): + +```bash +./scripts/run_test.sh +``` + +Run a single test file / test (tests import `common_dtensor` as a top-level module, so you must run from `test/` with `PYTHONPATH` including it): + +```bash +cd test +export PYTHONPATH=$(pwd):$PYTHONPATH +pytest dtensor/ragged_shard/test_norm.py -k test_ragged_shard_norm_dim0 -s +``` + +Lint/format (configured in `pyproject.toml`, line length 120; there is no pre-commit config or CI): + +```bash +ruff check . +black . +``` + +## Testing conventions + +- Tests subclass `DTensorTestBase` in `test/common_dtensor.py`, which extends torch's `MultiProcessTestCase`: each test **spawns `world_size` (4) processes** and requires the `@with_comms` decorator. Backend is NCCL if ≥2 CUDA GPUs are visible, else gloo on CPU (`test/dtensor/cpu_only/` tests are plain CPU-only `unittest` with a `ProcessPoolExecutor`). +- `scripts/run_test.sh` `pkill`s stray pytest processes between files because spawned distributed processes can hang the suite — if a local pytest run appears stuck, leftover workers are the likely cause. +- Multi-GPU tests are skipped (not failed) when GPUs are unavailable; a green CPU run does not exercise NCCL paths. + +## Architecture + +`vescale.dtensor` is a **minimally-patched fork of PyTorch 2.7's `torch.distributed.tensor`**, not a reimplementation. The guiding pattern throughout: subclass torch's class, override only what `RaggedShard` breaks, and delegate everything else back to torch. Each modified file documents its changes vs. torch in a module-level docstring comment — read those before editing. + +The extension points: + +- **`_api.py` — `DTensor(TorchDTensor)`**: carries a vescale `OpDispatcher` as the class attribute `_op_dispatcher`; `__torch_dispatch__` routes through it so all aten ops land in vescale code. Factory funcs (`distribute_tensor`, `from_local`, `redistribute`, `to_local`) are re-pointed to vescale's versions so they return vescale DTensors. +- **`_dispatch.py` — `OpDispatcher(TorchOpDispatcher)`**: merges torch's custom op handlers with vescale ones (fused AdamW/SGD handler, ragged norm handler, `found_inf` reduction for grad clipping). Adds `_cvt_dtensor` which runs before dispatch. +- **`_sharding_prop.py` — `ShardingPropagator(TorchShardingPropagator)`**: on init, merges torch propagator's `op_to_schema_info` / `op_to_rules` / `op_strategy_funcs` dicts with vescale's (vescale wins on conflict). Overrides `propagate_op_sharding_non_cached` to consult `is_ragged_shard` when adjusting shape/stride args. +- **`_redistribute.py` — `redistribute_local_tensor`**: fast path — if neither spec has a `RaggedShard`, calls torch's implementation directly. Otherwise handles three cases: ragged→ragged (single all-to-all via `RaggedShard._to_new_ragged_shard`), ragged→standard (all-gather via `_to_replicate_tensor`, then torch redistribute), standard→ragged (torch redistribute, then local `_split_tensor`). Ragged↔ragged is only supported within the same mesh dim. +- **`placement_types.py` — `RaggedShard(Placement)`**: shards a *contiguous flattened* tensor; `dims` = prefix dims collapsed together (must be `(0..k-1)`), `local_units` = relative element weights per rank (length == mesh dim size, sum must divide total numel). `_StridedRaggedShard` adds `split_factor` to compose with `Shard(0)` by recording a permutation (its `_to_replicate_tensor` is known-buggy per in-code TODO — avoid it in redistribute). +- **`vescale/utils/monkey_patch.py` — `patch_method`**: used to bolt `is_ragged_shard()` onto torch's own `Placement` base class so torch-side code can detect ragged placements without subclassing. + +Supporting machinery: + +- **`_ops/`** (pointwise/math/matrix/tensor): sharding rules registered via torch's `register_prop_rule` / `register_op_strategy` globals (re-exported in `_ops/utils.py`); vescale's `ShardingPropagator.__init__` picks them up, and `vescale/dtensor/__init__.py` force-imports `_ops` so registration happens on package import. Many files are near-copies of torch's, kept because the rules must output vescale `DTensorSpec`s and handle `RaggedShard`. +- **`vescale_utils/ragged_shard_utils.py`**: index arithmetic and spec helpers. Two idioms recur: `substitute_ragged_with_replicate` (swap the RaggedShard for `Replicate()` so torch's machinery handles the other mesh dims, then handle the ragged dim separately) and heavy `functools.lru_cache` on spec-derived computations (specs are hashable). +- **`vescale_utils/checkpoint.py`**: PyTorch DCP integration (`_break_ragged_box`, chunk/write-item builders) enabling communication-free checkpointing of ragged shards. + +Constraints enforced across the codebase: exactly one `RaggedShard` per placements tuple; RaggedShard does not support SymInts; cross-mesh redistribute is unimplemented. + +## Notes + +- `setup.py` reads the version from `vescale/__init__.py` (`__version__`); importing `vescale` prints the version. +- The root `Dockerfile` is a placeholder; `.clang-format` is a leftover for the legacy tree. The new `vescale/` code is pure Python. diff --git a/docs/texts/ndtimeline-design.md b/docs/texts/ndtimeline-design.md new file mode 100644 index 0000000..3a51e26 --- /dev/null +++ b/docs/texts/ndtimeline-design.md @@ -0,0 +1,273 @@ +# ndtimeline 设计文档 + +> **位置说明**：ndtimeline 模块位于 `legacy/vescale/ndtimeline/`，属于旧版 veScale（已整体迁入 `legacy/`，当前仓库根目录的活跃代码只有新版 DTensor）。本文档基于 legacy 代码分析。模块自带的简短 README 见 `legacy/vescale/ndtimeline/README.md`。 + +--- + +## 1. Motivation（设计动机） + +在千卡乃至万卡规模上训练 LLM 时，性能诊断面临三个传统工具解决不了的问题： + +1. **Straggler（掉队机器）难以定位**：个别机器发起 CUDA kernel 的速度变慢（"sink machines"），会拖慢整个集合通信，显著降低训练效率。要发现这类问题，必须**跨机器**对比同一时间窗口内各 rank 的行为。 +2. **传统工具只能看单机**：torch profiler / nsys 只采集单进程或单机内的信息，而训练的性能瓶颈往往发生在成百上千台机器之间的通信上，单机视图无法呈现全局画面。 +3. **全量 profiling 开销不可接受**：torch profiler 虽然细节丰富，但常开代价太高，且产生的 trace 文件巨大、难以分析，不适合生产环境长期开启。 + +因此需要一个满足以下目标的机制： + +- **跨机器联合采集**：所有 rank 的 span 能放在同一个时间坐标系下对比； +- **低开销、可常开**：采集路径对训练关键路径的影响要足够小； +- **trace 文件小且易分析**：只记录感兴趣的粗粒度 span（前反向计算、P2P 收发、集合通信），而不是全量 kernel trace； +- **与并行语义一致、易扩展**：采集器位于每个训练进程内，天然继承每个 rank 的并行拓扑语义（TP/DP/PP/DDP），新增并行方式时无需改动采集核心。 + +## 2. Problem Definition（问题定义） + +ndtimeline 要解决的核心问题可以形式化为： + +> 在一个由 N 个进程（分布在多台物理机上）、以 TP × DP × PP（× DDP）方式组织的分布式训练任务中，以可常开的低开销，为每个 rank 上选定的操作记录形如 `(metric_name, absolute_start_ts, duration, tags, step)` 的 span，并使得**不同机器上记录的 span 可以直接放到同一个时间轴上对齐可视化**，最终产出可用 Chrome Trace / Perfetto 打开的 trace 文件。 + +关键约束： + +- **时间对齐**：各机器墙钟（wall clock）无法精确同步，但产出的 span 必须能跨机器比较先后与重叠关系； +- **低开销**：记录的插桩点（start/stop）在训练热路径上，导出（事件同步 + 序列化 + 写盘）不能阻塞训练线程； +- **语义正确**：通信类操作实际运行在 NCCL 内部持有的 CUDA stream 上，事件必须记录到**同一条 stream** 才能量到真实耗时； +- **不破坏训练**：作为旁路观测系统，任何采集/导出异常都不得导致训练崩溃或 hang； +- **每机聚合**：同一台物理机上的多个 rank 的 span 应合并进同一个 trace 文件，减少文件数量并便于单机内对比。 + +系统对外提供两类 API： + +- **高层 API**（`ndtimeline.api`）：`init_ndtimers`（初始化）、`flush`（导出并校准全局时钟）、`wait`（优雅退出）、`set_global_step` / `inc_step`（维护全局 step）； +- **低层 API**：`NDTimerManager` 及 `ndtimeit` / `ndtimeit_p2p` / `ndtimer` 装饰器等插桩原语，供框架内部（如 `vescale.pipe`）使用。 + +## 3. Implementation Challenges（实现挑战） + +### 3.1 跨机器时间轴对齐（核心难点） + +机器间的 Unix 时钟存在漂移，NTP 精度不足以判断毫秒级的通信重叠关系。ndtimeline 的解法是**模拟全局时钟**（`timer.py: GlobalReferenceTime`）： + +- 初始化时做一次 `all_gather`，收集所有 rank 的当前 Unix 时间戳（ns），取**全局最小值** `initial_min_clock` 作为模拟时钟的纪元原点。由于 barrier + all_gather 本身是一个同步点，各 rank 时间戳的误差被限制在通信延迟内； +- 每个 rank 维护两个参考 CUDA event（round-robin 双缓冲）。`calibrate()` 时在默认 stream 上记录参考 event 并同步，记录此刻的 CPU 墙钟 `last_calibrated_at`； +- **CPU–GPU 时钟漂移修正**：相邻两次校准之间，比较两个参考 event 的 GPU `elapsed_time` 与 CPU `time.time()` 差值，得到残差系数 `gpu_clock_residual_coef = 1 - (gpu_time - cpu_time) / gpu_time`； +- 任意事件的绝对时间戳 = `gpu_time(ref → event) * gpu_clock_residual_coef + last_calibrated_at`。这样每个 GPU event 都被换算到"以全局最小时钟为原点、经漂移修正的墙钟"上，跨机器的 span 即可直接对齐。 + +挑战在于：校准本身需要同步（有开销），因此 `flush` 提供 `dynamic_calibrate`（至少每 30 分钟校准一次）与 `sequential_calibrate`（在主线程还是 flush 线程中校准）来控制精度与开销的平衡。 + +### 3.2 采集开销控制 + +- **CUDA Event 对象池**（`pool.py`）：`torch.cuda.Event` 的创建/销毁有成本且涉及 CUDA API。`CudaEventPool` 预分配 20 个 event（warmup 后入池），`get`/`release` 加锁复用；池耗尽时才新建； +- **异步 flush**（`timer.py: NDTimerManager.async_flush`）：flush 时把当前 timer 字典整体"换出"（用 meta 的拷贝重建一套新 timer 继续采集），旧 timer 交给 `ThreadPoolExecutor`（默认 3 worker）在后台做事件同步、耗时计算和 handler 派发。训练线程只做一次字典交换（期间临时 `gc.disable()` 避免 GC 触发 CUDA event 析构）； +- **带外进程聚合**（见 3.4）：写文件等重活完全移出训练进程； +- **级别过滤**（`NDMetricLevel`）：每个 timer 带级别（FRAMEWORK_INFO/USER_INFO/INFO/…/FRAMEWORK_TRACE/USER_TRACE/TRACE，有序枚举），低于 manager 级别的 metric 在 `start`/`stop` 入口直接返回，零开销； +- **未注册容错**：对未注册名字的 start/stop 只警告一次（记入 `_unregistered_timer_*` 列表后跳过），不反复打日志。 + +### 3.3 获取 NCCL 通信所用的 CUDA stream + +要度量一次 P2P/集合通信的真实 GPU 耗时，CUDA event 必须记录到该通信实际运行的 stream 上。但 PyTorch 原生**不暴露** NCCL 使用的 stream——`ProcessGroupNCCL` 内部维护一个 `ncclStreams_` 池，按需懒分配。 + +ndtimeline 的解法是**给 PyTorch 打补丁**（`legacy/patches/patched_pytorch_v2.2.1_rc3.patch`），在 `ProcessGroupNCCL` 上新增两个 C++ 接口并绑定到 Python： + +- `getP2PCudaStreamId(device_idx, peer, is_batched)`：batched 模式按 device key 查 `ncclStreams_`，非 batched 模式按 `getKeySendRecv(rank, peer)` 查；找到则返回 `stream.id()`，否则返回 -1； +- `getCollCudaStreamId(tensors)`：按 tensor 所在 device 的 key 查，返回 collective 使用的 stream id。 + +`stream.py` 中的 `get_nccl_p2p_stream` / `get_nccl_coll_stream` 把裸 stream id 包装成 `torch.cuda.Stream` 对象，并按 metric 名缓存到全局 `NCCL_STREAMS` 字典（每个 metric 固定对应一条 stream）。 + +**重要 caveat（代码与 README 均有强调）**：NCCL stream 是懒分配的，**第一次通信操作发生时 stream 才存在**，因此第一次通信必然查不到 stream（返回 -1，日志打印 `p2p stream is not available, skipped`），首个通信 span 会丢失，这是预期行为。这也意味着该补丁是 ndtimeline 通信度量的硬依赖——在未打补丁的 torch 上，`hasattr(nccl_backend, "get_p2p_cuda_stream_id")` 为假，通信类 metric 将全部退化为无 stream 记录。 + +### 3.4 单机聚合与训练进程减负 + +每个训练进程只承担"记录 + 序列化 + 发 socket"的轻活，聚合与落盘由独立的 **streamer 子进程**完成（`sock_streamer.py` + `handlers/sock_handler.py` + `binary_protocol.py`）： + +- 每台机器上由 `local_rank == 0` 的进程通过 `torch.multiprocessing.spawn`（spawn 方式、daemon）拉起一个 streamer 进程； +- streamer 进程在 Unix domain socket（默认 `/opt/tiger/tmp/ndtimeline/ndtimeline.sock`）上跑 `ThreadingUnixStreamServer`，每连接一线程收包，投递进容量 500000 的 `queue.Queue`，再由消费线程调用 post handler（默认 `LocalTimelineNDHandler`）； +- 训练进程侧，`SockNDHandler` 作为 pre handler 在 flush 线程中把每条 metric 序列化后 `sendall` 到该 socket；连接失败带退避重试，server 退出后置 `server_exited` 不再发送——采集失败绝不影响训练； +- 传输用自定义**二进制封包协议**（`binary_protocol.py`）： + + ``` + +---------------------------------------------------------------+ + | Magic Number 1Byte (0x9C) | Protocol Version 1Byte | Reserved 2Byte | + +---------------------------------------------------------------+ + | Payload Length 4Byte (little endian) | + +---------------------------------------------------------------+ + | Payload (pickle protocol 4, 上限 128MiB) | + +---------------------------------------------------------------+ + | EOF Symbol 1Byte ('\n') | + +---------------------------------------------------------------+ + ``` + + 接收端 `recv_and_validate` 逐字段校验 magic/version/reserved/长度/EOF，支持粘包处理（`preload_data` 缓冲跨包残留字节）。 + +### 3.5 Metric 与训练 step 的关联 + +一个 metric 在一次 flush 周期内可能产生多条 span（例如一个 step 内多次 microbatch 的 send），需要把 span 归属到正确的 global step： + +- **方式一（推荐）**：timer 注册时提供 `step_getter`，每条 span 记录时把当前 step 写进 tag（键 `_inner_global_step`）。`parse_record` 用 `itertools.groupby` 按 step 分组生成 record； +- **方式二**：不提供 `step_getter` 时，`flush` 传入 `step_range`，`parse_record` 要求 span 数量能被 step 数整除后**均分**到各 step（隐含"每 step 操作数相同"的假设，否则报 `NDHandlerError`，提示改用 `step_getter`）； +- 高层 API 还提供 `set_global_step` / `inc_step`，配合 manager 内部的 nature step getter 使用，免去用户写回调。 + +### 3.6 健壮性与资源管理 + +- 所有插桩入口（`start_timer`/`stop_timer`/`insert_record`）和 handler 调用都被 try/except 包裹，异常只记日志； +- `wait()` 先做 `dist.barrier()`，再等 flush futures（最多 10s），最后 sleep 3s 给 streamer 排空消息； +- 若干反直觉的资源处理均有注释说明：`_flush_timers` 末尾 `timer.meta = None` 防止 CudaTimer 因 meta 循环引用导致 event GC 失败；`LocalTimelineNDHandler` 切分缓冲时用 `record.copy()` 而非直接持有引用，同样是规避 GC 导致的内存泄漏；`loads`（pickle 反序列化）前后 `gc.disable/enable`； +- 双缓冲参考 event 使校准与正常记录互不踩踏。 + +## 4. Implementation Details（实现细节） + +### 4.1 代码结构 + +``` +legacy/vescale/ndtimeline/ +├── __init__.py # 包导出（try import _internal 兼容内部版） +├── api.py # 高层 API：init_ndtimers / flush / wait / set_global_step / inc_step +├── timer.py # 核心：GlobalReferenceTime、NDMetricLevel、DeviceTimerMeta、 +│ # DeviceTimer、NDTimerManager(Singleton)、ndtimeit 系列 +├── pool.py # CudaEventPool / DefaultEventPool（事件对象池） +├── stream.py # get_nccl_p2p_stream / get_nccl_coll_stream（依赖 patched torch） +├── sock_streamer.py # NDtimelineStreamer：streamer 子进程 + UDS server + 消费队列 +├── binary_protocol.py # 二进制封包/拆包协议 +├── world_info.py # WorldInfo = TopoInfo(rank/tp/dp/pp/ddp…) + TrainingInfo(role/trial/run) +├── predefined.py # 预定义 metric 名（forward-compute、cross-mesh-send 等） +├── variables.py # 常量：socket 路径/超时、event 池大小、tag 键、flush 特殊标记 +├── logger.py # NDTimelineLogger 单例（env VESCALE_NDTIMELINE_LOG_LEVEL 控制级别） +├── exceptions.py # ProtocolValidationError / NDHandlerError +├── is_internal.py # is_internal_vescale()：内部版/开源版能力门控 +├── fsdp_patch.py # patch_fsdp()（仅内部版有实现） +└── handlers/ + ├── handler_base.py # NDHandler 抽象基类（dispatch_key、ignore_metrics、call_impl） + ├── chrome_trace_event.py # Chrome Trace Event Format 数据类（X/B/E/s-t-f/C/M 事件） + ├── parser_handler.py # parse_record + DeviceTimerStreamRecord（span → 结构化 record） + ├── local_timeline_handler.py # 单机聚合 → trace_step{min}_{max} JSON 文件 + ├── sock_handler.py # SockNDHandler（训练进程侧 UDS 客户端，默认 pre handler） + ├── logging_handler.py # 日志摘要 handler + ├── local_raw_handler.py # 原始数据落盘（RotatingFileHandler，128MiB × 8） + └── do_nothing_handler.py # 空实现（占位/测试） +``` + +### 4.2 数据流（端到端） + +``` +训练热路径 flush 线程（ThreadPoolExecutor） streamer 子进程（local_rank 0 拉起） +───────── ────────────────────────────── ──────────────────────────────────── +ndtimeit_p2p(name, pg, peer) + ├─ get_nccl_p2p_stream ────▶ 取 NCCL 内部 stream（patched torch） + ├─ timer.start(): + │ event_pool.get() → event.record(stream) + └─ timer.stop(): + event.record(stream) + ……训练继续…… + api.flush(step_range) + └─ async_flush: + 换出新 timer 字典（gc.disable 保护） + submit _flush_timers ──────────────▶ + timer.elapsed(): + event.synchronize() + duration = start.elapsed_time(stop) + abs_ts = GlobalReferenceTime.elapsed_time(start) + event 归池；tag 按 legal_tags 过滤 + 对每个 handler 调用（dispatch_mode 过滤）: + SockNDHandler: + pickle + encode_package + sendall → UDS ──────────────▶ MsgHandler 收包校验 + q.put(args) + 消费线程: + LocalTimelineNDHandler + parse_record → record 入 rank 缓冲 + 发送 NDTIMELINE_FLUSH_SEPCIAL ───▶ 收到所有 local rank 的 special 后: + dump_records() → Chrome Trace JSON + （pid=dp_rank, tid=rank*10+类别序号, + ts/dur 秒→微秒） +``` + +### 4.3 关键类与机制要点 + +- **`GlobalReferenceTime`**（类级状态的单例式工具类）：`init` 预热两个参考 event 并做首次 `calibrate`（含 all_gather 取 `initial_min_clock`）；`elapsed_time(end_event)` 是 span 绝对时间戳的唯一来源。注意它假设所有 rank 同时调用 `calibrate`（首次有 barrier），因此校准只能通过 `flush` 路径集体触发。 +- **`DeviceTimer`**：一个 timer 可同时持有 GPU event pairs（`_event_pairs`）和 CPU 记录（`_extra_records`，`is_cpu_op=True` 时用 `time.time()`）。`elapsed()` 会把两类记录按起始时间**归并排序**成统一序列返回 `(total_elapsed, durations, since_starts, tags)`。带 stream 的 start 会先 `wait_stream(default_stream)`，保证事件捕获到从默认 stream fork 出的通信。 +- **`NDTimerManager.async_flush` 的换出机制**：`past_name2timer` 携带已有事件去后台处理；`fresh_name2timer` 由 meta 拷贝重建，训练线程立即切到新字典继续 start/stop——这是"记录与导出解耦"的关键。`next_iter_enabled` / `keep_timer_state` 决定换出后 timer 是否继续使能。 +- **Handler 派发**：`NDHandler.dispatch_key` 默认为类名；timer 的 `dispatch_mode="selected"` + `dst_names` 可让某 metric 只发给指定 handler；`specified_extra[dispatch_key]` 提供按 handler 区分的附加字段。`ignore_metrics` 默认忽略 `special` 标记，`LocalTimelineNDHandler` 特意清空该忽略列表以接收 flush 分界信号。 +- **`LocalTimelineNDHandler` 的分段落盘**：为每个 local rank 维护 `rank2steps`（deque of step-set）。每收到一个 rank 的 `special` 就追加一个新 set；当所有 rank 的 deque 长度 ≥ 2，说明上一段已收齐，把属于旧段的 record 落盘成 `trace_step{min}_{max}` 文件、缓冲中保留新段。落盘时按 metric 名查 `thread_sort_index` 表（main / pp → / pp ← / pp send-recv / tp collective / dp collective / other）决定 trace 中的线程分组，`tid = rank*10 + sort_index`（乘 10 给类别编号让位），`pid = dp_rank`。 +- **WorldInfo**：`TopoInfo` 携带 rank、local_rank、tp/pp/dp/ddp 的 rank 与 size、world_size、ip，并做非负/正数校验；`TrainingInfo` 携带 role_id（RL 多角色）、trial_id、run_id。所有 record 都带这些维度，这就是"N-Dimension"的含义——trace 可按任意并行维度过滤/对比。 +- **内外部版本门控**：`is_internal_vescale()`（能否 import `_internal`）决定 `fsdp` 模式、`MQNDHandler`（发到内部消息队列的 post handler）、`patch_fsdp` 等能力是否可用。开源版中 `init_ndtimers(mode="fsdp")` 直接抛 `NotImplementedError`；`fsdp` 模式的预置 timer（unshard-all-gather / grad-reduce-scatter / grad-all-reduce）在开源版也不会被注册。 +- **与训练框架的集成点**：`vescale/pipe/p2p_communication.py` 用 `ndtimeit_p2p` 包裹 `_communicate`（RECV_FORWARD/RECV_BACKWARD/SEND_FORWARD/SEND_BACKWARD/SEND_FORWARD_RECV_BACKWARD/SEND_BACKWARD_RECV_FORWARD）；`vescale/pipe/_schedules/pipedream_flush.py` 与 `zero_bubble_v.py` 度量 CROSS_MESH_RECV/SEND，并用 `ndtimer` 装饰器度量 FORWARD_COMPUTE/BACKWARD_COMPUTE。 + +### 4.4 使用方式 + +```python +# 必须在 torch.cuda.set_device 和默认进程组初始化之后调用 +from vescale.ndtimeline import init_ndtimers, flush, wait, inc_step + +init_ndtimers( + mode="hybrid", # 开源版仅支持 hybrid（从 VESCALE_DEVICE_MESH 读 TP/DP/PP） + step_getter=lambda: cur_step, # 或用 inc_step()/set_global_step() 维护 + enable_streamer=True, # fork streamer 子进程；默认 post handler 为 LocalTimelineNDHandler + level=NDMetricLevel.DEBUG, # 低于该级别的 metric 不记录 +) + +# ……训练循环……（框架插桩点自动记录；用户可用 ndtimeit/ndtimer 追加自定义 metric） +for step in range(n): + train_one_step() + inc_step() + if step % 10 == 0: + flush(step_range=range(step - 10, step)) # 异步导出 + 可选校准 + +wait() # 训练结束时等待 flush 完成 +``` + +产出的 `trace_step{min}_{max}` 文件是 Chrome Trace Event Format JSON，可直接拖入 https://ui.perfetto.dev/ 查看。若需跨机器汇总，可自行实现一个 post handler（如 MQHandler）把所有机器的 metric 送到中心存储后再生成 trace（内部版即通过 `MQNDHandler` 实现）。 + +## 5. How to Test This Feature（测试方法） + +### 5.1 现有单元测试（无需 GPU） + +单测位于 `legacy/test/ndtimeline/`： + +| 测试文件 | 覆盖内容 | +|---|---| +| `test_metric_level.py` | `NDMetricLevel` 有序枚举的比较语义（INFO/DEBUG/TRACE 三档及 FRAMEWORK_/USER_ 前缀的大小关系） | +| `test_parser_handler.py` | `parse_record`：带 tag 的正常输入生成 record 且 step 正确；三个数组长度不一致时抛 `NDHandlerError` | +| `test_local_raw_handler.py` | `LocalRawNDHandler` 的 RotatingFileHandler 行为（按 chunk 滚动、backup 数量上限） | + +运行方式（legacy 测试需从 `legacy/test/` 目录运行并把该目录加入 `PYTHONPATH`，因为测试间以顶层模块方式互相 import）： + +```bash +cd legacy +pip3 install -r requirements.txt && pip3 install -e . + +cd test +export PYTHONPATH=$(pwd):$PYTHONPATH + +# 运行整个 ndtimeline 单测目录 +pytest ndtimeline/ -s + +# 运行单个测试文件 / 单个用例 +pytest ndtimeline/test_parser_handler.py -s +pytest ndtimeline/test_parser_handler.py::test_normal_input_with_tags -s +``` + +注意：`test_local_raw_handler.py` 会在 `/opt/tiger/tmp/ndtimeline/`（`variables.py` 中的 `LOCAL_LOGGING_PATH`）下创建并清理滚动日志文件，需要该路径可写。 + +### 5.2 集成测试（需要多 GPU） + +`legacy/test/parallel/pipeline/instruction/test_schedule.py` 中的流水线调度测试真实调用了 `init_ndtimers → flush → wait` 全流程（在 4 卡 PP 场景中随 `ScheduleEngine` 运行）。按 legacy 套件的方式运行： + +```bash +cd legacy/test +export PYTHONPATH=$(pwd):$PYTHONPATH +export VESCALE_SINGLE_DEVICE_RAND=1 +pytest parallel/pipeline/instruction/test_schedule.py -s -k +``` + +（`legacy/scripts/run_test.sh` 会跑整个 legacy 测试集，并在每个文件前后 `pkill python3` 清理分布式测试的残留进程——本地跑多进程测试卡住时，残留进程通常是原因。） + +### 5.3 手工端到端验证 + +最小验证路径（需要至少 1 台多卡机器、patched torch 才能看到通信 span）： + +1. 在任意多进程训练脚本中（`torch.cuda.set_device` 与 `dist.init_process_group` 之后）调用 `init_ndtimers(enable_streamer=True)`，训练若干 step，周期性 `flush(step_range=...)`，结束时 `wait()`； +2. 确认日志出现 `ndtimeline streamer started` 与各 rank 的 `calibrate ...` 信息；预期看到若干 `p2p stream is not available, skipped`（首个通信 span 丢失，属预期）； +3. 在 cwd 下得到 `trace_step{min}_{max}` 文件，拖入 Perfetto，检查：每个 dp_rank 一个进程行、每个 rank 按 main/pp/tp/dp 分线程行、forward-compute 与 send/recv span 的时间先后关系符合调度预期； +4. 验证低开销路径：设置 `level=NDMetricLevel.INFO`（高于预置 timer 的默认 FRAMEWORK_DEBUG），确认对应 span 不再出现——级别过滤生效。 + +### 5.4 针对子系统的新测试建议 + +- **协议层**：可构造 payload 直接测试 `encode_package`/`recv_and_validate` 的粘包与坏 magic/坏 EOF/超长 payload 分支（纯 CPU，适合单测）； +- **时钟换算**：`GlobalReferenceTime` 依赖真实 CUDA，可在单卡上验证 `elapsed_time` 单调性与 `gpu_clock_residual_coef` 的合理范围； +- **handler 扩展**：新写一个 `NDHandler` 子类时，参照 `ParserNDHandler` 的单测方式直接构造 `call_impl` 入参验证，无需启动分布式环境。 diff --git a/docs/texts/ndtimeline-implementation-details.md b/docs/texts/ndtimeline-implementation-details.md new file mode 100644 index 0000000..7f08790 --- /dev/null +++ b/docs/texts/ndtimeline-implementation-details.md @@ -0,0 +1,707 @@ +# ndtimeline 实现细节(API、接入、Tutorial 与实现原理) + +> 本文档是 [`ndtimeline-design.md`](./ndtimeline-design.md) 的深入补充,聚焦**怎么用**和**怎么实现的**。设计文档侧重动机、问题定义与挑战;本文档侧重 API 参考、接入步骤、可运行的 PyTorch tutorial,以及逐组件的实现原理剖析。 +> +> 代码位置:`legacy/vescale/ndtimeline/`(属于旧版 veScale,当前仓库活跃代码为新版 DTensor)。 + +--- + +## 目录 + +1. [API 定义](#1-api-定义) +2. [如何接入 ndtimeline](#2-如何接入-ndtimeline) +3. [PyTorch Tutorial(可运行)](#3-pytorch-tutorial可运行) +4. [实现原理详解](#4-实现原理详解) +5. [局限与注意事项](#5-局限与注意事项) + +--- + +## 1. API 定义 + +ndtimeline 的 API 分两层:**高层 API**(`ndtimeline/api.py`,面向最终用户,封装好默认行为)和**低层 API**(`timer.py` 等,面向框架开发者,完全可控)。高层 API 的 `init_ndtimers` 依赖 veScale 的全局 device mesh(`VESCALE_DEVICE_MESH`),因此**纯 PyTorch 场景应直接使用低层 API**(见 §3)。 + +### 1.1 高层 API(`ndtimeline.api`) + +#### `init_ndtimers(...)` + +一次性完成:拉起 streamer 子进程、注入默认 handler、构造 `NDTimerManagerSingleton`、注册预置 timer。 + +```python +def init_ndtimers( + rank: Optional[int] = None, # 全局 rank,None 时用 dist.get_rank() + mode: Literal["fsdp", "hybrid"] = "hybrid", # 并行模式;fsdp 仅内部版支持 + wrapped_fsdp_module=None, # fsdp 模式用(内部版) + device_mesh=None, # fsdp 模式用(内部版) + mesh_shape: Optional[Tuple[int, ...]] = None, # fsdp 模式用(内部版) + local_rank: Optional[int] = None, # None 时用 VESCALE_DEVICE_MESH.get_local_rank() + step_getter: Optional[Callable[[], int]] = None, # 返回当前 global step 的回调;None 时用 set_global_step/inc_step + enable_streamer: bool = True, # 是否 fork streamer 子进程以启用 post handler + n_rank_per_host: Optional[int] = None, # 单机 GPU 数,None 时用 torch.cuda.device_count() + pre_handlers: Optional[List[NDHandler]] = None, # 训练进程内、flush 线程触发的 handler + post_handlers: Optional[List[NDHandler]] = None, # streamer 进程内触发的 handler + user_spcified_timers: Optional[List[DeviceTimerMeta]] = None, # 用户自定义 timer + level: NDMetricLevel = NDMetricLevel.DEBUG, # 级别阈值,低于它的 metric 不记录 + ip: str = "0.0.0.0", # 本机 IP + **kwargs, +) +``` + +**调用前提(硬性要求)**:必须在 `torch.cuda.set_device(...)` 和默认进程组 `dist.init_process_group(...)` **之后**调用。 + +**默认行为**: +- `mode="hybrid"`:从 `VESCALE_DEVICE_MESH` 读取 TP/DP/PP 的 size 与 rank,并断言 `tp*dp*pp == world_size`。 +- `mode="fsdp"`:仅内部版;会 `patch_fsdp()`,并从 process group / device_mesh / mesh_shape 推算 dp/ddp 拓扑。 +- `enable_streamer=True` 时: + - `local_rank == 0` 的进程拉起 streamer 子进程;若未给 `post_handlers`,默认注入 `LocalTimelineNDHandler(n_rank_per_host)`(在本机生成 Perfetto trace 文件)。 + - 若 `pre_handlers` 里没有 `SockNDHandler`,自动追加一个(负责把 metric 通过 Unix socket 发给 streamer)。 +- `hybrid` 模式注册的预置 timer:`forward-compute`、`backward-compute`、`forward/backward-recv`、`forward/backward-send`、`forward-send-backward-recv`、`backward-send-forward-recv`、`cross-mesh-recv/send`(均为 GPU op,`is_cpu_op=False`)。 +- `fsdp` 模式注册的预置 timer:`unshard-all-gather`、`grad-reduce-scatter`、`grad-all-reduce`、`forward-compute`、`backward-compute`。 + +#### `flush(...)` + +把已采集的 metric 异步导出到 handler,并可选地校准全局时钟。 + +```python +def flush( + step_range: Optional[range] = None, # 这批 metric 归属的 step 范围;给了 step_getter 时可忽略 + next_iter_enabled: bool = True, # flush 后 timer 是否继续使能 + submit2handler: bool = True, # False 表示丢弃 metric(只清空不导出) + dynamic_calibrate: bool = False, # 是否"至少每 30 分钟校准一次时钟" + keep_timer_state: bool = False, # True 时保持各 timer 现有 enable/disable 状态(忽略 next_iter_enabled) + sequential_calibrate: bool = True, # True=在主线程校准;False=放到 flush 后台线程校准 +) +``` + +#### `wait()` + +优雅退出:先 `dist.barrier()`,再等所有后台 flush future(最多 10s),最后 `sleep(3)` 给 streamer 排空消息。**应在训练结束时(`main.py` 末尾)调用**。 + +#### `set_global_step(step)` / `inc_step(step=1)` + +当 `step_getter=None` 时,用这两个 API 手动维护 manager 内部的 global step。每条 span 记录时会把当前 step 写进 tag。 + +### 1.2 低层 API(`timer.py`) + +#### `NDTimerManager` / `NDTimerManagerSingleton` + +`NDTimerManagerSingleton` 是 `NDTimerManager` 的单例子类,**所有 context manager(`ndtimeit` 等)都作用于这个单例**。 + +```python +NDTimerManagerSingleton( + world_info: WorldInfo, + handlers: Optional[List[Callable]] = None, # pre handler 列表 + max_workers: int = 3, # 后台 flush 线程池大小 + device_id: Optional[int] = None, # None 时用当前 CUDA device + init_cuda_dist: bool = True, # True 时初始化 event 池 + 全局时钟(含集合通信) + metric_level: NDMetricLevel = NDMetricLevel.TRACE, # 级别阈值 + is_nature_step: bool = True, # True 时,无 step_getter 的 timer 自动用 manager 内部 step +) +``` + +核心方法: + +| 方法 | 说明 | +|---|---| +| `register_timers(metas: List[DeviceTimerMeta])` | 从 meta 列表创建 `DeviceTimer`。**只能调用一次**(内部断言) | +| `start_timer(name, tag=None)` | 开始计时;`tag` 里可带 `NDTIMELINE_STREAM_KEY` 指定 stream | +| `stop_timer(name, tag=None)` | 停止计时 | +| `insert_record(name, start_ts, duration, tag=None)` | 直接插入一条已完成记录(不经 event),用于补录 | +| `async_flush(step_range, ...)` | 异步导出(见 §4.5) | +| `wait()` | 等待后台 flush 完成 | +| `clear()` | 丢弃当前所有未导出记录 | +| `disable_and_save()` / `recover_from_history()` | 临时禁用全部 timer 并恢复(期间关 GC) | +| `global_step`(property) | 读/写当前 step | +| `is_initialized()`(classmethod) | 单例是否已构造 | + +#### `DeviceTimerMeta`(timer 的声明式配置) + +```python +@dataclass +class DeviceTimerMeta: + name: str = "" + is_cpu_op: bool = False # True=用 time.time() 计 CPU 耗时;False=用 CUDA event + legal_tags: List[str] = [] # 允许保留进记录的 tag 键白名单 + step_getter: Optional[Callable] = None # 返回 step 的回调 + enabled: bool = True + level: NDMetricLevel = NDMetricLevel.FRAMEWORK_DEBUG + device_id: int = -1 # -1 时:GPU op 用当前 device,CPU op 用 0 + dispatch_mode: Literal["selected", "all"] = "all" # selected=只发给 dst_names 指定的 handler + dst_names: List[str] = [] # dispatch_mode="selected" 时的目标 handler dispatch_key + specified_extra: Dict[str, Any] = {} # 按 handler dispatch_key 区分的附加字段 + common_extra: Dict[str, Any] = {} # 对所有 handler 的附加字段 +``` + +#### `NDMetricLevel`(级别过滤) + +有序枚举,值越大越"啰嗦"。`timer.level > manager.metric_level` 时该 metric 在 `start/stop` 入口直接返回(零开销)。 + +``` +FRAMEWORK_INFO(2) < USER_INFO(3) < INFO(4) +< FRAMEWORK_DEBUG(12) < USER_DEBUG(13) < DEBUG(14) +< FRAMEWORK_TRACE(102) < USER_TRACE(103) < TRACE(104) +``` + +预置 timer 默认 `FRAMEWORK_DEBUG`;manager 默认阈值 `DEBUG`(低层)/`DEBUG`(高层),所以默认能收到预置 metric。 + +#### 插桩原语(context manager / 装饰器) + +```python +ndtimeit(name: str, tag: Optional[Dict] = None) + # 通用计时:在默认/当前 stream 上记录 start/stop event。可重入。 + +ndtimeit_p2p(name, nccl_pg, peer, is_batched=True, tag=None) + # P2P 通信计时:先通过 patched torch 取 NCCL p2p stream,把 stream 塞进 tag, + # 再 start/stop。事件记录到真实通信 stream 上,才能量到真实耗时。 + +ndtimeit_coll(name, pg, tensor, tag=None) + # 集合通信计时:取 NCCL coll stream(依赖 patched torch),同上。 + # 注意:此函数未在包级 __init__ 导出,需 from vescale.ndtimeline.timer import ndtimeit_coll + +ndtimer(metric: str, tags: Optional[Dict] = None) + # 函数装饰器,等价于用 ndtimeit 包住整个函数体。 +``` + +未初始化单例时,上述 context manager 都是**空操作**(直接 `yield`),因此插桩代码可以在不开启 ndtimeline 时安全保留。 + +### 1.3 Handler API(`handlers/`) + +#### `NDHandler`(抽象基类) + +```python +class NDHandler(ABC): + def __init__(self, designated_key="", ignore_metrics=None): + # dispatch_key 默认为类名;ignore_metrics 默认忽略特殊 flush 标记 NDTIMELINE_FLUSH_SEPCIAL + + def __call__(self, metric_name, elapsed, recent_elapsed_raw_parts, + recent_since_start_raw_parts, tags, step_range, world_info, extra): + # 过滤 ignore_metrics 后调用 call_impl + + @abstractmethod + def call_impl(self, metric_name, elapsed, recent_elapsed_raw_parts, + recent_since_start_raw_parts, tags, step_range, world_info, extra): ... +``` + +各参数含义: +- `metric_name`:metric 名(或特殊标记 `"special"`)。 +- `elapsed`:本周期该 metric 的总耗时(秒)。 +- `recent_elapsed_raw_parts`:每条 span 的耗时列表(秒)。 +- `recent_since_start_raw_parts`:每条 span 的**绝对起始时间戳**列表(秒,基于模拟全局时钟)。 +- `tags`:每条 span 的 tag(已经过 `legal_tags` 过滤)。 +- `step_range`:本批归属的 step 范围。 +- `world_info`:`WorldInfo`(拓扑 + 训练元信息)。 +- `extra`:timer 的附加字段。 + +#### 内置 handler + +| Handler | 运行位置 | 作用 | +|---|---|---| +| `SockNDHandler` | 训练进程(pre) | 把每条 metric 用二进制协议经 Unix socket 发给 streamer。**默认 pre handler** | +| `LocalTimelineNDHandler` | streamer 进程(post) | 聚合本机所有 rank,生成 Perfetto trace 文件 `trace_step{min}_{max}`。**默认 post handler** | +| `ParserNDHandler` | 任意 | 把原始入参解析成结构化 `DeviceTimerStreamRecord` 列表 | +| `LoggingNDHandler` | 任意 | 每个 metric 打一行日志(总耗时),最简单的验证手段 | +| `LocalRawNDHandler` | 任意 | 原始数据写滚动日志文件(128MiB × 8) | +| `DoNothingNDHandler` | 任意 | 空实现(占位/测试) | + +自定义 handler:继承 `NDHandler`,实现 `call_impl` 即可。若只想接收特定 metric,配合 timer 的 `dispatch_mode="selected"` + `dst_names=[你的 dispatch_key]`。 + +### 1.4 数据结构 + +#### `WorldInfo`(`world_info.py`) + +```python +WorldInfo( + rank, local_rank, # 必填 + dp_rank=0, ddp_rank=0, tp_rank=0, pp_rank=0, + dp_size=1, ddp_size=1, tp_size=1, pp_size=1, world_size=1, + ip="0.0.0.0", + role_id=0, run_id=0, trial_id=0, # 训练元信息(RL 多角色 / 试验 / 运行) + **extra_meta, +) +``` +内部拆成 `TopoInfo`(拓扑,带非负/正数校验)+ `TrainingInfo`(训练元信息)+ `extra_info`,支持 `world_info["rank"]` 式访问。 + +#### `DeviceTimerStreamRecord`(`handlers/parser_handler.py`) + +`parse_record` 的输出,一条结构化记录:`ts/rank/metric/step/avg_dur/start_ts[]/duration[]/tp_rank/pp_rank/dp_rank/ip/role_id/trial_id/run_id`。`to_dict()` 输出供下游(如内部 MQ)消费的字段。 + +### 1.5 关键常量(`variables.py`) + +| 常量 | 值 | 含义 | +|---|---|---| +| `SOCK_PARENT_DIR` | `/opt/tiger/tmp/ndtimeline` | streamer socket 与日志目录(**需可写**) | +| `SOCK_PATH` | `.../ndtimeline.sock` | Unix domain socket 路径 | +| `SOCK_TIMEOUT_CLIENT` | `2.0`s | socket 客户端超时 | +| `DEFAULT_CUDA_EVENT_POOL_SIZE` | `20` | event 池初始大小 | +| `NDTIMELINE_INNER_GLOBAL_STEP_KEY` | `_inner_global_step` | span 里 step 的 tag 键 | +| `NDTIMELINE_STREAM_KEY` | `stream_key` | tag 里指定 stream 的键 | +| `NDTIMELINE_FLUSH_SEPCIAL` | `special` | flush 分界特殊标记 | + +--- + +## 2. 如何接入 ndtimeline + +### 2.1 前提条件 + +1. **安装 legacy veScale**:`cd legacy && pip install -e .`。 + 注意:即使只用 ndtimeline,`import vescale.ndtimeline` 也会经由 `api.py` 传递性导入 `vescale.dtensor` 和 `vescale.devicemesh_api`,所以必须完整安装 legacy 包。 +2. **CUDA 环境**:计时核心基于 `torch.cuda.Event`,需要可用 GPU。 +3. **已初始化分布式**:默认进程组已建立(`dist.init_process_group`),且已 `torch.cuda.set_device`。全局时钟校准需要集合通信(barrier + all_gather)。 +4. **(可选,精确通信计时)patched PyTorch**:`ndtimeit_p2p`/`ndtimeit_coll` 依赖给 PyTorch 打的补丁(`legacy/patches/patched_pytorch_v2.2.1_rc3.patch`,新增 `ProcessGroupNCCL.get_p2p_cuda_stream_id`/`get_coll_cuda_stream_id`)。**未打补丁时**,这些函数取不到 NCCL stream,会退化为在默认 stream 上记录(量到的是"发起时刻"而非真实通信耗时),并打印 `p2p stream is not available, skipped`。**计算类计时不受影响**。 +5. **可写的 socket 目录**:默认 `/opt/tiger/tmp/ndtimeline` 需存在且可写(streamer 与 SockNDHandler 都用它)。 + +### 2.2 两条接入路径 + +**路径 A:legacy veScale 训练(hybrid/fsdp 模式)——用高层 API** + +veScale 的 pipeline / FSDP 框架内部已经在关键路径插好了桩(`pipe/p2p_communication.py` 用 `ndtimeit_p2p`,`pipe/_schedules/*.py` 用 `ndtimer`/`ndtimeit_p2p`)。用户只需: + +```python +# 1) 先初始化 veScale 全局 device mesh(含 DP/TP/PP 维度) +from vescale.devicemesh_api import VESCALE_DEVICE_MESH +VESCALE_DEVICE_MESH.init_device_mesh( + "cuda", mesh_shape=(dp, pp, tp), mesh_dim_names=("DP", "PP", "TP")) + +# 2) 初始化 ndtimeline(会自动读 device mesh、拉起 streamer、注册预置 timer) +from vescale.ndtimeline import init_ndtimers, flush, wait, inc_step +init_ndtimers(mode="hybrid", step_getter=lambda: global_step, enable_streamer=True) + +# 3) 训练循环中维护 step、周期性 flush +for step in range(n): + train_one_step() + inc_step() # 或依赖 step_getter + if step % FLUSH_INTERVAL == 0: + flush(step_range=range(step - FLUSH_INTERVAL + 1, step + 1)) + +# 4) 结束时优雅退出 +wait() +``` + +**路径 B:自定义 / 纯 PyTorch 框架——用低层 API** + +不依赖 veScale device mesh,自己构造 `WorldInfo` 和 `NDTimerManagerSingleton`,自己插桩。完整可运行示例见 §3。 + +### 2.3 插桩点选择 + +| 想度量 | 用什么 | 说明 | +|---|---|---| +| 一段 GPU 计算(forward/backward) | `with ndtimeit("forward-compute"):` | 在默认 stream 记录 | +| 一次 P2P send/recv | `with ndtimeit_p2p(name, pg, peer, is_batched):` | 需 patched torch 才精确 | +| 一次集合通信 | `ndtimeit_coll(name, pg, tensor)` | 需 patched torch 才精确 | +| 整个函数 | `@ndtimer("metric-name")` | 装饰器 | +| CPU 操作(数据加载等) | timer meta 设 `is_cpu_op=True`,再 `ndtimeit` | 用 `time.time()` 计时 | +| 已知起止时间的补录 | `manager.insert_record(name, start_ts, duration, tag)` | 不经 event | + +### 2.4 自定义 timer 与 handler + +```python +from vescale.ndtimeline import DeviceTimerMeta, NDMetricLevel +from vescale.ndtimeline.handlers import NDHandler + +# 自定义 timer:USER_DEBUG 级别,只发给指定 handler +metas = [DeviceTimerMeta( + name="my-metric", is_cpu_op=False, + level=NDMetricLevel.USER_DEBUG, + dispatch_mode="selected", dst_names=["MyHandler"], + common_extra={"source": "my_module"}, +)] +NDTimerManagerSingleton().register_timers(metas) # 注意:只能 register 一次 + +# 自定义 handler +class MyHandler(NDHandler): + def call_impl(self, metric_name, elapsed, parts, since_parts, tags, step_range, world_info, extra): + ... # 任意处理:写库、发 MQ、统计等 +``` + +### 2.5 关机顺序 + +`flush(...)`(把最后一批导出)→ `wait()`(barrier + 等后台 future + 给 streamer 3s 排空)。漏掉 `wait()` 可能丢失最后一批数据。 + +--- + +## 3. PyTorch Tutorial(可运行) + +下面是一个**自包含的 DDP 示例**,不依赖 veScale device mesh,直接用低层 API。它在 2 个进程(2 卡)上跑一个简单 MLP 的 DDP 训练,度量 forward / backward / 一次显式 allreduce,并在本机生成 Perfetto trace 文件。 + +### 3.1 完整代码 `ndtimeline_tutorial.py` + +```python +import os +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel as DDP + +# ---- ndtimeline 低层 API ---- +from vescale.ndtimeline import ( + NDTimerManagerSingleton, DeviceTimerMeta, NDMetricLevel, + ndtimeit, flush, wait, inc_step, +) +from vescale.ndtimeline.world_info import WorldInfo +from vescale.ndtimeline.handlers import SockNDHandler, LocalTimelineNDHandler +from vescale.ndtimeline.sock_streamer import NDtimelineStreamer + +# 预定义的 metric 名(也可直接用字符串) +FORWARD_COMPUTE = "forward-compute" +BACKWARD_COMPUTE = "backward-compute" +GRAD_ALLREDUCE = "grad-allreduce" +DATA_LOAD = "data-load" + + +def setup(): + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + return rank, local_rank + + +def init_ndtimeline(rank, local_rank): + world_size = dist.get_world_size() + n_rank_per_host = torch.cuda.device_count() + + # 1) 仅 local_rank 0 拉起 streamer 子进程,默认 post handler 生成本机 trace 文件 + NDtimelineStreamer.init( + local_rank, + [LocalTimelineNDHandler(n_rank_per_host=n_rank_per_host)], + ) + + # 2) 构造单例 manager。init_cuda_dist=True 会初始化 event 池 + 全局时钟 + # (内部含 barrier/all_gather,必须所有 rank 集体执行到这里) + world_info = WorldInfo( + rank=rank, local_rank=local_rank, + dp_rank=rank, dp_size=world_size, world_size=world_size, + ip="0.0.0.0", + ) + NDTimerManagerSingleton( + world_info, + handlers=[SockNDHandler()], # pre handler:经 socket 发给 streamer + init_cuda_dist=True, + metric_level=NDMetricLevel.DEBUG, # 阈值:收到 FRAMEWORK_DEBUG 及以上的 metric + ) + + # 3) 注册 timer(只能注册一次) + NDTimerManagerSingleton().register_timers([ + DeviceTimerMeta(FORWARD_COMPUTE, is_cpu_op=False, level=NDMetricLevel.FRAMEWORK_DEBUG), + DeviceTimerMeta(BACKWARD_COMPUTE, is_cpu_op=False, level=NDMetricLevel.FRAMEWORK_DEBUG), + DeviceTimerMeta(GRAD_ALLREDUCE, is_cpu_op=False, level=NDMetricLevel.FRAMEWORK_DEBUG), + DeviceTimerMeta(DATA_LOAD, is_cpu_op=True, level=NDMetricLevel.FRAMEWORK_DEBUG), # CPU op 示例 + ]) + + +class MLP(nn.Module): + def __init__(self, dim=64): + super().__init__() + self.net = nn.Sequential(nn.Linear(dim, dim), nn.ReLU(), nn.Linear(dim, dim)) + def forward(self, x): + return self.net(x) + + +def main(): + rank, local_rank = setup() + device = torch.device(f"cuda:{local_rank}") + init_ndtimeline(rank, local_rank) + + model = DDP(MLP().to(device), device_ids=[local_rank]) + optim = torch.optim.SGD(model.parameters(), lr=1e-3) + dim = 64 + num_steps = 10 + flush_every = 5 + + for step in range(num_steps): + inc_step() # 维护 manager 内部 global step(每条 span 会带上它) + + # ---- 度量一段 CPU 操作(数据加载)---- + with ndtimeit(DATA_LOAD): + x = torch.randn(32, dim, device=device) + + optim.zero_grad() + + # ---- 度量 GPU 计算 ---- + with ndtimeit(FORWARD_COMPUTE): + y = model(x) + loss = y.sum() + with ndtimeit(BACKWARD_COMPUTE): + loss.backward() + + # ---- 度量一次显式集合通信 ---- + # 注:精确的 allreduce 计时需 ndtimeit_coll + patched torch; + # 这里用通用 ndtimeit(在默认 stream 记录)演示流程。 + grad = model.module.net[0].weight.grad + with ndtimeit(GRAD_ALLREDUCE): + dist.all_reduce(grad, op=dist.ReduceOp.AVG) + + optim.step() + + # ---- 周期性 flush:把 [step-flush_every+1, step] 这批导出 ---- + if (step + 1) % flush_every == 0: + flush(step_range=range(step - flush_every + 1, step + 1)) + if rank == 0: + print(f"[step {step}] flushed") + + # ---- 优雅退出 ---- + wait() + if rank == 0: + print("done; look for trace_step*_*.json in the launch directory") + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() +``` + +### 3.2 运行 + +```bash +# 0) 安装 legacy veScale(只需一次) +cd legacy && pip install -e . + +# 1) 确保 socket 目录可写(默认 /opt/tiger/tmp/ndtimeline) +sudo mkdir -p /opt/tiger/tmp/ndtimeline && sudo chmod 777 /opt/tiger/tmp/ndtimeline +# (或改用一个你有权限的目录:见 §5 关于 SOCK_PARENT_DIR 的说明) + +# 2) 用 torchrun 起 2 个进程(2 卡) +torchrun --nproc_per_node=2 ndtimeline_tutorial.py +``` + +运行时会看到日志 `ndtimeline streamer started`、各 rank 的 `calibrate ...` 信息,以及若干 `p2p stream is not available, skipped`(因为没有 patched torch,且本例用的是通用 `ndtimeit`,属正常)。结束后在**启动目录**生成 `trace_step0_9` 之类的文件。 + +### 3.3 查看 trace + +把生成的 `trace_step{min}_{max}` 文件拖入 。你会看到: + +- 每个 dp_rank 一个 **process**(本例 dp_size=2,故 2 个 process); +- 每个 rank 按类别分成多个 **thread**:`main`(forward/backward-compute)、`dp collective`(grad-allreduce 等)、`other` 等; +- 每个 span 的宽度是耗时,横轴是**跨进程对齐的统一时间轴**(模拟全局时钟),因此可以直观比较两个 rank 的 forward 是否一前一后、allreduce 是否在等更慢的那个 rank。 + +### 3.4 最小冒烟测试(不想开 streamer / socket 时) + +如果只想确认"插桩 + 计时"跑通、暂不需要 trace 文件,可以把 handler 换成 `LoggingNDHandler` 并跳过 streamer: + +```python +from vescale.ndtimeline.handlers import LoggingNDHandler +NDTimerManagerSingleton(world_info, handlers=[LoggingNDHandler()], + init_cuda_dist=True, metric_level=NDMetricLevel.DEBUG) +NDTimerManagerSingleton().register_timers([...]) +# 训练 + flush;每个 metric 每次 flush 会打一行总耗时日志 +``` + +这条路不需要 `/opt/tiger/tmp` 可写,适合快速验证。 + +--- + +## 4. 实现原理详解 + +### 4.1 总体架构与数据流 + +``` +训练热路径 flush 线程(ThreadPoolExecutor) streamer 子进程(local_rank 0 拉起) +───────── ────────────────────────────── ──────────────────────────────────── +ndtimeit(name): + pool.get() → event.record() ……训练继续…… + ... 被测代码 ... + event.record() + flush(step_range) + └─ async_flush: + 换出新 timer 字典(gc.disable 保护) + submit _flush_timers ───────────────▶ + timer.elapsed(): + event.synchronize() + dur = start.elapsed_time(stop) # ms→s + abs = GlobalReferenceTime.elapsed_time(start) # 模拟全局时钟 + event 归池;tag 按 legal_tags 过滤 + 逐 handler 调用(dispatch 过滤): + SockNDHandler: + pickle + encode_package + sendall → Unix socket ─────────────▶ MsgHandler 收包+校验 + q.put(args) + 消费线程 internal_queue_consume: + LocalTimelineNDHandler + parse_record → record 入缓冲 + 发送特殊标记 "special" ────────────────▶ 收齐本机所有 rank 的 special: + dump_records() → Chrome Trace JSON + (pid=dp_rank, tid=rank*10+类别, + ts/dur 秒→微秒) +``` + +设计主线:**记录与导出解耦、导出与落盘解耦**。热路径只记 event;导出(同步+计算+序列化)在后台线程;落盘/聚合在独立子进程。三层各自异步,把对训练关键路径的干扰降到最低。 + +### 4.2 GlobalReferenceTime:模拟全局时钟(核心) + +这是整个系统最精巧的部分,解决"跨机器时间轴对齐"。位于 `timer.py`。 + +**问题**:各机器墙钟有漂移,NTP 精度不足以判断毫秒级通信重叠。需要一个"全局一致的时间坐标系"。 + +**三个关键量**: +- `initial_min_clock`(ns):全局时钟的**纪元原点**。 +- `last_calibrated_at`(ms):本 rank 最近一次校准时刻的墙钟。 +- `gpu_clock_residual_coef`:GPU 时钟相对 CPU 墙钟的漂移修正系数(≈1)。 + +**初始化 `init(world_sz, device)`**: +1. 在当前 device 上创建 **2 个参考 CUDA event**(双缓冲),先 record + synchronize 做 warmup。 +2. 调 `calibrate()` 完成首次校准。 +3. 置 `initialized=True`。 + +**校准 `calibrate()`**(round-robin 切换两个参考 event): +1. `ref_pointer` 指向下一个参考 event;`torch.cuda.synchronize()` 后在该 event 上 `record` 到默认 stream 并 `synchronize`,随即读 CPU 墙钟 `ts_ns = time.time_ns()`。这样这个 GPU event 就和一个 CPU 墙钟时刻绑定。 +2. **首次校准**(`not initialized`):先 `dist.barrier()` 同步所有 rank,然后 `all_gather` 收集各 rank 的 `ts_ns`,取**全局最小值**作为 `initial_min_clock`。 + > 为什么取 min 可行?barrier+all_gather 本身是一个同步点,各 rank 读到 `ts_ns` 的误差被限制在集合通信延迟内;取 min 作为共同原点,之后每个 rank 的绝对时间都是"相对这个原点的偏移",天然可比。 +3. `clock_diff = (ts_ns - initial_min_clock) / 1e6`(本 rank 相对原点的偏移,ms)。 +4. **CPU–GPU 漂移修正**:计算距上次校准的 CPU 墙钟间隔 `cpu_time`;若 `initialized` 且 `2e3 < cpu_time < 200000e3`(2 秒~200000 秒之间,避免间隔太短或太久导致测不准),则用两个参考 event 的 GPU `elapsed_time` 得到 `gpu_time`,算出 + `gpu_clock_residual_coef = 1 - (gpu_time - cpu_time) / gpu_time`。 + > 含义:GPU 时钟(`elapsed_time` 所基于)与 CPU 墙钟会有微小速率差,用这个系数把 GPU 时间换算成墙钟时间。 +5. 更新 `last_calibrated_at = ts`。 + +**换算 `elapsed_time(end_event)`**: +```python +gpu_time = ref_events[ref_pointer].elapsed_time(end_event) # 参考点→目标事件的 GPU 时长(ms) +return gpu_time * gpu_clock_residual_coef + last_calibrated_at # 修正后 + 参考点墙钟 = 目标墙钟(ms) +``` +即:任意 GPU event 的**绝对墙钟时间** = (它距最近参考点的 GPU 时长 × 漂移系数) + 参考点的墙钟时刻。这样所有 rank 的 event 都被换算到"以 `initial_min_clock` 为原点、经漂移修正"的统一墙钟上。 + +**何时校准**:`flush` 的 `dynamic_calibrate=True` 时,若距上次校准超过 30 分钟就触发;`sequential_calibrate=True` 在主线程校准,否则放到 flush 后台线程。首次校准在 `init` 里必然发生(含 barrier)。 + +### 4.3 CudaEventPool:event 对象池(`pool.py`) + +`torch.cuda.Event` 的创建/销毁有成本。`CudaEventPool` 预分配 `DEFAULT_CUDA_EVENT_POOL_SIZE=20` 个 event(`enable_timing=True, blocking=True`),每个先 `record()` 一次做 warmup 后入 `deque`。 + +- `get(tag)`:加锁从池里 `popleft`;池空才新建。把 tag 拷进 `event.tag`。 +- `release(event)`:加锁归还。 + +`DefaultEventPool` 是其类级单例包装。`DeviceTimer` 的每次 start/stop 都从池里取 event,用完归还,避免频繁创建。 + +### 4.4 DeviceTimer:单个 metric 的计时器(`timer.py`) + +状态:`_event_pairs`(GPU op 的 `[start_event, stop_event]` 列表)、`_extra_records`(CPU op 或补录的 `[start_ts, duration, tag]` 列表)、`_stream`、`_started`。 + +**`start(stream, tag, level)`**: +- 级别不够(`meta.level > level`)或未使能 → 直接返回。 +- 若 `is_cpu_op`:记录 `[time.time(), None, tag]` 到 `_extra_records`。 +- 否则从池取 start_event;若传了 `stream`,先 `stream.wait_stream(default_stream)`(保证 event 捕获到从默认 stream fork 出的通信),再 `event.record(stream=...)`。append `[start_event, None]`。 + +**`stop(tag, level)`**:CPU op 补上 duration;GPU op 取 stop_event record 并填进最近那对的末位。 + +**`elapsed(reset=True)`**(在 flush 线程调用,这是唯一会同步 GPU 的地方): +1. 对每对 event:`synchronize` 后, + - `single_elapsed = start.elapsed_time(stop) / 1e3`(耗时,秒); + - `single_since = GlobalReferenceTime.elapsed_time(start) / 1e3`(绝对起始,秒,见 §4.2); + - 合并 start/stop 的 tag,并**只保留 `legal_tags` 白名单内的键**; + - 两个 event 归还池。 +2. 若有 `_extra_records`(CPU/补录),按起始时间**归并插入**到上面两个序列里(GPU 与 CPU 记录合并成统一时间序列)。 +3. 返回 `(total_elapsed, durations[], since_starts[], tags[])`,`reset` 清空状态。 + +> `insert_record(start_ts, duration, tag)`:不走 event,直接塞一条已完成记录(用于把外部已知时间段补录进时间轴)。 + +### 4.5 NDTimerManager:注册、计时与异步 flush(`timer.py`) + +**注册 `register_timers(metas)`**: +- 每个 meta:`device_id==-1` 时,GPU op 用当前 device、CPU op 用 0;若 `step_getter is None and is_nature_step`,自动注入一个返回 `self._cur_global_step` 的 getter。 +- 由 meta 构造 `DeviceTimer`,存入 `_name2timer`。**只能注册一次**(断言)。 + +**计时入口 `start_timer/stop_timer/insert_record`**:按 name 查 timer;`KeyError`(未注册)只 warn 一次并记入 `_unregistered_timer_*`;其他异常 catch 后只打日志——**绝不让插桩异常冒泡到训练**。`tag` 里若有 `NDTIMELINE_STREAM_KEY` 会取出作为 stream 传给 `start`。 + +**异步 flush `async_flush(...)`**(关键机制): +1. 拷贝所有 meta;若 `not keep_timer_state`,把副本的 `enabled` 设为 `next_iter_enabled`。 +2. **换出**:把当前 `_name2timer` 作为 `past_name2timer`(带着已记录的 event),用 meta 副本新建一套 `fresh_name2timer`,然后在 `gc.disable()` 保护下把 `self._name2timer` 指向新字典。 + > 这一步是"记录与导出解耦"的核心:训练线程立即在新字典上继续 start/stop,旧字典(含未处理 event)交给后台慢慢同步。临时关 GC 是避免切换瞬间触发 CUDA event 析构。 +3. 回收已完成 future(打印其中异常)。 +4. 计算是否需要校准:`force_calibrate` 或(`dynamic_calibrate` 且距上次 >30min)。若 `sequential_calibrate`,现在就在主线程校准。 +5. 若 `submit2handler` 且有可导出的 timer,把 `_flush_timers(handlers, past_name2timer, step_range, world_info, require_calibrate)` 提交到 `ThreadPoolExecutor`,记入 `_futures`。 + +**后台导出 `_flush_timers(...)`**: +1. 若 `require_calibrate`(非顺序校准情形)→ `GlobalReferenceTime.calibrate()`。 +2. 逐个 timer:`elapsed_result = timer.elapsed()`;对每个 handler,按 `dispatch_mode`/`dst_names` 过滤后 `handler(name, *elapsed_result, step_range, world_info, extra)`(即展开成 `(name, elapsed, parts, since_parts, tags, step_range, world_info, extra)`)。handler 异常被 catch 打日志。 +3. `timer.meta = None`(防循环引用导致 event GC 失败)。 +4. 所有 timer 处理完后,给每个 handler 发一个特殊调用 `handler("special", 0.0, [], [], [], range(0,1), world_info, extra)`——作为"本批结束"的**分界信号**(`LocalTimelineNDHandler` 靠它判断何时落盘)。 + +**`wait()`**:有未完成 future 时,先 `dist.barrier()`,再 `wait(futures, timeout=10)`,收集异常,最后 `sleep(3)` 给 streamer 排空。 + +**单例**:`Singleton` metaclass + `NDTimerManagerSingleton`;`is_initialized()` 供所有 context manager 判断"没初始化就空转"。 + +### 4.6 NCCL stream 获取(`stream.py` + patched torch) + +要量真实通信耗时,event 必须记到 NCCL 实际使用的 stream 上。PyTorch 原生不暴露该 stream,故依赖补丁新增的 C++ 接口: + +- `getP2PCudaStreamId(device_idx, peer, is_batched)`:batched 按 device key 查 `ncclStreams_`,非 batched 按 `getKeySendRecv(rank, peer)` 查;命中返回 `stream.id()`,否则 `-1`。 +- `getCollCudaStreamId(tensors)`:按 tensor 所在 device 的 key 查。 + +`get_nccl_p2p_stream`/`get_nccl_coll_stream` 把裸 id 包成 `torch.cuda.Stream`,按 metric 名缓存到全局 `NCCL_STREAMS`。 + +**懒分配 caveat**:NCCL stream 在第一次通信时才创建,所以第一次通信查不到(`stream_id<0`),首个通信 span 必然丢失,日志 `p2p stream is not available, skipped` 属预期。未打补丁的 torch 上 `hasattr(..., "get_p2p_cuda_stream_id")` 为假,通信 metric 全部退化为默认 stream 记录。 + +### 4.7 Streamer 子进程 + 二进制协议 + +**为什么要独立进程**:写文件/聚合是重活,放训练进程会抖动;且一台机多个 rank 各写各的文件不便对比。于是每台机由 `local_rank==0` 拉起一个 streamer 进程统一聚合。 + +**`NDtimelineStreamer.init(local_rank, handlers)`**: +- 非 `local_rank 0` 直接返回(每机一个)。 +- 清理旧 socket,`torch.multiprocessing.spawn(start_method="spawn", daemon=True, join=False)` 起 1 个进程跑 `run`。 + +**`run(process_index, handlers)`**(在子进程内): +- 建 `queue.Queue(500000)`(在子进程里建,省主进程内存)。 +- 起消费线程 `internal_queue_consume`:不断 `q.get()`,把每条消息展开调用各 post handler;handler 抛 `NDHandlerError` 只 warn 继续。 +- 起 `ThreadingUnixStreamServer(SOCK_PATH, MsgHandler)` 监听 Unix socket。 + +**`MsgHandler.handle`**:每连接一线程,循环 `recv_and_validate(...)` 拆包 → `loads_fn`(pickle,前后关 GC)→ `q.put`。`BrokenPipeError`=客户端退出;`ProtocolValidationError`=坏包,跳过。 + +**二进制协议(`binary_protocol.py`)**: +``` +| Magic 1B (0x9C) | Version 1B (0x00) | Reserved 2B ("00") | +| Payload Length 4B (little-endian, 上限 128MiB) | +| Payload (pickle protocol 4) | +| EOF 1B ('\n') | +``` +`recv_and_validate` 逐字段校验,`read_or_recv`/`recv_to_buf` 用 `preload_data` 缓冲处理粘包/跨包残留。 + +**`SockNDHandler`**(训练进程侧客户端):`AF_UNIX` socket,`settimeout(2s)`;`try_to_connect` 带退避重试(`must=True` 时 50 次 ×0.8s),errno 106(已连接)视为成功;`call_impl` 把入参 dict 序列化 `sendall`。server 挂了置 `server_exited` 不再发——**采集失败绝不影响训练**。 + +### 4.8 Handler:step 归属 + 单机聚合 + Chrome Trace + +**`parse_record(...)`**(把原始入参 → `DeviceTimerStreamRecord`): +- 校验 `recent_elapsed_raw_parts / since_parts / tags` 三者等长,否则 `NDHandlerError`。 +- **两种 step 归属模式**: + 1. tag 里带 `_inner_global_step`(来自 `step_getter`):用 `itertools.groupby` 按 step 分组,每组一条 record(要求所有 tag 都带 step,否则报错)。 + 2. 无 step tag:按 `step_range` **均分**——要求 span 数能被 step 数整除,否则报错并提示改用 `step_getter`。 +- record 携带 rank/tp/pp/dp/ip/role/trial/run 等全维度信息("N-Dimension" 的由来)。 + +**`LocalTimelineNDHandler`**(聚合 + 落盘): +- 为每个 local rank 维护 `rank2buffer`(record 列表)和 `rank2steps`(deque of step-set,初始含 1 个空 set)。 +- 普通 metric:`parse_record` 后 append 到对应 rank 的 buffer,并把 step 记入当前 set。 +- 收到 `"special"`(flush 分界):给该 rank 的 deque append 新 set;**当所有 local rank 的 deque 长度 ≥2**,说明上一段收齐 → `dump_records()` 落盘,buffer 里只保留新段。 +- `dump_records()`: + - 遍历所有 rank 的 buffer,按 metric 名查 `thread_sort_index` 表(把 metric 映射到 `main` / `pp ->` / `pp <-` / `pp send/recv` / `tp collective` / `dp collective` / `other` 等线程类别); + - `tid = rank*10 + sort_index`(乘 10 给类别编号让位),`pid = dp_rank`; + - 每条 span → `CompleteEvent(ts=start_ts*1e6, dur=duration*1e6, ...)`(秒→微秒,Chrome Trace 单位); + - 追加 `ThreadMetadataEvent` / `ProcessMetadataEvent`(命名线程/进程); + - 所有 event `.to_objects()` 后 `json.dump` 到 `trace_step{min}_{max}`。 + - 切分缓冲时用 `record.copy()` 而非直接持引用,规避 GC 导致的内存泄漏。 + +**`chrome_trace_event.py`**:Chrome Trace Event Format 的数据类——`CompleteEvent`(ph=X)、`BeginEvent`(B)/`EndEvent`(E)、`FlowEvent`(s/t/f,画跨线程/进程连线,会自动修正"结束早于开始"以免不渲染)、`CounterEvent`(C)、`ProcessMetadataEvent`/`ThreadMetadataEvent`(M)。产出的 JSON 可直接被 Perfetto / `chrome://tracing` 打开。 + +### 4.9 级别过滤与 tag 白名单 + +- **级别**:`DeviceTimer.start/stop/insert_record` 入口判 `meta.level > level`(level 取自 manager 的 `metric_level`)则直接返回。低于阈值的 metric **连 event 都不记**,零开销。 +- **tag 白名单**:只有 `meta.legal_tags` 里的键会保留进最终记录。注册时若给了 `step_getter`,会自动把 `_inner_global_step` 和 `stream_key` 加进白名单。 + +### 4.10 健壮性与资源管理细节 + +- 所有插桩入口、所有 handler 调用都 try/except,异常只记日志。 +- 未注册 name 只 warn 一次(记入列表去重)。 +- 多处 `gc.disable()` 保护:timer 字典切换、`pickle.loads`、`disable_and_save/recover_from_history`——避免 GC 触发 CUDA event 析构带来开销或失败。 +- `timer.meta = None`、`record.copy()` 都是为了切断引用、帮助 event/record 正确回收。 +- 双缓冲参考 event 让"校准"与"正常记录"互不踩踏。 +- `wait()` 的 barrier + 超时 + sleep,保证退出时数据尽量不丢。 + +--- + +## 5. 局限与注意事项 + +1. **依赖 patched PyTorch 才能精确测通信**:无补丁时 `ndtimeit_p2p`/`ndtimeit_coll` 退化为默认 stream 计时(量的是发起时刻),并打印 skip 日志;计算计时不受影响。 +2. **首个通信 span 必丢**:NCCL stream 懒分配所致,属预期。 +3. **高层 `init_ndtimers` 绑定 veScale device mesh**:`hybrid` 模式需 `VESCALE_DEVICE_MESH` 已按 DP/TP/PP 初始化;`fsdp` 模式仅内部版。纯 PyTorch 请用低层 API(§3)。 +4. **import 即依赖整个 legacy 包**:哪怕只用 ndtimeline,也要 `pip install -e .` 安装 legacy。 +5. **socket 目录硬编码** `/opt/tiger/tmp/ndtimeline`,需可写;换目录要同步改 `variables.py` 的 `SOCK_PARENT_DIR`(streamer 与客户端都读它)。 +6. **开源版只有单机聚合**:跨机器汇总需自实现一个 post handler(如 MQHandler)把各机 metric 送中心存储;内部版的 `MQNDHandler`/`SMon`/what-if 分析未开源。 +7. **不自动告警**:ndtimeline 产出 trace 供人/下游分析,本身不做"最慢 N 个 rank"的自动判定与修复——那是 Megatron/NeMo StragglerDetector 或编排层的职责(详见 [`straggler-detection-survey.md`](./straggler-detection-survey.md))。 +8. **全局时钟校准是集体操作**:`init`/`flush` 触发的校准含 barrier/all_gather,所有 rank 必须同步到达;不要只在部分 rank 调用 `flush`。 +9. **`register_timers` 只能调一次**;step 用 `step_getter` 或 `inc_step/set_global_step`,二选一保持一致。 diff --git a/docs/texts/straggler-detection-survey.md b/docs/texts/straggler-detection-survey.md new file mode 100644 index 0000000..01a5d31 --- /dev/null +++ b/docs/texts/straggler-detection-survey.md @@ -0,0 +1,161 @@ +# 慢节点(Straggler)检测方案调研总结 + +> 背景:veScale 的 `legacy/ndtimeline` 是一个基于 CUDA Event + 模拟全局时钟的跨机器分布式 tracing 系统,核心动机之一是检测大规模 LLM 训练中的 straggler(慢节点)。本文档调研业界与学术界的其他实现方案,并将 ndtimeline 放入整体技术地图中对比定位。 +> +> 注:调研中核查发现,部分流传的系统名("Astro MLSys'24""Barometer""TicToc")**查无实据**,已用确实存在的工作(FALCON、GREYHOUND、Mycroft、NURD 等)替换。另:ByteDance OSDI'25 论文中描述的 "NDTimeline" 正是本仓库 `legacy/ndtimeline` 的内部同源系统。 + +--- + +## 一、先厘清:三个不同的问题 + +"检测慢节点"其实是三件不同的事,不同方案各管一段: + +| 问题 | 含义 | 代表方案 | +|---|---|---| +| **Detection(检测)** | 有没有慢节点?是哪个 rank/节点? | Megatron StragglerDetector、NeMo、FALCON、Minder | +| **Diagnosis(诊断)** | 它为什么慢?时间花在哪条依赖链上? | **ndtimeline**、MegaScale tracer、Mycroft、Flight Recorder | +| **Mitigation(缓解)** | 慢已经发生了,怎么不让它拖垮训练? | MapReduce backup task、Oobleck、Bamboo、Malleus、StragglAR | + +**ndtimeline 本质上是 Diagnosis 工具**(跨机器统一时间轴可视化),而不是自动 Detection 工具——它产出 trace 让人(或下游分析)看出谁慢、为什么慢,本身不做"排名告警"。理解这点,才能看清它和其他方案的关系是**互补**而非竞争。 + +## 二、检测类方案:按"信号来源"分层 + +所有检测方案都遵循同一个范式:**本地采集 → 跨 rank 对比("谁慢只有比了才知道")**。区别在于采什么信号、在哪一层采。 + +### 2.1 应用/框架层:比"每 rank 的迭代/计算耗时" + +| 方案 | 机制 | 粒度 | 侵入性 | 特点 | +|---|---|---|---|---| +| **Megatron-LM StragglerDetector** | CUDA event + `perf_counter` 双计时 + NVML 温度/功耗,`all_gather` 到 rank0 排名 | rank 级 | 低(context manager) | 开箱即用、开销极低、可常驻;只到 rank 不到算子 | +| **NeMo / nvidia-resiliency-ext** | 自动 wrap train step,算"相对 + 个体历史"双性能分,低于阈值(0.7)告警,可自动停训 | rank+节点级 | 极低(callback) | 产品化最好,有 WandB 可视化和自动终止 | +| **MegaScale**(ByteDance, NSDI'24) | 轻量 CUDA event 分段计时 → 集中库 → **step latency 热力图**(慢 10% 即显色)+ 自动健康检查驱逐 | per-rank/段 | 需插桩 | 12288 卡生产验证,检测+定位+自愈闭环 | +| **FALCON / GREYHOUND**(HKUST) | shim 层拦截 NCCL,自相关函数(ACF)提周期 + 贝叶斯变点检测(BOCD)区分真退化与抖动;定位时注入 marker 把 ring 拆成 P2P 握手测链路 | per-GPU+per-link | shim(环境变量注入) | >99% 精度,检测+定位+4 级缓解全闭环 | + +补充: +- **MegaScan / MegatronApp**(arXiv:2507.19845):CUDA-event-driven 的**算子粒度**在线慢节点检测 + 根因定位,通过 collective 边界做跨机器时钟校准——与 ndtimeline 思路最接近的学术工作之一。 +- **DeepSpeed**:**没有专门的 straggler 检测**。相关能力是 FlopsProfiler(submodule 级 latency/FLOPs)和 `wall_clock_breakdown`(forward/backward/update 分阶段计时),均为单 rank 本地指标,无跨 rank 对比。 + +### 2.2 通信层:比"每个 collective 的完成时间/状态" + +| 方案 | 机制 | 粒度 | 特点 | +|---|---|---|---| +| **PyTorch Flight Recorder** | per-rank 环形缓冲记录每个 collective 生命周期+调用栈,超时经 TCPStore 全量 dump,`fr_trace` 跨 rank 对齐找不匹配 | collective 级 | 零代码、Meta 万卡常驻;**定位"谁没来"(hang/desync),不定位"谁来得慢"** | +| **NCCL Inspector**(NVIDIA, NCCL 2.23+) | 官方 profiler 插件,记录 per-collective 执行时间/带宽,区分 NVLink vs IB 流量,出 Prometheus | collective×rank | **实时**,能分"GPU 内部 vs 网络"瓶颈 | +| **Mycroft**(ByteDance+Harvard, SOSP'25) | NCCL 内部插桩追 collective 的控制/数据依赖,空间对比各 rank 状态 | collective 内部依赖 | 90% 情况 15s 内发现异常 | +| **SysOM-AI** | eBPF uprobe 记 NCCL host 入口时间,超 μ+kσ 判 straggler | per-rank/collective | 无需改代码、跨 CPU/GPU/NCCL 层 | +| **Google CoMMA** | NCCL profiler plugin 收集 per-communicator/per-collective 指标 | collective 级 | 开源,配合 Cluster Director 使用 | + +补充:PyTorch 的 `TORCH_DISTRIBUTED_DEBUG=DETAIL` 与 `TORCH_NCCL_DESYNC_DEBUG=1` 能定位 desync 的"罪魁祸首 rank",但开销显著,不适合生产常驻;注意 NCCL 核心库本身**没有**应用级 watchdog/heartbeat,这些机制在 PyTorch 的 `ProcessGroupNCCL` 层实现。 + +### 2.3 硬件层:看"GPU 物理真相" + +| 方案 | 机制 | 粒度 | 特点 | +|---|---|---|---| +| **NVML 降频位掩码** | `HwSlowdown/HwThermalSlowdown/HwPowerBrake` 任一位触发 = **2 倍以上降频**,是近乎确定的 straggler 信号 | 单 GPU+降频原因 | 零开销、最底层;但只查得出硬件降频,查不出软件/数据问题 | +| **DCGM + dcgm-exporter** | 被动健康监控(温度/功耗/ECC/NVLink/XID)+ 主动诊断(L1~L4,从 <2.5s 到 2.25h)+ 降频原因分解,出 Prometheus/Grafana | 单 GPU 子系统级 | 覆盖最全,能回答"为什么慢"(热降频?电源?内存?),Policy Engine 可自动修复 | + +### 2.4 网络层:验证"是节点慢还是链路慢" + +- **nccl-tests**(`busbw` 对比理论峰值,Together AI 要求达 ~92%) +- **perftest**(`ib_write_bw/lat` 逐对测链路带宽/延迟,`--data_validation` 检测数据损坏) +- **UFM / Cumulus NetQ**(fabric 健康、拥塞监控、端口级告警) +- **Meta RoCE 网络设计**(SIGCOMM'24):DCQCN 禁用、接收方驱动准入控制、增强 ECMP——从网络架构层面消除 straggler 根因 + +这类是**主动基准/被动监控**,用于训练前验收或定位网络根因,不是训练内实时检测。 + +### 2.5 平台/编排层:自动化与因果定位 + +| 方案 | 机制 | 亮点 | +|---|---|---| +| **Google Cloud Cluster Director** | 被动观察网络流量构建**通信图**,图遍历追溯"是等别人慢,还是自己就是源头" | 因果分析(非阈值)、零配置常驻、分钟级定位 | +| **Meta**(Fleetscanner/Ripple/Hardware Sentinel + 开源 GCM) | 定期微基准(45-60 天覆盖)+ 共置测试(数天)+ 内核异常分析(实时),二分搜索隔离,影子节点验证 | 全栈覆盖,Hardware Sentinel 实时且无需额外资源 | +| **K8s 生态**(AKS NPD、阿里 ACK GPU 级 fencing、AWS EKS NMA、HyperPod HMA) | 解析 XID/降频/链路 → Node Condition → 自动 cordon/drain/替换 | **GPU 级隔离**(ACK)、自动闭环修复 | +| **Minder**(NSDI'25)、**Hound**(SIGMETRICS'18) | 故障机器检测 / 离线因果推断根因 | 学术代表 | + +## 三、缓解/容忍类方案(简述) + +这类不检测、而是让训练对慢/故障节点**不敏感**: + +- **经典**:MapReduce speculative/backup task(复制落后任务并发跑,先完成者保留)、SkewTune(数据倾斜动态重分区)——**LLM 训练不直接适用**(pipeline 串行依赖,无法简单复制)。 +- **弹性容错**:Oobleck(SOSP'23,pipeline 模板+(f+1) 副本)、Bamboo(NSDI'23,pipeline bubble 里塞冗余计算抗抢占)、Varuna(EuroSys'22)、GEMINI(SOSP'23,内存检查点快速恢复)、Parcae(NSDI'24,主动规避抢占)。 +- **straggler-aware 重配/容忍**:Malleus(SIGMOD'25,按 per-GPU 性能自适应重排四维并行,2.6-5.3× 效率提升)、StragglAR(AllReduce 时让快的 GPU 先做 ReduceScatter 掩盖慢者)、PipeMorph/Adaptra(自适应 pipeline 调度)。 +- **调度规避**:Gandiva(OSDI'18,progress rate 估计做 placement)、NURD(MLSys'22,负无标注学习预测 straggler 概率)。 + +## 四、ndtimeline 的定位 + +把 ndtimeline 放进上面这张地图: + +**它属于"应用/框架层计时对比"这一类,但走的是其中少有人走的一条路——构建跨机器统一时间轴。** + +| 对比维度 | ndtimeline | Megatron/NeMo(哨兵型) | FALCON/GREYHOUND(统计型) | Mycroft(依赖型) | +|---|---|---|---|---| +| 跨机器时钟 | **模拟全局时钟**(collective 边界校准,构建统一时间轴) | 不做,事后 all_gather 比数值 | 不做,ACF 提周期 | 不做,空间对比 | +| 输出 | Perfetto 可交互 timeline(全量、可回放) | 文本日志(最慢 N 个 rank) | 自动告警+缓解 | 依赖图 | +| 定位粒度 | 算子/段 + collective,能看 wait vs compute | rank 级 | GPU+链路级 | collective 内部 | +| 是否自动判定 | **否,需人/下游分析** | 是 | 是 | 是 | +| 回答的问题 | **为什么慢、慢在哪条依赖链** | 有没有、是哪个 | 是哪个+怎么救 | 谁先变慢(根因) | + +### ndtimeline 的独特价值 + +1. **模拟全局时钟**:不同于简单的集中式时间收集,ndtimeline 通过周期性时钟同步构建了一个虚拟的全局时间轴,使得跨机器的事件可以在同一时间坐标系下精确对齐——这在万卡级别是关键挑战。多数统计型检测器直接绕开了这个问题(只比相对耗时,不构建统一时间轴)。 +2. **可回放的现场**:多数检测方案只给"谁慢"的统计结论;ndtimeline 给的是可回放的现场——能直观看到 pipeline bubble、通信等待、跨机器 collective 的时间错位。这正是 MegaScale 的 "distributed tracer" 和 MegaScan 在做的事。 +3. **粗粒度低开销**:把多个 GPU 执行打包为 forward/backward compute marker,在保持可观测性的同时最小化性能影响;CUDA event 异步记录 + 后台线程导出 + 独立 streamer 子进程落盘。 +4. **内部版还叠了自动分析**:OSDI'25 论文证实,内部版 NDTimeline 在采集+可视化之上还有 **SMon 自动热力图 dashboard** 和 **what-if 反事实模拟**(重建"没有这个 straggler 会快多少",量化影响)。开源仓库只有前端的采集+可视化部分。 + +### ndtimeline 的短板 + +- 不自动排名/告警(要自己写分析,或配 Megatron/NeMo 那种哨兵); +- 不触达硬件根因(得配 DCGM/NVML 才知道是热降频还是真故障); +- 不含修复/隔离(要配 K8s 生态或 FALCON 那种闭环); +- 需要插桩 + trace 数据量大,更适合"深度诊断"而非"7×24 哨兵"; +- 依赖 patched PyTorch(NCCL stream 接口),多机汇总需自行实现 MQHandler。 + +## 五、给实践的结论 + +1. **没有银弹,生产系统都是分层组合**:哨兵(Megatron/NeMo,常驻、低开销、答"有没有")+ 诊断(ndtimeline/Flight Recorder,按需、答"为什么")+ 硬件(DCGM/NVML,答"硬件根因")+ 编排(K8s/Cluster Director,答"自动隔离修复")。 +2. **ndtimeline 该和谁搭**:把 Megatron/NeMo 当"报警器"常开,报警后拉 ndtimeline 的 trace 做深度复盘,再用 DCGM 历史指标确认是不是硬件问题——这是一条很顺的链路。 +3. **趋势是"因果化"和"闭环化"**:从"比阈值"走向"追根因"(Cluster Director 的因果图、Mycroft 的依赖追踪、what-if 模拟),从"告警"走向"检测→定位→隔离→恢复"自动闭环。ndtimeline 的模拟全局时钟+what-if 思路正好踩在前一个趋势上。 + +## 六、分层组合参考 + +| 阶段 | 推荐方案组合 | +|---|---| +| 集群验收 / Burn-in | DCGM Diagnostics L3-4 + perftest 全对测试 + nccl-tests(目标 ~92% 理论峰值) | +| 训练前预检 | DCGM Level 2(<3min)+ NVML 降频检查 | +| 训练中实时哨兵 | Megatron/NeMo StragglerDetector + dcgm-exporter + NCCL Inspector → Prometheus/Grafana | +| 训练中异常响应 | K8s NPD/HMA 自动隔离 + NCCL watchdog timeout + Flight Recorder dump | +| 深度根因诊断 | **ndtimeline trace** + Flight Recorder 跨 rank 对齐 + DCGM 历史指标 + NetQ/UFM 网络事件 | + +--- + +## 参考资料 + +**框架内嵌检测** +- Megatron StragglerDetector: [README_STRAGGLER.md](https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/README_STRAGGLER.md) · [API Docs](https://docs.nvidia.com/megatron-core/developer-guide/latest/apidocs/core/core.utils.html) +- NeMo Resiliency: [文档](https://docs.nvidia.com/nemo-framework/user-guide/24.12/resiliency.html) · [NVRx API](https://docs.nvidia.com/nemo/megatron-bridge/0.3.0/apidocs/bridge/bridge.training.nvrx_straggler.html) +- PyTorch Flight Recorder: [教程](https://docs.pytorch.org/tutorials/intermediate/debug_hangs_with_flight_recorder.html) · [博客](https://pytorch.org/blog/flight-recorder-a-new-lens-for-understanding-nccl-watchdog-timeouts/) · [NCCL 环境变量](https://docs.pytorch.org/docs/stable/torch_nccl_environment_variables.html) +- DeepSpeed: [Flops Profiler](https://www.deepspeed.ai/tutorials/flops-profiler/) · [Monitor](https://deepspeed.readthedocs.io/en/latest/monitor.html) + +**基础设施 / 硬件 / 网络层** +- NVIDIA DCGM: [文档](https://docs.nvidia.com/datacenter/dcgm/latest/user-guide/feature-overview.html) · [GitHub](https://github.com/nvidia/dcgm) · [dcgm-exporter](https://github.com/nvidia/dcgm-exporter) · [Diagnostics](https://docs.nvidia.com/datacenter/dcgm/latest/user-guide/dcgm-diagnostics.html) +- NVML: [ClocksThrottleReasons](https://docs.nvidia.com/deploy/nvml-api/group__nvmlClocksThrottleReasons.html) +- K8s GPU 健康检查: [Azure AKS](https://learn.microsoft.com/en-us/azure/aks/gpu-health-monitoring) · [阿里云 ACK](https://www.alibabacloud.com/help/en/ack/ack-managed-and-ack-dedicated/user-guide/anomaly-diagnosis-in-gpu-accelerated-nodes) · [AWS EKS NMA](https://docs.aws.amazon.com/eks/latest/userguide/node-health-nma.html) · [SageMaker HyperPod HMA](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-eks-resiliency-health-monitoring-agent.html) · [node-problem-detector](https://github.com/kubernetes/node-problem-detector) +- NCCL Inspector: [博客1](https://developer.nvidia.com/blog/enhancing-communication-observability-of-ai-workloads-with-nccl-inspector/) · [博客2](https://developer.nvidia.com/blog/real-time-performance-monitoring-and-faster-debugging-with-nccl-inspector-and-prometheus/) +- 网络: [nccl-tests](https://github.com/nvidia/nccl-tests) · [perftest](https://github.com/linux-rdma/perftest) · [Meta RoCE (SIGCOMM'24)](https://engineering.fb.com/2024/08/05/data-center-engineering/roce-network-distributed-ai-training-at-scale/) · [Together AI 实践](https://www.together.ai/blog/a-practitioners-guide-to-testing-and-running-large-gpu-clusters-for-training-generative-ai-models) +- 平台: [Google Cluster Director](https://cloud.google.com/blog/products/compute/stragglers-in-ai-a-guide-to-automated-straggler-detection) · [CoMMA](https://github.com/google/CoMMA) · [Meta 硬件可靠性](https://engineering.fb.com/2025/07/22/data-infrastructure/how-meta-keeps-its-ai-hardware-reliable/) · [Meta GCM](https://github.com/facebookresearch/gcm) + +**学术 / 研究系统** +- MegaScale (NSDI'24): [arXiv:2402.15627](https://arxiv.org/abs/2402.15627) +- Understanding Stragglers / What-if Analysis (OSDI'25): [arXiv:2505.05713](https://arxiv.org/abs/2505.05713) · [开源代码](https://github.com/ByteDance-Seed/StragglerAnalysis) +- MegaScan / MegatronApp: [arXiv:2507.19845](https://arxiv.org/html/2507.19845v1) +- FALCON: [arXiv:2410.12588](https://arxiv.org/abs/2410.12588) +- GREYHOUND (ATC'25): [USENIX](https://www.usenix.org/conference/atc25/presentation/wu-tianyuan) · [GitHub](https://github.com/wutianyuan1/Greyhound) +- Mycroft (SOSP'25): [arXiv:2509.03018](https://arxiv.org/abs/2509.03018) +- Malleus (SIGMOD'25): [arXiv:2410.13333](https://arxiv.org/abs/2410.13333) +- StragglAR: [arXiv:2505.23523](https://arxiv.org/abs/2505.23523) · [GitHub](https://github.com/arjundevraj/stragglar) +- Minder (NSDI'25): [arXiv:2411.01791](https://arxiv.org/abs/2411.01791) +- Oobleck (SOSP'23): [arXiv:2309.08125](https://arxiv.org/abs/2309.08125) · Bamboo (NSDI'23): [arXiv:2204.12013](https://arxiv.org/abs/2204.12013) · Varuna (EuroSys'22): [arXiv:2111.04007](https://arxiv.org/abs/2111.04007) · GEMINI (SOSP'23): [PDF](https://zhuangwang93.github.io/docs/Gemini_SOSP23.pdf) · Parcae (NSDI'24): [arXiv:2403.14097](https://arxiv.org/abs/2403.14097) +- Gandiva (OSDI'18): [USENIX](https://www.usenix.org/conference/osdi18/presentation/xiao) · NURD (MLSys'22): [arXiv:2203.08339](https://arxiv.org/abs/2203.08339) · Hound (SIGMETRICS'18): [ACM DL](https://dl.acm.org/doi/abs/10.1145/3292040.3219641) +- MapReduce (OSDI'04): [PDF](https://static.googleusercontent.com/media/research.google.com/en//archive/mapreduce-osdi04.pdf) · SkewTune (SIGMOD'12): [PDF](https://homes.cs.washington.edu/~magda/papers/kwon-sigmod12.pdf) +- 综述: Tails in the Cloud — [Springer](https://dl.acm.org/doi/abs/10.1007/s11227-020-03241-x) +- Llama 3 训练报告(故障统计): [arXiv:2407.21783](https://arxiv.org/abs/2407.21783)