vLFS:基于学习的AI芯片虚拟化调度 论文总结

发布于:2025-03-19 ⋅ 阅读:(11) ⋅ 点赞:(0)

摘要

AI芯片的虚拟化由于在计算资源上是独立的,导致了对同一调度器时间片的竞争。为此提出vLFS,一种基于深度强化学习的调度算法,通过结合任务负载特性和底层芯片的运行时状态以及主机与设备之间的协同调度,提高AI芯片的利用率。

介绍

由于现有的启发式算法通常是在调度完成某个实例的所有任务后,再继续为其他实例调度任务。当某个实例的任务执行的很快,导致该实例的计算资源会闲置下来,造成资源浪费。为例能够自动调整待分配任务的数量,提出了vLFS,可以根据工作负载自动调整分配给每个实例的任务数量。核心思想是将AI芯片中的调度器设计为一个强化学习RL代理,使其能够自主学习最优调度策略。此外还设计了一种主机与设备之间的协调调度方法,以确保程序顺序和执行顺序的一致性。

背景和动机

A. AI芯片利用率不足

现有的虚拟化技术是将所有虚拟实例使用同一个调度器来执行任务分配,并且以轮询方式访问这些实例,会造成调度时间片的隐性竞争,即使某个实例里有在等待执行的任务且资源为空闲,但是因为必须等待调度器的服务而无法执行。

B. 调度不及时

现有的调度策略要么分配相同时间片要么根据启发式算法调整时间片大小,但都不是最优的。

例如a,实例B的3个任务只能在实例A的任务调度后才能调度执行,导致实例B的前两个任务完成后有一段资源空闲(idle那一块),图b就是根据强化学习策略能够预测出要适当减少实例A的时间片,使调度器在实例B的任务即将完成时切换到调度实例B,从而及时分派新任务,避免实例B的芯片资源闲置。启发式算法难以准确估算并行任务的工作负载总执行时间。

C. 基于学习的调度

启发式调度方法难以适应底层设备和工作负载的变化,但是强化学习可以,能够自动分析工作负载的特征信息。vLFS使用网络模型预测最优任务调度策略,根据任务的实际执行结果动态调整网络模型的参数,从而提高芯片利用率。

vLFS的设计

摘要

vLFS是一种基于强化学习(RL)的在线虚拟化调度设计,旨在优化AI芯片的任务分配效率,其核心是将调度器表示为RL代理,通过主机与设备的协同机制实现动态调整。主机端部署RL模型于驱动API,捕获用户提交任务的多级特征(任务级、实例级、全局级)并收集运行时信息,经神经网络分析后预测动作概率分布,通过随机采样选择最优调度策略;设备端接收预测的任务数量,由驱动程序指导任务执行,并收集执行结果计算奖励函数,反馈至主机优化RL模型。这种闭环设计集成于AI软件栈(包括深度学习框架、加速库、运行时和驱动程序)的任务分发关键路径,确保任务有效检测与记录,从而实现兼容性与执行效率的双重优化,同时提升芯片资源利用率。

主机/设备协同调度

vLFS通过主机与设备的协同调度机制优化AI芯片虚拟化中的任务分配效率,其中主机端部署强化学习(RL)模型,分析任务的关键特征(如任务类型和参数),并将这些特征输入神经网络预测需执行的任务数量,生成调度动作信号发送至设备,而不直接拦截任务以避免增加处理延迟;设备端则负责实际的任务拦截与分派,根据主机信号动态调整任务数量和执行顺序,利用设备原有能力维护程序顺序与执行顺序的一致性,无需对主机进行大量代码修改即可实现高效调度;任务执行后,设备反馈调度结果至主机,主机据此从特征向量中移除已完成任务,确保两端数据同步更新,从而在减少整体延迟的同时提升资源利用率和系统性能。

状态空间与动作空间

状态空间

在每个时间步长t,vLFS会收集实例和任务的状态特征,组成一个观测向量。我们通过特征选择确定vLFS观测向量中包含的最优状态特征。表I列出了选定的六个特征。

动作空间

在每个时间步长t,给定状态s下,vLFS从所有可能的动作中选择一个动作。在AI芯片虚拟化中,一个可能的动作是确定从某个实例中需要分派的任务数量。

特征提取

由于实例和任务数量随环境变化而波动,vLFS通过设置任务视野(tv)和实例视野(iv)阈值来调节可见范围,确保在边界内的信息被详细分析,超出边界的则通过加权整合处理,从而解决状态信息向量无法适应任意规模的问题。

策略网络

vLFS利用策略梯度算法训练,通过基于训练中观察到的奖励对神经网络参数执行梯度下降来学习;其策略网络采用LeakyReLU激活函数,包含四层全连接层(神经元数分别为64、32、16、8),输出层使用sigmoid函数归一化动作,直接输出各可能动作的概率分布。

奖励函数

vLFS旨在预测实例资源闲置时机,并通过及时调度任务提升芯片利用率,其核心是通过最小化任务等待时间确保用户任务及时执行,防止资源浪费,因此将任务等待时间和芯片利用率纳入奖励函数关键因素。在每个时间步长t,vLFS监控实例芯片利用率Utilt并计算利用率奖励Ru:当Utilt超过60时,资源高效利用,代理获高奖励以维持或提升利用率;若低于60,资源未充分利用,则施加惩罚激励优化调度策略,从而提高整体芯片性能。

为了提升系统整体性能,vLFS计算所有实例平均利用率奖励Ru(记为Rū)作为整体芯片利用率指标。在每个时间步长t,系统记录实例任务等待时间并计算变异系数(cv):当cv低于0.01时,任务等待时间高度平衡,调度策略合理,代理获丰厚奖励以强化高效行为;cv在0.01到0.15之间时,等待时间相对平衡,属可接受范围,代理获适度奖励以维持或优化策略;若cv超0.15,等待时间极不平衡,调度失效导致资源浪费,则施加惩罚,降低该动作概率,激励代理探索更优调度策略。

当主机预测的调度动作传至设备执行时,模型预测的任务数量可能与实际执行数量不符,尤其在非法动作情况下,因此设备通过发送包含执行结果及实例信息的确认信号反馈给主机;主机接收后精确计算已完成任务数,与预测数比较得出偏差,该偏差作为惩罚项纳入奖励函数,以降低代理选择非法动作的概率,从而提升学习效率和模型预测准确性。

三者结合是

评估

评估方法

vLFS在Intel主机处理器上通过mlpack库嵌入RL模型,并以Cambricon AI芯片(通过PCIe连接)作为设备处理器进行性能评估,实验配置见下表,但因软件栈复杂性,当前实现暂不支持同步原语和异步内存分配,未来将优化并扩展至其他芯片。评估以Cambricon现有调度器为基线,vLFS作为附加组件仅修改驱动层代码,以禁用vLFS的性能为基准。工作负载基于Driver API构建,包括深度学习中常见的计算密集型(如矩阵乘加)和内存密集型(如memcopy)任务,分为稀疏实例(SparseInst)和密集实例(DenseInst),以模拟多样化场景并验证调度效果。

性能分析

本文通过芯片利用率、任务平均等待时间和平均总时间评估vLFS与启发式算法性能,测试显示vLFS在不同工作负载下有效提升效率。图5分析表明,在密集与稀疏任务实例竞争的多轮调度中,vLFS与启发式方法均保持密集任务实例高芯片利用率,但对稀疏任务实例,vLFS从一开始即调度任务,避免芯片闲置,优于启发式方法;一轮调度后,vLFS对密集任务维持高利用率,对稀疏任务在20微秒内完成调度,较启发式方法的50微秒节省一半以上时间,显著提高效率。

在密集与稀疏任务实例竞争时,启发式算法下密集任务的等待时间(Wt)和总时间(Tt)波动大且不稳定,稀疏任务因阻塞导致Wt激增及资源闲置;而vLFS算法下,密集与稀疏任务的Wt和Tt均较小且稳定,表明vLFS能更公平、高效地调度多实例,优化芯片资源利用。

开销分析

vLFS的开销测试表明,即使任务视野(tv)达200,开销仍可接受;在重负载下,其开销不超过内核平均等待时间的10%,可被等待时间掩盖,且主要源于特征预处理,模型推理和训练开销小,适合长期在线学习。