在深度學(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)對所有可用GPU資源的充分調(diào)度，不受制于供應(yīng)商限制。

本文將深入探討如何混合AMD/NVIDIA GPU集群以支持PyTorch分布式訓(xùn)練。通過建立CUDA與ROCm的技術(shù)橋接，我們將闡述如何實(shí)現(xiàn)以下目標(biāo)：

• 無需重構(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ù)譜系，每個層級在分布式機(jī)器學(xué)習(xí)和高性能計(jì)算環(huán)境中均需采取差異化的管理策略。這些集群主要依賴GPU作為核心計(jì)算加速器，同時可能在CPU架構(gòu)、內(nèi)存配置、存儲系統(tǒng)及網(wǎng)絡(luò)互連技術(shù)方面存在差異。本章重點(diǎn)分析GPU異構(gòu)性對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加速器之間的代際差異。在此類場景中，這些GPU共享基礎(chǔ)架構(gòu)、驅(qū)動程序和通信庫，使PyTorch等框架能夠通過抽象層有效管理這些差異。

輕度異構(gòu)集群面臨的挑戰(zhàn)：

• 計(jì)算能力不平衡： 老舊GPU架構(gòu)在處理新型模型時性能滯后，形成系統(tǒng)瓶頸。
• 內(nèi)存容量不匹配： VRAM容量較小的設(shè)備限制了可處理的批量大小。
• 互連性能變化： PCIe Gen3與NVLink/NVSwitch技術(shù)在數(shù)據(jù)傳輸速率上存在顯著差異。

解決方案：

• 參數(shù)服務(wù)器分布式策略 實(shí)現(xiàn)更具容錯性的分布式工作負(fù)載架構(gòu)
• 動態(tài)負(fù)載均衡： 實(shí)施智能工作負(fù)載分配機(jī)制，跟蹤設(shè)備利用率，將較小批次任務(wù)分配至性能較低的GPU。
• 梯度壓縮技術(shù)： 減少帶寬受限節(jié)點(diǎn)的通信開銷。
• 容器化部署： 使用針對特定GPU架構(gòu)優(yōu)化的CUDA/ROCm版本構(gòu)建Docker鏡像，提高兼容性。

由于NVIDIA的NCCL或AMD的RCCL等供應(yīng)商專用庫針對各自生態(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分別為其專有硬件平臺設(shè)計(jì)，采用不同的指令集、內(nèi)存管理機(jī)制和驅(qū)動程序接口。這種缺乏統(tǒng)一基礎(chǔ)架構(gòu)的情況導(dǎo)致依賴共享通信后端的負(fù)載均衡策略和基于統(tǒng)一內(nèi)存語義的全分片數(shù)據(jù)并行(FSDP)技術(shù)無法跨平臺運(yùn)行。

目前業(yè)界尚未形成標(biāo)準(zhǔn)化解決方案來應(yīng)對強(qiáng)度異構(gòu)環(huán)境帶來的挑戰(zhàn)。這一技術(shù)缺口需要開發(fā)策略，在最小化AMD與NVIDIA GPU間通信開銷的同時，實(shí)現(xiàn)混合供應(yīng)商集群的透明利用，并達(dá)到接近原生性能水平。這一目標(biāo)可定義為：

• 透明資源利用： 執(zhí)行分布式訓(xùn)練任務(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)有集體通信庫來解決強(qiáng)度異構(gòu)環(huán)境帶來的挑戰(zhàn)。

NCCL與RCCL的兼容性探索

NCCL (NVIDIA) 和 RCCL (AMD) 是經(jīng)過 GPU優(yōu)化的 集體通信庫，集成了直接利用GPU Direct RDMA以及必要時使用套接字傳輸?shù)牡讓觾?yōu)化機(jī)制。

在研究RCCL變更日志時，我發(fā)現(xiàn)的首個積極信號是—與NCCL <version>的兼容性注釋。無論采用何種版本配置或應(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é)同工作。這些庫依賴特定硬件集成，且基于不同的內(nèi)核級優(yōu)化、內(nèi)存層次結(jié)構(gòu)和IPC機(jī)制，使真正的兼容性實(shí)現(xiàn)變得極為困難。

解決這一問題需要高效的通信中間件技術(shù)。

統(tǒng)一通信框架技術(shù)

在尋找適當(dāng)解決方案的過程中，我發(fā)現(xiàn)了統(tǒng)一通信框架(UCF)—一個由工業(yè)界、研究機(jī)構(gòu)和學(xué)術(shù)界組成的聯(lián)盟，其使命是統(tǒng)一高性能計(jì)算通信標(biāo)準(zhǔn)。

具有前景的解決方案—統(tǒng)一集體通信(UCC)—是一個開源項(xiàng)目，為高性能計(jì)算、人工智能、數(shù)據(jù)中心和I/O領(lǐng)域提供集體通信操作的API和庫實(shí)現(xiàn)。該項(xiàng)目旨在通過網(wǎng)絡(luò)拓?fù)涓兄惴?、簡化軟件方法和網(wǎng)絡(luò)內(nèi)計(jì)算硬件加速，提供高效且可擴(kuò)展的集體通信操作。

UCC與傳輸層中間件—統(tǒng)一通信X(UCX)協(xié)同工作，利用其高性能點(diǎn)對點(diǎn)通信原語和功能組件。UCX的設(shè)計(jì)汲取了多個項(xiàng)目的經(jīng)驗(yàn)，包括Mellanox的HCOLL和SHARP、華為的UCG、開源Cheetah及IBM的PAMI Collectives。最為關(guān)鍵的是—UCC和UCX均實(shí)現(xiàn)了對ROCm和CUDA的全面支持。

UCC作為實(shí)驗(yàn)性后端已集成到PyTorch分布式模塊中。它可以作為PyTorch分布式模塊的直接后端使用，也可以作為OpenMPI中集體通信操作的后端。為此需要從源代碼構(gòu)建支持MPI的torch庫，并使用mpirun啟動器執(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ā)揮硬件投資價值，將分散的計(jì)算資源整合為高效統(tǒng)一的高性能計(jì)算環(huán)境。

異構(gòu)Kubernetes集群實(shí)現(xiàn)

在企業(yè)級基礎(chǔ)設(shè)施管理中，組織面臨著如何高效配置資源支持團(tuán)隊(duì)需求的挑戰(zhàn)。同時還需要支持各種規(guī)模的機(jī)器學(xué)習(xí)工作負(fù)載的快速高效運(yùn)行，包括小型實(shí)驗(yàn)和長期訓(xùn)練萬億參數(shù)級大模型的場景。

Kubernetes因其強(qiáng)大的資源編排能力以及最大化資源利用率和協(xié)調(diào)多樣化硬件的能力，已成為分布式機(jī)器學(xué)習(xí)的基礎(chǔ)平臺。

要在Kubernetes上調(diào)度分布式訓(xùn)練任務(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ú)部署每個Pod。

VolcanoJob作為Volcano調(diào)度器提供的Kubernetes CRD，簡化了這一流程。Volcano專為高性能批處理工作負(fù)載設(shè)計(jì)，提供組調(diào)度(gang scheduling)功能確保分布式任務(wù)的原子執(zhí)行（即所有必需資源可用時所有Pod同時啟動，否則全部不啟動），并提供插件自動化基礎(chǔ)設(shè)施配置。與Kubeflow的Training Operators（如MPIOperator）強(qiáng)制所有worker使用統(tǒng)一資源模板不同，Volcano允許按任務(wù)定義Pod，從而實(shí)現(xiàn)對異構(gòu)資源的精確控制。

在異構(gòu)Kubernetes集群上部署混合GPU分布式訓(xùn)練工作負(fù)載，需配置以下VolcanoJob CRD功能：

自動SSH配置

ssh插件生成包含預(yù)共享SSH密鑰的ConfigMap，實(shí)現(xiàn)Pod間無密碼認(rèn)證。每個容器中的sshd設(shè)置利用這些密鑰，無需手動證書管理。

worker pod DNS解析

svc插件創(chuàng)建無頭服務(wù)(headless service)，分配可預(yù)測的DNS主機(jī)名。Pod通過Volcano注入的環(huán)境變量（如VC_MPIWORKER_AMD_HOSTS）識別對等節(jié)點(diǎn)，主Pod利用這些變量構(gòu)建mpirun主機(jī)列表。

資源特定任務(wù)組

每個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個master + 2個worker）同時調(diào)度，防止異構(gòu)集群中的資源死鎖。

任務(wù)完成定義

帶有CompleteJob動作鍵的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庫環(huán)境，以及將分布式模塊鏈接到這些庫的PyTorch構(gòu)建。啟動器僅需UCC、UCX和OpenMPI支持，由于其集體操作不涉及GPU特定處理，因此不需要GPU感知構(gòu)建。此環(huán)境配置需要從源代碼構(gòu)建相關(guān)庫和PyTorch。

通過在Kubernetes平臺上啟用混合GPU集群，組織能夠?qū)⒎稚⒌挠布Y源轉(zhuǎn)化為統(tǒng)一的創(chuàng)新平臺。這種方法有效消除了成本高昂的供應(yīng)商鎖定，最大化現(xiàn)有投資價值并提升GPU資源利用率。 無論是擴(kuò)展萬億參數(shù)模型訓(xùn)練還是整合具有不同基礎(chǔ)設(shè)施的團(tuán)隊(duì)，異構(gòu)環(huán)境使團(tuán)隊(duì)能夠以更高效、智能的方式進(jìn)行模型訓(xùn)練，無需徹底更換硬件平臺。

技術(shù)局限性

缺乏RDMA驗(yàn)證：由于AWS EFA對g4ad實(shí)例的支持狀態(tài)尚不明確，適當(dāng)?shù)腞DMA兼容性尚未得到完全測試。UCX同樣缺乏針對零拷貝RDMA操作的官方AWS EFA兼容性認(rèn)證，因此當(dāng)前實(shí)現(xiàn)主要依賴TCP傳輸。

次優(yōu)通信性能：僅使用TCP傳輸層會顯著降低通信帶寬和增加延遲，這一點(diǎn)已通過OSU基準(zhǔn)測試結(jié)果得到驗(yàn)證。

機(jī)器學(xué)習(xí)框架集成要求：盡管PyTorch和Horovod支持用于集體操作的MPI后端，但Horovod在本實(shí)現(xiàn)中尚未進(jìn)行全面測試。此外，大多數(shù)框架需要顯式的MPI后端集成，而這種集成并非在所有框架中普遍可用。

PyTorch中有限的MPI后端支持：PyTorch的MPI風(fēng)格集體后端功能集相對有限，優(yōu)先支持NCCL/Gloo后端，且僅完全支持分布式數(shù)據(jù)并行(DDP)模式。全分片數(shù)據(jù)并行(FSDP)等高級策略依賴于allgather_base等操作，這些操作在PyTorch的MPI后端中尚未實(shí)現(xiàn)。

總結(jié)

對于尋求在機(jī)器學(xué)習(xí)和深度學(xué)習(xí)工作負(fù)載中實(shí)現(xiàn)快速擴(kuò)展的組織而言，在多供應(yīng)商GPU集群上執(zhí)行分布式訓(xùn)練的能力提供了極具戰(zhàn)略價值的技術(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ù)靈活性。