前言

Cloud Native

Koordinator 是阿里云基于過去我們建設(shè)的統(tǒng)一調(diào)度系統(tǒng)中積累的技術(shù)和實踐經(jīng)驗，對外開源了新一代的調(diào)度系統(tǒng)。Koordinator 支持 Kubernetes 上多種工作負(fù)載的混部調(diào)度。它的目標(biāo)是提高工作負(fù)載的運行時效率和可靠性（包括延遲敏感型負(fù)載和批處理任務(wù)）。Koordinator 不僅擅長混部場景，也同樣支持大數(shù)據(jù)、AI 訓(xùn)練等任務(wù)調(diào)度場景。本文分享了使用 Koordinator 支持異構(gòu)資源管理和任務(wù)調(diào)度場景的實踐經(jīng)驗。

AI/LLMs 帶來新機遇和新挑戰(zhàn)

Cloud Native

從 2022 年 11 月 ChatGPT 發(fā)布到現(xiàn)在，ChatGPT 所引起的關(guān)注、產(chǎn)生的影響可能已經(jīng)超越了信息技術(shù)歷史上的幾乎所有熱點。眾多業(yè)界專家都被它征服，比如阿里云 CEO 張勇的看法是：“所有行業(yè)、應(yīng)用、軟件、服務(wù)，都值得基于大模型能力重做一遍?！盢VIDIA CEO 黃仁勛稱它帶來了 AI 的 iPhone 時刻。ChatGPT 開啟了新的時代，國內(nèi)外的企業(yè)和科研機構(gòu)紛紛跟進(jìn)，幾乎每周都有一個甚至多個新模型推出，從自然語言處理、計算機視覺到人工智能驅(qū)動的科學(xué)研究、生成式 AI 等，應(yīng)用百花齊放；大模型成為業(yè)務(wù)提效和打開下一個增長點的關(guān)鍵。同樣對于云計算、基礎(chǔ)設(shè)施、分布式系統(tǒng)的需求也撲面而來。

為支撐百億級、千億級別參數(shù)量的大模型訓(xùn)練需求，云計算和基礎(chǔ)設(shè)施需要提供更強大、可擴展的計算和存儲資源。大模型訓(xùn)練依賴的的核心技術(shù)之一是分布式訓(xùn)練，分布式訓(xùn)練需要在多個計算節(jié)點之間傳遞大量的數(shù)據(jù)，因此需要一個帶寬更高、延遲更低的高性能網(wǎng)絡(luò)。

為了發(fā)揮計算、存儲和網(wǎng)絡(luò)資源的最佳效能，保障訓(xùn)練效率，調(diào)度和資源管理系統(tǒng)需要設(shè)計更合理的策略。在此基礎(chǔ)上，基礎(chǔ)設(shè)施還需要在可靠性上持續(xù)增強，具備節(jié)點故障治愈和容錯能力，確保訓(xùn)練任務(wù)的持續(xù)運行。 大模型訓(xùn)練離不開異構(gòu)計算設(shè)備，典型的就是我們熟知的 GPU。

在 GPU 領(lǐng)域，NVIDIA 仍然占據(jù)著主導(dǎo)地位，其他廠商如 AMD 和國內(nèi)的芯片制造商的機會在努力追趕。以 NVIDIA 為例，其強大的產(chǎn)品設(shè)計能力、扎實的技術(shù)實力和靈活的市場策略使其能夠快速推出更優(yōu)秀的芯片，但產(chǎn)品間的架構(gòu)差異較大，例如 NVIDIA A100 型號和 NVIDIA H100 型號的系統(tǒng)架構(gòu)差異十分明顯，使用方式上也存在許多需要注意的細(xì)節(jié)，這給上層的調(diào)度系統(tǒng)和資源管理系統(tǒng)帶來了不小的挑戰(zhàn)。

Koordinator+KubeDL 的強強聯(lián)合

Cloud Native

我們在阿里云支撐的大模型訓(xùn)練場景中，使用了 Koordinator 來解決基本的任務(wù)調(diào)度需求和異構(gòu)設(shè)備資源管理需求。同時，使用 KubeDL 管理訓(xùn)練作業(yè)生命周期和訓(xùn)練作業(yè)排隊調(diào)度需求。

Koordinator 不僅擅長混部調(diào)度場景，還針對大數(shù)據(jù)、AI 模型訓(xùn)練場景，提供了包括彈性 Quota 調(diào)度、Gang 調(diào)度等通用的任務(wù)調(diào)度能力。此外，它還具備精細(xì)化的資源調(diào)度管理能力，不僅支持中心化分配 GPU，還能感知硬件系統(tǒng)拓?fù)浞峙滟Y源，同時支持 GPU&RDMA 的聯(lián)合分配和設(shè)備共享能力。

我們選擇使用 KubeDL 來管理訓(xùn)練作業(yè)生命周期，是因為它不僅在支撐了內(nèi)部大量 AI 領(lǐng)域相關(guān)場景，而且得益于其優(yōu)秀的設(shè)計和實現(xiàn)都十分優(yōu)秀，可運維性、可靠性和功能擴展性都非常出色，自身是一個統(tǒng)一的 controller，可以支持多種訓(xùn)練工作負(fù)載，如 TensorFlow、PyTorch、Mars 等。

此外，它還可以適配不同調(diào)度器提供的 Gang 調(diào)度能力，可以幫助已經(jīng)使用 KubeDL 項目的存量場景平滑的切換到 Koordinator；KubeDL 還內(nèi)置了一個通用的作業(yè)排隊機制，可以有效解決作業(yè)自身的調(diào)度需求。 Koordinator 和 KubeDL 的強強聯(lián)合，可以很好的解決大模型訓(xùn)練的調(diào)度需求。

Job 調(diào)度

Cloud Native

Job 是一種更高層次的抽象，通常具有特定的計算任務(wù)或操作。它可以分割成多個子任務(wù)并行完成，也可以拆分成多個子任務(wù)協(xié)作完成。通常 Job 不會依賴其他的工作負(fù)載，可以獨立的運行。而且 Job 比較靈活，在時間維度、空間維度、或者資源方面的約束都比較少。

?

Job 排隊

Job 同樣需要經(jīng)過調(diào)度程序調(diào)度，這也就意味著 Job 同樣在調(diào)度時需要排隊。那為什么需要排隊呢？或者說我們可以通過排隊解決哪些問題？ 是因為系統(tǒng)中的資源有限的，我們的預(yù)算也是有限的，而 Job 的數(shù)量和計算需求往往是無限的。

如果不進(jìn)行排隊和調(diào)度，那些計算需求較高或者執(zhí)行時間較長的 Job 就會占用大量的資源，導(dǎo)致其他 Job 無法獲取到足夠的資源進(jìn)行計算，甚至可能導(dǎo)致集群系統(tǒng)崩潰。 因此，為保證各個 Job 能夠公平的獲得資源，避免資源爭奪和沖突，就需要對 Job 進(jìn)行排隊和調(diào)度。

我們使用 KubeDL提供的通用的 Job 排隊和調(diào)度機制解決這個問題。KubeDL 因為本身就內(nèi)置支持了多種訓(xùn)練工作負(fù)載，因此它天然支持按照 Job 粒度進(jìn)行調(diào)度；并且它具備多租戶間的公平性保障機制，減少 Job 間的資源爭奪和沖突，排隊和調(diào)度的過程中，KubeDL 根據(jù) Job 的計算需求、優(yōu)先級、資源需求等因素進(jìn)行評估和分配，確保每個 Job 都能夠得到合適的資源進(jìn)行計算。KubeDL 支持多種擴展插件，如 Filter 插件，Score 插件等，可以進(jìn)一步擴展其功能和特性滿足不同場景的需求。

?

彈性 Quota

Job 排隊要解決的核心問題之一是資源供給的公平性，一般在調(diào)度系統(tǒng)中都是通過彈性 Quota 機制來解決。 彈性 Quota 機制要解決的幾個核心問題：首先是保障公平性，不能讓某一些任務(wù)的資源需求過高導(dǎo)致其他任務(wù)被餓死，應(yīng)盡量讓大部分任務(wù)都能得到資源；其次需要有一定的彈性能力，能夠把空閑的額度共享給當(dāng)下更需要資源的任務(wù)，同樣還要能夠在需要資源時，把共享出去的資源拿回來，這意味還需要提供具備靈活的策略滿足不同場景的需求。



Koordinator 實現(xiàn)了彈性 Quota 調(diào)度能力，可以保障租戶間的公平性。我們在設(shè)計之初就考慮兼容 scheduler-plugins 社區(qū)項目中定義的 ElaticQuota CRD，這樣方便存量的集群和用戶可以平滑的過度到 Koordinator。 另外，我們不僅是兼容 ElasticQuota 原有按照 Namespace 管理 Quota 的能力，還支持按照支持按照樹形結(jié)構(gòu)進(jìn)行管理，可以跨 Namespace。

這樣的方式可以很好的支持一個復(fù)雜的組織的額度管理需求，比如一家公司里多個產(chǎn)品線，每個產(chǎn)品線的預(yù)算和使用情況都不一樣，都可以轉(zhuǎn)為 Quota 進(jìn)行管理，并借助彈性 Quota，把暫時沒有用到的空閑資源通過額度的形式臨時共享給其他部門使用。

?

Coscheduling

當(dāng)一個 Job 經(jīng)過排隊被調(diào)度后，Job Controller 會創(chuàng)建出一批子任務(wù)，對應(yīng)到 K8s，就是一批 Pod。這些 Pod 往往需要協(xié)調(diào)一致的啟動運行。這也就要求調(diào)度器在調(diào)度時一定要按照一組 Pod 分配資源，這一組 Pod 一定都可以那可以申請到資源或者一旦有一個 Pod 拿不到資源都認(rèn)為是調(diào)度失敗。

這也就是調(diào)度器需要提供的 All-or-Nothing 調(diào)度語義。 如果我們不這樣按照一組調(diào)度，會出現(xiàn)因為多個作業(yè)在資源調(diào)度層面出現(xiàn)爭搶，是有可能出現(xiàn)資源維度的死鎖，即至少兩個 Job 會出現(xiàn)拿不到資源的情況，即使原本空閑資源足夠其中一個 Job 運行的，也會拿不到。

比如下圖中，Job A 和 Job B 同時創(chuàng)建一批 Pod，如果不在中間的 Scheduling Queue 進(jìn)行排序而是隨意的調(diào)度，就會出現(xiàn) Job A 和 Job B 的 Pod 各持有了一部分節(jié)點的部分資源，如果此時集群資源緊張，很有可能 Job A 和 Job B 都可能拿不到資源。但如果排序后，我們嘗試先讓其中一個 Job 的 Pod 先一起嘗試優(yōu)先分配資源，那么至少保障一個 Job 可以運行。



當(dāng)一個 Job 切分的一組 Pod 非常大時，而集群內(nèi)的資源又不是十分充足，或者 Quota 不是很多時，可以把這樣的一組 Pod 切分成更多個子組，這個切割的大小以能運行任務(wù)為基礎(chǔ)，假設(shè)一個 Job 要求最小切割粒度是每組 3 個 Pod，那么這個最小粒度，一般在調(diào)度域中稱為 min available。

具體到 AI 模型訓(xùn)練領(lǐng)域，一些特殊的 Job 比如 TFJob，它的子任務(wù)有兩種角色，這兩種角色在生產(chǎn)環(huán)境中，也是需要設(shè)置不同的 min available 的。這種不同角色的區(qū)分的場景還有可能要求每個角色的 min available 都滿足時才可以認(rèn)為符合 All-or-Nothing 語義。

Koordinator 內(nèi)置了 Coscheduling 調(diào)度能力，它兼容社區(qū)的 scheduler-plugins/coscheduling 定義 PodGroup CRD，還支持把多個 PodGroup 聯(lián)合調(diào)度，這樣就可以支持按角色設(shè)置 min available 場景。Koordinator 實現(xiàn)了一個 KubeDL Gang Scheduler 插件，這樣就可以和 KubeDL 做集成一起支撐這類調(diào)度場景。

精細(xì)化設(shè)備管理

Cloud Native

K8s 設(shè)備管理的局限性

K8s 是通過 kubelet 負(fù)責(zé)設(shè)備管理和分配，并和 device plugin 交互實現(xiàn)整套機制，這套機制在 K8s 早期還是夠用的，其他廠商如 AMD 和國內(nèi)的芯片制造商也抓住機會努力追趕。

kubelet 與 device plugin 協(xié)作流程 首先 K8s 只允許通過 kubelet 來分配設(shè)備，這樣就導(dǎo)致無法獲得全局最優(yōu)的資源編排，也就是從根本上無法發(fā)揮資源效能。比如一個集群內(nèi)有兩臺節(jié)點，都有相同的設(shè)備，剩余的可分配設(shè)備數(shù)量相等，但是實際上兩個節(jié)點上的設(shè)備所在的硬件拓?fù)鋾?dǎo)致 Pod 的運行時效果差異較大，沒有調(diào)度器介入的情況下，是可能無法抵消這種差異的。

其次是不支持類似 GPU 和 RDMA 聯(lián)合分配的能力。大模型訓(xùn)練依賴高性能網(wǎng)絡(luò)，而高性能網(wǎng)絡(luò)的節(jié)點間通信需要用到 RDMA 協(xié)議和支持 RDMA 協(xié)議的網(wǎng)絡(luò)設(shè)備，而這些設(shè)備又和 GPU 在節(jié)點上的系統(tǒng)拓?fù)鋵用媸蔷o密協(xié)作的，比如下面這張圖是 NVIDIA 的 A100 型號機型的硬件拓?fù)鋱D，我們可以看到，PCIe Switch 下掛了 GPU 和高性能網(wǎng)卡，我們需要就近分配這兩個設(shè)備，才能做到節(jié)點間通信的低延遲。

而且這里比較有意思的是，當(dāng)如果需要分配多個 GPU 時，如果涉及到了多個 PCIe Switch，就意味著需要分配多個網(wǎng)卡，這就和 K8s 的另一個限制有關(guān)系，即聲明的資源協(xié)議是定量的，而不是隨意變化的，也就是說用戶實際上也不知道這個 Pod 需要多少支持 RDMA 的網(wǎng)卡，用戶只知道要多少個 GPU 設(shè)備，并期望就近分配 RDMA 的網(wǎng)卡而已。

而且 kubelet 也不支持設(shè)備的初始化和清理功能，更不支持設(shè)備的共享機制，后者在訓(xùn)練場景一般用不到，但在線推理服務(wù)會用到。在線推理服務(wù)本身也有很明顯的峰谷特征，很多時刻并不需要占用完整的 GPU 資源。

K8s 這種節(jié)點的設(shè)備管理能力一定程度上已經(jīng)落后時代了，雖然現(xiàn)在最新的版本里支持了 DRA 分配機制（類似于已有的 PVC 調(diào)度機制），但是這個機制首先只在最新版本的 K8s 才支持，但實際情況是還有大量存量集群在使用，并且升級到 K8s 最新版本也并不是一個小事情，所以我們得想其他辦法。

Koordinator 精細(xì)化設(shè)備管理機制

我們在 Koordinator 中提出了一種方案，可以解決這些問題，做到精細(xì)化的資源調(diào)度。

Koordinator 精細(xì)化設(shè)備管理機制 從上面的圖中可以看到，用戶創(chuàng)建的一個 Pod，由 koord-scheduler 調(diào)度器根據(jù) koordlet 上報的 Device CRD 分配設(shè)備，并寫入到 Pod Annotation 中，再經(jīng) kubelet 拉起 Sandbox 和 Container，這中間 kubelet 會發(fā)起 CRI 請求到 containerd/docker，但在 Koordinator 方案中，CRI 請求會被 koord-runtime-proxy 攔截并轉(zhuǎn)發(fā)到 koordlet 內(nèi)的 GPU 插件，感知 Pod Annotation 上的設(shè)備分配結(jié)果并生成必要的設(shè)備環(huán)境變量等信息返回給 koord-runtime-proxy，再最終把修改后的 CRI 請求轉(zhuǎn)給 containerd/docker，最終再返回給 kubelet。這樣就可以無縫的介入整個容器的生命周期實現(xiàn)自定義的邏輯。 Koordinator Device CRD 用來描述節(jié)點的設(shè)備信息，包括 Device 的拓?fù)湫畔?，這些信息可以指導(dǎo)調(diào)度器實現(xiàn)精細(xì)化的分配邏輯。

Koordinator Device 對象

Future: NRI 模式

前面提到了 Koordinator 單機側(cè)依靠 koord-runtime-proxy 協(xié)作完成設(shè)備信息注入，我們自己也意識到，koord-runtime-proxy 這種方式其實不太好在大家的集群內(nèi)落地。這涉及到修改 kubelet 的啟動參數(shù)問題。 所以 Koordinator 社區(qū)后續(xù)會引入 NRI/CDI 等機制解決這個場景的問題。這塊工作正在和 Intel 相關(guān)團(tuán)隊共建。 NRI/CDI 是 containerd 支持的一種插件化機制。其部署方式有點類似于大家熟悉的 CNI，支持在啟動 Sandbox/Container 前后獲得機會修改參數(shù)或者實現(xiàn)一些定制邏輯。這相當(dāng)于是 containerd 內(nèi)置的 runtimeproxy 機制。

?

GPU&RDMA 按照硬件拓?fù)渎?lián)合分配

前面也提到，大模型訓(xùn)練不僅僅只用到了 GPU，還依賴 RDMA 網(wǎng)絡(luò)設(shè)備。要確保 GPU 和 RDMA 之間的延遲盡可能的低，否則會因為設(shè)備間的延遲放大到整個分布式訓(xùn)練網(wǎng)絡(luò)中，拖慢整體的訓(xùn)練效率。 這就要求在分配 GPU 和 RDMA 時需要感知硬件拓?fù)?，盡可能就近分配這種設(shè)備。嘗試按照同 PCIe，同 NUMA Node，同 NUMA Socket 和 跨 NUMA 的順序分配，延遲依次升高。 而且我們還發(fā)現(xiàn)同一個硬件廠商的不同型號的 GPU，它們的硬件系統(tǒng)拓?fù)涫遣灰粯拥?，這就要求我們的調(diào)度器需要感知這些差異。比如下圖是 NVIDIA A100 型號的 System Topology 和 NVIDIA H100 的一個簡單的設(shè)備連接圖。

NVIDIA A100 System Topology NVIDIA A100 GPU 之間的 NVLINK 聯(lián)通方式和 NVIDIA H100 型號就不一樣，NVSwitch 的數(shù)量也不一樣，這種差異就會給使用方式帶來很大的差異。

NVIDIA H100

NVIDIA-based system 在多租模式下的差異

NVIDIA H100 GPU 在多租 VM 場景下的特殊之處，多個 GPU 之間聯(lián)通需要操作 NVSwitch 才可以實現(xiàn)。 在多租場景中，NVIDIA 為保障安全，會通過 NVSwitch 管理 NVLink 的隔離狀態(tài)，并且要求只能由授信的軟件操作 NVSwitch。這個授信軟件是可以自定義的。

NVIDIA 支持多種模式，一種是 Full Passthrough 模式，這種模式把 GPU 和 NVSwitch 都直通到 VM 的 Guest OS，這樣做的好處是使用起來很簡單，但代價是當(dāng) GPU VM 多了，NVLINK 的帶寬會減少（原文：Reduced NVLink bandwidth for two and four GPU VMs）。

另一種稱為 Shared NVSwitch 多租戶模式，它只要求把 GPU 直通到 Guest OS，然后通過一個特殊的 VM，稱為 Service VM 管理 NVSwitch，并通過 ServiceVM 調(diào)用 NVIDIA Fabric Manager 激活 NVSwitch 實現(xiàn) GPU 間通信。這種模式就不會出現(xiàn)因為 Full Passthrough 模式的弊端，但使用方式明顯要更復(fù)雜一些。

這種特殊的硬件架構(gòu)和使用方式，還導(dǎo)致在分配 GPU 時有一些額外的要求。NVIDIA 定義了哪些 GPU 設(shè)備實例可以組合一起分配，比如用戶申請分配 4個 GPU，那必須是按照規(guī)定的 1,2,3,4 號一起或者 5,6,7,8 一起分，否則就會導(dǎo)致 Pod 無法運行。這種特殊的分配方式的背后原因我們不得而知，但分析這些分配約束可以發(fā)現(xiàn)，廠商規(guī)定的這種組合關(guān)系正好符合硬件系統(tǒng)拓?fù)浣Y(jié)構(gòu)，也就是可以滿足前面講到的 GPU&RDMA 聯(lián)合分配期望的分配結(jié)果。

NVIDIA H100 System Topology

審核編輯：劉清