基于机器学习的GPU L2缓存管理优化框架AutoScratch：面向推理工作负载的性能提升

作者与发表信息 本研究由来自NVIDIA的Yaosheng Fu、Evgeny Bolotin、Aamer Jaleel、Gal Dalal、Shie Mannor、Jacob Subag、Noam Korem、Michael Behar和David Nellans共同完成。该研究论文发表于2023年的第6届MLSys会议（Proceedings of the 6th MLSys Conference），会议地点为美国迈阿密海滩。

学术背景与研究动机 本研究的核心领域是计算机体系结构，特别是针对深度学习的GPU（图形处理器）内存子系统优化。随着深度学习（DL）训练和推理（inference）对计算需求的爆炸式增长，GPU的计算单元（如NVIDIA Tensor Core）性能提升迅速，但其内存带宽（尤其是DRAM带宽）的提升速度未能同步，导致内存带宽日益成为GPU性能，特别是能效比（performance per watt）的关键瓶颈。此外，DRAM的功耗在GPU总功耗中占比显著，这对于面向推理场景的被动散热、低功耗GPU（如NVIDIA T4、L4，其热设计功耗仅70-72瓦）构成了严峻挑战。

为缓解内存墙问题，一种常见方法是增加片上缓存（如SRAM）容量。然而，GPU传统上采用硬件管理的缓存，虽然易于使用，但无法对关键数据提供精确的控制。相比之下，许多专用深度学习加速器采用软件显式管理的暂存器内存（scratchpad memory），能以编程方式精细管理宝贵的片上存储资源，但增加了编程复杂度。为了兼顾易用性与控制精度，NVIDIA在其Ampere架构GPU中引入了L2缓存驻留控制（L2 residency control）机制，允许程序员将特定数据标记为“L2驻留”，使其免受默认硬件缓存替换策略的驱逐，从而像暂存器一样被“钉”在L2中。

然而，手动为复杂的深度学习应用（可能包含数百上千个GPU内核）确定最优的L2驻留数据配置极其困难，且配置错误可能导致性能下降。因此，本研究旨在开发一个自动化框架，以消除人工干预，自动为深度学习推理应用发现并优化L2缓存驻留配置，弥合硬件管理缓存与软件管理暂存器之间的鸿沟。

详细工作流程 本研究提出并实现了一个名为AutoScratch的机器学习优化框架。其核心目标是将应用的共享激活缓冲区（shared activation buffer）划分为多个内存切片（memory slice），并自动学习一个最优的二进制配置向量，指示哪些切片应设为L2驻留，哪些由硬件管理，以最大化应用性能（通常以减少DRAM流量和缩短执行时间为目标）。

1. 框架集成与问题建模： 首先，研究将AutoScratch与NVIDIA的TensorRT SDK集成。TensorRT是一个用于高性能深度学习推理的优化编译器。研究修改了TensorRT，暴露了L2驻留控制的应用程序接口（API），使得任何经TensorRT优化的推理应用都能无缝接入AutoScratch，无需修改应用代码。优化过程针对的是深度学习推理中常见的共享激活缓冲区。该缓冲区在推理迭代中被各层复用，其大小与批次大小（batch size）成正比。研究将这一缓冲区均匀划分为N个等大小的内存切片（实验中采用1MB粒度），从而将L2驻留配置问题转化为一个在{0,1}^N空间中的搜索问题，其中“1”表示对应切片为L2驻留，“0”则表示由硬件管理。对于一个100MB的缓冲区，若L2驻留容量上限为30MB，穷举搜索的组合数将高达C(100,30)≈10^25，故需借助机器学习进行高效搜索。

2. 两种机器学习优化器的设计与实现： 研究开发并比较了AutoScratch的两种实现，分别基于不同的机器学习范式：

• AutoScratch-RL（基于强化学习）：

  • 环境（Environment）： 定义为在特定GPU平台上运行的特定深度学习推理工作负载。
  • 状态（State）： 定义为当前的L2驻留配置，表示为一个N维比特向量。
  • 动作（Action）： 定义为对单个内存切片进行“提升”（设为L2驻留）或“无操作”。动作表示为一个N+1维的one-hot向量。
  • 智能体（Agent）： 采用近端策略优化（Proximal Policy Optimization， PPO）算法。智能体包含演员（actor）和评论家（critic）两个神经网络，均为具有两个隐藏层（每层64个神经元）的多层感知机（MLP）。
  • 奖励（Reward）： 定义为执行单个推理迭代的时间减少量，即 rt = t(st) - t(st+1)。奖励信号直接来源于在真实GPU硬件上运行应用所测量的性能提升。
  • 训练流程： 训练在真实GPU（NVIDIA L4）上进行（on-policy）。每个训练回合（episode）从全零配置（无L2驻留）开始。智能体根据当前状态选择动作，应用新配置后运行工作负载，获得奖励，并转移到新状态。当连续多个状态未发生变化时，回合终止。整个训练过程包含数十万步，最终保留在整个训练过程中性能最佳的配置。

• AutoScratch-EA（基于进化算法）：

  • 算法： 采用一种正则化进化算法（Regularized Evolutionary Algorithm）。
  • 种群（Population）： 维护一个包含M个L2配置向量的种群（实验中M=100）。
  • 初始化： 随机生成M个配置向量，并在目标GPU上评估其适应度（即推理执行时间）。
  • 进化循环： 重复K步（实验中K=20,000，含提前终止机制）。每一步中，从当前种群中随机选择一个大小为S的子集（S=25），评估其中个体的适应度，并选择适应度最佳（执行时间最短）的个体作为父代。
  • 变异（Mutation）： 对父代配置向量执行随机单比特翻转操作，生成子代。
  • 评估与更新： 评估子代的适应度，将其加入种群，并移除种群中最“年长”（最早加入）的个体（而非最差个体），以管理训练噪声。
  • 输出： 进化结束后，选择整个过程中适应度最佳的配置向量作为最终方案。

3. 评估方法学： * 工作负载： 使用MLPerf推理基准测试套件进行评估，包括ResNet50、SSD-ResNet34、SSD-MobileNet、3D-UNet、BERT、DLRM和RNNT等多个代表性模型，涵盖了计算机视觉和自然语言处理任务。 * 硬件平台： 主要评估平台为NVIDIA L4 GPU。这是一款面向推理优化的GPU，具有48MB L2缓存，其中36MB可通过软件API配置为驻留模式。其72瓦的TDP和300GB/s的内存带宽代表了典型的推理场景约束。 * 对比基线： 包括：1) 基线：禁用L2驻留控制，完全由硬件管理缓存；2) 随机搜索（RS）：随机生成20,000个配置并选取最佳；3) 人工设计启发式方法（HDH）：NVIDIA专家设计并集成于TensorRT（v8.5）中的启发式方法，需要程序员指定L2驻留预算（0到1.0）。 * 评估指标： 主要衡量脱片DRAM流量减少和端到端推理性能（每秒推理次数，IPS）提升。同时比较两种AutoScratch优化器的训练开销。

主要研究结果 1. DRAM流量与性能提升： 在NVIDIA L4 GPU上，AutoScratch取得了显著效果。 * DRAM流量： AutoScratch-RL和AutoScratch-EA分别实现了29% 和27% 的几何平均DRAM流量降低（最高可达44%）。其中，写流量（WR）的减少通常比读流量（RD）更显著，因为L2驻留有效防止了激活数据写回DRAM。 * 性能： AutoScratch-RL和AutoScratch-EA分别带来了9.3% 和8.8% 的几何平均性能加速，在BERT模型上最高可获得22% 的性能提升。性能提升与DRAM流量减少趋势基本一致，但并非完全线性，因为不同算子对内存带宽的敏感度不同（例如，MLP-based的BERT和DLRM比卷积网络获得更高加速）。

2. 优化器比较与训练开销： * 效果对比： 两种优化器（RL和EA）在最终优化效果上非常接近，AutoScratch-RL略优但差异不大。 * 训练时间： AutoScratch-EA在训练效率上具有巨大优势。其训练时间仅为数分钟，而AutoScratch-RL需要数小时。具体而言，AutoScratch-EA的平均训练时间比AutoScratch-RL快58倍。这使得AutoScratch-EA更具实用性，可以集成到像TensorRT这样的编译框架中，作为可选的基于配置文件的优化（PGO）步骤。

3. 与替代优化方法的对比： * 随机搜索（RS）： 在搜索空间较小的工作负载（如ResNet50）上表现尚可，但在搜索空间大的工作负载（如3D-UNet）上完全失败，甚至可能劣于基线，其几何平均加速接近0。 * 人工设计启发式（HDH）： 优于随机搜索，在最佳预算设置下可获得2.6%的几何平均加速。然而，最佳L2驻留预算因工作负载而异，需要人工为每个应用单独调优。而AutoScratch能够自动在给定的L2容量预算内发现最优配置，无需额外参数调整，且性能提升显著优于HDH。

4. 敏感性与鲁棒性分析： * 批次大小敏感性： 在ResNet50和BERT上进行的批次大小扫描实验表明，AutoScratch在较宽的批次大小范围内都能有效工作。当批次大小过小时，大部分数据本就适合缓存，优化空间小；当批次大小过大时，可驻留的数据比例有限。AutoScratch有时能将最佳性能点推向更大的批次大小，因为它缓解了DRAM带宽瓶颈。 * 训练稳定性： 多次随机种子训练的结果显示，两种AutoScratch版本都具有高度一致性，性能变异系数在0.02以内。

5. 模拟器对比实验（补充材料）： 为了与硬件缓存替换策略进行理论对比，研究还进行了模拟实验，比较了AutoScratch与签名命中预测器（SHIP）、软件辅助的SHIP（SHIP-SW）以及理论最优的Belady算法。 * AutoScratch在DRAM流量减少上优于SHIP策略（额外减少11%），且缩小了与理论最优Belady算法的差距（仍有26%的流量减少差距）。 * 软件辅助的SHIP-SW比原始SHIP略有改进（额外2%），但无法达到AutoScratch的水平。这证明了基于机器学习的软件缓存控制是一种有前景的新技术。

结论与价值 本研究成功设计并实现了AutoScratch，一个基于机器学习的框架，用于自动优化GPU的L2缓存驻留控制，以提升深度学习推理应用的性能。通过在最新的推理优化型NVIDIA L4 GPU上进行硅片实测，AutoScratch在MLPerf推理基准测试中实现了平均29%的DRAM流量降低和9%的性能提升。研究证明，利用进化算法（EA）的AutoScratch-EA在取得与强化学习（RL）版本相近优化效果的同时，训练速度大幅提升（58倍），使其更易于实际部署。

该研究的价值在于： 1. 科学价值： 提出并验证了一种软硬件协同设计的新思路，即通过机器学习自动化地利用硬件提供的精细控制机制（如L2驻留控制），将硬件管理缓存的易用性与软件管理暂存器的性能潜力结合起来。这为未来计算机体系结构设计提供了启示：通过设计更多可由软件/ML优化的硬件“钩子”（hooks），可以解锁更高的性能与能效。 2. 应用价值： AutoScratch通过集成到TensorRT，为广大深度学习开发者提供了一种全自动、黑盒式的性能优化工具，无需人工干预或针对特定应用进行复杂调优，即可显著提升推理效率，降低功耗，这对于大规模云服务部署至关重要。 3. 方法论贡献： 系统性地比较了强化学习与进化算法在解决同一体系结构优化问题上的效能与效率，为后续研究在选择优化方法时提供了重要参考。

研究亮点 1. 问题新颖性与实用性： 针对Ampere架构GPU新引入的L2驻留控制机制，首次系统性地提出了基于机器学习的自动化优化方案，解决了该机制在实际应用中难以手动配置的痛点。 2. 端到端自动化框架： 实现了从模型编译（集成TensorRT）到在线/离线优化的完整自动化流程，并对业界标准基准测试套件进行了全面评估，证明了其普遍有效性。 3. 双优化器对比研究： 不仅提出了基于RL的解决方案，还创新性地引入并适配了进化算法，并通过实验证实EA在取得相近效果的前提下，训练效率远高于RL，指出了更实用的技术路径。 4. 在真实硬件上训练与评估： 整个机器学习优化循环直接在目标GPU硅片上运行（on-policy），避免了模拟器不准确或离线策略学习可能带来的问题，结果更具说服力。 5. 超越传统硬件策略： 通过模拟对比，证明了AutoScratch学习的软件策略能够超越先进的硬件缓存替换策略（如SHIP），并逼近理论最优性能，展示了软件定义缓存管理的潜力。

其他有价值内容 论文在讨论部分指出，AutoScratch主要适用于那些性能受DRAM带宽/能耗限制、并配备大容量L2缓存的GPU配置（如L4）。在高端的A100 GPU上，虽然AutoScratch也能显著减少DRAM流量，但由于A100的HBM带宽本身不是瓶颈，因此流量减少并未转化为明显的性能加速，但能提升能效（performance per watt）。此外，论文展望了未来研究方向，例如将AutoScratch与批次大小等参数联合优化，或者探索预训练的RL策略在不同应用间的泛化能力。