前言

Cloud Native

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

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

Cloud Native

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

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

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

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

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

Cloud Native

我們?cè)诎⒗镌浦蔚拇竽Ｐ陀?xùn)練場(chǎng)景中，使用了 Koordinator 來解決基本的任務(wù)調(diào)度需求和異構(gòu)設(shè)備資源管理需求。同時(shí)，使用 KubeDL 管理訓(xùn)練作業(yè)生命周期和訓(xùn)練作業(yè)排隊(duì)調(diào)度需求。

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

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

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

Job 調(diào)度

Cloud Native

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

?

Job 排隊(duì)

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

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

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

?

彈性 Quota

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



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

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

?

Coscheduling

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

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

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



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

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

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

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

Cloud Native

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

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

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

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

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

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

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

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

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

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

Koordinator Device 對(duì)象

Future: NRI 模式

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

?

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

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

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

NVIDIA H100

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

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

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

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

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

NVIDIA H100 System Topology

審核編輯：劉清