本文发表于第33届国际高性能并行与分布式计算研讨会（HPDC ‘24），作者为Zichao Yang、Hao Guo（共同第一作者，中国科学院大学）、Heng Wu、Yuewen Wu、Hua Zhong、Wenbo Zhang（通讯作者，中国科学院软件研究所）、Chuan Zhou（中央民族大学）以及Yan Liu（浪潮软件科技有限公司）。
一、研究背景与研究目的
随着深度学习在计算机视觉和自然语言处理等领域的广泛应用，云数据中心需要为海量的深度学习模型训练任务提供高效的硬件资源调度。实现高质量调度决策的关键，在于准确预测深度学习模型在不同硬件环境下的训练迭代时间（iteration time）。现有的预测方法主要分为基于算子（operator-based）和基于图（graph-based）两类，但它们都因未能提取运行时执行的序列特征（sequential features）而存在明显缺陷。基于算子的方法往往忽略了深度学习框架（如PyTorch、TensorFlow）在调用GPU执行每个算子时产生的开销，本文观察发现这部分开销占迭代总时间的比例高达4%到68%，导致预测精度不足。基于图的方法虽然直接利用模型计算图进行预测，但面对从小型模型（如仅含5个算子的MLP）到大型模型（如包含超过1000个算子的Transformer）的巨大规模差异，预测模型要么结构臃肿、收敛缓慢，要么需要压缩特征而牺牲精度。为了解决上述问题，本文作者提出了ETS方法，旨在通过一种新颖的执行追踪滑动窗口（execution trace sliding window）技术，同时实现高精度预测和模型快速收敛。
二、研究工作流程
ETS的核心思想源于一个关键观察：尽管深度学习模型被定义为有向无环计算图，但模型在运行时的算子执行是严格顺序的，并且在算子执行间隙穿插着深度学习框架开销。基于此，研究工作分为两个主要阶段：离线训练与在线推理。
在离线训练阶段，研究者首先构建了一个大规模的数据集。他们收集了超过14,000个深度学习模型配置，这些配置来源于两个部分。第一部分是超参数优化（HPO）数据集，包含18种真实世界模型通过变化超参数（如批次大小、输入形状、卷积核大小）生成的6500余个配置，涵盖视觉（如ResNet、EfficientNet）、自然语言处理（如Transformer）和语音识别（如LSTM）等领域。第二部分是RandWire数据集，包含7500余个通过神经架构搜索（Neural Architecture Search）随机生成的模型配置。所有配置均使用NVIDIA Nsight Systems工具进行性能分析，以获取详细的运行时执行追踪信息，其中包括每个算子的类型、训练阶段（前向传播forward pass、反向传播backward propagation、参数更新parameter update）、超参数、张量（tensor）大小、计算量（FLOPs）以及算子间GPU空闲的框架开销时间。
在数据获取后，ETS的核心方法——执行追踪滑动窗口被应用于数据处理。对于一个给定的深度学习模型，其完整的执行追踪被视为一个包含算子执行时间和框架开销的序列。研究团队设计了一种滑动窗口机制，以一个固定的窗口宽度（最终调优为12）和步长（stride，最终调优为4）在该序列上滑动切片。每个窗口内的算子序列及其对应的执行时间和开销构成一个定长的训练样本。这一方法解决了两个关键挑战：首先，窗口内的序列数据明确包含了算子执行顺序和框架开销，这是准确预测这些开销并提升整体预测精度的关键特征。其次，通过滑动窗口将不同规模的模型执行追踪截断并规范化为等长的样本，使得预测模型无需处理变化巨大的输入尺寸，从而能够设计得更轻量，实现快速收敛，并且滑动窗口本身还起到了数据增强（data augmentation）的作用。
模型训练时，每个样本的输入为窗口内所有算子的特征向量序列。这些特征向量是研究者精心设计的手工特征编码（handcrafting feature encoding），包括算子的类型（采用独热编码和频率编码结合的双层方法以缩减特征尺寸）、所处训练阶段、超参数（如卷积层核大小、步长）、各类张量大小的对数值、FLOPs计算量、批次大小（batch size）及优化器（optimizer）信息。模型的输出则为该窗口内每个算子的预测执行时间及其后跟随的框架开销。研究者遵循最佳实践，测试了多种顺序模型结构，包括LSTM、GRU、RNN、Transformer和图卷积网络（GCN）。通过综合比较准确率，最终选择了四层双向LSTM（Bi-LSTM）作为预测模型的基础结构，并采用平均绝对误差（MAE）作为损失函数。训练前，研究者还使用Z-score标准化器对特征进行了归一化处理，特别是对不同算子类型的超参数进行了独立归一化。
在线推理阶段，对于一个待预测的新模型，ETS首先进行极高效的剖析（profiling）（每个独特结构在三秒内即可完成）以获取其确定的算子执行顺序。然后，同样利用滑动窗口从该模型的执行追踪中切分出样本，输入训练好的LSTM模型进行推理。由于滑动窗口之间存在重叠区域，每个算子及其后续开销可能会被多个窗口的预测结果覆盖。ETS通过平均这些重叠区域的预测值来获得最终的预测执行追踪，并汇总得到整个模型的迭代时间。该预测时间最终被传递给集群调度器（cluster scheduler），以辅助其做出更优的调度决策。
三、主要实验结果
在涵盖未见过的模型（如GoogLeNet、Swin Transformer、VGG等）的HPO测试集和RandWire测试集上，ETS展现了卓越的预测性能。在整体评估中，ETS实现了5.9%的平均相对误差和58.6毫秒的均方根误差（RMSE），显著优于基于算子的方法Habitat和基于图的方法DNNPerf。在对RandWire数据集中按算子数量（0至1200个）分组的模型进行测试时，ETS的预测误差保持稳定，而DNNPerf在处理大模型时因图特征提取中的信息丢失导致精度下降，证明了ETS在处理不同规模模型时的良好泛化能力。预测误差的分解分析显示，ETS对构成迭代时间主体的重要算子（如Conv2D、BatchNorm2D）的执行时间和框架开销预测非常准确；对占比较小的算子和开销的预测误差较高，但对整体精度影响甚微。
在收敛速度方面，得益于滑动窗口对特征尺寸的规制，ETS使用的轻量级LSTM模型在总体数据集上仅需约11分钟即可完成训练，而基于图的DNNPerf则需要超过100分钟。超参数影响分析证实，滑动窗口的宽度设为12、步长设为4是平衡精度与推理延迟的最佳选择。使用带重叠的滑动窗口并通过平均策略能进一步提升预测的稳定性与精度。
四、研究结论与应用价值
本研究的结论是，ETS作为一种基于执行追踪滑动窗口的迭代时间预测方法，成功克服了现有方法无法兼顾深度学习框架开销与模型规模差异的难题。其科学价值在于揭示并利用了深度学习模型运行时的顺序执行特性，提出了一种新颖的序列特征提取方法，为性能建模领域提供了新思路。应用价值则体现在，将ETS集成到集群调度系统中能够带来显著的调度质量提升。通过在阿里PAI和微软Philly两个大规模GPU集群的作业追踪数据上，模拟Allox和Hydra两种调度策略的实验结果表明，相比于其他基线预测方法，ETS能将平均作业完成时间（JCT）降低多达22%，将截止日期延迟（Tardiness）的增幅控制在仅4%，其调度效果非常接近使用真实迭代时间（即预测误差为0%的理想情况）的水平，有效提升了云数据中心的资源利用率和作业执行效率。
五、研究亮点
本研究的亮点可归纳为三点。第一，创新性地提出了“执行追踪滑动窗口”方法，首次从模型运行时执行追踪中提取包含算子顺序和框架开销的序列化特征，解决了预测中框架开销被忽略的核心痛点。第二，该方法巧妙地利用固定大小的滑动窗口处理任意规模的模型，将特征尺寸标准化，从而极大地加快了预测模型的收敛速度，同时滑动窗口的重叠与平均机制还增强了预测鲁棒性和准确性。第三，研究构建了包含14000余个多样化模型配置的大规模数据集，并通过严谨的实验在各类模型结构上验证了方法的有效性和泛化能力，展示了其在真实集群调度场景中的巨大潜力和实用价值。此外，ETS的轻量化设计使其易于通过迁移学习等技术适配到新的硬件执行环境中。