在深度學習的背景下，NVIDIA的CUDA與AMD的ROCm框架缺乏有效的互操作性，導致基礎設施資源利用率顯著降低。隨著模型規模不斷擴大而預算約束日益嚴格，2-3年更換一次GPU的傳統方式已不

具可持續性。但是Pytorch的最近幾次的更新可以有效利用異構計算集群，實現對所有可用GPU資源的充分調度，不受制于供應商限制。

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

• 無需重構訓練代碼即可充分利用異構硬件資源
• 維持高性能集體通信操作—如all_reduce和all_gather—通過UCC和UCX技術框架高效聚合和傳輸AMD與NVIDIA GPU節點間的數據（如梯度），實現同步化訓練
• 在采用AWS g4dn (NVIDIA T4)和g4ad (AMD Radeon V520)實例構建的異構本地及Kubernetes集群上部署分布式PyTorch訓練任務

集群異構性分析

集群異構性呈現從輕度到強度的連續譜系，每個層級在分布式機器學習和高性能計算環境中均需采取差異化的管理策略。這些集群主要依賴GPU作為核心計算加速器，同時可能在CPU架構、內存配置、存儲系統及網絡互連技術方面存在差異。本章重點分析GPU異構性對HPC集群的影響，包括單一供應商生態系統內的輕度差異及多供應商環境下的顯著差異。

輕度異構環境

輕度異構環境主要涉及同一供應商生態系統內的技術差異，如NVIDIA V100與A100或AMD MI50與MI250X加速器之間的代際差異。在此類場景中，這些GPU共享基礎架構、驅動程序和通信庫，使PyTorch等框架能夠通過抽象層有效管理這些差異。

輕度異構集群面臨的挑戰：

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

解決方案：

• 參數服務器分布式策略 實現更具容錯性的分布式工作負載架構
• 動態負載均衡： 實施智能工作負載分配機制，跟蹤設備利用率，將較小批次任務分配至性能較低的GPU。
• 梯度壓縮技術： 減少帶寬受限節點的通信開銷。
• 容器化部署： 使用針對特定GPU架構優化的CUDA/ROCm版本構建Docker鏡像，提高兼容性。

由于NVIDIA的NCCL或AMD的RCCL等供應商專用庫針對各自生態系統進行了深度優化，因此集體通信在輕度異構環境中仍能保持較高效率。

強度異構環境

強度異構環境涉及來自不同供應商的GPU設備組成的混合集群（如NVIDIA與AMD）。

NVIDIA CUDA與AMD ROCm分別為其專有硬件平臺設計，采用不同的指令集、內存管理機制和驅動程序接口。這種缺乏統一基礎架構的情況導致依賴共享通信后端的負載均衡策略和基于統一內存語義的全分片數據并行(FSDP)技術無法跨平臺運行。

目前業界尚未形成標準化解決方案來應對強度異構環境帶來的挑戰。這一技術缺口需要開發策略，在最小化AMD與NVIDIA GPU間通信開銷的同時，實現混合供應商集群的透明利用，并達到接近原生性能水平。這一目標可定義為：

• 透明資源利用： 執行分布式訓練任務無需重寫模型代碼或按供應商分割集群。
• 接近原生的性能表現： 最小化AMD與NVIDIA GPU間的通信開銷，接近NCCL/RCCL原生性能，并利用支持RDMA的集體通信和GPU P2P通信實現高效分布式計算。

在后續內容中，我將詳細闡述為在AWS G4dn實例(配備NVIDIA T4 GPU)和AWS G4ad實例(配備AMD Radeon V520 GPU)上啟用PyTorch分布式訓練的集體通信所進行的技術探索。重點將置于利用現有集體通信庫來解決強度異構環境帶來的挑戰。

NCCL與RCCL的兼容性探索

NCCL (NVIDIA) 和 RCCL (AMD) 是經過 GPU優化的 集體通信庫，集成了直接利用GPU Direct RDMA以及必要時使用套接字傳輸的底層優化機制。

在研究RCCL變更日志時，我發現的首個積極信號是—與NCCL <version>的兼容性注釋。無論采用何種版本配置或應用何種優化調整，系統始終返回：

NCCL WARN NET/Socket: message truncated: receiving X bytes instead of Y.

這一探索最終證實為技術瓶頸，因為盡管RCCL是NCCL的移植版本，但底層架構差異阻礙了RCCL與NCCL在異構集群中的協同工作。這些庫依賴特定硬件集成，且基于不同的內核級優化、內存層次結構和IPC機制，使真正的兼容性實現變得極為困難。

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

統一通信框架技術

在尋找適當解決方案的過程中，我發現了統一通信框架(UCF)—一個由工業界、研究機構和學術界組成的聯盟，其使命是統一高性能計算通信標準。

具有前景的解決方案—統一集體通信(UCC)—是一個開源項目，為高性能計算、人工智能、數據中心和I/O領域提供集體通信操作的API和庫實現。該項目旨在通過網絡拓撲感知算法、簡化軟件方法和網絡內計算硬件加速，提供高效且可擴展的集體通信操作。

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

UCC作為實驗性后端已集成到PyTorch分布式模塊中。它可以作為PyTorch分布式模塊的直接后端使用，也可以作為OpenMPI中集體通信操作的后端。為此需要從源代碼構建支持MPI的torch庫，并使用mpirun啟動器執行支持OpenMPI的分布式任務。

這一發現是技術突破的關鍵是成功確定可行配置，并使用PyTorch和MPI成功運行了多節點分布式數據并行訓練任務。

AWS G4ad (AMD GPU)和G4dn (NVIDIA GPU)實例上運行的分布式數據并行PyTorch訓練任務。

通過采用UCC和UCX技術框架，異構GPU集群不再是遙不可及的理想，而是可實現的技術目標。這一突破有望使組織能夠充分發揮硬件投資價值，將分散的計算資源整合為高效統一的高性能計算環境。

異構Kubernetes集群實現

在企業級基礎設施管理中，組織面臨著如何高效配置資源支持團隊需求的挑戰。同時還需要支持各種規模的機器學習工作負載的快速高效運行，包括小型實驗和長期訓練萬億參數級大模型的場景。

Kubernetes因其強大的資源編排能力以及最大化資源利用率和協調多樣化硬件的能力，已成為分布式機器學習的基礎平臺。

要在Kubernetes上調度分布式訓練任務，無論使用Kubeflow MPI Operator還是PyTorch operator，任務清單都需要使用AMD或NVIDIA設備插件提供的特定資源定義：

# NVIDIA
 resources:
   limits:
     nvidia.com/gpu: 1
 
 # AMD
 resources:
   limits:
     amd.com/gpu: 1

配置強度異構任務需要自定義資源定義(CRD)或變更準入控制器(mutating webhook handler)，以統一資源命名（如heterogenous.com/gpu: 1），或者手動單獨部署每個Pod。

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

在異構Kubernetes集群上部署混合GPU分布式訓練工作負載，需配置以下VolcanoJob CRD功能：

自動SSH配置

ssh插件生成包含預共享SSH密鑰的ConfigMap，實現Pod間無密碼認證。每個容器中的sshd設置利用這些密鑰，無需手動證書管理。

worker pod DNS解析

svc插件創建無頭服務(headless service)，分配可預測的DNS主機名。Pod通過Volcano注入的環境變量（如VC_MPIWORKER_AMD_HOSTS）識別對等節點，主Pod利用這些變量構建mpirun主機列表。

資源特定任務組

每個task定義唯一Pod模板：

—mpimaster協調訓練過程，使用MPI和UCX參數優化GPU通信。

—mpiworker-nvidia和mpiworker-amd分別指定不同resources和供應商特定容器鏡像。

組調度機制

minAvailable: 3確保所有Pod（1個master + 2個worker）同時調度，防止異構集群中的資源死鎖。

任務完成定義

帶有CompleteJob動作鍵的policies字段允許配置將任務標記為完成狀態的事件。此處為mpimaster任務的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

運行PyTorch分布式任務需要具備特定GPU類型感知的UCC、UCX和MPI庫環境，以及將分布式模塊鏈接到這些庫的PyTorch構建。啟動器僅需UCC、UCX和OpenMPI支持，由于其集體操作不涉及GPU特定處理，因此不需要GPU感知構建。此環境配置需要從源代碼構建相關庫和PyTorch。

通過在Kubernetes平臺上啟用混合GPU集群，組織能夠將分散的硬件資源轉化為統一的創新平臺。這種方法有效消除了成本高昂的供應商鎖定，最大化現有投資價值并提升GPU資源利用率。 無論是擴展萬億參數模型訓練還是整合具有不同基礎設施的團隊，異構環境使團隊能夠以更高效、智能的方式進行模型訓練，無需徹底更換硬件平臺。

技術局限性

缺乏RDMA驗證：由于AWS EFA對g4ad實例的支持狀態尚不明確，適當的RDMA兼容性尚未得到完全測試。UCX同樣缺乏針對零拷貝RDMA操作的官方AWS EFA兼容性認證，因此當前實現主要依賴TCP傳輸。

次優通信性能：僅使用TCP傳輸層會顯著降低通信帶寬和增加延遲，這一點已通過OSU基準測試結果得到驗證。

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

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

總結

對于尋求在機器學習和深度學習工作負載中實現快速擴展的組織而言，在多供應商GPU集群上執行分布式訓練的能力提供了極具戰略價值的技術機遇。由于主流機器學習框架缺乏原生支持，目前實現這一目標仍需投入大量工程資源。

開放、標準化實現的發展將有助于實現異構硬件生態系統的民主化訪問，從而在不犧牲性能的前提下提供經濟高效的技術靈活性。