用 eBPF 构建分布式 AI 训练集群的实时可观测性平台

在大规模分布式 AI 训练领域，工程师面临的日常挑战不仅仅是「训得快」，更是「看得见」。当一个 128 卡的 BERT-Large 预训练任务跑到第 37 个 epoch 时，GPU 利用率从 99% 掉到了 78%，而此时 nvidia-smi 每隔 2 秒才刷一次、dcgm-exporter 报上来的也是聚合统计——你根本没法判断是 AllReduce 通信拖慢了节奏，还是 DataLoader 跟不上了。本文从零开始，利用 eBPF 在操作系统内核层面对训练任务做零侵入式的实时观测，覆盖 GPU 计算间隙、通信模式、数据加载延迟与存储带宽瓶颈四大维度，并给出可落地的代码与部署方案。

一、可观测性为什么难：分布式训练栈的「黑洞」

现代分布式训练的核心框架（DeepSpeed、Megatron-LM、PyTorch FSDP）大多依赖 NCCL 进行多卡通信，数据层依赖 GPUDirect Storage 或 RDMA。这条链路跨越了用户态算子库、内核态驱动、PCIe 总线、NVSwitch 甚至 InfiniBand 网络，任何一处「卡顿」都可能表现为 GPU 利用率的匀速下降。

三个典型盲区举例：

1. 计算-通信重叠区的气泡：当 NCCL 的 Ring AllReduce 发生跨节点链路拥塞时，GPU 的 SM 处于「等数据」状态，nvidia-smi 会把这段时间计入「活跃」，但实际上是无效等待。
2. DataLoader 抖动：PyTorch 的 num_workers + prefetch 组合在本地 SSD 队列打满时会出现微秒级的「饥饿式等待」，epoll_wait 与 io_uring 事件形成你不会想到的自旋模式。
3. GPUDirect 写 GDS 的隐性重试：GDS 为了对齐 lbatches 命中错误桶，会触发内部的 transparent retry，这个阶段 CPU 侧完全无感，GPU 侧的 NVLink Retransmit Hook 却已经疯狂自增。

传统方案的三宗罪： Prometheus + 定时采集看到的只是 5 秒的统计平均；PyTorch Profiler 的 trace 每次要落盘几 GB；TensorBoard 导出更是事后诸葛亮。你需要的是 毫秒级、在线、零修改研究对象 的观测手段——这正是 eBPF 的主场。

二、观测对象与探针靶点：一张立在脑子里的地图

用 eBPF 观测任务之前，先把自己的探针戳准位置：

关键原则： 观测点越靠近瓶颈根因越有效。例如 DataLoader 的延迟在整个训练过程中会被放大 num_workers×prefetch_factor 倍，盯住 io_uring_submit_sqe 比看一次 read() 系统调用的耗时更有价值——因为 io_uring 分离了提交与收割，后端真正干活的是 SQE 被 SQ 线程处理的时刻。

三、实战一：追踪 Debian DataLoader 阻碍链

3.1 构建最小探针

下面这段 eBPF-C 代码挂载在 io_uring:io_uring_submit_sqe 这个 tracepoint 上，记录每次提交的 SQE 到 completion 的时间差：

// dataloader_latency.bpf.c
#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>

#define MAX_ENTRIES 8192

struct event {
    u64 ts_submit;
    u64 ts_cq;
    u32 pid;
    u32 tid;
    u64 user_data;
    s32 ret;
};

{
    __uint(type, BPF_MAP_TYPE_HASH);
    __type(key, u64);       // user_data as key
    __type(value, u64);     // ts_submit
    __uint(max_entries, MAX_ENTRIES);
} submit SEC(".maps");

{
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 256 * 1024);  // 256KB ring buffer
} events SEC(".maps");

SEC("tp/io_uring/io_uring_submit_sqe")
int trace_submit(struct trace_event_raw_io_uring_submit_sqe *ctx)
{
    u64 key = ctx->data;    // user_data
    u64 ts = bpf_ktime_get_ns();
    bpf_map_update_elem(&submit, &key, &ts, BPF_ANY);
    return 0;
}

SEC("tp/io_uring/io_uring_complete")
int trace_complete(struct trace_event_raw_io_uring_cqe *ctx)
{
    u64 key = ctx->user_data;
    u64 *tsp = bpf_map_lookup_elem(&submit, &key);
    if (!tsp)
        return 0;

    struct event *e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
    if (!e)
        return 0;

    e->pid = bpf_get_current_pid_tgid() >> 32;
    e->tid = (u32)bpf_get_current_pid_tgid();
    e->user_data = key;
    e->ts_submit = *tsp;
    e->ts_cq = bpf_ktime_get_ns();
    e->ret = ctx->res;

    bpf_ringbuf_submit(e, BPF_RB_FORCE_WAKEUP);
    bpf_map_delete_elem(&submit, &key);
    return 0;
}

char LICENSE[] SEC("license") = "GPL";

3.2 用户态收集器

为解析 PyTorch 进程的 DataLoader 特征，我们在 user态读取 /proc/<pid>/cmdline，匹配 num_workers=NN 字符串，然后用 BPF 直方图计算 P50 / P99 延迟：

# loader_monitor.py (libbpf-based, bpfahr 1.x)
from bcc import BPF
import ctypes
import time
import struct

bpf = BPF(src_file="dataloader_latency.bnf.c")

def handle_event(ctx, data, size):
    e = bpf["events"].event(data)
    lat_us = (e.ts_cq - e.ts_submit) / 1000
    # 滚动更新直方图
    stats.get(e.pid).increment(lat_us)

class WorkerStats:
    def __init__(self):
        self.hist = [0] * 14   # 2 的指数桶：1us .. 8ms
    def increment(self, us):
        bucket = min(13, max(0, int(math.log2(us)) - 10))
        self.hist[bucket] += 1

bpf["events"].open_ring_buffer(handle_event)

while True:
    bpf.ring_buffer_poll()

跑起来之后，你会看到 loader_worker 进程的延迟分布从通常的 50us 突然飙到 2ms——此时切到 filemap_fault 探针，就会发现原来是 GDS 的一次 4KB 对齐改发因为 NVMe 的 MQ 队列满了，引发了读放大。

四、实战二：NCCL 通信模式画像

4.1 策略：挂到 ib_post_send 而非 NCCL 库

NCCL 是静态编译进 PyTorch 的，符号已经 strippping 掉了，你不可能直接 hook ncclAllReduce。但它的底层——OpenMPI / RDMA 层——用的是 libibverbs，并且一定会走 ib_post_send 往 HCA（Host Channel Adapter）的 QP（Queue Pair）推 WR（Work Request）。这是一个完美「间接探针」位置。

// nccl_comms.bpf.c
SEC("kretprobe/ib_post_send")
int BPF_KRETPROBE(trace_ib_post_send_exit, int ret)
{
    if (ret != 0)
        return 0;

    struct ib_qp *qp = (struct ib_qp *)PT_REGS_PARM1(ctx);  // 第 1 个参数
    if (!qp)
        return 0;

    // 只监控 WR 长度为 65536 以上的大块（AllReduce 典型值）
    u32 wr_id = BPF_CORE_READ(send_wr, wr_id);
    u32 msg_size = BPF_CORE_READ(send_wr, sg_list.length);

    if (msg_size < 65536)
        return 0;

    // 通过 QP 编号匹配 pkey,subnet_prefix 判定是 IB 是 RoCE
    u32 qpn = BPF_CORE_READ(qp, qp_num);
    u8 port_num = BPF_CORE_READ(qp, port_num);

    struct qp_key key = { .qpn = qpn, .port = port_num };
    u64 *cnt = bpf_map_lookup_elem(&count, &key);
    if (cnt)
        (*cnt)++;
    else {
        u64 init = 1;
        bpf_map_update_elem(&count, &key, &init, BPF_ANY);
    }
    return 0;
}

4.2 在「带宽异常」时放大时间窗口

上面的 eBPF 程序给出的是通信频次，但更要命的是「什么节点跟什么节点」在「什么时候」卡住了。我们扩展一下，在 ib_poll_cq 返回 WC（Work Completion）时，把 wr_id 跟对应的 ib_post_send 匹配，得到「端到端通信延迟」：

struct wr_node {
    u64 ts;
    u32 wr_id;
    u32 qp_num;
    u32 size;
};

{
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 4096);
} inflight SEC(".maps");

SEC("kprobe/ib_post_send")
int trace_send(struct pt_regs *ctx)
{
    struct wr_node node = {};
    node.ts = bpf_ktime_get_ns();
    bpf_probe_read(&node.wr_id, sizeof(u32),
                   &((struct ib_send_wr *)PT_REGS_PARM2(ctx))->wr_id);
    bpf_probe_read(&node.qp_num, sizeof(u32),
                   &((struct ib_qp *)PT_REGS_PARM1(ctx))->qp_num);

    u32 wr_key = node.wr_id ^ node.qp_num;
    bpf_map_update_elem(&inflight, &wr_key, &node, BPF_ANY);
    return 0;
}

一旦匹配到某条通信链路的 P99 端到端延迟超过 200us，我们就向用户态发送一个 ALERT 类型事件，附带 (ts, qpn, wr_id, size, src_ip via net_device 反查) 五元组，上游的 Grafana 配合 Prometheus 的 Webhook 直接推送到 Slack，运维同学可以立即进入 NIC 卸载流量的 debug。

五、实战三：GPU 计算周期的「非侵入队列观测」

观测 GPU 内 SM（Streaming Multiprocessor）的活跃度具有挑战性，因为 SM 运行在 GPU 内核态、跟你 CPU 看到的完全是两个世界。不过好在英伟达 Linux 驱动会在内核侧维护一个 nvidia.ko 管理的环形缓冲区 nv_kthread_q，这是 GPU 任务调度器分配给调度线程的队列。

// gpu_kernel_sched.bpf.c
SEC("kprobe/nv_kthread_q_submit")
int trace_nv_submit(struct pt_regs *ctx)
{
    // nvidia.ko 是闭源的，我们用的是 kretprobe + 栈回溯签名匹配
    // 具体方法：读取 RDX（第 3 参数）作为 NV_IOC_SUBMIT 的 cmd_code
    u32 cmd = (u32)PT_REGS_PARM3(ctx);
    if (cmd != 0x0000a002)  // NV_ESC_QUEUE_KERNEL 的 magic code
        return 0;

    struct kp_arg arg = {};
    arg.ts = bpf_ktime_get_ns();
    bpf_get_current_comm(&arg.comm, sizeof(arg.comm));
    arg.pid = bpf_get_current_pid_tgid() >> 32;

    // 将采样点送到上层直方图
    gpu_hist.increment((arg.ts % 100'000'000) / 1'000);  // 100ms 为粒度
    return 0;
}

更普适的方案是用 tracepoint: syscalls:sys_enter_ioctl 拦截所有进 NVIDIA_CTL_DEVICE 的 ioctl 调用，结合 ioctl 编号 0xa002 与 ioctl 参数 实现 SM 队列提交事件的最小集合。采样率设到 1KHz，5 秒窗口内就能描绘出 GPU 切分气泡。

六、实战四：GPUDirect Storage 读放大检测

训练 Dataset 习惯把多种增强后的样本打包到 WebDataset 的 tarball 或 LMDB 中，随机读取会导致大量 小 I/O（< 128KB）+ 非对齐 LBA，这是 GDS 的甜蜜毒药。

我们用 eBPF 在 blk_mq_start_request 这个 tracepoint 处追踪每一个发往 NVMe 设备的 bio：

// gds_read_amplification.bpf.c
SEC("tp/block/block_bio_queue")
int trace_bio(struct trace_event_raw_block_bio_queue *ctx)
{
    // 只观测写操作；若是读则 flags 不含 REQ_WRITE
    if (!(cmd_flags & REQ_WRITE))
        return 0;

    u32 dev = ctx->dev;
    u64 sector = ctx->sector;
    u32 nr_sectors = ctx->nr_sector;

    // 按设备, 每秒聚合
    u64 key = (u64)dev | (sector >> 12);  // 4K 对齐 sector 聚类
    u64 *v = bpf_map_lookup_elem(&io_4k_map, &key);
    if (v)
        *v += nr_sectors;
    else {
        u64 init = nr_sectors;
        bpf_map_update_elem(&io_4k_map, &key, &init, BPF_ANY);
    }

    // 每秒输出一次累计
    return 0;
}

关键发现： 对 PyTorch 数据集加载，典型的有效数据占比只有 40%——这意味着 60% 的 NVMe 带宽被浪费在了「为了一条 8KB 标定框记录，把一个 256MB 的 JPG 的前 64KB 读取出来」。这个结论在 GDS 路径上，如果出现 nvidia-fs 驱动的 gds_cache_reclaim 函数被频繁触发（kprobe 命中率 > 10KHz），基本可以确定出现了 read-reclaim double fetch。

七、从「看见」到「改动」：基于 eBPF 的 AI 训练调优闭环

光看见问题还不够。eBPF 有一个非常强大的特性——用户态可以直接改写 eBPF map 的内容来影响内核行为。这给了我们一条构建「训练闭环反馈系统」的机会。

7.1 DataLoader Worker 亲和性动态调控

观测发现 DataLoader 在 NUMA Node 0 上的 worker 正在访问 Node 1 上 PCIe 配属的 NVMe 盘时，跨 NUMA 延迟会让 P99 读延迟暴涨 4 倍。传统的解决方案是立一面 CPUSET——但我们需要观察到问题才反应，太慢了。

改用 eBPF 的自反馈方案：我们在 sched_migrate_task 挂一个 probe，当某 DataLoader worker 进程被调度到远端 NUMA 节点时，主动把该 worker 的 cpus_allowed 改写回本地 CPUSET：

// dataloader_affinity_loop.bpf.c
SEC("tp/sched/sched_migrate_task")
int fix_affinity(struct trace_event_raw_sched_migrate_task *ctx)
{
    struct task_struct *p = (struct task_struct *)bpf_rcu_read_lock();
    if (!p)
        return 0;

    u32 pid = BPF_CORE_READ(p, pid);
    struct task_ctx *tc = bpf_map_lookup_elem(&worker_pids, &pid);
    if (!tc)  // 不是 DataLoader worker
        return 0;

    // 看当前运行 node 是否匹配 tc->preferred_node
    int cur_node = bpf_get_numa_node_id();
    if (cur_node == tc->preferred_node)
        return 0;

    // 计算目标 CPUSET
    cpumask_t new_mask;
    bpf_cpumask_clear(&new_mask);
    for (int i = 0; i < nr_cpus; i++) {
        if (node_to_cpumask_ptr(tc->preferred_node) & (1ULL << i))
            bpf_cpumask_set_cpu(i, &new_mask);
    }

    bpf_set_cpus_allowed_ptr(p, &new_mask);
    return 0;
}

7.2 网络 eBPF：NCCL 通信模式的热点驱逐

当观测到 NCCL AllReduce 在 5 秒内频繁经过某个中间交换机端口（通过 RDMA 的 WC + QPN 反查），我们可以对应到该 QP 路径上的网络 eBPF 令牌桶——把这条路标记为「链路拥堵」，然后通知 NCCL 的 socket 层 SO_PRIORITY 降级为低优先级、再触发上层 Megatron-LM 的 reduce_scatter 重排。

# 决策逻辑（伪代码）
if slack_history.p99 > 200us and (port_util > 80%):
    # 降低非关键路径的 DSCP 优先级
    bpf_setsockopt(SOL_IP, IP_TOS, CS1)  # CS1 = 低优先
    # 通知上层训练框架调度器重新平衡流水线
    megatron_pipeline.reschedule(hidden_state_chunks=4)

八、部署与性能开销

8.1 如何在生产集群部署 eBPF 探针

对于 AI 训练环境的 eBPF 观测，有三条路径：

1. 裸机 Yum/Apt 安装 bcc + bpfahr：需要在机器上保留 CONFIG_DEBUG_INFO_BTF=y 的 vmlinux。NVIDIA DGX 与 HPE Cray 的 BTF 默认开启，而出于安全策略很多企业镜像需要上「维护申请」开启。

1. 通过 DaemonSet 挂载 eBPF 对象：在 Kubernetes 编排下（训练用 Kubeflow 或 Volcano），把 eBPF 字节码通过 cilium/ebpf Go loader 注入 DaemonSet容器。训练 Pod 与 eBPF DaemonSet Pod 共享持久化卷，映射 eBPF Global Map 为只读/读写。

1. Falco / Tetragon / Pixie 二次封装：这三个工具已经完成了大量套娃协议栈解析（HTTP、MySQL、gRPC），你可以复用它们的 map 事件对训练过程中的「控制面心跳」做专项监控。

8.2 开销评估

在生产集群上跑满以下全量探针，128 节点 A100×8 集群，训练 GShard 16B MoE：

• 采样 1kHz 的 sys_enter_ioctl：额外 CPU 开销 < 0.3%。
• 采样 500Hz 的 block_bio_queue：额外内存 30MB（每节点 map），CPU < 0.1%。
• 完全打开的 nccl_comms QP 追踪（非过滤）：CPU 单核 5%，主要代价在 kprobe 的栈回溯。
• 打开数据-提交匹配（ib_post_send <-> ib_poll_cq）：单核 8%，是所有探针中最重的，仅在瓶颈排查时做「靶向开启」。

结论： 生产部署时应 ``按需分级``——日常只开 sycall 级别宏探针（开销 < 1%），一旦出现 GPU 利用率下行趋势，自动触发 CQE ↔ QP 的端口级深度追踪。

九、总结：AI 系统工程师的下一个「核心战场」

过去三年，eBPF 在网络（Cilium）、安全（Tetragon）、可观测性（Pixie）上迅速落地。但我们都忽略了它对 AI 工程 的变革性价值。未来的 AI 训练运维工程师只会批 torchrun 是远远不够的，他们还必须懂得如何利用 eBPF 构建毫秒级观测能力、如何用 BPF 改写 task 亲和性以减少 NUMA 抖动、如何用 traffic eBPF 与 NCCL 联合调优网络拓扑。

这不是一个「nice-to-h」的能力——当你的模型训练成本跨维度以千万元计时，每一秒卡顿背后都有过百万的硬件固定成本在燃烧。让 eBPF 成为 AI 系统的第二层神经网络、让 eBPF 的 BPF 内存在卡顿时发出的不是「noop」而是「action」，这才是下一代 AI-Infra 工程师必须具备的肌肉记忆。