Megatron Timers 源码精读——分布式训练性能打点

本文精读 Megatron-Core 的 megatron/core/timers.py（约 300 行）。目标：让你面试时能讲清"训练里怎么打点测一段算子耗时、为什么必须 torch.cuda.synchronize()、为什么分布式打点要跨 rank 聚合、log_level/DummyTimer 的开销控制怎么做"。这是个"小而完整"的子系统，特别适合讲清"分布式时序测量的陷阱"。

一、为什么需要专门的 Timer

训练几百亿参数、几千卡，一句"前向多快"没法用 time.time() 包一下就答——它有四个坑：

  1. GPU 异步：PyTorch 的 CUDA op 是异步 launch，CPU 侧 time.time() 测的是"把 kernel 丢进队列"的时间，不是 GPU 真跑完的时间。必须 torch.cuda.synchronize() 强制等 GPU 空闲再取时间。
  2. 多 rank 不同步：上千张卡的耗时不同，rank 0 的前向可能是 5ms，rank 1000 可能是 8ms（某段通信慢），单看一个 rank 会漏掉 straggler。要 all-gather 各 rank 耗时取 min/max。
  3. 打点本身有开销：到处插 synchronize 会让训练变慢，且 collect barrier 又会引入人为同步。要能用 log_level 关掉无关打点，DummyTimer 零开销。
  4. 聚合与输出：上千卡把每个 rank 的时间都打出来太长，需要 max/minmax/all 几种聚合策略。

timers.py 就是围绕这四点设计的：Timer（单计时器 + cuda sync）、DummyTimer（零开销占位）、Timers（一组计时器 + 跨 rank 聚合 + 日志/TensorBoard 输出）。

实际用法（源码 schedules.py 等）

1
2
3
4
5
6
7
config.timers('forward-compute', log_level=2).start()
...   # 前向
config.timers('forward-compute').stop()

config.timers('forward-backward', log_level=1).start(barrier=config.barrier_with_L1_time)
...   # 整个 forward+backward
config.timers('forward-backward').stop()

config.timers 是一个全局 Timers 实例，按名字取/建计时器，log_level 控制是否真正计时。在 schedules.py、attention.py、distrib_optimizer.py、finalize_model_grads.py、p2p_communication.py 等关键路径都有打点——覆盖 forward / backward / 通信 / 优化器 / 梯度同步全链路。

二、Timer 类：单计时器，cuda.synchronize 是灵魂

Timer 继承抽象基类 TimerBase，核心字段：

1
2
3
4
5
6
7
class Timer(TimerBase):
    def __init__(self, name):
        self._elapsed = 0.0       # 累计已计时段(可 reset)
        self._active_time = 0.0   # 累计活跃时间(reset 不清,用于总开销统计)
        self._started = False      # 是否正在计时
        self._barrier_group = None # 可指定 barrier 的 process group
        self._start_time = time.time()

2.1 start / stop：必先 cuda.synchronize

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def start(self, barrier=False):
    assert not self._started
    if barrier:
        torch.distributed.barrier(group=self._barrier_group)  # 先跨 rank 同步
    torch.cuda.synchronize()        # 再等 GPU 空闲
    self._start_time = time.time()
    self._started = True

def stop(self, barrier=False):
    assert self._started
    if barrier:
        torch.distributed.barrier(group=self._barrier_group)
    torch.cuda.synchronize()
    elapsed = time.time() - self._start_time
    self._elapsed += elapsed
    self._active_time += elapsed
    self._started = False

两个关键设计：

  1. torch.cuda.synchronize() 不可省。PyTorch 的 CUDA kernel 是异步提交的——start 时如果不 sync，time.time() 取的是"launch 完 kernel"的 CPU 时间，此时 GPU 上可能还堆着一堆没跑完的 kernel，测出来的时间偏小且不可靠。sync 强制 CPU 等 GPU 把队列里所有 kernel 跑完，才取时间戳。stop 同理——等这段 GPU 活真做完才记结束。

  2. barrier 与 barrier_group：可选地先做一次 torch.distributed.barrier(group=...)，让所有 rank 起跑线一致。barrier 后再 sync + 取时间，能减少"不同 rank 起点错位"带来的噪声。但 barrier 本身有开销、且要求所有 rank 都调（否则 hang），所以默认 barrier=False，只在需要精确对齐的场景（如测整个 forward-backward）才开——源码里 config.barrier_with_L1_time 就是控制 L1 级（forward-backward）要不要 barrier。

面试金句：“stop/start 里 torch.cuda.synchronize() 是分布式计时的命门——不 sync 测的是 CPU launch 时间不是 GPU 计算时间，数据完全不可信。sync 本身有开销，所以打点要用 log_level 控制、DummyTimer 跳过。barrier 是可选的跨 rank 对齐，要所有 rank 一起调否则 hang。”

2.2 elapsed 与 active_time

1
2
3
4
5
6
def elapsed(self, reset=True, barrier=False):
    # 如果正在计时先 stop, 取出 _elapsed, 按 reset 决定是否清零, 再按需 restart
    ...

def active_time(self):
    return self._active_time   # 累计活跃时间, reset 不清
  • elapsed(reset=True)：取当前累计耗时，默认清零（下一次从 0 计）。如果计时正在进行会先 stop 再 start（无缝续测）。这个 reset 语义让"一个 step 累计、log 时取出并清零"成为惯用法。
  • active_time()：不随 reset 清零，统计这个 timer 总共活跃了多久——用来评估"打点/某段代码总占比"。区分 _elapsed（可清零的窗口值）和 _active_time（不可清零的累计值）是个细节但有用的设计。

2.3 reset

1
2
3
4
def reset(self):
    # 注意: 不清 _active_time
    self._elapsed = 0.0
    self._started = False

只清窗口值不清累计值——保证 active_time 永远单调递增反映真实总活跃时长。

三、Timers 类：分组管理 + log_level + 跨 rank 聚合

Timers 是一组 timer 的管理者，持有全局 config.timers 实例。三个核心能力：log_level 控制、DummyTimer 零开销、跨 rank 聚合输出。

3.1 call：按名字取/建 timer + log_level 过滤

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def __init__(self, log_level, log_option):
    self._log_level = log_level            # 当前启用的日志级别
    self._log_option = log_option          # 'max'/'minmax'/'all'
    self._timers = {}                       # name -> Timer
    self._log_levels = {}                   # name -> 该 timer 的 log_level
    self._dummy_timer = DummyTimer()
    self._max_log_level = 2

def __call__(self, name, log_level=None):
    if name in self._timers:                # 已建过: 校验 log_level 一致后返回
        ...
        return self._timers[name]
    if log_level is None:
        log_level = self._max_log_level    # 默认用最高级
    # 若该 timer 的 log_level > 当前启用的 _log_level, 返回 DummyTimer(零开销)
    if log_level > self._log_level:
        return self._dummy_timer
    self._timers[name] = Timer(name)
    self._log_levels[name] = log_level
    return self._timers[name]

log_level 机制：每个 timer 创建时带一个 log_level（1 或 2），若它高于全局 _log_level（运行时由 --log-level 参数设），就返回 DummyTimer——它的 start/stop/reset 全是空函数，调用 elapsed 直接抛异常提醒"这个 timer 没启用，别拿它算时间"。

这就是"打点零开销"的关键：生产训练用低 log_level（只开 L1 的 forward-backward 等粗粒度），调试用高 log_level（开 L2 的 forward-compute/backward-compute 等细粒度）。源码里：

  • log_level=1：粗粒度（forward-backward 整段），默认开。
  • log_level=2：细粒度（forward-compute、backward-compute 单段），需要 --log-level 2 才启用。

DummyTimer 让上层代码不用 if 判断——config.timers('foo', log_level=2).start() 在 log_level=1 时退化成空操作，开销近乎为零。

3.2 _get_elapsed_time_all_ranks：跨 rank all-gather

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def _get_elapsed_time_all_ranks(self, names, reset, barrier):
    if barrier:
        torch.distributed.barrier()
    world_size = torch.distributed.get_world_size()
    rank = torch.distributed.get_rank()
    # [world_size, len(names)] 的张量, 放在 GPU 上
    rank_name_to_time = torch.zeros((world_size, len(names)), dtype=torch.float,
                                    device=torch.cuda.current_device())
    for i, name in enumerate(names):
        if name in self._timers:
            rank_name_to_time[rank, i] = self._timers[name].elapsed(reset=reset)
    # all-gather 各 rank 的时间到所有卡
    dist_all_gather_func(rank_name_to_time.view(-1), rank_name_to_time[rank, :].view(-1))
    return rank_name_to_time

关键设计：

  • 先把本 rank 各 timer 的 elapsed 填进自己那一行，再 all_gather 把所有 rank 的时间广播到所有卡。于是每个 rank 都拿到完整的 [world_size, len(names)] 矩阵。
  • 张量放 GPU（torch.cuda.current_device()），all-gather 走 NCCL（快）；数据量小（world_size × timer 数），走 GPU collective 也无妨。
  • 源码注释解释了为啥用 all-gather 而不是 gather：早期 PyTorch 没 gather_base，且数据量小，all-gather 简单。现在虽可优化但沿用。
  • 强假设：所有 rank 必须用相同的 names 列表调这个函数，否则 all-gather 形状对不上 → hang。源码注释明说。
  • 这里取 elapsed 时不再传 barrier——因为本函数开头已经 barrier 过一次，避免不同 timer 有不同 barrier_group 的坑（注释也提到）。

3.3 聚合：max / minmax / all

拿到 [world_size, len(names)] 后，按 log_option 三种聚合：

1
2
3
4
5
6
7
def _get_global_min_max_time(self, names, reset, barrier, normalizer):
    rank_name_to_time = self._get_elapsed_time_all_ranks(names, reset, barrier)
    for i, name in enumerate(names):
        rank_to_time = rank_name_to_time[:, i]
        rank_to_time = rank_to_time[rank_to_time > 0.0]   # 过滤掉没打点的 rank(0.0)
        if rank_to_time.numel() > 0:
            name_to_min_max_time[name] = (rank_to_time.min()/norm, rank_to_time.max()/norm)
  • max：只报所有 rank 的最大值——反映 straggler（最慢那张卡），是最该看的指标（训练速度被最慢的卡拖累）。
  • minmax：报 (min, max)——min 是最快卡（理想值）、max 是最慢卡（实际拖累值），区间反映不均衡度。
  • all：逐 rank 报所有值——调试某张卡异常时用，但上千卡输出爆炸。

> 0.0 过滤把"没打这个点的 rank"（值为 0）剔除，避免拉低 min。normalizer/1000.0 把秒转毫秒输出。

面试金句：“聚合优先看 max——分布式训练是木桶效应，最慢的 rank 决定整体吞吐。minmax 能看不均衡度，all 用于定位 straggler。注意聚合前要 all-gather 各 rank 时间，且所有 rank 必须传相同 names 列表否则 hang。”

四、输出：log 打到 stdout、write 到 TensorBoard

4.1 log：打到指定 rank 的 stdout

1
2
3
4
5
6
def log(self, names, rank=None, normalizer=1.0, reset=True, barrier=False):
    output_string = self.get_all_timers_string(names, normalizer, reset, barrier)
    if rank is None:
        rank = torch.distributed.get_world_size() - 1   # 默认打到最后一个 rank
    if rank == torch.distributed.get_rank() and output_string is not None:
        print(output_string, flush=True)
  • 只让指定 rank（默认最后一个）打印，避免上千卡各打印一遍刷屏。
  • normalizer：归一化因子。惯用法：训练每 N 步打一次，normalizer=N 就得到"每步平均耗时"。源码注释的例子就是"用 logging interval 做 normalizer 算每步均值"。
  • reset=True：打完清零，开始下一个统计窗口。
  • 输出格式：(min, max) time across ranks (ms): 或 max time across ranks (ms):，每个 timer 一行，名字用 . 补齐到 48 字符对齐——'foo '.ljust(48, '.')。

4.2 write：写 TensorBoard，只记 max

1
2
3
4
5
6
def write(self, names, writer, iteration, normalizer=1.0, reset=True, barrier=False):
    name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
    if writer is not None:
        for name in name_to_min_max_time:
            _, max_time = name_to_min_max_time[name]
            writer.add_scalar(name + '-time', max_time, iteration)
  • TensorBoard 只记跨 rank 的 max_time（一个标量），因为标量图看趋势，max 最能反映真实步时。
  • 注释解释了为什么不用 add_scalars（复数）——它会给每个 timer 单独建一个 run，污染 runs 列表，所以逐个 add_scalar。

五、踩坑与最佳实践

源码里几处 assert 和注释，其实是踩过的坑，面试可讲：

  1. assert not self._started / assert self._started：防止 start 嵌套或 stop 没 start。Timer 不可重入，嵌套会丢时间。要嵌套用不同名字的 timer。
  2. barrier 必须所有 rank 一起调：注释明说"It is up to the user to make sure all the ranks in barrier_group call it otherwise, it will result in a hang."。条件性 barrier（只有部分 rank 走）会 hang。
  3. _get_elapsed_time_all_ranks 的 names 必须全局一致：否则 all-gather 张量形状对不上，hang。这是分布式打点最常见的死锁源。
  4. elapsed 取值时不再 barrier：因为聚合函数开头已 barrier，避免重复 barrier 和不同 barrier_group 冲突。
  5. sync 开销：每次 start/stop 都 sync 两次（start 一次、stop 一次），细粒度打点会显著拖慢。所以细粒度 timer 都标 log_level=2，生产 --log-level 1 时它们变 DummyTimer 零开销。
  6. _active_time 不清零：reset 只清窗口值，要算"这个 timer 累计花了多久"（含所有步）用 active_time()，算"这步花了多久"用 elapsed(reset=True)。

怎么用（给训练调优者的速记）

  • 跑基线：--log-level 1，看 forward-backward、optimizer、data-loaders 等粗粒度时间，找出占比最大的段。
  • 定位 straggler：--log-level 2 --log-timers all，对比各 rank 的 forward-compute/backward-compute，找出明显慢的卡。
  • 调通信：关注 forward-backward 里的 P2P/all-reduce 时间占比，判断是不是通信瓶颈（配合 NCCL profiler）。
  • TensorBoard：write 把 max-time 写图，看趋势有没有随着 batch/seq 变长而非线性恶化。
  • 注意：打点开 sync 会让数据比真实训练慢，绝对值偏高但相对占比仍可信——调优看占比、不看绝对值。

六、面试速答清单

Q1：Megatron 的 Timer 为什么 start/stop 里必须有 torch.cuda.synchronize()？

PyTorch 的 CUDA op 异步 launch，CPU 侧 time.time() 测的是"把 kernel 提交进队列"的时间，不是 GPU 真跑完的时间，数据不可靠。synchronize() 强制 CPU 等 GPU 队列排空再取时间戳，才能测到真实 GPU 耗时。代价是 sync 有开销，所以细粒度打点用 log_level 控制、生产关掉。

Q2：log_level 和 DummyTimer 怎么控制开销？

每个 timer 带一个 log_level（1 粗 2 细），全局 _log_level 由 --log-level 设。timer 的 log_level 高于全局时，Timers.__call__ 返回 DummyTimer——它的 start/stop/reset 是空函数、调 elapsed 抛异常，于是细粒度打点在生产里零开销。这样上层代码不用 if 判断、调试时开高级别即可。

Q3：跨 rank 怎么聚合时间？为什么看 max？

_get_elapsed_time_all_ranks 把本 rank 各 timer 的 elapsed 填进 [world_size, len(names)] 张量（放 GPU），用 all-gather 广播到所有卡。聚合有 max（只看最慢卡）、minmax（min+max 看不均衡）、all（逐 rank，调试用）。优先看 max——分布式训练是木桶效应，最慢 rank 决定整体吞吐，max 反映 straggler。

Q4：barrier 和 barrier_group 是干什么的？有什么坑？

可选地在 start/stop 前做 torch.distributed.barrier(group=barrier_group)，让所有 rank 起跑线一致减少噪声，默认 None 走全局组。坑：barrier 必须所有相关 rank 一起调，否则 hang；所以默认 barrier=False，只在需要精确对齐的粗粒度（如 forward-backward）用 config.barrier_with_L1_time 开。

Q5：elapsed(reset=True) 和 active_time() 区别？

elapsed 是当前窗口累计值，默认 reset 清零，用于"一个 step 累计、log 时取出清零"的窗口统计。active_time 是 timer 自创建以来所有活跃时间总和，reset 不清，用于算"这段代码总共占用了多久"的累计开销。两者分离让窗口值和累计值各取所需。

Q6：用 Timer 定位训练慢点应该怎么操作？

先 --log-level 1 看粗粒度（forward-backward/optimizer/data-loader）找占比最大段；再 --log-level 2 --log-timers all 看细粒度并对比各 rank 找 straggler；通信段占比高则配合 NCCL profiler 看是不是带宽瓶颈。注意打点开 sync 绝对值偏高，看占比不看绝对值。TensorBoard 用 write 记 max-time 看趋势。

七、一张图收口

1
2
3
4
5
6
7
8
9
10
11
12
13
上层:  config.timers('forward-compute', log_level=2).start()  ...  .stop()
                │
                ▼  (log_level 高于全局? 返回 DummyTimer 零开销)
Timer:  start: barrier(可选) → cuda.synchronize → time.time()
        stop:  barrier(可选) → cuda.synchronize → _elapsed += dt, _active_time += dt
        elapsed(reset): 取窗口值(可清零)    active_time: 取累计值(不清零)

聚合:  _get_elapsed_time_all_ranks
        [world_size, len(names)] 张量(GPU) → all-gather(所有 rank 名单必须一致否则 hang)
        → max / minmax / all 三种聚合(过滤 0 值的没打点 rank)

输出:  log(normalizer=N 算每步均值, reset 清窗口, 只 last rank 打印)
        write(TensorBoard, 只记跨 rank max_time 一个标量)

主线一句话：Megatron Timer = cuda.synchronize 保证测的是 GPU 真时间 + log_level/DummyTimer 控制打点开销 + all-gather 跨 rank 取 max 反映 straggler + max/minmax/all 三种聚合输出。 这套小机制把"分布式训练哪里慢"这件事做成了可分级、零开销、跨 rank 可比的工具——读懂它就读懂了"怎么给一个异构异步系统做可信时序测量"。

参考资料

  • Megatron-Core 源码：megatron/core/timers.py
  • 打点使用点：pipeline_parallel/schedules.py、tensor_parallel/layers.py、optimizer/、distributed/finalize_model_grads.py
  • PyTorch CUDA async launch 与 torch.cuda.synchronize 文档
  • 与本文 Megatron-LM 源码精读、显存计算法则篇交叉对照