意外的联盟：AMD 和 NVIDIA GPU 用于分布式机器学习任务，一起 (AI生成的图像使用 Ideogram 2.0)
到2024年，根据美国乔治城大学安全与新兴技术中心（CSET）的一份报告，科技巨头如苹果、微软和英特尔已经收购了60家AI和机器学习公司。这些收购旨在推动创新并整合可扩展的机器学习方面的专业知识，其中高效集成GPU对于高性能AI基础设施来说至关重要。然而，收购竞争对手或与合作伙伴合并往往会带来混合硬件集群。一家组织可能最初使用的是NVIDIA的GPU，但在收购了依赖于AMD GPU集群的团队后，会面临各种挑战。克服供应商绑定并统一异构集群对于实现经济高效且分布式的训练解决方案来说至关重要。
供应商锁定加强了集成难题，把集群困在孤立的生态系统中。NVIDIA的CUDA和AMD的ROCm很少互相兼容，这使得团队不得不将工作负载分开，导致现有设施的利用率不高。随着模型变得越来越大，预算越来越紧张，每2到3年更换一次GPU变得越来越不可行。解决办法是利用异构集群，充分利用所有可用的GPU资源，不论供应商。
本文是关于MLOps中分布式处理系列的第三部分。我们将探讨如何在混合AMD/NVIDIA GPU集群的AWS上配置设备以在PyTorch中进行分布式训练。通过桥接CUDA和ROCm，我们学习相关方法。
- 使用混合硬件而无需重写训练代码。
- 保持高性能的集体操作——如
all_reduce和all_gather——这些操作能有效地聚合和共享数据（例如梯度），以便使用UCC和UCX进行同步训练。 - 在异构的本地和Kubernetes集群上运行分布式PyTorch作业，使用AWS g4dn（NVIDIA T4）和g4ad（AMD Radeon V520）实例。
查看本系列的其他文章：
第1部分 — 使用分布式计算加速MLOps以实现可扩展的机器学习让您的MLOps平台能够更快地训练更大规模的模型——因为您的模型需要“团队合作”精神medium.com 第二部分 — 如何在MLOps中高效使用GPU进行分布式机器学习实践行业巨头如何大幅降低机器学习训练成本并实现5.6倍的速度提高在我的代码仓库中，您可以找到云基础设施的搭建、代码以及详细指南。
GitHub — RafalSiwek/distributed-mlops-overview: 关于 MLOps 的分布式处理方案探索 — RafalSiwek/distributed-mlops-overviewgithub.com 集群异质性现象集群异质性范围从轻微到强烈，每个级别都需要在分布式机器学习和高性能计算环境中采取不同的管理措施。这些集群主要依赖GPU作为加速器，但在CPU、内存配置、存储和互连技术方面也可能存在差异。本章主要讨论具有GPU异质性的HPC集群，包括单一供应商生态系统内的轻微差异及多供应商环境中的显著差异。
轻度异质性轻微的异构性涉及同一供应商生态系统内的差异，例如NVIDIA的V100和A100，或是AMD的MI50和MI250X。在这种情况下，这些GPU具有共同的架构、驱动程序和通信库，使得像PyTorch这样的框架可以通过抽象层次来处理这些不同。
在轻度混合集群中遇到的挑战包括：
- 计算不平衡: 较旧的GPU落后于较新型号，形成瓶颈。
- 内存不匹配: 具有较小VRAM的设备限制了批处理量。
- 互连性差异: PCIe Gen3与NVLink/NVSwitch影响数据传输速度。
解决方法:.
- 利用参数服务器分布式策略，实现更容错的分布式工作负载分配。
- 动态负载均衡: 实现智能工作负载，跟踪设备利用率，将较小批次分配给性能较弱的GPU。
- 梯度压缩技术: 减少带宽受限的节点的通信开销。
- 容器化: 使用针对特定GPU代的定制化Docker镜像（包含CUDA/ROCm版本），以确保兼容。
在适度的异质性下，集体通信仍然有效，这得益于如 NVIDIA 的 NCCL 和 AMD 的 RCCL 等特定于供应商的库，这些库专门针对各自的生态系统进行了优化。
很强的多样性高度的异质性包括来自不同供应商（例如 NVIDIA 和 AMD）的不同类型 GPU 设备在集群中的整合。
NVIDIA CUDA 和 AMD ROCm 是专门为其硬件设计的，使用不同的指令集、内存管理和驱动接口。这种缺乏共同基础的现状导致依赖共享通信后端的负载均衡功能以及依赖统一内存语义的 FSDP 在这两个平台上也无法正常运行。
目前尚不存在标准化的解决方案来解决强异构性的挑战。这一缺口要求提出相应的策略，以实现透明地使用来自不同供应商的集群，同时通过减少AMD和NVIDIA GPU之间的通信开销，达到接近原生的性能。这确立了以下目标：
- 透明利用: 无需修改模型代码或根据供应商分割集群即可进行分布式训练。
- 接近原生性能: 减少AMD和NVIDIA GPU之间的通信延迟，达到接近NCCL/RCCL的性能，并利用RDMA支持的集合操作和GPU间直接通信，实现高效的分布式技术。
在以下部分中，我将详细说明我在使用 PyTorch 在 AWS G4dn 实例（配备 NVIDIA T4 GPU）和 AWS G4ad 实例（配备 AMD Radeon V520 GPU）上进行分布式训练时，实现群组通讯的努力。这里主要会讲的是利用现有的群组通讯库来解决强烈异构环境中的问题。
RCCL是NCCL的版本吗？NVIDIA 和 AMD的RCCL 是 针对GPU优化的 集体通信库（简称CC库），集成了底层优化，直接利用GPU Direct RDMA技术，并在必要时使用套接字传输技术。
我遇到的第一个积极迹象是在研究RCCL 变更日志时看到的——「与NCCL <版本> 的兼容性」。无论我使用什么版本配置或应用了什么优化更改，结果总是：
NCCL 警告 NET/Socket: 消息被截短：实际上只收到了 X 字节，而不是 Y 字节这被证明是一个死胡同，因为看起来尽管 RCCL 是 NCCL 的移植，许多底层更改使得 RCCL 和 NCCL 在集群中难以实现异构使用。这些库依赖于底层硬件集成，并且依赖于内核级优化、内存层次结构以及 IPC 机制，这会使兼容性变得相当复杂。
高效的消息中间层是必不可少的。
统一通信平台,在寻找合适方案的过程中，我偶然遇到了统一通信框架（UCF）——一个由业界、实验室和学术界组成的联盟，其使命是统一高性能计算中的通信。
激发希望的产品——统一集体通信（UCC），这是一个开源项目，它提供了一种用于高性能计算、人工智能、数据中心和I/O中的群组通信操作的API和库。它的目标是通过考虑网络拓扑、简单软件方法以及网络内计算的硬件加速，提供高效且可扩展的集体通信操作。
UCC的组成部件（参考：[UCC](https://github.com/openucx/ucc））
UCC 使用传输层中间件——统一通信X（UCX），利用其快速的点对点通信例程和工具集。UCX的设计吸取了多个项目的经验，包括Mellanox的HCOLL和SHARP、华为的UCG、Cheetah开源以及IBM的PAMI Collectives。最重要的是，UCC和UCX都支持ROCm和CUDA。
UCX 的组件（参考自 UCX）
UCC 实现为一个实验性的后端变体，作为 PyTorch 分布式模块中的一个实验后端风味。它可以被直接指定为 PyTorch 分布式模块的后端，也可以作为 OpenMPI 集体操作的后端支持。为此，需要从源代码构建 MPI 支持的 torch，并且必须使用 mpirun 启动器来执行这些基于 OpenMPI 的分布式任务。
这被证明是一个突破；在UCF团队的指导下（来自GitHub的链接，关于如何配置：https://github.com/openucx/ucc/issues/1043），我成功确定了一个可行的配置并运行了一个使用PyTorch和MPI的多节点分布式数据并行任务（来自GitHub的示例代码：https://github.com/RafalSiwek/distributed-mlops-overview/blob/main/examples/mpi/multinode_ddp_mpi.py）。
在 AWS G4ad（AMD GPU）和 G4dn（NVIDIA GPU）实例类型上，进行 PyTorch 的分布式数据并行训练。
使用UCC和UCX后，混合GPU集群似乎不再是遥远的梦想，而是一个可以实现的现实。这一突破可能促使组织释放其硬件投资的限制，将分散的资源整合成一台协同工作的超级计算机。
这里详细记录了所有关于环境配置的底层细节和研究和测试的过程。
GitHub - RafalSiwek/troubleshoot-heterogenous-distributed-operations: 我在不同类型的计算资源组成的集群上运行分布式训练任务时遇到困难的仓库，记录了我的努力和挑战……github.com 实现异构的Kubernetes集群在管理企业规模的基础设施时，组织面临着提供所需资源和支持团队的挑战。他们还需要支持任何规模的快速高效机器学习工作负载，包括小型实验和万亿参数量级的长期训练任务。
Kubernetes 成为分布式机器学习编排的基石，凭借其高效地最大化资源利用率和协调不同硬件资源。
要在Kubernetes上调度分布式训练任务，无论是使用Kubeflow MPI Operator还是Pytorch操作符，任务配置文件都需要使用由AMD或NVIDIA提供的设备插件资源。
# NVIDIA (分配给NVIDIA GPU的资源)
resources:
limits:
nvidia.com/gpu: 1
# AMD (分配给AMD GPU的资源)
resources:
limits:
amd.com/gpu: 1指定一个高度异构的任务需要使用自定义资源定义（CRD）或 mutating webhook 处理程序，以便将资源统一使用一个通用名称，例如 异构.com/gpu: 1，或者手动为每个 Pod 单独部署。
VolcanoJob（https://volcano.sh/en/docs/vcjob/），由 Volcano 调度器 提供的 Kubernetes CRD，简化了这一过程。Volcano 被设计用于高性能批处理工作负载，提供 团伙调度 以确保分布式任务的原子执行（即，当所有所需资源都可用时，所有 Pod 同时启动，否则一个也不启动），并通过插件来自动化基础设施的搭建。与 Kubeflow 的 Training Operators（如 MPIOperator）强制所有工人都使用统一资源模板不同，Volcano 允许 每个任务的 Pod 定义，从而实现对异构资源的精确控制。
要在异构的Kubernetes集群上启动混合GPU分布式训练任务，需要配置以下VolcanoJob CRD。
- 自动SSH设置
ssh插件生成包含预共享SSH密钥的ConfigMap，实现Pod间无密码身份验证。每个容器中的sshd使用这些密钥，消除了手动证书管理的需要。 - 工作Pod的DNS解析
svc插件创建了一个无头服务，分配了可预测的DNS主机名。Pod通过Volcano注入的环境变量解析同伴（例如VC_MPIWORKER_AMD_HOSTS），主Pod利用这些变量构建mpirun主机列表。 - 特定资源的任务组
每个task定义了一个独特的Pod模板：
—mpimaster负责训练，使用MPI和UCX标志来优化GPU通信。
—mpiworker-nvidia和mpiworker-amd指定了不同的资源和特定供应商的容器镜像。 - 集群调度
minAvailable: 3确保所有Pod（1个主Pod + 2个工作Pod）一起调度，防止混合集群中的资源死锁。 - 作业完成条件定义
policies字段中的CompleteJob动作键允许配置触发作业完成状态的事件。在这种情况下，由mpimaster任务的TaskCompleted事件触发。
apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
name: mpi-training-heterogeneous
namespace: volcano-job-training
spec:
minAvailable: 3 # 所有3个Pod必须被分配
plugins:
ssh: [] # 通过ConfigMap自动生成SSH密钥
svc: [] # 为稳定主机名创建无头服务
schedulerName: volcano
tasks:
- name: mpimaster
policies:
- action: CompleteJob # 当启动任务成功完成后，作业即完成
event: TaskCompleted
replicas: 1
template:
spec:
containers:
- command:
- /bin/bash
- -c
# 为插件注入无密码配置而创建SSH目录
- mkdir -p /var/run/sshd; /usr/sbin/sshd;
# Volcano 使用VC_MPIWORKER_*_HOSTS注入工作主机:
MPI_HOST=${VC_MPIWORKER_AMD_HOSTS},${VC_MPIWORKER_NVIDIA_HOSTS};
NUM_WORKERS=$(($(echo ${MPI_HOST} | tr -cd ',' | wc -c) + 1));
# 使用mpirun启动训练作业，并推送提取的MPI_HOST和NUM_WORKERS内容
mpirun -np ${NUM_WORKERS} --allow-run-as-root --host ${MPI_HOST} -x MASTER_ADDR=${VC_MPIWORKER_AMD_HOSTS} -x MASTER_PORT=29603 \
# 配置OpenMPI使用UCC作为集体操作后端
-mca pml ucx -mca coll_ucc_enable 1 -mca coll_ucc_priority 100 \
# 配置UCX传输层和UCC集体操作参数，以支持g4ad实例
-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 \
# 指向具有MPI感知能力的PyTorch作业执行代码
/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; # 启动并保持sshd服务运行
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专用GPU
restartPolicy: OnFailure
- name: mpiworker-amd
replicas: 1
template:
spec:
containers:
- command:
- /bin/bash
- -c
- mkdir -p /var/run/sshd; /usr/sbin/sshd -D; # 启动并保持sshd服务运行
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专用GPU运行 PyTorch 分布式作业需要具有 GPU 类型感知的 UCC、UCX 和 MPI 库的 worker 环境，这些库与链接了这些库的 PyTorch 构建版本一起使用。启动器仅需要 UCC、UCX 和 OpenMPI，因为其集体操作不涉及 GPU 特定处理，所以不需要 GPU 感知构建。对于此设置，库和 PyTorch 必须从源代码构建。Dockerfile 配置在我的仓库中提供：我的仓库。此外，特定 worker 实例的预构建镜像在我的公共 DockerHub 注册表中可供下载。
- 搭载AMD GPU（ROCm 6.2.2和PyTorch 2.5.1）的AWS g4ad实例
- 搭载NVIDIA GPU（CUDA 12.4和PyTorch 2.5.1）的AWS g4dn实例
- 纯CPU标准的MPI启动器实例
通过在 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集群上执行分布式训练，为寻求在机器学习和深度学习工作负载中的快速扩展的组织提供了有吸引力的机会。然而，但要实现这一点，在现阶段则需要付出大量的工程努力，因为主流机器学习框架缺乏原生支持。
通过开放标准的实现和发展，可以让普通用户也能方便地访问异构硬件生态系统，同时保持低成本灵活性和高性能。
我想感谢Edgar Gabriel和我们一起讨论的UCC的实现问题。
你可以通过我的社交媒体联系我