尽管全世界都疯狂追逐NVIDIA GPU，如H100，但Databricks却采取了战略性的举措，选择利用AMD GPU提升他们的LLM训练能力，并且这对公司来说效果非常好。

去年，Databricks与AMD达成合作，使用了其第三代EPYC Instance处理器。今年六月，Databricks收购了MosaicML，后者正在使用AMD MI250 GPU训练AI模型。这家企业软件公司看到了AMD的潜力，并且期待着这家芯片制造公司下一代MI300X的发布，以在生成式AI领域崭露头角。

AMD GPU似乎在社区中获得了广泛的采用，证明了它们在人工智能领域的实力。包括Lamini和Moreh在内的知名人工智能初创公司已经采用了AMD MI210和MI250系统来调优和部署定制的LLM模型。仅在一周前，Lamini透露他们的LLM模型是在AMD的Instinct GPU上运行的。

Moreh成功使用1200个AMD MI250 GPU训练了一个具有惊人的2210亿个参数的语言模型。AMD最近还向Moreh注资了2200万美元的B轮基金。AI2的开源LLM模型OLMo也采用了大量集群的AMD GPU进行训练。

Databricks宣布他们获得了AMD加速器云(AAC)的一组重要的多节点MI250集群的提前访问权限。该集群由32个节点组成，每个节点装有4个AMD Instinct MI250 GPU，并且具有800Gbps的互连。这个设置非常适合在AMD硬件上进行规模化的LLM训练。

现在该公司已经扩展到使用128个MI250 GPU。

原因很简单——软件

我们都知道NVIDIA真正独占的是其软件CUDA，这是其所有人工智能能力的背后。AMD也意识到了这一点，并一直在软件创新方面走在前列，特别是AMD Radeon Open eCosystem (ROCm)软件平台，这是AMD替代CUDA的解决方案。

近期，AMD人工智能高级副总裁Vamsi Boppana表示，ROCm是该公司目前的头等大事。“我们在软件方面有更多的人力资源，AMD首席执行官Lisa Su也明确表示，她希望在软件方面进行重大而持续的投资，”他说。

ROCm经历了重要升级，从5.4版本升级到5.7版本。此外，Databricks强调，FlashAttention的ROCm内核已经升级到FlashAttention-2，带来了大幅的性能提升。

Lamini还表示，AMD的ROCm已经可以投入生产，并声称“它具有巨大的潜力，可以加速LLM微调和更多领域的人工智能发展，与CUDA相当甚至更好”。

Databricks还赞扬了AMD在OpenAI的Triton编译器中的积极参与。这种贡献使得机器学习工程师能够开发在包括NVIDIA和AMD系统在内的各种硬件平台上高效运行的定制内核，从而简化了移植定制模型代码和内核的过程，无需使用ROCm特定的内核。

对于Databricks来说，AMD GPU的表现如何？

Databricks在采用ROCm 5.7和FlashAttention-2进行训练时，训练性能有了显著的改善，提升了1.13倍，在与ROCm 5.4和FlashAttention的之前结果进行对比时。此外，Databricks展示了强大的扩展性，性能从单个节点上的166 TFLOP/s/GPU提升到32个节点上的159 TFLOP/s/GPU，同时保持了一致的全局训练批量大小。

在64个MI250 GPU上成功训练了具有10亿和30亿参数的MPT模型后，训练过程保持稳定，最终模型的评估指标与著名的开源模型，如Cerebras-GPT-1.3B和Cerebras-GPT-2.7B相当。

为了进行训练，Databricks利用了基于PyTorch的开源训练库，如LLM Foundry，建立在Composer、StreamingDataset和PyTorch FSDP之上。这得益于PyTorch对CUDA和ROCm的支持，使得在NVIDIA和AMD GPU上的无缝操作成为可能，无需进行代码修改。

Databricks表示，他们对即将发布的下一代AMD Instinct MI300X GPU充满期待。他们预计他们基于PyTorch的软件栈将继续无缝运行和有效扩展。

此外，AMD和Triton的整合将简化移植定制模型代码和内核的过程，消除对ROCm特定内核的需要。

Databricks的研究员Abhi Venigalla表示：“H100仍然位居榜首，但我们期待着尽快测试AMD的新MI300X，我们相信它将非常具有竞争力！”

Lamini也急切等待MI300X的推出，该产品配备了192GB的高带宽内存(HBM)，这将使他们的模型表现更好。

总结起来，Databricks转向AMD GPU在LLM训练领域迈出了重要一步。这一发展也证明了AMD在GPU领域逐渐获得的地位。