Kimi 大模型训推混部的稳定性与资源优化实战

演讲嘉宾｜黄维啸

编辑 ｜Kitty

策划 ｜QCon 全球软件开发大会

在 InfoQ 举办的 QCon 全球软件开发大会上，月之暗面系统工程师黄维啸分享了“Kimi 稳定高效的 LLM 基础设施构建之道”，他介绍了月之暗面在训推混部集群中的实践经验，重点探讨如何快速定位并隔离故障，实现任务的高效恢复，从而提升系统整体稳定性。同时，他还分享了如何在资源有限的情况下最大化利用率，避免浪费，进一步将该思路应用于强化学习任务的训练中。给关注大规模模型训练的同行们提供了一些可靠的技术思考与实践参考。

预告：将于 10 月 23 - 25 召开的 QCon 上海站设计了「大模型推理的工程实践」专题，聚焦于大模型推理的工程实践，无问芯穹、百度、昆仑芯探讨如何通过技术创新、架构优化和工程实践，提升大模型推理的效率和可靠性等。敬请关注！

以下是演讲实录（经 InfoQ 进行不改变原意的编辑整理）。

今天我给大家介绍一下 Kimi 稳定高效的 LLM 基础设施的一些工作，我会分为四个主题来讲。第一个是关于大规模训推集群的挑战。第二个是全链路稳定性的提升。第三个是资源的高效利用，尤其是面对大量资源时。第四个是强化学习方面，我们如何充分利用训练和推理的资源，从而让强化学习效率变得非常高。

大规模训推集群的挑战

在大规模训推集群方面，我们遇到了以下挑战：首先，当资源规模庞大时，故障频次会显著增加。当机器资源较少时，故障率的影响可以忽略不计，但当机器资源非常多时，故障几乎每隔几小时就会出现一次。如果任务规模庞大，任务会频繁中断，这会导致训练效率大幅下降。因此快速监测并恢复故障是非常关键的。

其次，资源的使用不够高效。许多团队共同使用这些资源，尤其是云服务相关的用户，会发现开发机和文件存储存在大量浪费的情况。例如，为用户分配了一台 GPU 开发机，如 H800，但很多时候 GPU 的占用率会出现短暂峰值后就闲置，这造成了巨大的成本浪费。此外，用户的文件存储和模型数据管理混乱，不清楚哪些数据可以删除，哪些需要保留，这涉及数据治理问题。如果仅根据用户用量进行自动清理，可能不是最优方案，我们需要建立高效的资源使用范式。

第三，大模型应用以及许多线上应用呈现出明显的潮汐效应。例如，在深夜时间，Kimi 的使用率会很低，线上服务的流量存在明显波动。我们需要根据时间动态分配资源。与 CPU 任务不同，Kimi 的服务存在 PD 分离这样的负载，我们需要高效地应对潮汐效应，同时确保 PD 分离不会出现问题。

最后，在今年年初发布的 Kimi k1.5 论文中，我们涉及了一些强化学习的工作。在强化学习中，训练和推理两种工作负载都会在机器上运行。我们需要平衡训练和推理的效率，使二者能够高效利用资源，并确保训练效率达到最优，这也是一个需要解决的重要问题。

全链路稳定性提升

我们针对稳定性问题，包括机器故障率高的问题，采取了一系列措施来解决并提高资源利用率。

首先，我们实施了任务全生命周期的监控。演讲 PPT 的标题是“Varys”，如果大家看过《权力的游戏》，应该知道里面有一个情报总管叫 Varys，绰号“八爪蜘蛛”。他在《权力的游戏》的每一个角落都能布置眼线，收集各种信息和情报。因此我们的系统也希望像“八爪蜘蛛”一样，建立全面的监控体系，对训练任务的全生命周期进行监控。具体措施包括以下几点：

1. Precheck 与周期巡检：我们会进行任务前的预检查（Precheck），以及定期的周期性巡检。这些都是比较通用的做法。

2. 调用栈全链路监控：我们会对调用栈进行全链路监控。在 CPU 端，特别是在 Python 侧和 torch 侧，我们会进行调用栈监控。在 GPU 端，我们会插入一些 CUDA 事件，对 CUDA 操作进行监控，确保整个调用链路的稳定性。

3. 智能日志分析：这一部分，许多友商也有相应的解决方案。我们会进行离线日志分析，同时结合用户使用的 TensorBoard，对日常异常指标进行检测，通过智能分析提高问题发现的效率。

4. 连续异步 checkpoint 机制：我们设计了一种连续异步 checkpoint 机制，后续我会详细说明这一机制的具体内容和优势。

PreCheck、周期巡检

我们会进行 Precheck 和周期巡检工作。Precheck 是在任务启动之前进行的，主要做一些基本的故障检测。这些检测有不同的级别，对于一些希望快速启动的任务，可以跳过部分检测。我们通常会进行一些标准的集合通信性能检测，例如 NVIDIA 的 all reduce 以及其他集合通信检测。但这些检测不一定能完全暴露机器的问题。因此，我们会测试一些 GPU 的 d2h（设备到主机）相关的硬件指标，因为这些指标可能无法通过巡检暴露出来。对于一些重要任务，需要检查 GPU 的 d2h 是否出现 PCIe 降速等问题。

当使用训练框架进行 Precheck 时，框架会创建许多通信 group，这些 group 之间会进行各种集合通信。利用这些 group 进行检测是很自然的，尤其是当标准集合通信检测无法完全暴露机器问题时。例如，NVIDIA 的通信机制中，如果进行全局 all reduce，可能会采用 0 号卡进、7 号卡出，再 7 号卡进、0 号卡出的策略。这种机制无法检测 0 号到 7 号卡之间其他卡（如 1、2、3、4、5、6 号卡）的互通情况。但在大模型训练中创建的 group 里，可能会让同号卡之间进行通信。这种情况下，标准集合通信检测就无法监测到问题，因此我们使用训练中创建的 group 进行 group 通信检测，避免任务被调度到故障机器上。如果出现问题，会告警给用户，用户可以根据任务的容忍程度决定是重启任务还是继续运行。

我们还会进行周期巡检和故障维护。对于 GPU 故障和网线故障，特别是 ROCE 网线故障，这种情况在集群刚交付时会比较频繁。这是因为新交付的机器可能存在网线批次问题，所以一开始需要进行大量网线故障检测，但随着时间推移，故障会逐渐减少。经过这些检测后，我们可以筛选出一批健壮的机器，用于更好地服务大模型的预训练。对于发现的故障机器，我们会立即进行维护，等待任务排空后尝试重启自愈。如果故障仍未修复，我们会通过自动化流程通知供应商。整个过程中基本不需要人工介入，一旦机器出现故障，供应商会立即收到通知，并安排 On Call 人员进行修复。修复完成后，机器会自动进入自检流程，然后重新加入集群。这些就是一些基础的 Precheck 和周期巡检工作。

调用栈全链路监控

在监控 CPU 调用栈时，可以获取很多信息，例如 Python 调用栈可以告诉我们代码卡在了哪一行，以及是否有 worker 在等待其他 worker。这些信息可以在 CPU 上初步判断出来，我们会周期性地记录 Python 调用栈，并在任务崩溃后记录 Torch 调用栈。然后，我们会对这些调用栈进行聚类，方便用户快速找到问题。

如下图所示，纵轴表示 worker 的数量，横轴表示聚类类型。聚类类型本身没有具体意义，只是用于区分不同的聚类结果。图中可以看到三个蓝点，左上角的蓝点表明大部分任务卡在了这里，正在等待。它主要等待的是中间和右下角的两个 worker。具体是在等待什么内容呢？通过日志可以查看，点击日志可以看到问题发生在 Python 代码的某一行，从而快速定位是代码问题还是机器问题。因为前面提到的故障检测并不能完整地暴露机器故障信息，有时甚至会遇到某台机器的某张显卡在运行 Flash Attention V3 时卡住，这种情况无法通过常规监控发现，只能通过调用栈追踪手段来发现。

通过聚类信息定位到 Python 代码的某一行后，如何知道是哪个 kernel 调用或哪个通信卡住了呢？这就需要 GPU 调用栈监控。我们会分级别插入 CUDA 事件，并周期性地导入事件信息，将其集合成一个时序图。例如，可以看到一些 PP（Pipeline Parallelism）的 rank 在完成 forward 操作后会进行 send 和 receive 操作，而一个组的 rank 会执行 forward 和 backward 操作。通过这个时序图，我们可以清晰地看出是哪个 DP（Data Parallelism）卡住了，以及是 forward 还是 backward 操作中出现了 GPU 硬件上的卡顿。通过这种方式，我们可以比较完善地找到是代码问题还是其他地方出现的卡顿，这种方法非常高效。

智能日志分析

智能日志分析的核心工作是从任务的各种日志中查询和匹配特殊模式，并且结合之前获取的调用栈信息以及机器故障情况，进行综合分析。通过这种方式，我们可以生成一份非常完善且详细的任务分析报告。

例如，当一台机器出现“CUDA out of memory error”（CUDA 内存不足错误）时，仅从这个错误信息本身，我们很难明确具体原因。但通过进一步分析，我们可能会发现 NVIDIA Collective Communication Library（NCCL）通讯超时的情况。对于每一个异常点，我们都可以查看详细的日志来深入了解问题。最终，如果发现任务中出现了 XID（eXception ID，异常 ID）事件，我们就能清晰地判断出是这台机器的 GPU 出现了故障，从而导致了任务的异常。

此外，对于一些 Python 中的致命错误（fatal Python error），我们也可以通过点击详细日志来查看，明确问题发生在代码的哪一行，从而快速定位问题所在。这种综合分析方法对于任务排查非常有效，能够帮助我们快速定位问题根源，无论是代码问题还是硬件故障。

连续异步 checkpoint

异步 checkpoint 技术大家应该并不陌生，这项技术已经出现很久了。我们开发这项功能也已经有段时间了，大概从去年年初开始开发并投入使用，一直用到现在。我们实现了一个名为“save continuous”的功能，即针对一些非常重要的预训练任务，开发了一个近乎无损的异步 checkpoint 服务。

我们希望任务在存完一个 checkpoint 后，能紧接着存下一个。虽然从理论上讲应该如此操作，但实际中存在诸多挑战。例如，如果频繁保存 checkpoint，存储系统可能难以承受，尤其是存储容量方面。因此，我们需要一套 checkpoint 备份同步系统，针对短期内新增大量 checkpoint 的场景，实现快速删除和采样备份等功能。这样，我们就能很好地运行整套 pipeline。任务会时刻保留前几个 checkpoint，并对其进行采样备份。一旦任务失败，能迅速从最新 checkpoint 加载。

目前，我们在存 checkpoint 和 load checkpoint 方面做了很多分布式优化。存 checkpoint 时无需进行通信，只需执行一次设备到主机（D2H）的操作，然后异步存储即可。加载时，每台机器各自加载自己的部分。这样一来，存和加载的速度都非常快。整体来看，这套系统运行起来后，对于任务，尤其是大规模节点的预训练任务，收益是非常显著的。

高效资源利用

前面我主要讲了预训练任务中的故障检测与隔离手段，接下来主要讲如何提高资源利用率，避免资源浪费。

全方位 LLM 开发效率提升动态申请云上开发资源

我们会动态申请云上的开发资源。基于云的 ECS（Elastic Compute Service），我们为用户提供开发机，并在 ECS 上实现容器化功能，以支持用户使用开发机。之所以不直接使用 ECS，是因为 ECS 不支持自定义镜像功能，而我们希望用户能够使用 Docker 镜像来自定义开发机。这样，用户可以使用任意 Docker 镜像启动开发机，启动后通过远程方式交互，提交 GPU 资源用于调试。需要注意的是，这种场景主要用于开发机，不推荐用于大规模任务。用户拥有一台 CPU 开发机后，通过远程启动操作，就能高效利用 GPU 资源。例如，用户可以自动启动一些交互式的 GPU worker，直接连接到这些 worker 上，且文件系统完全互通。任务完成后，我们会自动清理或用户可以通过 Ctrl+C 释放 GPU worker，避免一直占用 GPU 资源但不使用而造成资源浪费。对于 CPU worker，由于开发机资源有限，可能需要较大的 CPU 资源，例如进行模型结构转换或数据处理时。此时，用户可以提交交互式的 CPU worker，利用云的弹性功能，自动扩容一批按量付费的节点，任务完成后自动回收，实现良好的成本管控。

任意级目录用量统计

我们实现了一个任意级别目录用量统计功能。虽然这不是什么新功能，但由于我们对接的系统众多，涉及多种云存储，而并非所有云存储都能实现任意级目录用量查询。因此，我们开发了一个名为 FScounter 的工具，能够实时扫描所有目录。尽管扫描过程耗时，但我们通过缓存机制避免每次都进行全局扫描，从而实现对任意级目录的容量、节点数和文件数的监控。我们将其整合成一个饼图，用户可以直观地看到文件系统中资源的使用情况，例如某个文件占用较多空间。这就像电脑上的实时监控助手，帮助用户了解资源使用情况，方便用户管理。

模型异步 evaluation

我们实现了模型的异步 evaluation 功能，这与前面提到的异步 checkpoint 技术相辅相成。在保存 checkpoint 后，有些 checkpoint 需要备份并进行 evaluation。我们会监控这些 checkpoint，并利用空闲资源进行 evaluation。这样，用户在训练任务后不久就能看到 evaluation 结果，对 GPU 资源的利用也非常高效。

跨区域、多实验、灵活对比功能

我们实现了一套跨区域、多实验、灵活对比的功能。我们有一个统一托管的 TensorBoard，并优化了 TensorBoard 的数据读取，支持任意实验之间的对比，包括不同 set 和不同 room 之间的实验。例如，用户在一个 set 中用一种机型做实验，在另一个 set 中用另一种机型做实验，也能灵活对比。此外，我们还支持子实验的合并以及实验之间的互相对比等功能。通过这些功能，用户对平台的使用体验以及对 GPU 资源的利用效率都得到了显著提升。

跨机房推理模型分发

我们实现了一个跨机房推理模型分发功能。由于我们的服务部署在多个地方，包括自建机房和云厂商等，我们需要将一个模型分发到这些不同地点。为此，我们采用了一套跨机房推理分发方案。在实现过程中，我们复用了开源项目 Dragonfly 的部分方案，并对其进行了配置和改造。我们构建了一个基于中心化的模型分发系统，通过暴露一个 HTTP 服务器，其他所有节点可以基于点对点（P2P）模式拉取对应的模型，并高效缓存模型的权重。通过这种方式，我们基本可以将模型拉取时间控制在几十秒级别。如果模型已经被缓存，拉取速度还会更快。

训推多级潮汐系统

我们实现了一套训练多级潮汐系统，这套系统的开发时间相对较早，大概在去年上半年。当时，Kimi 刚刚开始受到广泛关注，推理服务迅速占据了大量 GPU 资源，导致训练任务几乎没有可用的卡。因此，我们急需开发一套训练多级潮汐系统来解决这一问题。目前，我们将资源分为四个等级：

1. 在线推理服务：这是我们的核心服务，必须保证稳定运行，避免用户出现卡顿等问题。

2. 不可抢占的训练任务：这些任务优先级较高，我们希望它们不受干扰。

3. Spot 训练任务：这些任务可以被潮汐系统抢占。当流量高峰时，Spot 训练资源会被挤占；当流量低谷时，这些资源会被释放出来。

4. 低优先级的离线推理任务：例如异步 evaluation 等，这些任务优先级最低。

系统的设计如下：

• 在线推理服务（用红线表示）占用的资源数量是固定的，必须保证其稳定运行。

• Spot 训练资源（用绿线表示）在流量高峰时会被挤占，低谷时释放。未被 Spot 训练任务用满的资源会被离线任务（用橙色块表示）利用。

• 不可抢占的训练任务（用蓝色表示）虽然优先级高，但也不会永远占满资源，空闲时离线任务也会利用这些资源（图中未画出）。

整个系统资源基本处于满负荷状态，通过潮汐机制高效利用 GPU 资源。然而，这种潮汐机制也存在一定风险，尤其是在参数与数据分离（PD 分离）的场景下。我们需要让调度系统能够感知到具体的参数（P）和数据（D），并进行相应的分离操作。有时还需要进行成对的潮汐操作，因此需要考虑多个优先级。例如，根据流量情况和用户优先级，动态调整推理和解码任务的资源分配，以满足当前负载需求。这与传统的 CPU 资源管理有所不同。

强化学习中的混合部署强化学习 Infra 的挑战

如何在强化学习中混合部署训练和推理，尤其是训练和推理之间存在交互时，如何高效利用资源、避免资源浪费。强化学习的基础设施面临诸多挑战，这里主要提到了两个比较典型的方面，当然实际挑战不止这些。

首先，强化学习的计算流程非常复杂，其训推 pipeline 流程比预训练和单纯的 PD 分离推理都要复杂。训练和推理的计算流对硬件的要求不同，例如训练可能主要在 H800 上进行，而推理可以在 H800、H20 或国产硬件等多种异构资源上完成。这就需要我们考虑如何充分发挥这些异构资源的优势，以实现高效的强化学习流程。

其次，在 rollout 过程中，用户的请求回复长度各不相同。一开始可以设置较大的 batch size，但如果到了快结束时，集群可能会出现闲置，因为一些长请求在长链式推理情况下会占用大量资源，导致整体强化学习的训练速度变慢。此外，训练和推理是两种不同的模式。训练框架和推理框架运行的是两套不同的代码，各有各的优化方案。不能简单地直接用训练框架进行推理，因为虽然推理主要是 forward 流程，但训练框架中只实现了预训练和 SFT 中的一些重要优化，而推理框架中有许多技术栈需要优化。这就使得将训练和推理框架融合在一起变得非常困难。

以 Kimi K1.5 论文中的情况为例，我们有 Megatron 训练框架和 vLLM 推理框架，如何将这两套框架融合是个难题。而且训练和推理任务的 checkpoint 格式完全不同，训练时为了高效存储和读取 checkpoint 需要特定格式，而推理一般使用 Hugging Face 格式，如何在这两种格式之间高效转换也是个问题。此外，训练和推理任务相互切换 GPU 资源会造成一些闲置和浪费，且它们的并行模式也不一样。例如，一些 Moe 模型在训练时可能开启了较大的 EP，而在推理时可能既有 EP 又有 TP，将这两种模式融合起来比较困难。

k1.5 RL System Overview and Partial Rollout

针对这些挑战，我们设计了一套较为完善的 IO 系统。在 Kimi K1.5 论文中有一张图展示了整个系统的模式。系统中有一个 master 节点，它负责准备训练数据，然后将这些数据发送给 Trainer worker。Trainer worker 会进行梯度更新，更新完成后，它会将权重发送给 roll out worker，也就是推理流程的节点。roll out worker 开始执行 roll out，得到数据后进行打包，同时进行 evaluation，然后将数据发送回 master。master 节点有一个 Replay buffer，用于记录尚未结束的 roll out worker，以及将结束的数据重新发送给训练环节，从而将整个 pipeline 运转起来。

在 rollout 过程中，由于请求回复长度不一致，导致了一些长尾问题。为了解决这个问题，我们设计了一个优化方案，称为 partial rollout。在长链式推理场景下，我们会设定一个阈值。如果某些请求达到这个阈值，它们可能会被截断，被截断的部分会被保存到 Replay buffer 中。在下一个迭代时，Replay buffer 会优先处理这些被截断的请求，并且这些请求的优先级会高于其他请求。这就需要我们在虚拟机层面和调度层面建立一套优先级系统，以实现这个 Replay buffer 机制。通过这种方式，我们能够较好地缓解长尾问题。虽然长尾问题无法完全解决，但通过这个方案，我们能够显著提升长链式推理场景下 GPU 的利用率。

k1.5 Hybrid Deployment

Kimi K1.5 采用了 Hybrid deployment，基于容器体系。具体来说，每个机器上有一个 Pod，每个 Pod 包含两个 sidecar 容器：一个名为 Megatron，用于训练；另一个名为 vLLM，用于推理。训练和推理被分在两个不同的 sidecar 中，这样做的主要目的是因为训练和推理镜像中分别进行了针对各自任务的优化，甚至 CUDA 版本也可能不同。通过使用不同的镜像和 sidecar 进行隔离，系统变得更加可扩展，且可以灵活替换镜像。

在训练过程中，会加载训练所需的权重，训练完成后，训练权重会被 offload，同时将其转换为 Hugging Face 格式。然后，Hugging Face 格式的模型会被注册到共享内存和全局注册表中。接下来，通过 RDMA 技术，基于 mooncake 引擎，将模型直接发送给 vLLM，vLLM 会上传其权重并开始 rollout。rollout 完成后，可以选择终止 vLLM 进程。

终止有两种方式：一种是直接杀死进程，另一种是进行 offload。具体选择哪种方式取决于显存的使用情况。如果显存非常紧张，可能只能选择直接杀死进程；如果显存情况较好，则可以选择 offload，尽管 offload 并不能完全卸载所有资源，例如一些与 NCCL 等相关的资源很难完全 offload，但是也能节省大部分的 GPU 显存

整个流程能够高效地利用资源。与直接启动 vLLM 并将 checkpoint 保存到磁盘相比，通过内存、RDMA 进行传输的速度要快得多。主要优势在于利用 sidecar 隔离环境，训练和推理共享资源，避免资源闲置，并且通过 RDMA 和内存传输（包括显存直接传输）来提高效率。 Greedy Rollout

我们能够兼容训练和推理的不同并行策略。但这里有一个问题：假设推理所需的节点比训练更多，特别是在长链式推理场景下，样本数非常多，我们希望达到很大的 batch size，但训练并不需要这么多机器，那该怎么办呢？我们有一个解决方案，叫做 greedy rollout。这个方案主要解决的是，当推理的 batch size 需求很大，但现有机器无法满足时，我们可以通过额外启动一批称为 isolated 的 Pod 来解决。这些 isolated Pod 只运行推理任务，且基于相同的框架，通过 RDMA 从 Pod 中获取权重并启动。

整个流程如下：scheduler 通过启动 Pod，然后通过 greedy 方式启动一堆 Pod。这些 Pod 启动后，会进行 rollout，rollout 会直接请求到一个 Hybrid rollout proxy 上，这个 proxy 会进行负载均衡。这样可以确保即使有大量的并发请求，系统也有足够的资源来处理。处理完成后，如果 Isolated Pod 发现数据量减少，它会自动退出。通过这种方式，我们可以最大化利用空闲资源和异构资源，提高 rollout 的速度，尤其是 Isolated Pod 可以部署到其他硬件上，如 H20 或国产硬件，从而更好地实现整个系统的吞吐能力。

这个流程中存在一些重点问题。例如，在进行多机推理时，特别是 PD 分离或多机推理场景下，RDMA 的使用可能会出现冲突，这种场景下 RDMA 的配置需要特别注意，以避免吞吐层面的冲突。此外，rollout 的请求需要进行动态负载均衡，我们通过 proxy 实现了这一点，避免某些机器启动后没有收到请求。同时，Isolated Pod 和 Core Pod 可能是异构的，这意味着我们可以使用不同的硬件，如 H20 或国产芯片，并在这些硬件上启动优化后的 Pod，从而在系统层面最大化吞吐量。基于这些措施，我们可以高效地实现强化学习，更高效地利用 GPU 资源。

嘉宾介绍

黄维啸，月之暗面系统工程师。毕业于清华大学，7 年 AI Infra 系统经验。目前在月之暗面负责 Infra 平台、系统优化相关工作。曾在旷视科技公司主导公司 AI 平台 Brain++ 从 0 到 1 的研发工作。

会议推荐

10 月 23 - 25 日，QCon 上海站即将召开，现在大会已开始正式报名，可以享受 8 折优惠，单张门票立省 1360 元（原价 6800 元），详情可联系票务经理 18514549229 咨询。