字节跳动机器学习系统云原生落地实践

2022 年 2 月 18 日 专知

分享嘉宾：辛涛字节跳动高级工程师

编辑整理：赵冬月河北农业大学

出品平台：DataFunTalk

导读：作为一家高度重视人工智能的公司，字节跳动内部构建了一套丰富的产品矩阵，涵盖数据处理、模型开发、离线训练到在线推理等机器学习研发的全流程，以支持抖音、头条等产品的高速发展。亿级的用户规模和不断深研的业务场景，对字节机器学习平台从研发体验、训练时效、任务编排、资源运维等方面不断提出新的要求挑战，以 K8s 为核心的云原生理念正是为解决以上问题提出，并在业界取得了广泛应用。本文主要对字节跳动机器学习平台的云原生化的改造工作进行介绍，期望可以带给读者一些实践经验。

今天的介绍主要围绕以下三点展开：

机器学习系统云原生化动机
机器学习系统云原生化落地
未来展望

机器学习系统云原生化动机

首先和大家简单分享下字节内部机器学习训练场景训练框架特点，及在实际生产环境中遇到的问题痛点。

1. 推广搜场景

抖音、西瓜、火山、头条等作为字节的核心C端产品，推广搜在其中扮演着极其重要的角色。Deep-Wide 是其主流训练模型，亿级别用户规模和 item（视频、商品）特征表决定了它必须是采用 PS-Worker 的训练框架。PS-Worker 框架包含两种角色：

PS（ParameterServer）: 主要存放模型参数，扩展能力相对比较弱，单个异常会导致整个任务的失败，木桶效应强烈，网络带宽大；
Worker: 迭代计算模型参数梯度并返回给PS，单个异常不影响任务，可以弹性伸缩，慢 worker 的梯度会被丢弃。

PS-Worker 框架特点，在工程实践时要求：

PS 必须保证是一个高优资源，保障其稳定性，避免受干扰；
PS 资源尽量保持同质化，例如相同的机型（CPU型号）；
PS 和 Worker 必须考虑网络拓扑的亲和性。

2. CV/NLP等通用场景

字节的产品多以视频和文字的形式展现，CV、NLP 等通用的深度学习技术同样有着广泛应用，这类的场景模型参数能够单机塞满，因此使用 Ring AllReduce 框架。AllReduce 框架没有中心式的节点，只有 Worker 一种角色，有以下特点：

Worker 存储完整模型，具备故障容忍和弹性能力；
Worker 木桶效应强烈，算力要求高，网络带宽需求大。

因此在工程实践上面临以下挑战：

Worker 资源必须同质化（相同GPU）；
Worker 必须考虑网络拓扑的亲和性；
Worker 故障/扩缩需要让相关 Worker 感知和更新通讯拓扑；
对弹性资源的支持需要业务、框架、编排调度相互协调。

3. 历史痛点

受一些历史技术因素影响，字节的 AI 离线场景基本全部构建在 Yarn 生态上，在线微服务和推理则运行在 K8s 生态上，在离线场景技术体系比较割裂。一方面，平台技术的割裂对用户研发体验不佳，同时也不利于机器学习研发流程 DevOps/ AIOps 的演进。另一方面，在离线资源池的隔离阻碍了资源的高效流转，不利于资源利用率的提升。

另一方面，字节复杂多样的业务场景和高速变化的业务需求，对当前的机器学习系统在深度定制、产品多样上也不断提出新的要求。

基于以上业务场景和问题需求，字节开启了机器学习系统的云原生化改造工作。

4. 云原生

那么什么是云原生，云原生为什么可以解决、怎么解决以上问题呢？先从云原生的概念讲起。

① 云原生CNCF定义

概念：在公有云、私有云和混合云等新型动态环境中，构建和运行可弹性扩展的应用。

技术：包括 Kubernetes、容器、服务网格、微服务、不可变基础设施和声明式 API。

上图是 CNCF 官网上云原生产品的全景图，可以看到它的生态其实是非常丰富的。

② 云原生操作系统 Kubernetes

在云原生领域，Kubernetes 作为云原生操作系统，是最核心的组件之一，它有着以下几个特点：

强大的抽象能力
良好的扩展和分布式特性
高度统一的规范标准

K8s 整体架构如上图所示，包括 Master 和 Kubelet 两大部分：

Master 包括 ETCD、API server、scheduler 和 controller manager。

ETCD 是 KV 数据库存储，存储 K8s 中的所有的资源对象。
API server 在 ETCD 的基础上封装了一层缓存，提供了 restful api 支持，是集群所有对象增删改查的唯一入口。
Scheduler 是资源调度中心，用于将工作负载分发到物理节点上。
Controller Manager 是对象控制中心，驱动 K8s 内置对象，直到满足用户指定的终态，它是实现 K8s 声明式语义的一个重要组件。

Kubelet 是单机 agent，负责单机节点上容器生命周期的管理和物理资源的管控上报。

K8s 在以上架构实现之上提供了丰富的可扩展机制：

Scheduler 调度策略可以以 plugin 的方式进行扩展。
Opertor 提供了业务依据业务场景定制 controller 的能力。比如定义一个简单的 job crd 并按照规范实现对应的 job controller ，在用户声明需要一个需要2个 worker 的 job cr 后，job controller 会直接创建2个 worker 。
CRI/CNI/CSI 分别定义了容器运行时、容器网络、容器存储的扩展标准。
DevicePlugin 机制可以很方便的扩展系统硬件资源，像使用 CPU、Mem 一样进行资源的上报、调度和分配，常见的 devicepluin 有 GPU、RDMA 、高性能NIC 等。

标准化的扩展能力让 K8s 不止非常友好的支持了微服务，同时在大数据和机器学习领域也愈加流行，像开源的 Kubeflow 就在 K8s 之上构建了机器学习场景的完整工具集。字节机器学习系统的云原化工作同样是围绕 K8s 生态进行开展。

机器学习系统云原生化落地

1. 云原生机器学习系统概览

上图是字节机器学习系统的整体概览图，自上而下来看，包括以下几部分：

一站式机器学习平台：提供了包括实验环境，沙盒调试，特征工程，离线训练和在线推理等完整 AI 研发生命周期的支持；
统一 Operator：在线推理主要是使用 K8s 原生的 deployment 做支持，离线训练场景则是自研的 operator 进行支持；
统一调度和资源并池：通过在离线统一调度和资源池合并池化，加速资源流传，满足训练推理不同的资源诉求；
异构资源/设备支持：通过 K8s deviceplugin 框架扩展实现 GPU、Habana、RDMA 等异构资源设备的感知上报和调度分配；
通用基础运维组件：日志、指标、Webshell 远程登录、Quota 资源管理等基础运维能力均进行统一形式进行提供。

2. 在线推理场景

在线推理业务场景时延敏感，对稳定性的要求比较高，此部分工作主要围绕以下四个方面展开。

① 共享 GPU 优化资源利用率

构建了容器化形态下的 GPU 共享能力，将 GPU 资源的申请的最小单位从卡级别降低到0.1卡级别；
调度器支持 GPU 资源多维度调度以及多种灵活的调度策略，将多个 Pod 调度到同一张 GPU 卡上，资源维度包括算力、显存、视频编解码，调度策略包括重显存模型和重算力模型的亲和性反亲和性、binpacking 等；
共享模式支持时间片和 MPS 两种并行模式。

② 微拓扑感知资源分配策略优化单机性能

支持 NUMA 亲和性调度，保证为容器分配的 CPU、Memory、GPU 在同一个 NUMA 节点上；
支持 GPU 拓扑感知调度，保证为容器分配的 GPU 资源拓扑最优，通过 NV-Link 加速 GPU 设备之间的数据通信。

③ 高可用MPS

字节内部大量使用了 MPS 来提高 GPU 设备的空间利用率，通过以 Pod Sidecar 的形式实现了 MPS 容器化，将 MPS 的部署粒度从整机所有卡降低到 POD 级别，以减小 MPS 故障影响面。

④ 监控与周边工具

基于 DCGM 工具和 NVML 库，从维度和粒度两个方面提高指标的覆盖程度。维度方面，将指标从单纯显存和算力监控，扩展到空间利用率、频率、温度、PCIE 字节数等。粒度方面，将算力、显存、编解码等利用率指标从卡粒度，精细到服务和进程级别。
实现了内部版 nvidia-smi 工具，解决容器内无法通过 nvidia-smi 看到 GPU 进程的痛点，在保证安全隔离性的前提下，帮助业务调试和诊断进程信息。

离线训练部分是通过统一的 operator 进行支持。

3. 离线训练场景

离线训练场景通过统一的 Operator 进行支持，自研 TrainingJob Operator 既支持公司自研 Lagrange 等训练框架，同时也支持 tensflow、pytorch 开源产品等。

一个 operator 支持多种框架，一方面降低了接入成本，优化了用户体验；
operator 与框架做一个深度的联动，框架可以通过感知 worker 数量变化来动态调整训练的学习率；
目前字节内部微服务以常态混部作为资源优化的主要手段，但对于 socket 和 GPU 服务来说，仍会以 HPA 弹性出让的方式为主，统一的 operator 有利于做在离线一体的弹性策略的支持。

4. 在离线统一调度

字节在线 K8s 集群节点的量级在万台规模，Yarn 集群节点可支撑五万+规模，为支持超大规模的集群调度，自研分布式调度器支持在离线的统一调度。

调度器主要包括三个组：Dispatcher、Scheduler、Binder。Dispatcher 是一个单实例，主要负责监听集群中的 pending pod，然后把这些 pod 基于某些特定的分配策略，分配给不同的 Scheduler 去执行；Scheduler 多实例，各个实例之间是乐观并发工作，负责具体的调度计算工作，同时把计算调度的结果传递给 binder；binder 单实例，单点收敛 scheduler 乐观并发的结果，做一些冲突检查，最终提交调度结果。

调度策略上丰富了对离线场景调度功能的支持，例如 DRF、Gang 调度、优先级抢占、Fairshare 等功能。Gang 调度是离线训练场景基础功能 feature 之一，调度时要求对一批实例（pod）要么全部成功，要么全部失败，不存在部分成功的中间状态。

调度上的统一同时需要要求资源 Quota 统一，在此通过把资源抽象成 Guatanteed 和 BestEffort 两大类资源类型，配合通过 min/max 语义支持资源的弹性超售分配。

5. 异构微拓扑调度

在上文中可以看到字节内部存在多种异构资源，这部分主要通过扩展 K8s deviceplugin框架进行支持；微拓扑感知调度，例如网卡亲和 numa 亲和，可以很大程度提升训练速度，整体流程如上图所示：

Collector 收集静态物理拓扑信息和节点实时已分配信息；
CNR 是自定义 K8S CRD 资源，存储节点拓扑信息；
Scheduler Plugin 根据 CNR 和 Pod 请求信息进行节点的绑定；
NumaAffinity DevicePlugin 完成最终亲和性分配和绑定。

6. 在离线一体的弹性训练

对于在线服务潮汐现象明显，相反离线训练资源短缺的情况，在离线统一由 K8s 做资源编排的背景下，基于 Virtual Kubernetes 实现了在线跨集群的资源整合，以支持离线弹性训练需求。

Virtual-Kubernetes 基于开源的 Virtual-Kubelet（简称VK）项目演进而来，提供了跨集群整合零散资的能力。在 K8s 中 Kubelet 负责节点的资源管控和 pod 启停，Virtual-Kubelet 模拟了这一行为，只是启动的节点是一个不在本集群中的虚拟节点。举例来说，在离线分属不同的 K8s 集群，其中离线的弹性训练实例就可以通过 VK 部署到在线集群中的空闲节点中，上图右下角虚框标识的 Node' 节点实际上就是一个虚拟节点。

以 Horovod 举例，Horovod 是 Uber 开源的一个分式训练框架，使用的是 Ring Allreduce 来进行通信。首先通过 VK 的方式，提供跨集群的潮汐资源。AutoScaler是一个弹性扩缩控制器，监控潮汐资源的状态，通过策略触发 training job 扩缩配置的调整。TrainingJobOperator 提供 worker replica 调整能力，与训练框架相互配合，完成数据的 checkpoint 和 failover 机制等。

未来展望

最后一部分是对未来的展望。

1. 字节云原生践行