对于传统的 HPC 工作负载，AMD 的 MI250X 在双精度浮点运算方面仍然很强大。通过一些人工智能模型，AMD 优先考虑 HPC 的决定就显而易见了。当然，除非您碰巧已经拥有 37,888 个可供您使用的。
橡树岭国家实验室的 1.69 exaflops“Frontier”超级计算机就是这种情况，该计算机刚刚使用仅 3,072 个 MI250X GPU 的分区训练了一个万亿参数模型。
对于人工智能训练和推理中常用的混合精度工作负载，MI250X 是一个有趣的计算引擎。在 HPC 工作负载中，该芯片是冠军，可产生 96.7 teraflops 的 FP64 矩阵性能。但众所周知，双精度对于 AI 工作负载来说是无端的杀伤力，您可以用四分之一 (FP16) 的精度换取高 4 倍的原始吞吐量。
不幸的是，在 MI250X 上下降到 FP16 每台设备只能获得 383 teraflops，比 Nvidia A100 好一点，但与 H100 在该精度下的 989 teraflops 性能相差甚远（没有稀疏性支持，它会翻倍至不到 2 petaflops）如果你有一个稀疏矩阵，其中数学技巧使它看起来更密集）。然而，383 teraflops 仍然不容小觑，ORNL 确实拥有大量 MI250X GPU。
要解决的问题是如何在市场中的数百个节点上有效地训练大型模型，其中这些工作负载几乎完全使用针对 CUDA 优化的工具在 Nvidia 硬件上运行。
AMD 凭借其ROCm 运行时在这方面取得了很大进展，现已发展到第六代。在上个月的 AI 活动中，该芯片制造商声称，由于 vLLM、HIP Graph 和 Flash Attention 的优化，性能提升了 1.3 倍到 2.6 倍。虽然前景光明，但仍然有很多代码针对 CUDA 进行了优化。
ORNL 团队在尝试优化 Frontier 以训练万亿参数变压器模型时必须克服其中一些障碍。最值得注意的是，在寻找模型并行技术的最佳组合时，研究人员采用了 Megatron-DeepSpeed 来优化管道并行性。研究人员指出，该框架已成为在人工智能部署中实现高度并行性的标准代码库。
如果 Frontier 是使用 Nvidia GPU 构建的，那么利用 Megatron-DeepSpeed 会相对简单，但事实并非如此。要使其在 AMD 硬件上运行，需要与 AMD 开发人员合作，将该项目移植到 ROCm。
不用说，这并不像运行 HIPIFY 将代码转换为 AMD 的异构计算接口以实现可移植性 (HIP) 运行时那么简单。无论芯片制造商多少次表示他们可以将 CUDA 代码无缝转换为某种与供应商无关的格式，在这种规模下，事情很少那么简单，但情况正在变得更好。
研究人员遇到的令人头疼的问题之一是 DeepSpeed 的操作是在执行训练管道时构建的。不幸的是，这种特殊的细微差别不能很好地与 ROCm 配合，需要研究人员禁用即时编译并预构建它们。
即便如此，研究人员仍需要 AMD 开发人员的帮助来填补 ROCm 运行时的空白。也就是说，必须构建某些基本 CUDA 包的 ROCm 等效项。其中包括 APEX 库，Megatron-DeepSpeed 使用该库进行混合精度计算。
该团队还改编了 ROCm 的 FlashAttention 和 FlashAttention2 实现，以便与 Frontier 上提供的编译器一起使用。后者似乎是一个聪明的做法，因为实验室认为 FlashAttention2 将吞吐量提高了 30%。
至于张量并行性，ORNL 发现，由于调用的“AllReduce”操作数量巨大，尝试跨节点扩展会导致延迟瓶颈。通过将张量并行性限制在具有 8 个 GPU 的单个节点，可以获得最佳结果。请记住，每个“Aldebaran”MI250X实际上是两个 GPU 芯片融合在一起，每个芯片带有 64 GB HBM2e。对于软件来说，它看起来像是两个 GPU，这就是测试。（顺便说一下，后续的“Antares”MI300X看起来不像八个GPU，而是一个，尽管它有八个小芯片，因为它们的互连和缓存耦合更紧密。）
最后，该团队实施了 ZeRO-1 优化器来减少内存开销，并实施了 Amazon Web Services 的 ROCm 集体通信库 (RCCL) 插件（这允许 EC2 开发人员使用 libfabric 作为网络提供商）来提高 Frontier 节点之间的通信稳定性。
在效率方面，该团队发现，对于每个处理器给定的问题大小（也称为弱扩展），数据并行训练的效率为 100%。换句话说，解决问题的 GPU 越多，能解决的问题就越大。
ORNL 发现收益递减的地方是针对固定问题规模进行扩展。直观上，您可能会认为，如果 500 个 GPU 可以在 X 时间内训练一个模型，那么 1,000 个 GPU 将在 X/2 时间内完成。事实上，扩大规模会带来各种瓶颈，这一点在 ORNL 的测试中得到了证实。
通过将全局批量大小设置为 8,000 并改变处理器数量，团队发现在使用 1,024 个 GPU 的 1750 亿参数模型测试中能够实现 89.9% 的效率，在使用 3,072 个 GPU 的 1 万亿参数模型测试中实现 87.05% 的效率。
虽然该团队发现重新分配 Frontier 任务来训练大型 AI 模型取得了成功，但他们强调，需要做更多的工作来提高这些工作负载在 AMD 硬件上的性能。大多数现代训练框架并不是为非 Nvidia 硬件构建的，并且对 ROCm 平台的支持仍然很少。
尽管如此，橡树岭国家实验室仍然希望从这次实验中汲取的经验教训可以作为其他设施的蓝图，例如运行非 Nvidia、非 CUDA 系统的阿贡国家实验室。
