位元組跳動機器學習系統雲原生落地實踐

導讀：作為一家高度重視人工智慧的公司，位元組跳動內部構建了一套豐富的產品矩陣，涵蓋資料處理、模型開發、離線訓練到線上推理等機器學習研發的全流程，以支援抖音、頭條等產品的高速發展。億級的使用者規模和不斷深研的業務場景，對位元組機器學習平臺從研發體驗、訓練時效、任務編排、資源運維等方面不斷提出新的要求挑戰，以 K8s 為核心的雲原生理念正是為解決以上問題提出，並在業界取得了廣泛應用。本文主要對位元組跳動機器學習平臺的雲原生化的改造工作進行介紹，期望可以帶給讀者一些實踐經驗。

今天的介紹主要圍繞以下三點展開：

機器學習系統雲原生化動機

機器學習系統雲原生化落地

未來展望

01機器學習系統雲原生化動機

首先和大家簡單分享下位元組內部機器學習訓練場景訓練框架特點，及在實際生產環境中遇到的問題痛點。

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 生態進行開展。

02

機器學習系統雲原生化落地

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. 位元組雲原生踐行

先總結一下位元組內部對雲原生的理解。位元組對雲原生踐行主要體現在效率和成本兩方面。

①效率方面

透過 CNI/CSI 完成對網路儲存的抽象，透過 DevicePlugin 完成對異構資源的抽象等，以此來達到基礎設施的標準化;

有狀態服務，透過擴充套件 Statefulset 進行支援，離線訓練透過自研的 operator 支援，完成業務框架的標準化;

面向服務或者模型的 CI/CD 流程完成流程規範的標準化;

基於容器和映象提供了標準一致的執行環境，進而達到交付形態的標準化。

② 成本方面

容器秒級彈性特徵，讓業務可以放心按需分配資源，無需考慮常態部署下的N倍容災;

各種混部方式，比如在離線混、計算儲存混、01 出借和常態混部等，有效提升了資源配置的利用率，降低了總體資源需求。

2. 未來展望

目前的編排排程能力在叢集聯邦層次相對比較薄弱，後續會進一步加強叢集聯邦編排排程能力，在跨任務編排時更多的考慮 Quota 分配、任務排隊、依賴的資料地理親和性和網路頻寬等等因素。

任務/資料編排方面，可能會嘗試以 Argo(開源的編排元件)去構建雲原生的 DAG 編排能力，支援機器學習研發訓練各種流程中的 pipeline 場景。資料編排可能會考慮引入如 Alluxio 等 cache 能力。

最後更長期持續跟進的事項是構建在離線一體的雲基礎設施能力。目前在離線 GPU 資源分屬不同的 K8s 叢集，雖然透過 VK 一定程度上解決了資源的隔離，但是長遠來看，還是希望進一步探索深度並池的可能，支援 CPU 與 GPU 或者在離線 GPU 之間的常態混部。