在深度學(xué)習(xí)的背景下，NVIDIA的CUDA與AMD的ROCm框架缺乏有效的互操作性，導(dǎo)致基礎(chǔ)設(shè)施資源利用率顯著降低。隨著模型規(guī)模不斷擴(kuò)大而預(yù)算約束日益嚴(yán)格，2-3年更換一次GPU的傳統(tǒng)方式已不具可持續(xù)性。但是Pytorch的最近幾次的更新可以有效利用異構(gòu)計(jì)算集群，實(shí)現(xiàn)對(duì)所有可用GPU資源的充分調(diào)度，不受制于供應(yīng)商限制。
本文將深入探討如何混合AMD/NVIDIA GPU集群以支持PyTorch分布式訓(xùn)練。通過建立CUDA與ROCm的技術(shù)橋接，我們將闡述如何實(shí)現(xiàn)以下目標(biāo)：
- 無(wú)需重構(gòu)訓(xùn)練代碼即可充分利用異構(gòu)硬件資源
- 維持高性能集體通信操作—如all_reduce和all_gather—通過UCC和UCX技術(shù)框架高效聚合和傳輸AMD與NVIDIA GPU節(jié)點(diǎn)間的數(shù)據(jù)（如梯度），實(shí)現(xiàn)同步化訓(xùn)練
- 在采用AWS g4dn (NVIDIA T4)和g4ad (AMD Radeon V520)實(shí)例構(gòu)建的異構(gòu)本地及Kubernetes集群上部署分布式PyTorch訓(xùn)練任務(wù)
集群異構(gòu)性分析
集群異構(gòu)性呈現(xiàn)從輕度到強(qiáng)度的連續(xù)譜系，每個(gè)層級(jí)在分布式機(jī)器學(xué)習(xí)和高性能計(jì)算環(huán)境中均需采取差異化的管理策略。這些集群主要依賴GPU作為核心計(jì)算加速器，同時(shí)可能在CPU架構(gòu)、內(nèi)存配置、存儲(chǔ)系統(tǒng)及網(wǎng)絡(luò)互連技術(shù)方面存在差異。本章重點(diǎn)分析GPU異構(gòu)性對(duì)HPC集群的影響，包括單一供應(yīng)商生態(tài)系統(tǒng)內(nèi)的輕度差異及多供應(yīng)商環(huán)境下的顯著差異。
輕度異構(gòu)環(huán)境
輕度異構(gòu)環(huán)境主要涉及同一供應(yīng)商生態(tài)系統(tǒng)內(nèi)的技術(shù)差異，如NVIDIA V100與A100或AMD MI50與MI250X加速器之間的代際差異。在此類場(chǎng)景中，這些GPU共享基礎(chǔ)架構(gòu)、驅(qū)動(dòng)程序和通信庫(kù)，使PyTorch等框架能夠通過抽象層有效管理這些差異。
輕度異構(gòu)集群面臨的挑戰(zhàn)：
- 計(jì)算能力不平衡：老舊GPU架構(gòu)在處理新型模型時(shí)性能滯后，形成系統(tǒng)瓶頸。
- 內(nèi)存容量不匹配：VRAM容量較小的設(shè)備限制了可處理的批量大小。
- 互連性能變化：PCIe Gen3與NVLink/NVSwitch技術(shù)在數(shù)據(jù)傳輸速率上存在顯著差異。
解決方案：
- 參數(shù)服務(wù)器分布式策略實(shí)現(xiàn)更具容錯(cuò)性的分布式工作負(fù)載架構(gòu)
- 動(dòng)態(tài)負(fù)載均衡：實(shí)施智能工作負(fù)載分配機(jī)制，跟蹤設(shè)備利用率，將較小批次任務(wù)分配至性能較低的GPU。
- 梯度壓縮技術(shù)：減少帶寬受限節(jié)點(diǎn)的通信開銷。
- 容器化部署：使用針對(duì)特定GPU架構(gòu)優(yōu)化的CUDA/ROCm版本構(gòu)建Docker鏡像，提高兼容性。
由于NVIDIA的NCCL或AMD的RCCL等供應(yīng)商專用庫(kù)針對(duì)各自生態(tài)系統(tǒng)進(jìn)行了深度優(yōu)化，因此集體通信在輕度異構(gòu)環(huán)境中仍能保持較高效率。
強(qiáng)度異構(gòu)環(huán)境
強(qiáng)度異構(gòu)環(huán)境涉及來自不同供應(yīng)商的GPU設(shè)備組成的混合集群（如NVIDIA與AMD）。
NVIDIA CUDA與AMD ROCm分別為其專有硬件平臺(tái)設(shè)計(jì)，采用不同的指令集、內(nèi)存管理機(jī)制和驅(qū)動(dòng)程序接口。這種缺乏統(tǒng)一基礎(chǔ)架構(gòu)的情況導(dǎo)致依賴共享通信后端的負(fù)載均衡策略和基于統(tǒng)一內(nèi)存語(yǔ)義的全分片數(shù)據(jù)并行(FSDP)技術(shù)無(wú)法跨平臺(tái)運(yùn)行。
目前業(yè)界尚未形成標(biāo)準(zhǔn)化解決方案來應(yīng)對(duì)強(qiáng)度異構(gòu)環(huán)境帶來的挑戰(zhàn)。這一技術(shù)缺口需要開發(fā)策略，在最小化AMD與NVIDIA GPU間通信開銷的同時(shí)，實(shí)現(xiàn)混合供應(yīng)商集群的透明利用，并達(dá)到接近原生性能水平。這一目標(biāo)可定義為：
- 透明資源利用：執(zhí)行分布式訓(xùn)練任務(wù)無(wú)需重寫模型代碼或按供應(yīng)商分割集群。
- 接近原生的性能表現(xiàn)：最小化AMD與NVIDIA GPU間的通信開銷，接近NCCL/RCCL原生性能，并利用支持RDMA的集體通信和GPU P2P通信實(shí)現(xiàn)高效分布式計(jì)算。
在后續(xù)內(nèi)容中，我將詳細(xì)闡述為在AWS G4dn實(shí)例(配備NVIDIA T4 GPU)和AWS G4ad實(shí)例(配備AMD Radeon V520 GPU)上啟用PyTorch分布式訓(xùn)練的集體通信所進(jìn)行的技術(shù)探索。重點(diǎn)將置于利用現(xiàn)有集體通信庫(kù)來解決強(qiáng)度異構(gòu)環(huán)境帶來的挑戰(zhàn)。
NCCL與RCCL的兼容性探索
NCCL (NVIDIA)和RCCL (AMD)是經(jīng)過GPU優(yōu)化的集體通信庫(kù)，集成了直接利用GPU Direct RDMA以及必要時(shí)使用套接字傳輸?shù)牡讓觾?yōu)化機(jī)制。
在研究RCCL變更日志時(shí)，我發(fā)現(xiàn)的首個(gè)積極信號(hào)是—與NCCL
的兼容性注釋。無(wú)論采用何種版本配置或應(yīng)用何種優(yōu)化調(diào)整，系統(tǒng)始終返回：
NCCL WARN NET/Socket: message truncated: receiving X bytes instead of Y.
這一探索最終證實(shí)為技術(shù)瓶頸，因?yàn)楸M管RCCL是NCCL的移植版本，但底層架構(gòu)差異阻礙了RCCL與NCCL在異構(gòu)集群中的協(xié)同工作。這些庫(kù)依賴特定硬件集成，且基于不同的內(nèi)核級(jí)優(yōu)化、內(nèi)存層次結(jié)構(gòu)和IPC機(jī)制，使真正的兼容性實(shí)現(xiàn)變得極為困難。
解決這一問題需要高效的通信中間件技術(shù)。
統(tǒng)一通信框架技術(shù)
在尋找適當(dāng)解決方案的過程中，我發(fā)現(xiàn)了統(tǒng)一通信框架(UCF)—一個(gè)由工業(yè)界、研究機(jī)構(gòu)和學(xué)術(shù)界組成的聯(lián)盟，其使命是統(tǒng)一高性能計(jì)算通信標(biāo)準(zhǔn)。
具有前景的解決方案—統(tǒng)一集體通信(UCC)—是一個(gè)開源項(xiàng)目，為高性能計(jì)算、人工智能、數(shù)據(jù)中心和I/O領(lǐng)域提供集體通信操作的API和庫(kù)實(shí)現(xiàn)。該項(xiàng)目旨在通過網(wǎng)絡(luò)拓?fù)涓兄惴ā⒑?jiǎn)化軟件方法和網(wǎng)絡(luò)內(nèi)計(jì)算硬件加速，提供高效且可擴(kuò)展的集體通信操作。
UCC與傳輸層中間件—統(tǒng)一通信X(UCX)協(xié)同工作，利用其高性能點(diǎn)對(duì)點(diǎn)通信原語(yǔ)和功能組件。UCX的設(shè)計(jì)汲取了多個(gè)項(xiàng)目的經(jīng)驗(yàn)，包括Mellanox的HCOLL和SHARP、華為的UCG、開源Cheetah及IBM的PAMI Collectives。最為關(guān)鍵的是—UCC和UCX均實(shí)現(xiàn)了對(duì)ROCm和CUDA的全面支持。
UCC作為實(shí)驗(yàn)性后端已集成到PyTorch分布式模塊中。它可以作為PyTorch分布式模塊的直接后端使用，也可以作為OpenMPI中集體通信操作的后端。為此需要從源代碼構(gòu)建支持MPI的torch庫(kù)，并使用mpirun啟動(dòng)器執(zhí)行支持OpenMPI的分布式任務(wù)。
這一發(fā)現(xiàn)是技術(shù)突破的關(guān)鍵是成功確定可行配置，并使用PyTorch和MPI成功運(yùn)行了多節(jié)點(diǎn)分布式數(shù)據(jù)并行訓(xùn)練任務(wù)。
AWS G4ad (AMD GPU)和G4dn (NVIDIA GPU)實(shí)例上運(yùn)行的分布式數(shù)據(jù)并行PyTorch訓(xùn)練任務(wù)。
通過采用UCC和UCX技術(shù)框架，異構(gòu)GPU集群不再是遙不可及的理想，而是可實(shí)現(xiàn)的技術(shù)目標(biāo)。這一突破有望使組織能夠充分發(fā)揮硬件投資價(jià)值，將分散的計(jì)算資源整合為高效統(tǒng)一的高性能計(jì)算環(huán)境。
異構(gòu)Kubernetes集群實(shí)現(xiàn)
在企業(yè)級(jí)基礎(chǔ)設(shè)施管理中，組織面臨著如何高效配置資源支持團(tuán)隊(duì)需求的挑戰(zhàn)。同時(shí)還需要支持各種規(guī)模的機(jī)器學(xué)習(xí)工作負(fù)載的快速高效運(yùn)行，包括小型實(shí)驗(yàn)和長(zhǎng)期訓(xùn)練萬(wàn)億參數(shù)級(jí)大模型的場(chǎng)景。
Kubernetes因其強(qiáng)大的資源編排能力以及最大化資源利用率和協(xié)調(diào)多樣化硬件的能力，已成為分布式機(jī)器學(xué)習(xí)的基礎(chǔ)平臺(tái)。
要在Kubernetes上調(diào)度分布式訓(xùn)練任務(wù)，無(wú)論使用Kubeflow MPI Operator還是PyTorch operator，任務(wù)清單都需要使用AMD或NVIDIA設(shè)備插件提供的特定資源定義：
# NVIDIA
resources:
limits:
nvidia.com/gpu: 1
# AMD
resources:
limits:
amd.com/gpu: 1
配置強(qiáng)度異構(gòu)任務(wù)需要自定義資源定義(CRD)或變更準(zhǔn)入控制器(mutating webhook handler)，以統(tǒng)一資源命名（如heterogenous.com/gpu: 1），或者手動(dòng)單獨(dú)部署每個(gè)Pod。
VolcanoJob作為Volcano調(diào)度器提供的Kubernetes CRD，簡(jiǎn)化了這一流程。Volcano專為高性能批處理工作負(fù)載設(shè)計(jì)，提供組調(diào)度(gang scheduling)功能確保分布式任務(wù)的原子執(zhí)行（即所有必需資源可用時(shí)所有Pod同時(shí)啟動(dòng)，否則全部不啟動(dòng)），并提供插件自動(dòng)化基礎(chǔ)設(shè)施配置。與Kubeflow的Training Operators（如MPIOperator）強(qiáng)制所有worker使用統(tǒng)一資源模板不同，Volcano允許按任務(wù)定義Pod，從而實(shí)現(xiàn)對(duì)異構(gòu)資源的精確控制。
在異構(gòu)Kubernetes集群上部署混合GPU分布式訓(xùn)練工作負(fù)載，需配置以下VolcanoJobCRD功能：
自動(dòng)SSH配置
ssh插件生成包含預(yù)共享SSH密鑰的ConfigMap，實(shí)現(xiàn)Pod間無(wú)密碼認(rèn)證。每個(gè)容器中的sshd設(shè)置利用這些密鑰，無(wú)需手動(dòng)證書管理。
worker pod DNS解析
svc插件創(chuàng)建無(wú)頭服務(wù)(headless service)，分配可預(yù)測(cè)的DNS主機(jī)名。Pod通過Volcano注入的環(huán)境變量（如VC_MPIWORKER_AMD_HOSTS）識(shí)別對(duì)等節(jié)點(diǎn)，主Pod利用這些變量構(gòu)建mpirun主機(jī)列表。
資源特定任務(wù)組
每個(gè)task定義唯一Pod模板：
—mpimaster協(xié)調(diào)訓(xùn)練過程，使用MPI和UCX參數(shù)優(yōu)化GPU通信。
—mpiworker-nvidia和mpiworker-amd分別指定不同resources和供應(yīng)商特定容器鏡像。
組調(diào)度機(jī)制
minAvailable: 3確保所有Pod（1個(gè)master + 2個(gè)worker）同時(shí)調(diào)度，防止異構(gòu)集群中的資源死鎖。
任務(wù)完成定義
帶有CompleteJob動(dòng)作鍵的policies字段允許配置將任務(wù)標(biāo)記為完成狀態(tài)的事件。此處為mpimaster任務(wù)的TaskCompleted事件。
apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
name: mpi-training-heterogeneous
namespace: volcano-job-training
spec:
minAvailable: 3 # Gang scheduling: All 3 pods must be allocated
plugins:
ssh: [] # Auto-generates SSH keys via ConfigMap
svc: [] # Creates headless service for stable hostnames
schedulerName: volcano
tasks:
- name: mpimaster
policies:
- action: CompleteJob # The job is completed when the launcher task completes successfully
event: TaskCompleted
replicas: 1
template:
spec:
containers:
- command:
- /bin/bash
- -c
# Create SSH directories for the plugin to inject passwordless configuration
- mkdir -p /var/run/sshd; /usr/sbin/sshd;
# Volcano injects worker hosts via VC_MPIWORKER_*_HOSTS:
MPI_HOST=${VC_MPIWORKER_AMD_HOSTS},${VC_MPIWORKER_NVIDIA_HOSTS};
NUM_WORKERS=$(($(echo ${MPI_HOST} | tr -cd ',' | wc -c) + 1));
# Launch the training job with mpirun and push the extracted MPI_HOST and NUM_WORKERS content
- mpirun -np ${NUM_WORKERS} --allow-run-as-root --host ${MPI_HOST} -x MASTER_ADDR=${VC_MPIWORKER_AMD_HOSTS} -x MASTER_PORT=29603 \
# Configure OpenMPI to use UCC for collective operation backend
-mca pml ucx -mca coll_ucc_enable 1 -mca coll_ucc_priority 100 \
# Fine-tune UCX transport layer and UCC collectives parameters to support g4ad instances
-x UCX_ROCM_COPY_D2H_THRESH=0 -x UCX_ROCM_COPY_H2D_THRESH=0 \
-x UCC_EC_ROCM_REDUCE_HOST_LIMIT=0 -x UCC_EC_ROCM_COPY_HOST_LIMIT=0 \
-x OMPI_MCA_mpi_accelerator_rocm_memcpyD2H_limit=0 -x OMPI_MCA_mpi_accelerator_rocm_memcpyH2D_limit=0 \
# Point on the MPI-aware PyTorch job execution code
/opt/conda/envs/py_3.12/bin/python 1000 1000 --batch_size 500
/mpijob/main.py --backend=mpi 1000 1000 --batch_size 500
image: docker.io/rafalsiwek/opmpi_ucx_simple:latest
name: mpimaster
ports:
- containerPort: 22
name: mpijob-port
restartPolicy: OnFailure
- name: mpiworker-nvidia
replicas: 1
template:
spec:
containers:
- command:
- /bin/bash
- -c
- mkdir -p /var/run/sshd; /usr/sbin/sshd -D;
image: docker.io/rafalsiwek/g4dn_distributed_ml:1.0_pytorch_mpi_opperator
name: mpiworker
ports:
- containerPort: 22
name: mpijob-port
- containerPort: 29603
name: torch-port
resources:
limits:
nvidia.com/gpu: "1" # NVIDIA-specific GPU
restartPolicy: OnFailure
- name: mpiworker-amd
replicas: 1
template:
spec:
containers:
- command:
- /bin/bash
- -c
- mkdir -p /var/run/sshd; /usr/sbin/sshd -D;
image: docker.io/rafalsiwek/g4ad_distributed_ml:1.0_pytorch_mpi_opperator
name: mpiworker
ports:
- containerPort: 22
name: mpijob-port
- containerPort: 29603
name: torch-port
resources:
limits:
amd.com/gpu: "1" # AMD-specific GPU
運(yùn)行PyTorch分布式任務(wù)需要具備特定GPU類型感知的UCC、UCX和MPI庫(kù)環(huán)境，以及將分布式模塊鏈接到這些庫(kù)的PyTorch構(gòu)建。啟動(dòng)器僅需UCC、UCX和OpenMPI支持，由于其集體操作不涉及GPU特定處理，因此不需要GPU感知構(gòu)建。此環(huán)境配置需要從源代碼構(gòu)建相關(guān)庫(kù)和PyTorch。
通過在Kubernetes平臺(tái)上啟用混合GPU集群，組織能夠?qū)⒎稚⒌挠布Y源轉(zhuǎn)化為統(tǒng)一的創(chuàng)新平臺(tái)。這種方法有效消除了成本高昂的供應(yīng)商鎖定，最大化現(xiàn)有投資價(jià)值并提升GPU資源利用率。無(wú)論是擴(kuò)展萬(wàn)億參數(shù)模型訓(xùn)練還是整合具有不同基礎(chǔ)設(shè)施的團(tuán)隊(duì)，異構(gòu)環(huán)境使團(tuán)隊(duì)能夠以更高效、智能的方式進(jìn)行模型訓(xùn)練，無(wú)需徹底更換硬件平臺(tái)。
技術(shù)局限性
缺乏RDMA驗(yàn)證：由于AWS EFA對(duì)g4ad實(shí)例的支持狀態(tài)尚不明確，適當(dāng)?shù)腞DMA兼容性尚未得到完全測(cè)試。UCX同樣缺乏針對(duì)零拷貝RDMA操作的官方AWS EFA兼容性認(rèn)證，因此當(dāng)前實(shí)現(xiàn)主要依賴TCP傳輸。
次優(yōu)通信性能：僅使用TCP傳輸層會(huì)顯著降低通信帶寬和增加延遲，這一點(diǎn)已通過OSU基準(zhǔn)測(cè)試結(jié)果得到驗(yàn)證。
機(jī)器學(xué)習(xí)框架集成要求：盡管PyTorch和Horovod支持用于集體操作的MPI后端，但Horovod在本實(shí)現(xiàn)中尚未進(jìn)行全面測(cè)試。此外，大多數(shù)框架需要顯式的MPI后端集成，而這種集成并非在所有框架中普遍可用。
PyTorch中有限的MPI后端支持：PyTorch的MPI風(fēng)格集體后端功能集相對(duì)有限，優(yōu)先支持NCCL/Gloo后端，且僅完全支持分布式數(shù)據(jù)并行(DDP)模式。全分片數(shù)據(jù)并行(FSDP)等高級(jí)策略依賴于allgather_base等操作，這些操作在PyTorch的MPI后端中尚未實(shí)現(xiàn)。
總結(jié)
對(duì)于尋求在機(jī)器學(xué)習(xí)和深度學(xué)習(xí)工作負(fù)載中實(shí)現(xiàn)快速擴(kuò)展的組織而言，在多供應(yīng)商GPU集群上執(zhí)行分布式訓(xùn)練的能力提供了極具戰(zhàn)略價(jià)值的技術(shù)機(jī)遇。由于主流機(jī)器學(xué)習(xí)框架缺乏原生支持，目前實(shí)現(xiàn)這一目標(biāo)仍需投入大量工程資源。
開放、標(biāo)準(zhǔn)化實(shí)現(xiàn)的發(fā)展將有助于實(shí)現(xiàn)異構(gòu)硬件生態(tài)系統(tǒng)的民主化訪問，從而在不犧牲性能的前提下提供經(jīng)濟(jì)高效的技術(shù)靈活性。
本文的源代碼可以在這個(gè)項(xiàng)目中找到：
https://avoid.overfit.cn/post/41b87700f05642b0b0cbd4729274ed1a
作者：Rafa? Siwek
熱門跟貼