随着人工智能的进步不断要求云端进行新的创新,我们发现自己正处于这样一个时刻:硬件与软件的共同优化对于优化人工智能基础设施以实现最佳性能、可扩展性和可互换性至关重要。
在Hot Chips 2024上,微软分享了 Maia 100 的规格,这是微软第一代定制 AI 加速器,
做额为一颗专为部署在 Azure 中的大规模 AI 工作负载而设计。Maia 100 系统采用垂直集成以优化性能并降低成本,其平台架构采用定制服务器主板和量身定制的机架,以及旨在提高性能和成本效率的软件堆栈,用于 Azure OpenAI Services 等服务上的高级 AI 功能。
Maia 100 加速器专为各种基于云的 AI 工作负载而设计。该芯片尺寸约为 820 平方毫米,采用台积电的 N5 工艺和 COWOS-S 中介层技术。Maia 100 的标线大小 SoC 芯片配备大型片上 SRAM,结合四个 HBM2E 芯片,可提供每秒 1.8 TB 的总带宽和 64 GB 的容量,以满足 AI 规模的数据处理要求。
Maia 100 设计支持高达 700W 的 TDP,但配置为 500W,可提供高性能,同时根据其目标工作负载有效管理电源。
Maia 100 的架构根据现代机器学习需求而量身定制,体现了对人工智能系统进行深思熟虑的研究的应用,以实现最佳计算速度、性能和准确性。
1、高速张量单元可提供快速的训练和推理处理,同时支持多种数据类型,包括低精度数据类型,如微软于 2023 年通过 MX 联盟首次推出的 MX 数据格式。该张量单元构建为 16xRx16 单元。
2、矢量处理器是一个松散耦合的超标量引擎,采用定制指令集架构 (ISA) 构建,支持多种数据类型,包括 FP32 和 BF16。
3、直接内存访问 (DMA) 引擎支持不同的张量分片方案。
4、硬件信号量支持在 Maia 系统上进行异步编程。
Maia 加速器采用低精度存储数据类型和数据压缩引擎设计,以减少大型推理作业所需的带宽和容量要求,而大型推理作业通常因数据移动而成为瓶颈。为了进一步提高数据利用率和能效,大型 L1 和 L2 暂存器由软件管理,以实现最佳数据利用率和能效。
2023 年,微软牵头成立了超级以太网联盟,帮助业界使用专为超高带宽计算而设计的基于以太网的互连。Maia 100 支持高达 4800 Gbps 的全收集和分散减少带宽,以及 1200 Gbps 的全到全带宽。这种以太网互连采用定制的 RoCE 类协议,提供增强的可靠性和平衡性。Maia 的后端网络协议支持 AES-GCM 加密,也使其成为机密计算的理想选择。Maia 100 还由统一的后端网络支持,用于扩展和扩展工作负载,从而提供灵活性以支持直接和交换机连接。
Maia 100 的硬件和软件架构从头开始设计,旨在更高效地运行大规模工作负载,它垂直整合了我们在云架构每一层所学到的知识——从先进的冷却和网络需求到允许快速部署模型的软件堆栈。Maia 软件开发工具包 (SDK) 允许用户快速将用 PyTorch 和 Triton 编写的模型移植到 Maia。
Maia SDK 为开发人员提供了一套全面的组件,以便将模型快速部署到 Azure OpenAI 服务:
1、框架集成:一流的 PyTorch 后端,支持 Eager 模式和图形模式;
2、开发人员工具:用于调试和性能调整模型的工具,例如调试器、分析器、可视化器以及模型量化和验证工具;
3、编译器:我们有 2 种针对 Maia 的编程模型和编译器 - Triton 编程模型提供灵活性和可移植性,而 Maia API 适合最高性能。
4、内核和集体库:使用编译器,我们开发了一组高度优化的 ML 计算和通信内核,让您能够快速开始使用 Maia。还支持编写自定义内核。
5、Maia 主机/设备运行时:主机设备运行时层带有硬件抽象层,负责内存分配、内核启动、调度和设备管理。
Maia 编程模型利用信号量进行异步编程以实现同步,从而实现计算与内存和网络传输的重叠。它使用两个执行流进行操作:控制处理器通过队列发出异步命令,硬件线程执行这些命令,从而通过基于信号量的同步确保高效的数据处理。
要对 Maia 进行编程,开发人员可以从两种编程模型中进行选择:Triton 是一种流行的开源领域特定语言 (DSL),用于深度神经网络 (DNN),可简化编码并可在 GPU 和 Maia 上运行;Maia API 是一种 Maia 特定的自定义编程模型,旨在通过更详细的控制实现最高性能。Triton 需要的代码行数较少,并可自动处理内存和信号量管理,而 Maia API 则需要程序员编写更多代码并进行明确管理。
Maia 使用基于 Gather 的方法进行大型分布式通用矩阵乘法 (GEMM),而不是基于 All-Reduce 的方法。这种方法有几个优点:通过将后 GEMM 激活函数(如 GELU)直接融合在 SRAM 中来提高处理速度和效率;通过将计算与网络通信重叠来减少空闲时间,并通过网络发送量化数据来减少延迟,从而加快数据传输速度并提高整体系统性能。
此外,我们在集群级别利用静态随机存取存储器 (SRAM) 来缓冲激活和中间结果。网络读写也直接由 SRAM 提供,从而可以直接访问 CSRAM。这显著减少了 HBM 读取,从而改善了延迟。
我们通过跨集群并行计算并利用片上网络 (NOC) 进行片上激活收集来进一步提高性能。
Maia 100 可互换性的关键在于它能够使用一行代码针对 Maia 执行 PyTorch 模型。这由 PyTorch 后端支持,该后端既可以在 Eager 模式下运行以获得最佳开发人员体验,也可以在图形模式下运行以获得最佳性能。利用 PyTorch 和 Triton,开发人员可以在硬件后端之间实现完全的可移植性和灵活性,从而优化工作负载性能,而不会牺牲效率和针对 AI 工作负载的能力。
Maia 100 凭借其先进的架构、全面的开发人员工具以及与 Azure 的无缝集成,正在彻底改变微软管理和执行 AI 工作负载的方式。通过算法协同设计或硬件与软件、为模型开发人员和自定义内核作者提供的内置硬件可选性以及垂直集成设计以优化性能并提高能效同时降低成本,Maia 100 为在微软 AI 基础架构上运行先进的基于云的 AI 工作负载提供了新选择。