如何在新时代为人工智能/机器学习共同设计软件/硬件架构?

美妆时尚专家的秘诀:帮助你在新时代为人工智能/机器学习共同设计软件/硬件架构

设计高效AI/ML架构的整体视角

最近在计算机视觉、自然语言处理等领域,先进的生成式AI技术取得了突破性的研究成果,包括稳定扩散、神经渲染(NeRF)、文本转3D、大型语言模型(LLM)等创新模型架构。这些先进技术需要更复杂的AI网络,并且需要数量级更多的计算资源和能效架构。

为了满足上述架构要求,必须在硬件(HW)和软件(SW)/算法设计之间建立桥梁。协同设计涉及到对HW和SW组件进行设计、测试和优化的迭代过程,直到系统满足所需的性能要求。然而,传统上SW和HW是独立设计的,因为SW程序员很少需要考虑要在哪个HW上运行,而HW通常被设计为支持广泛的SW。对于支持AI工作负载的高效SW/HW协同设计的研究非常有限。

Source: Image by the author.

我们必须设计高效的‘SW/算法感知’HW和‘HW感知’SW/算法,使它们紧密结合以充分利用有限的计算资源。我们应该如何做呢?以下是一种不断发展的方法论,可供您参考,如果您正在为AI架构设计新的HW/SW特性。

1. 确定代理AI工作负载。

为了启动分析,我们需要代理AI模型,并定义一个优先级列表以进行进一步调查。您可以参考多个资源,包括最新的研究论文(CVPR、Siggraph、大型科技公司的研究实验室)和开源代码、客户反馈或请求、行业趋势等。基于您的专业判断,筛选出几个具有代表性的模型。这一步骤至关重要,因为您将使用它们来设计‘未来的架构’

2. 彻底的模型架构分析。

务必全面调查模型架构,以理解其功能、创新,并尽可能详细地将其分解为不同的部分。当前技术堆栈中是否支持新的未支持的操作符?哪些是计算密集型层?这是一个数据传输(内存)密集型模型吗?需要哪种数据类型?可以应用哪种量化技术而不损失准确性?模型的哪个部分可以进行HW加速,哪里存在潜在的性能优化?

例如,在神经渲染中,模型要求同时进行渲染和计算(矩阵乘法),您需要检查当前的SW堆栈是否支持渲染/计算同时进行。在LLMs中,与输入序列长度相关的键值(KV)缓存大小正在增加,了解存储器要求以及潜在的数据传输/存储器层次结构优化以处理大型KV缓存是至关重要的。

Left: Memory layout on NVIDIA A100 when serving an LLM with 13B parameters; Right: Memory usage and Throughput v.s. Batch Size. Figure source: Efficient Memory Management for Large Language Model Serving with Paged Attention [1]

3. SW使能和原型设计

下载在第2步中确定的模型的开源代码,并在目标SW框架/HW上运行。这一步骤并不简单,特别是对于新的/具有颠覆性的模型来说。由于目标是为性能分析实现可行的解决方案,不必在此阶段交付产品质量的代码。可以对SW进行简单修复而不进行性能调优,以继续第4步。一个重要的步骤是将开发框架(Pytorch)中的预训练模型转换为目标新框架所需的新格式。

torch.onnx.export(model, dummy_input, "resnet50.onnx", verbose=False, input_names=input_names, output_names=output_names, export_params=True)

然而,通常有很多情况需要进行大量的支持工作。例如,运行可微分渲染模型需要支持自动求导。新框架很可能尚未准备好这个功能,需要开发团队数月的努力。另一个例子是针对LLMs的GPTQ量化,在推理框架中可能不被支持。架构师可以选择在英伟达系统上运行负载进行性能分析,而不是等待工程团队。因为英伟达是学术开发的硬件选择。这使得在SW启用过程中,基于观察到的差距形成了一份SW需求清单。

4. 性能分析和架构创新。

有许多指标可以衡量AI模型的性能。以下是我们应考虑的主要指标。

4.1 FLOPs(浮点运算)和 MACs(乘加运算)。

这些指标通常用于计算深度学习模型的计算复杂性。它们提供了一种快速简便的方法,了解所需的算术操作数量。可以通过纸上分析、Vtune报告、或者诸如flops-counter.pytorch和pytorch-OpCounter等工具来计算FLOPs。

4.2 内存占用和带宽 (BW)

内存占用主要包括权重(网络参数)和输入数据。例如,一个具有13B参数的Llama模型在FP16下约消耗13*2(FP16=2字节)= 26GB内存(输入可忽略,因为权重占用更多的空间)。LLMs的另一个关键因素是KV高速缓存的大小。KV高速缓存占总内存的30%,并且是动态的(请参考步骤2中的图片)。大模型通常受内存限制,因为速度取决于从系统内存快速转移数据到本地内存,或者从本地内存转移到本地缓存/寄存器的速度。可用的内存带宽远比峰值计算TOPS更好地预测推理延迟(LLMs的令牌生成时间)。一个性能指标是内存带宽利用率(MBU),其定义为实际带宽/峰值带宽。理想情况下,MBU接近100%表示内存带宽得到充分利用。

仅仅有足够的内存是不够的!

Nvidia is increasing Flops by orders of magnitude but not memory BW. Source: Substack/SemiAnalysis

由于内存是一个瓶颈,需要在先进模型压缩和内存/缓存技术方面进行探索。以下是一些先驱性的工作:

  • MemGPT:它利用不同级别的内存层次结构资源,例如快速而小型的RAM和大型而慢速的存储存储器。信息必须在它们之间明确传输。[2]
  • 低精度量化(GPTQ、AWQ、GGML)以减少模型的内存占用
  • 内存计算(PIM):通过消除对数据移动的需求来降低功耗并提高性能。

4.3 延迟/吞吐量。

在计算机视觉中,延迟是生成一帧图像的时间。在LLMs的上下文中,它是从第一个令牌到下一个令牌生成的时间。吞吐量是每秒的令牌/帧数。延迟是衡量AI系统性能的关键指标,是SW/HW性能的综合因素。有各种优化策略可供考虑,以下是其中一些:

  • 优化带宽受限的操作,如归一化、逐点操作、SoftMax和ReLU。估计归一化和逐点操作所需的运行时间比矩阵乘法多约40%,而FLOPS只比矩阵乘法少250倍和700倍。为解决这个问题,可以使用内核融合来融合多个操作符以节省数据传输成本,或者将昂贵的操作符(SoftMax)替换为轻量级的操作符(ReLU)。

<!–运算符类别的比例。来源:数据传输就是你所需要的

  • 专用的硬件架构。集成专用的硬件(AVX、GPU、TPU、NPU)可以显著提高速度和节约能源,这对于需要在资源受限设备上进行实时处理的应用程序尤其重要。例如,英特尔 AVX 指令可以比原生 Python 代码提高多达 60,000 倍的速度。
优化性能工程的程序在两个 4096×4096 矩阵相乘时的加速比例。来源:在摩尔定律之后,顶尖的计算机性能将受到什么推动?[3]

英伟达图形处理器上的张量核(V100,A100,H100 等)可以在一个时钟周期内乘和加两个 FP16 和/或 FP32 矩阵,而 Cuda 核每个周期只能执行1次操作。然而,张量核的利用率非常低(端到端训练为 3% – 9%),导致能源成本高,性能低下。目前正在积极研究改善 systolic array 利用率的方法(FlexSA,多方向 SA 等),我将在下一系列文章中详细介绍。

Systolic array. 来源:Telesens

另外,由于内存和数据传输对于大型人工智能模型始终是一个瓶颈,因此探索考虑更大更高效的芯片内存的先进架构至关重要。一个例子是 Cerebras 核心内存设计,其中内存可以独立地对每个核进行寻址。

Cerebras 内存架构
  • 还有许多其他优化:并行处理、LLM 的 KV 缓存量化、稀疏激活和端到端优化 – 我将在接下来的文章中进行更详细的解释

4.4 功率和能效

功率是我们需要考虑的另一个因素,特别是对于低功耗用户场景。性能和功率之间始终存在权衡。如下图所示,内存访问操作的能耗比计算操作要高几个数量级。减少内存传输是为了节省功耗而强烈要求的。

计算和内存的能耗。来源:斯坦福的宋 H [4]

结论

以上是我们需要衡量AI模型性能的主要指标,使用诸如Vtune、Nsight或其他建模工具的性能分析器,架构师可以深入了解下一层的性能细节(逐层的计算/内存流量、计算效率、带宽利用率等),并利用追踪技术来确定热点。由于软件效率低下,性能往往会出乎意料地低。根据代理模型进行AI软件/硬件的协同设计是一个迭代的过程。

这是一个合作努力,旨在对新兴的AI模型进行彻底的架构分析。调查周期可能相当长,并且产品进度和分析深度之间总是存在权衡。同时需要在硬件、软件和AI技术方面具备强大的专业知识,以设计高效的AI解决方案。

参考文献

[1] Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, Ion Stoica, 采用分页注意力实现大型语言模型服务器的高效内存管理,2023,arxiv

[2] Charles Packer, Vivian Fang, Shishir G. Patil, Kevin Lin, Sarah Wooders, Joseph E. Gonzalez, MemGPT: 将LLM作为操作系统,2023,arxiv

[3] Charles Leiserson,Neil Thompson,Joel Emer,Bradley Kuszmaul,Butler Lampson,Daniel Sanchez和Tao Schardl,There’s plenty of room at the Top: What will drive computer performance after Moore’s law? 2020,Science

[4] Song Han,Jeff Pool,John Tran,William J. Dally,为高效神经网络同时学习权重和连接,2015,arxiv