基于深度强化学习的NVIDIA MIG GPU多任务智能调度优化研究

为解决NVIDIA MIG（Multi-Instance GPU）环境下多任务调度效率低下的问题，研究人员开展了一项基于深度强化学习（DRL）的智能调度策略研究。通过设计状态-动作空间和奖励机制，该研究实现了动态配置和任务分配的协同优化，实验表明其调度方案相比现有方法（如FAR、MISO）显著降低了完成时间（makespan），为异构计算资源的高效利用提供了新思路。

　　

随着人工智能和高性能计算需求的爆发式增长，GPU已成为关键计算资源。NVIDIA推出的多实例GPU（Multi-Instance GPU, MIG）技术允许将物理GPU划分为多个独立实例，从而支持多任务并行执行。然而，MIG的复杂配置选项（如1-1-1-1-1-1-1、2-1-1-3等）和任务执行时间的异构性（如超线性、近线性或次线性扩展），使得如何动态调度任务并实时调整GPU配置以最小化总完成时间（makespan）成为一个NP难问题。传统启发式方法（如FAR、MISO）无法充分适应动态工作负载，而静态配置（如FixBest、NoMIG）则可能导致资源利用率低下。

为解决这一问题，研究团队提出了一种基于深度强化学习（DRL）的智能调度框架。该框架将MIG调度问题建模为马尔可夫决策过程（MDP），通过设计紧凑的状态表示（包括GPU配置和任务批处理信息）、动作空间（包含推进时间、重新配置GPU和执行任务）以及以负时间进展为奖励的机制，引导智能体学习最小化makespan的策略。研究使用近端策略优化（PPO）算法进行训练，并采用动作掩码（action mask）技术处理无效动作，显著提升了学习效率。

关键技术方法包括：

1. 1.

   状态表示：将GPU状态（当前配置和各切片剩余时间）和任务批状态（任务在各实例大小下的执行时间向量）编码为固定维度的输入，支持神经网络处理。

2. 2.

   动作设计：包含时间推进、GPU重新配置和任务执行三类动作，确保状态转移符合MIG约束。

3. 3.

   奖励函数：以负时间进展值作为奖励，鼓励智能体快速完成任务。

4. 4.

   训练环境：使用合成工作负载（如PoorScaling、GoodScaling、MixScaling等）和真实基准测试（Rodinia、Altis）生成多样化数据集，并通过离散化时间值（M=14级）降低状态空间复杂度。

5. 5.

   评估指标：采用相对最优下限的百分比（p_opt）衡量调度效果。

研究结果：

• •

  状态与动作空间的有效性：通过规范状态表示（如按任务时间向量排序）和动作掩码，减少了无效探索，加速了收敛（图10）。

• •

  时间离散化的优势：对比连续时间表示，离散化（M=14）在训练效率和最终性能上均更优（图11），最终p_opt可达12.65%（GoodScaling负载）。

• •

  批处理大小的影响：实验表明批大小N=14在收敛速度和性能间达到最佳平衡（图12），过小（N=7）或过大（N=28）均会降低效率。

• •

  与基线方法的对比：在多种合成和真实工作负载下，DRL调度器显著优于FAR（平均22.71%）、MISO（26.72%）、FixBest（75.37%）和NoMIG（101.52%）等基线（表4），尤其在混合扩展负载（MixScalingExtreme）中，p_opt降至19.56%。

• •

  在线调度验证：在线场景中，DRL智能体通过动态调整配置和任务顺序，实现了与离线接近的性能，证明了其在实际系统中的适用性。

研究结论与意义：

该工作首次将深度强化学习应用于NVIDIA MIG环境的多任务调度问题，提出了一种可处理高维状态和动作空间的通用框架。通过精心设计的状态表示、动作约束和奖励机制，智能体能够学习到接近最优的调度策略，在合成和真实工作负载上均表现出色。相比传统启发式方法，DRL调度器平均将makespan优化了20%以上，为解决GPU资源动态分配问题提供了新范式。未来工作可扩展至多GPU集群调度，并进一步探索迁移学习在跨工作负载泛化中的应用。论文发表于《Future Generation Computer Systems》，为异构计算系统的自动优化奠定了重要基础。