亚马逊AWS官方博客
高性能大规模机器学习分布式训练 – EFA(Elastic Fabric Adapter)最佳实践
背景
大语言模型 LLM 和生成式人工智能进一步推动了新一代机器学习的大规模普及和应用,同时也有越来越多的团队选择构建自己的大(语言)模型,以实现更好的生成效果。但是,高效地训练这些包含数十亿甚至数千亿参数的模型是一个非常大的挑战,因为模型越大,理论上所需要的计算资源就越多,如何设计一个高性能的基础设施来支持这些密集型的分布式工作负载,已经成为一个关键性的问题。
为了解决这个难题,在亚马逊云科技平台上可以使用面向大规模机器学习(以及高性能计算 HPC)场景下的 Elastic Fabric Adapter(下文称 EFA),它是一种 Amazon EC2 实例上创新的网络接口,可以显著提高机器学习和高性能计算应用程序的网络性能。EFA 利用了 OS-Bypass 技术,如下图所示,可以实现更低的延迟和更高的吞吐。
![]() |
图片来自 https://www.linkedin.com/pulse/why-aws-removed-tcp-from-operating-system-srd-protocol-singh/
从上图上也可以看到,EFA 底层借助了亚马逊云科技自己创新实现的 Scalable Reliable Datagram(SRD)网络传输协议,最小化网络抖动并确保对网络拥塞作出最快反应,下图是来自于这篇博客中的 TCP 和 SRD 的对比测试。更多有关 SRD 的信息,可以参考这篇论文,里面探讨了 SRD 的设计思路、解决方案和实际效果。
![]() |
EFA 可以无缝地与亚马逊云科技的加速计算实例集成(更多支持实例类型,可以参考文档),如 GPU 类型的 P4de 和 P5。在训练大模型时,使用 NCCL 通信库(NVIDIA Collective Communications Library)进行应用程序可以扩展到数千个 CPU 或 GPU,并支持 NVIDIA GPUDirect RDMA 功能。
![]() |
同样,EFA 也支持 Trainium 类型的 Trn1 和 Trn1n 等实例(更便捷地使用配备了 Trainium 和 Inferentia 自研芯片的实例类型,可以参考博客通过 Optimum Neuron 加速 Llama 2 和 SDXL 在 Inf2 和 Trn1 上部署推理)。
![]() |
用户可以通过 EFA 来创建一个低延迟的高速网络,连接数千个 GPU 或者加速芯片,实现模型并行训练。相比传统的以太网和 TCP/IP 协议,EFA 提供了更高效的点对点通信,从而大大缩短了分布式训练的时间,对于希望构建自己大模型的团队来说,EFA 无疑可以成为基础设施架构的最关键的组件之一。目前 EFA 已经发展到了第二代,在硬件和软件层面都有很多性能增强,更多信息可以参考这篇博客。
对于想使用 EFA 进行大规模训练的团队来说,为了能很好地接入 NCCL 这套分布式通信体系,亚马逊云科技开源了一套 OFI NCCL 插件(OFI – Open Fabrics Interfaces,也称为 Libfabric),它使得机器学习开发者能够在运行基于 NCCL 的应用程序时使用 Libfabric 作为网络提供程序,是 NCCL 和 EFA 之间的一个翻译层。
上面的文字和图片中提到了很多概念,为了简化理解,这里按照模块化的逻辑进行了分层整理,对应结构如下图所示:
![]() |
这时候可能已经关注到上图中标星号“*”的模块,为了方便开发者使用 EFA,亚马逊云科技非常贴心地提供了 EFA Installer 一键部署的程序帮助到用户,安装涵盖了上述标星号“*”的模块。
开启 EFA
想要在亚马逊云科技平台上使用 EFA 非常简单,只需要在创建 ENI(Elastic network interfaces)网卡的时候,选择开启 EFA 就可以,或者在命令行/程序中使用 InterfaceType=efa 即可。
![]() |
注:需要注意的是如果是使用 P5 实例,要特别关注 NetworkCardIndex 和 DeviceIndex 的设置。其中,Network Card 的概念区别于 ENI,也并不常见,因为亚马逊云科技平台上大多数实例类型都只支持一个 Network Card,只有这里列出的特定实例类型(其中 GPU 和 Trainium 对应的实例当然也涵盖在内)才具有多个 Network Card,这些实例能够提供更高的网络性能,包括超过 100 Gbps 的带宽能力和提升的数据包处理性能。每个 ENI 都需要连接到一个 Network Card 上,主 ENI 必须在 Network Card 上分配索引为 0。如果在支持多 Network Card 的实例上启用了 EFA,每个 Network Card 可以分配最多一个 EFA。对于 P5 来说,建议 EFA 的使用遵循以下规则:
- 对于 NetworkCardIndex 0,指定 DeviceIndex 为 0
- 对于 NetworkCardIndex 1-31,指定 DeviceIndex 为 1
参考的命令行如下所示:
如果想要大规模使用 EFA,就需要用户提前进行一些必要的规划和设计,然后在使用中注意如何达到更高的性能水平,并进行持续的监控。所以,本文通过设计、性能、工具和监控四个方面介绍下使用 EFA 的最佳实践。
1. 设计篇
优秀的架构需要提前进行设计,我们需要了解 EFA 的特点和限定因素,来规划出一个合理的基础设施架构,因为 EFA 主要是网络相关的内容,那我们就从网络方面开始入手。
1.1 网络
1.1.1 子网规划
就目前而言,EFA 之间的数据面流量只能在同一个 VPC 子网内发送和接受,暂时还不能跨子网路由,这意味着启用 EFA 的实例需要在同一个子网内。相对地,EFA 的正常 TCP/IP 流量仍然可以跨子网发送,可用于实例之间的控制面通信。
对于子网 IP 地址的规模也需要提前计算清楚,像 p5.48xlarge 和 trn1n.32xlarge 这样的大实例,分别至少需要挂载 32 和 16 个 ENI 接口,这样没台实例就会分别消耗 32 和 16 个 IP 地址。这样在一个 100 台规模的集群情况下,至少要有 3200 和 1600 个地址需要预留出来,并且这些 IP 地址不会在 EFA 的通信中(非 TCP/IP)使用。因此,这里建议采用大实例类型的客户选择 IPv6 地址,数量丰富不容易受容量的限制,但是一定要注意的是对应的 libfabric、NCCL 等工具也要做好 IPv6 的适配,采用实际的工作负载进行测试验证。
如果使用 IPv4 地址的话,就更加需要做好地址池的规划,因为单个 VPC CIDR 最大支持/16 位的掩码,理论上有 65536 个 IPv4 地址,所以在子网划分上要提前计算并准备,也要为未来扩展集群做好充分的准备。
1.1.2 安全规则组
启用了 EFA 的 ENI 和更普遍使用的 ENI(ENA ,Elastic Network Adapter)一样,也需要配置安全规则组,EFA 所在的安全组务必配置允许从本安全组入站以及出站到本安全组的规则,即自身到自身的设置。EFA 使用 RDMA 技术进行节点之间的直接内存访问通信,需要安全组规则允许所有入站和出站(即所有端口号之间)的流量。另外,由于 EFA 通信依赖实例间直接连接,尽管通常不会调整和配置 VPC 的 ACL,最好还是确保网络 ACL 等不会阻止 EFA 实例之间的访问。
如果安全组没有按照上述要求设置,通常会在通信时看到以下错误提示,进行相应的调整即可。
1.1.3 EFA 与 ENA
这里另外需要注意的一点是,EFA 和 ENA 在实例 OS 上看到的是两个不同的 PCIe 设备,具有不同的 ID。但是,如果同一个 Network Card 上如果既有 EFA 也有普通的 ENA,那么它们仍然是共享对应 Network Card 的网络带宽,例如 100 Gbps。
![]() |
通过 lspci 命令可以查看对应的 PCIe 信息,下面是一个参考示例,可以很明显看到在操作系统上它们是不同的设备。
1.2 计算
1.2.1 Placement Group
对启用 EFA 的实例来说,Placement Group 并不是绝对必需的,但是建议将启用 EFA 的实例运行在 Placement Group 集群模式中,因为这可以将实例启动到单个可用区中的低延迟组内。更多关于 Placement Group 的信息,可以参考这里的文档。
![]() |
1.2.2 容器化
容器技术已经在机器学习领域的训练过程中被越来越广泛地采用,这主要是因为容器带来的一致性、可复制性、资源效率、部署灵活和协作便利等优势。容器可以通过封装例如 NCCL、CUDA、Neuron SDK 等库以及应用程序代码和运行环境,保证了不同实例上训练环境的一致,尽可能避免“测试的时候可用,线上运行就报错“的问题。此外,通过容器平台调度,GPU、Trainium 等加速硬件资源也可以动态、有效地分配,大幅提升资源利用效率。容器镜像实现了对训练好的模型及环境的打包,也便于快速分发和部署,多人协作时可以避免环境不一致等问题。目前,容器已经是机器学习领域的首选训练方式,相信其采用率会持续增长。
亚马逊云科技加速计算团队目前已经为用户准备好了一些容器镜像,可以参考这个链接,容器内已经封装好了 EFA、NCCL、CUDA、cuDNN、OpenMPI 等组件,在测试单容器功能或者构建自己的镜像时可以参考。以下是一个可以参考的单容器启动命令,注意这里可以通过 device 的方式将 EFA 挂载到容器上,EFA 会以/dev/infiniband/uverbs 路径体现。
此外,在测试多容器之间的通信时,如通过 mpirun 命令测试 NCCL,也需要在容器内开启免密 SSH 登录,可以参考以下命令:
1.2.3 Kubernetes 和 EKS
用户的容器编排平台通常会选择事实标准的 Kubernetes 环境,集成 EFA 到 Kubernetes 或者 Amazon EKS 集群中的应用可以无缝地利用性能超强的加速计算实例,减少大规模分布式训练的完成时间。在 Kubernetes 或者 Amazon EKS 集群中使用 EFA,需要使用像 GPU 一样的 Device Plugin,亚马逊云科技提供了在 EKS 上使用 EFA 的开源项目供用户参考。
在使用 EFA 时,只需要在 spec 中的 resources 中的 limits 和 requests 中引用对应资源即可。
注:可以看到上面的资源中还涉及了内存巨页 Huge Pages,理论上内存巨页并不是完全必须的,但是会带来微弱的性能损失(具体差异还需要在实际工作负载中进行测试)。这里需要提到一个 Libfabric 环境变量参数 FI_EFA_FORK_SAFE,Neuron 文档中提到旧版本的 Linux 内核<5.15 需要设置环境变量 FI_EFA_FORK_SAFE 为 1,以使 Libfabric 正常工作。然而,根据 Libfabric 文档说明,设置 FI_EFA_FORK_SAFE 会启用 fork()支持,但是这可能会对性能有轻微影响,建议只有在需要时才应设置,此外需要内存巨页和 fork 是不兼容的。因此,上述 hugepages 参数并不是完全必要的,请根据实际情况考虑设置。
1.2.4 SageMaker
另外,如果想要寻求一个可以快速启动带有加速硬件的计算资源时,也可以考虑使用 Amazon SageMaker 服务。SageMaker 默认内置提供了与 EFA 设备的集成,允许用户在 SageMaker 环境中运行分布式训练作业时利用 EFA 设备。用户可以将 EFA 集成添加到 SageMaker 现有 Docker 容器中,关于如何配置容器,请参考文档。
![]() |
需要注意的是,如果深度用户在 SageMaker 环境内通过 ip 或者 ifconfig 等工具查看网卡信息,会看到网卡和 EFA 数量并不是匹配的,这是因为容器环境设备挂载和网络模式决定的,完全不影响 EFA 使用。但这里通常不需要用户额外关注,如果使用 SageMaker 环境建议更多关注大模型本身并对环境的运行做好监控,持续关注是否达到性能预期。
2. 性能篇
在完成了前期网络和计算规划后,用户就可以开始对性能进行管理和调优,在早期完成对 EFA 性能的了解和测试。
2.1 TCP 和 RDMA 流量
再次强调一个点,ENA 和 EFA 在同一个 Network Card,在操作系统层面会显示为两个不同的 PCIe 设备如下图所示,但是它们是共享 Network Card 上的网络带宽的,例如 100 Gbps。
![]() |
下图是对于两种流量划分的测试,通过下载 EFA Installer 安装包来模拟 TCP 流量,通过 EFA Installer 中带的 efa_test.sh 脚本(利用了 fi_pingpong 命令)来模拟 RDMA 流量。然后,通过/sys/class/net/<ensX>/statistics/和/sys/class/infiniband/<efa_X>/ports/1/hw_counters/路径分别来获取 TCP 和 RDMA 流量,可以明显看到两种流量的区分,后续在监控时也可以同时关注两种流量。
![]() |
2.2 NCCL Test
想要查看 GPU 分布式训练的基础网络性能,NCCL Test 是必备测试,以查看集群间的带宽情况。执行 NCCL Test 测试很方便,在 EC2 实例上运行可以参考以下命令。
在容器或者 Kubernetes 环境中可以参考以下配置。
执行测试过程中,可以注意查看日志,以 P4d 为例下面的日志应该被视为正确检测实例并生成拓扑结构:
注意 P4d 以上使用了 GPUDirect RDMA,日志中应该查看到以下日志:
最终我们要查看运行结果,NCCL Test 结果中提供了两种带宽:1)算法带宽 algorithm bandwidth 和 2)总线带宽 bus bandwidth。通常来说,总线带宽更有用于估计网络性能,这里需要注意的是带宽的单位是 GB/s而不是 Gb/s。
算法带宽和总线带宽是具有固定的数值关系的,针对不同 NCCL 运算的总线带宽到算法带宽的校正系数为:
- AllReduce 是 2*(n-1)/n
- ReduceScatter 是 (n-1)/n
- AllGather 是 (n-1)/n
- Broadcast 是 1
- Reduce 是 1
其中 n 是 rank 的数量,总线带宽需要乘以校正系数才能得到算法带宽。一个参考的 NCCL Test 的结果如下所示(这里数值仅为示意,请根据实际情况进行测试):
![]() |
建议开启 4 台以上的跨机实例进行测试,如果只有 2 台示例,会收到实例内 NVLink 的影响,导致数值并不准确,甚至超过实例标称最大带宽值。
2.3 Perftest
如果需要进阶地了解 EFA 在 RDMA 中的特定操作的 micro-benchmark 性能,可以考虑使用 Perftest 工具,它是一组针对 uverbs 编写的性能基准测试,这些测试可用于硬件或软件调优。第一代 EFA 为 NCCL 通信添加了 RDMA read 语义,更进一步在 EFAv2 中添加了更完整的 RDMA 接口,允许更复杂的语义。
![]() |
uverbs(userspace verbs)在之前挂载设备的时候已经见过(/dev/infiniband/uverbs0),它是 Linux 内核的一个组件,会创建名如上的设备节点以支持 InfiniBand RDMA 网络(SRD 协议也是受 InfiniBand 启发,从设备名称中也可以看到这点,关于 Infiniband 和 SRD 的对比可以参考 Deep-Dive Into 100G networking & Elastic Fabric Adapter on Amazon EC2),该组件允许用户空间应用通过 verbs 接口直接访问 RDMA 硬件。
2.4 实际业务
在经过各种 Benchmark 测试后,强烈建议在小规模集群情况下(>=4)使用真实的业务模型进行验证,以查看线上业务在 EFA 上进行分布式训练时的运行情况。如果遇到问题,可以及时定位性能瓶颈点,如果一切顺利,再向大规模集群逐渐扩展开。
3. 工具篇
在完成性能验证和调优后,还需要一些趁手的工具对 EFA 进行查看,这里列举了一些查看 EFA 及其对应设备的工具集,希望对用户有所帮助。
3.1 设备信息
在 Infiniband 环境中,通常会使用 show_gids 命令来查看 RDMA 信息汇总表,其中列出了设备、端口、索引 GID、IP 地址和设备名称等,这些信息本身暂时不包含在 EFA Installer 安装程序中,不过可以通过参考这个脚本来进行编写,因为前文提到 EFA 也参考了 Infiniband 协议,同样在系统路径/sys/class/infiniband/下保存了设备信息。
![]() |
除此之外,EFA Installer 中涵盖了查看 fabric 的工具 fi_info,更多关于 fi_info 的参数可以查看链接。
另外,EFA Installer 中也预装了 ibv_devices 和 ibv_devinfo,可以对 EFA 设备进行更多信息的查看。
3.2 网卡映射
另外,如果希望记录 EFA 设备和对应网卡,可以参考 ibdev2netdev 工具进行相应地修改,这里提供了一个 EFA 的参考版本,这里利用了 fi_info 命令中地址信息进行的对应。
3.3 镜像构建
无论使用虚拟机还是容器, efa-nccl-baseami-pipeline 这个开源项目都提供了基础镜像构建的参考模版,也包括了构建规范的示例,可以将其与 CodeBuild/CodePipeline 等服务集成以进行自动构建。此外,这个项目还提供了监控相关的参考实现,参见下一节监控篇。
4. 监控篇
线上真实运行大规模 EFA 环境后,务必要对 EFA 的使用进行密切监控,可以根据自身业务和模型的特点,对监控指标设定警告阈值,这样遇到问题及时响应并处理。
4.1 系统指标
用户可以通过/sys/class/infiniband/<efa_X>/ports/1/hw_counters/这个路径来查看对应的流量信息:
4.2 VPC Flow Log
可以通过 VPC Flow Log 来捕获 EFA 的进出流量,为 EFA 创建 Flow Log 的方式与为 ENI 创建 Flow Log 一样,注意在 Flow Log 日志条目中,EFA 流量由 srcAddress 和 destAddress 标识,都被格式化为 MAC 地址形式,以下所示:
4.3 CloudWatch 集成
前文提到,efa-nccl-baseami-pipeline 这个开源项目也提供了和 CloudWatch 集成的监控模版,包括了对 GPU 和 EFA 的监控,如下图所示:
![]() |
其中,对 EFA 指标的抓取如下图所示,用户也完全可以参考这些模版进行相应的调整,实现自己的监控体系:
![]() |
综上,本文通过设计、性能、工具和监控四个部分介绍了在大规模机器学习训练场景下使用 EFA 的最佳实践,希望能帮助到有需求的用户。在真实环境中,相信一定还会遇到更多的问题需要处理,可以试着静下心来采用 Problem Formulation 的方式逐一解决。
小结
EFA 是亚马逊云科技推出的一项创新网络加速技术,专门针对分布式机器学习训练进行了优化,能够大幅减少实例间通信延迟并提高吞吐,能够显著提升机器学习训练的性能。通过 EFA,用户可以实现数千数万张加速芯片集群的大规模分布式训练,同时保留了云上的弹性和灵活性优势。EFA 扩展了机器学习训练的规模和性能,是广大机器学习用户实现大规模分布式训练的绝佳选择。
参考资料
https://ieeexplore.ieee.org/document/9167399
https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/p5-efa.html
https://github.com/aws-samples/aws-efa-nccl-baseami-pipeline
https://www.hpcworkshops.com/08-efa/00-efa-basics.html