MLKAPS：基于机器学习的高性能计算内核自动调优框架

1. 项目概述与核心价值在搞高性能计算HPC的朋友估计都遇到过同一个头疼的问题同一个计算内核换了个输入矩阵大小或者换台机器跑性能就可能天差地别。为了榨干硬件的每一分算力我们往往需要为不同的输入场景和硬件环境手动准备一大堆“优化配方”——也就是内核的设计参数配置。这个过程不仅极度依赖专家经验耗时耗力而且人工探索总有盲区很难保证在所有情况下都达到最优。今天要聊的MLKAPS就是一个试图用自动化和机器学习把我们从这种繁琐且不完美的调优工作中解放出来的框架。简单来说MLKAPS是一个面向HPC内核的自动调优工具。它的目标很直接你给它一个黑盒内核比如一个矩阵分解函数以及这个内核有哪些输入参数比如矩阵的行数n、列数m和可调的设计参数比如线程数T、分块大小block_size它就能自动帮你找到一个“决策树”。这个决策树在程序运行时能根据实际的输入参数瞬间选出当前情况下性能最好的那套设计参数配置。它的核心价值在于将原本需要海量手动测试和专家经验的调优过程转化为一个数据驱动的、可自动扩展的优化流程并且能在庞大到人力无法穷举的参数空间里找到连专家都可能遗漏的“性能盲点”。2. MLKAPS框架设计思路拆解2.1 问题定义输入参数 vs. 设计参数要理解MLKAPS在做什么首先得厘清两个核心概念输入参数和设计参数。输入参数这是用户调用内核时给定的或者由运行环境决定的、不可更改的条件。比如你要分解一个矩阵这个矩阵的行数n和列数m就是输入参数再比如当前系统可用的内存大小、CPU核心数也属于环境决定的输入参数。它们是问题的“已知条件”。设计参数这是内核内部为了实现不同算法变体、适应不同硬件特性而暴露出来的可调节“旋钮”。比如一个OpenMP并行循环用多少线程num_threads、矩阵计算时如何分块tile_size、选择哪种算法分支等。设计参数的不同组合构成了一个庞大的配置空间。调优的本质就是为每一组特定的输入参数在这个配置空间里找到一个能让目标通常是执行时间最短最优的配置。传统的、高度优化的库如Intel MKL内部就维护着这样的映射关系通常以硬编码的决策树或查找表形式存在。但构建和维护这套映射极其困难尤其是当输入空间和设计空间都很大时会遭遇“维度灾难”——可能的组合数量呈指数级增长穷举搜索完全不现实。2.2 核心思路从黑盒采样到决策树生成MLKAPS的整个流水线可以概括为四个阶段其核心思想是用相对廉价的机器学习模型预测替代昂贵的真实内核执行评估从而高效探索巨大的参数空间。采样与建模既然无法测试所有配置就需要智能地选择一部分配置进行实际运行采样用这些数据训练一个代理模型。这个模型MLKAPS主要用梯度提升决策树GBDT学习从“输入参数设计参数”到“性能指标如运行时间”的映射关系。采样的策略直接决定了模型的质量也是MLKAPS的创新重点。优化有了一个可以快速预测性能的代理模型后我们就能以很低的成本在模型上为大量不同的输入点在输入空间上布的一个网格运行优化算法如遗传算法为每个输入点找到一个“模型认为”最优的设计配置。决策树生成上一步我们得到了一个映射表(输入点1 - 最优配置A), (输入点2 - 最优配置B)...。最后一步就是基于这个映射表为每一个设计参数训练一个回归或分类决策树。这个决策树的输入是实际的输入参数输出就是最优的设计参数值。最终这个决策树可以编译成轻量级的C代码嵌入到目标内核中实现运行时的快速决策。整个流程的关键在于昂贵的真实内核执行只发生在最初的“采样”阶段而且通过智能采样策略可以用尽可能少的采样次数获得一个对优化区域预测足够准确的代理模型。后续的优化和决策树生成都是在模型上完成的成本极低。2.3 为何选择决策树作为最终输出你可能会有疑问既然已经有了GBDT代理模型为什么不直接用这个模型做预测原因在于部署开销和可解释性。低运行时开销一个浅层的决策树进行前向推断只需要几次整数比较和跳转速度极快几乎不增加运行时负担。这对于HPC内核来说至关重要调优决策本身消耗的时间必须远小于它带来的性能收益。可嵌入性决策树可以很容易地转换成if-else语句或switch-case结构的纯C代码无需引入复杂的机器学习库依赖便于集成到现有的高性能库中。可解释性与深度神经网络等“黑箱”模型相比决策树的决策路径清晰可见。库的维护者可以查看和理解自动生成的调优逻辑这增加了信任度也便于后续的调试和融合专家知识。3. 自适应采样策略如何聪明地“问问题”采样是MLKAPS流水线的基石也是最体现其智能的地方。目标是用有限的“提问次数”采样预算尽可能问出对找到最优配置最有价值的问题。MLKAPS实现并对比了多种策略其核心都是在探索了解未知区域和利用聚焦已知优势区域之间做权衡。3.1 基础策略空间填充采样这是一种偏向“探索”的策略目的是让采样点均匀覆盖整个参数空间确保代理模型有一个全局的、基础的了解。随机采样最简单但容易留下空白区域或产生聚类。拉丁超立方采样这是MLKAPS中使用的默认引导方法。它确保每个参数的每个区间都被采样一次能比纯随机采样更好地覆盖多维空间为后续建模提供一个质量更高的初始数据集。注意单纯的空间填充采样在优化场景下效率不高因为它平等对待所有区域包括那些性能很差的“荒原”。我们需要把宝贵的采样预算更多地投入到可能产出“宝藏”的区域。3.2 基于方差的采样Hierarchical Variance SamplingHVS策略试图更智能地分配采样预算。它首先用LHS进行初始采样然后根据现有数据用一个决策树对参数空间进行划分。计算每个划分区域的“体积”和其内部样本性能的“方差”。下一轮采样时倾向于在体积 × 方差乘积大的区域投放更多样本。其逻辑是方差大的区域说明模型对这个区域的预测很不确定需要更多数据来“学习”同时区域体积大代表这是一个广泛的未充分探索的区域。这种方法能有效刻画复杂的目标函数曲面。实操心得在HPC内核调优中某些极差的配置例如线程数设置严重不合理导致大量冲突会产生异常长的运行时间这会在目标函数中制造出极高的方差区域。原始的HVS会浪费大量采样预算去刻画这些“性能悬崖”的陡峭下坡但这对于寻找最优配置并无帮助。MLKAPS在实践中为性能指标设置了一个上限截断将异常值限制在一个合理范围内有效防止了HVS在坏上的过度采样这是一个非常实用的工程改进。3.3 优化驱动的采样GA-Adaptive这是MLKAPS论文中提出的新策略其思想更加激进代理模型不需要学好整个空间它只需要在包含优秀配置的区域学得极其精准即可。GA-Adaptive直接模拟了MLKAPS最终的优化阶段形成了一个“采样-建模-优化”的微循环。它的算法流程可以概括为初始化使用LHS采集一批初始样本建立初始代理模型。迭代循环在每一轮中 a. 用当前所有样本训练/更新GBDT代理模型。 b.利用随机选取一部分输入点使用遗传算法在代理模型上为这些输入点寻找“模型认为”的最优配置然后将这些配置加入真实采样队列。 c.探索使用一种子采样器如HVSr在现有样本分布的基础上补充一些侧重于探索的样本点。 d. 执行真实采样将新数据加入样本集。动态平衡通过一个线性变化的参数ε控制“利用”和“探索”的比例。开始时ε较小侧重探索快速构建全局模型随着采样进行ε逐渐增大越来越侧重对优势区域的利用和精细刻画。为什么有效自我修正如果模型对某个配置过于乐观GA会选中它真实采样后会得到一个很差的性能数据这个数据会修正模型避免后续优化再掉入这个陷阱。局部精修在表现好的配置附近反复采样可以降低该区域预测的噪声帮助模型区分性能相近的候选配置这对于最终找到精确的最优解至关重要。从论文中的实验结果图可以看到在追求全局预测精度时HVS表现最好但在追求对最优配置的局部预测精度上GA-Adaptive策略显著优于其他方法。而这正是自动调优最关心的指标。4. 从模型到部署优化与决策树生成实操4.1 基于代理模型的优化当采样完成我们获得了一个相对可靠的代理模型后就进入了“纸上谈兵”的优化阶段。这一步完全在模型上进行成本极低。构建优化网格在输入参数空间例如n和m都在1000到5000之间上创建一个均匀的网格。网格密度是一个权衡太粗可能错过细节太细则优化点太多计算量增大。论文中发现在16x16的网格上已经能取得很好效果继续增加到24x24收益不明显。并行优化为网格上的每一个输入点独立启动一个优化器实例MLKAPS使用pymoo库的NSGA-II算法。每个优化器的任务是在设计参数空间中寻找使代理模型预测值如时间最小的配置。由于是在模型上运行遗传算法可以轻松进行成千上万次评估快速收敛。输出映射表优化完成后我们得到了一张表记录了每一个网格输入点所对应的、模型推荐的最优设计参数配置。关键点由于采样阶段并不受这个网格限制模型在优化时能够利用网格点周围邻域的信息进行推断。这意味着最终为某个网格点找到的优化配置并不是基于该单一点的信息而是基于其周围一片区域的采样数据学习得到的这增强了方案的鲁棒性和泛化能力避免了过拟合。4.2 决策树训练与代码生成现在我们有了一系列(输入 最优设计参数)的配对数据。对于每一个设计参数比如线程数T、分块大小B我们单独训练一个决策树模型。模型选择对于连续值参数如分块大小使用回归树对于类别参数如算法变体选择使用分类树。MLKAPS使用scikit-learn库来实现。树深度控制这是性能与开销权衡的关键。一棵很深的树可以做出非常精细、局部的决策可能更接近最优但运行时需要经过更多判断分支开销大。一棵浅树则相反它可能用一个配置覆盖一大片输入区域虽然牺牲了一点最优性但决策速度极快。需要根据内核本身的执行时间来决定树的复杂度——如果内核本身运行需要几秒钟那么一个稍复杂的决策树开销可以忽略如果内核是毫秒级的微内核则必须使用极浅的树。代码生成训练好的决策树会被序列化保存同时MLKAPS会将其转换为等价的C代码函数。这个函数就是最终交付物可以被直接编译进目标HPC库中。生成的代码通常是纯逻辑判断没有外部依赖确保了部署的便捷性和高性能。5. 实战评估以Intel MKL dgetrf为例理论再好还得看疗效。论文选择Intel MKL中的dgetrfLU分解和dgeqrfQR分解内核作为测试对象这两个内核本身就经过了Intel专家团队的深度手工优化挑战性极大。5.1 实验设置与挑战规模硬件使用了两种差异显著的Intel架构——Knights Mill和Sapphire Rapids涵盖了众核和现代服务器CPU的不同特性。内核与参数以dgetrf为例输入参数是矩阵维度n和m范围1000-5000。设计参数有8个包括线程数、分块方式等。这构成了一个巨大的联合搜索空间。搜索空间论文中提到dgetrf的设计空间有约4.6×10^13种可能配置而感兴趣的输入区域有约1.6×10^7种可能。这是一个典型的“维度灾难”场景穷举完全不可行。5.2 MLKAPS性能提升结果MLKAPS仅采样了约整个设计空间的一百亿分之一的配置就取得了显著效果对专家调优的超越在Sapphire Rapids平台上MLKAPS为85%的输入改进了Intel MKL原有手工调优的性能取得的几何平均加速比达到1.30倍即平均加速30%。这意味着即使在经过高度优化的工业级库中自动调优方法依然能发现人工忽略的“盲点”。自适应能力在两个不同的硬件架构上MLKAPS找到的最优配置和带来的加速效果是不同的这证明了其能自动适应不同的硬件环境。采样效率GA-Adaptive采样策略在“对最优配置的预测精度”这一关键指标上显著优于传统的LHS和HVS方法验证了其优化驱动采样的有效性。5.3 与现有自动调优工具对比论文还将MLKAPS与两个知名的自动调优框架进行了对比对比Optuna在dgeqrf内核上MLKAPS获得的平均性能比Optuna高36%。Optuna使用了TPE和CMA-ES等高级优化算法但MLKAPS通过其“采样-建模-优化-决策树”的完整流程特别是在采样阶段针对目标进行的优化整体上更胜一筹。对比GPTuneGPTune是另一个基于贝叶斯优化和高斯过程模型的先进调优器。MLKAPS在GPTune自己的ScaLAPACK用例和Intel MKLdgetrf用例上都表现出了更好的性能。同时MLKAPS的流水线在处理大量样本时开销相对更低展现了更好的可扩展性。6. 常见问题、挑战与避坑指南在实际应用MLKAPS或类似自动调优框架时会遇到一些典型问题和挑战。6.1 采样预算与收敛判断问题应该采样多少配置才算够怎么知道模型已经足够好可以停止采样进入优化阶段应对这是一个经验与指标结合过程。可以监控代理模型在预留的验证集上的预测误差如MAE是否收敛。更重要的指标是用当前模型找到的“最优配置”在真实评估中的性能提升是否趋于稳定。实操中建议设置一个总时间预算然后观察随着采样进行最佳性能的提升曲线。当连续多轮采样带来的最佳性能提升小于一个阈值如1%时可以考虑停止。6.2 目标函数的噪声与稳定性问题HPC内核的执行时间并非确定值会受到系统负载、缓存状态等因素干扰存在测量噪声。噪声会干扰模型学习特别是方差大的区域。应对重复测量对同一个配置进行多次运行如5-10次取中位数或最小值作为该配置的性能值。中位数对异常值更鲁棒。设置性能上限如前所述对异常长的运行时间进行截断避免它们过度影响采样策略和模型训练。使用鲁棒的模型GBDT本身对噪声有一定的容忍度。在模型评估时关注其在最优配置区域的预测一致性而非全局的绝对误差。6.3 决策树的深度与过拟合问题决策树深度设置多少合适太深容易过拟合训练网格太浅则区分能力不足。应对使用验证集将优化网格的数据按比例分成训练集和验证集用验证集上的误差来选择树深度。正则化利用scikit-learn决策树提供的参数如min_samples_split节点分裂所需最小样本数、min_samples_leaf叶节点最小样本数来限制树的生长防止过拟合。业务逻辑约束有时设计参数有内在逻辑。例如分块大小可能是缓存行大小的整数倍。在决策树生成后可以后处理其输出使其符合这些约束。6.4 高维类别参数的处理问题如果设计参数中包含类别变量例如算法变体A/B/C且类别很多如何有效处理应对GBDT可以原生处理类别特征。在采样时需要确保对每个类别都有足够的覆盖。对于决策树生成使用分类树。一个技巧是对于有序的类别如性能由低到高的几种算法可以尝试将其编码为有序整数有时能帮助模型学习。6.5 从研究到生产的落地挑战论文中的实验是在受控环境下针对单个内核进行的。真实生产环境可能是多个内核链式调用或者需要同时优化多个目标如时间、能耗、精度。思路多目标优化MLKAPS的优化阶段可以使用NSGA-II这类多目标遗传算法寻找帕累托最优解集。最终需要根据业务需求选择一个折衷点或生成一个更复杂的决策机制。工作流集成将MLKAPS生成的决策树集成到大型HPC库中需要设计清晰的接口。通常可以封装成一个独立的“调优模块”在库初始化时加载决策树函数。持续学习硬件和软件栈会更新。可以设计一个轻量级的在线评估机制当发现当前决策树在某个新输入上性能不佳时触发小范围的重新采样和模型更新实现渐进式优化。MLKAPS框架展示了一条将机器学习系统化应用于HPC性能调优的有效路径。它通过将昂贵的全局搜索问题分解为“智能采样构建代理模型”和“基于模型的廉价优化”两个阶段并创新性地提出了面向优化目标的GA-Adaptive采样策略最终输出易于部署的决策树在实用性和性能上都取得了很好的平衡。对于库的开发者而言它提供了一个强大的工具来构建或增强其自适应调优能力对于用户而言它意味着无需手动干预就能让计算任务更充分地利用硬件资源。尽管在应对极端高维问题、多目标权衡和在线适应方面仍有探索空间但MLKAPS无疑为自动化、智能化的HPC性能优化提供了一个坚实且富有前景的范本。