📚 课程目录第4章 DeepSpeed内置监控工具

4. DeepSpeed内置监控工具

训练大模型时，最怕的不是显存不够，而是“不知道发生了什么”。DeepSpeed 在提供 ZeRO、Offload、混合精度等优化能力的同时，也内置了一套贯穿训练全过程的监控与日志体系。理解这套体系，才能把吞吐、显存、通信、损失曲线这些关键信号变成可观测、可回溯、可告警的数据。

一、核心原理：DeepSpeed 监控的分层结构

DeepSpeed 的监控能力并不是单一模块，而是由三层组成：最底层是 DeepSpeed Monitor，负责在训练步中采集标量指标；中间层是 日志系统，负责把指标按 rank、按时间窗口聚合后输出到控制台或文件；最上层是 外部集成，通过 TensorBoard、WandB 或自定义回调把指标推送到可视化平台。三者关系可以用下面的结构图表示。

二、DeepSpeed Monitor：指标采集的入口

DeepSpeed Monitor 通过 DeepSpeedMonitor 类实现，核心是 monitor.write_events() 与 monitor.log()。它支持两种典型用法：一是直接调用 engine.monitor 记录自定义指标；二是在配置文件里开启 tensorboard 或 wandb 输出。下面是一个最小示例：

import deepspeed

# 在 DeepSpeed 配置中开启监控输出
ds_config = {
    "train_batch_size": 64,
    "tensorboard": {
        "enabled": True,
        "output_path": "./logs/tb",
        "job_name": "ds-monitor-demo"
    },
    "wandb": {
        "enabled": True,
        "project": "llm-training",
        "group": "zero3"
    }
}

model, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=ds_config
)

for step, batch in enumerate(dataloader):
    loss = model(batch)
    model.backward(loss)
    model.step()

    # 自定义指标写入 Monitor
    model.monitor.write_events([
        ("train/loss", loss.item()),
        ("train/lr", optimizer.param_groups[0]["lr"]),
        ("train/grad_norm", model.get_global_grad_norm())
    ])

注意 write_events 接收的是 (name, value) 列表，name 建议使用 train/、eval/ 前缀，便于在 TensorBoard 中分组。Monitor 默认只在 rank0 写入，避免多进程重复写文件。

三、日志系统：从控制台到结构化文件

DeepSpeed 的日志由 deepspeed.utils.logging 统一管理，支持 logging.INFO、logging.DEBUG 等级别。训练时常见输出包括：step 耗时、吞吐、显存占用、ZeRO stage 状态、通信量等。若需要结构化日志，可开启 steps_per_print 控制打印频率：

ds_config = {
    "train_batch_size": 64,
    "steps_per_print": 50,
    "wall_clock_breakdown": True
}

wall_clock_breakdown 会输出前向、反向、优化器 step、通信等各阶段耗时，是定位“训练慢在哪”的第一手材料。日志默认输出到 stdout，也可通过 deepspeed.utils.logging.configure_logging 重定向到文件。

四、TensorBoard 与 WandB 集成

TensorBoard 集成只需在配置中写 tensorboard.enabled=True，DeepSpeed 会自动创建 SummaryWriter，把 Monitor 中的标量写入 event 文件。WandB 集成类似，但需要提前 wandb login，并在配置中提供 project、group、name 等字段。两者对比见下表：

维度	TensorBoard	WandB
配置方式	`tensorboard.enabled`	`wandb.enabled`
依赖	tensorboard / torch.utils.tensorboard	wandb 账号与
输出形式	本地 event 文件，可离线查看	云端 run，支持协作与对比
适用场景	单机调试、内网训练	多实验对比、团队共享
自定义指标	通过 Monitor 写入	通过 Monitor 或 wandb.log

五、自定义回调：把监控接到业务告警

当内置输出不够用时，可以继承 DeepSpeedMonitor 或直接包装 engine.step()，在每一步后检查指标并触发回调。例如当 grad_norm 超过阈值时发邮件或写告警文件：

class AlarmCallback:
    def __init__(self, threshold=10.0):
        self.threshold = threshold

    def on_step_end(self, engine, step):
        grad_norm = engine.get_global_grad_norm()
        if grad_norm is not None and grad_norm > self.threshold:
            print(f"[ALARM] step={step} grad_norm={grad_norm:.2f} 超过阈值")
            # 这里可以接入邮件、钉钉、Prometheus pushgateway

alarm = AlarmCallback(threshold=10.0)

for step, batch in enumerate(dataloader):
    loss = model(batch)
    model.backward(loss)
    model.step()
    alarm.on_step_end(model, step)

实践建议：把 Monitor 的指标名统一成 train/loss、train/grad_norm、perf/step_time、mem/peak 四类前缀，TensorBoard 和 WandB 都能自动分组，后期排查效率会高很多。

六、常见误区与复盘

误区一：所有 rank 都写日志。多卡训练时若每个 rank 都写 TensorBoard，event 文件会互相覆盖或膨胀。应确保只在 rank0 写入，或使用 deepspeed.comm.get_rank() 判断。
误区二：只记录 loss，不记录吞吐和显存。loss 正常不代表训练高效。建议同时记录 samples/s、step_time、peak_memory，才能发现通信瓶颈或显存泄漏。
误区三：WandB 未登录导致训练中断。在离线环境使用 WandB 时，应设置 WANDB_MODE=offline，或改用 TensorBoard。
误区四：忽略 wall_clock_breakdown。它是定位“前向慢还是通信慢”的关键开关，建议在性能调优阶段始终开启。

注意：DeepSpeed Monitor 的写入是同步的，若每步都写大量指标，可能拖慢训练。建议通过 steps_per_print 或自定义采样间隔控制写入频率，通常每 10～50 步写一次即可。

掌握 DeepSpeed 内置监控工具后，训练过程就不再是黑盒。Monitor 负责采集，日志系统负责落地，TensorBoard 与 WandB 负责可视化，自定义回调负责告警。四者配合，才能把“调优”从经验驱动变成数据驱动。