mGPU 技术揭秘： mGPU 算力和显存隔离底层方案

AI 大模型的涌现，不仅推动了 AI 技术的边界，也对计算资源提出了前所未有的挑战。大模型的多样性要求计算资源具备高度灵活性，同时也促使企业追求更高的成本效益。GPU 作为高性能计算的核心引擎，结合 GPU 共享技术可以显著提高计算资源的利用率和灵活性，为 AI 技术及其他高性能计算领域的进一步突破提供强有力的支持。

是火山引擎基于内核虚拟化隔离 GPU 并结合自研调度框架提供的容器共享 GPU 方案。在保证性能和故障隔离的前提下，mGPU 支持多个容器共享一张 GPU 显卡，支持算力与显存的灵活调度和严格隔离，降低 GPU 使用成本，为充分利用 GPU 硬件资源进行训练和推理提供有效保障。

mGPU 是针对容器场景下 NVIDIA GPU 资源 QoS 的一个解决方案，能够为 NVIDIA GPU 提供显存与算力资源的切分服务。在容器场景下，mGPU 对 NVIDIA GPU 资源的切分是通过设备驱动劫持来实现，包括了用户态和内核态的组件。对上层软件而言，业务不重编、CUDA 库不替换、业务无感。

mGPU 底层驱动部分的整体架构如下图所示：

mGPU Kernel Module：mGPU 内核模块，用于劫持容器对 NVIDIA GPU 驱动程序的调用 （open / close / mmap / ioctl  / poll），对容器实例进行显存和算力控制。

mGPU Container Runtime Hook：用于劫持 nvidia-container-runtime 对 nvidia-container-runtime-hook / nvidia-container-toolkit 的 prestart hook 调用，从环境变量中解析 mGPU 容器配置相关信息，通过 RPC 向 mGPU Daemon 发起 mGPU 容器创建请求。

mGPU Daemon：作为 RPC server，接受 mGPU Container Runtime Hook 发起的容器创建请求，通过 mGPU 内核模块提供的 ioctl 接口对 mGPU 容器进行注册。

◇ 基本概念

首先我们看一下 GPU 负载是如何在 NVIDIA GPU 硬件上工作起来的。

• GPU task 借助 GPU kernel driver 在调度系统中完成初始化操作。此过程会创建一个 pushbuffer（本质上是一个 queue），接着以 channel 的形式对任务进行封装。所有的 channel 都会被整合进 Time Slice Group（TSG）里，随后将这个 TSG 插入到 GPU 的某个 runlist 中，如此一来，该 TSG 就能够被调度了。最后，会把用户可访问的、指向 push buffer 的指针传递至用户空间，供用户态驱动使用。

• 当 task 需要 GPU 工作时，用户态驱动会把相关 command 入队到 stream 里，而这些 stream 会被映射到 push buffer。commands 可能包括 kernel launch、copy commands 以及各种其他操作。如果没有明确使用 stream，则会构建一个 default stream。由于 channel push buffers 被映射到用户可写内存中，所以 command 入队操作不需要经过系统调用，非常高效。

• commands 在 push buffer 中会一直累积，直到 GPU hardware scheduler（Host Interface）选择一个 channel 进行调度。scheduler 如何选择 channel 呢？会分两个级别进行：GPU 在任务的时间片组（TSG）之间进行全局轮询，并且会循环扫描当前活动的 TSG 中的所有 channel 以获取命令。

• 一旦 Host Interface 选定了一个有待处理命令的 channel，它就会利用 PBDMA 单元将 commands 从 CPU 拉到 GPU 里，这些 commands 会被解析并传递给某个 engine 进行工作。

概括来说，runlist 中有多个 entry，每个 entry 是一个 TSG 或一个 channel， 一个 TSG 由至少一个 channel 组成，每个 TSG 都具有不同的时间片，以便在不同 GPU task 之间对 GPU 进行时间片共享。每当一个 TSG 内的所有工作都已完成，或者由于 TSG 的 timeslice 到期需要被抢占时，当前正在运行的 TSG 就会进行上下文切换。因此，Host Interface 将从 runlist 选择下一个 TSG，开始调度与下一个 TSG 相关的工作负载。

◇ mGPU 实现

基于硬件时间片的算力调度器：通过拦截原生的用户态驱动对 GPU 硬件时间片设置的 ioctl ，调整拦截到的参数后再转发给原生的内核态驱动， 从而实现硬件时间片按比例缩放调整；

基于软件时间片的算力调度器：mGPU 为每张显卡创建一个内核线程，根据该显卡上配置的各个容器的算力权重，对该卡上各个容器对应 channel 进行动态 enable / disable  切换， 从而实现精准的时间片算力 QoS。

通过分析可知， CUDA 显存管理相关的 API，会在 CUDA 用户态驱动 （UMD）统一汇聚为对 nvidiactl 字符设备的 ioctl 操作。利用这一特点， mGPU 为容器创建了 GPU 虚拟卡并挂载到容器内部，从而可以在 mGPU 内核模块里对显存分配、释放、查询接口等实现劫持、统计、并转发。

• 当内核模块拦截到显存分配请求时，判断请求分配的 size 是否超过该容器对应的该卡的限额，如果超过则返回 OOM，反之则记录进程分配显存的相关信息，并把请求转发给 NVIDIA kernel driver；

• 当内核模块拦截到显存释放请求后，先查询分配记录来进行释放，接着再转发至 NVIDIA kernel driver；

• 当内核模块拦截到显存查询请求后，会返回容器隔离限制之下的显存相关信息。

这里我们以 V100/32G 显存的服务器做 ResNet-50 推理为例，测试 mGPU 使用前后的性能差异。如下图所示，使用了 mGPU 功能之后， GPU 负载打满时性能几乎无损：

目前，生成式 AI 正在推动数据中心对高性能 AI 芯片的需求增长。根据 Gartner 的最新预测，2024 年，全球 AI 芯片收入将增长 33％。而借助共享 GPU 技术的弹性能力，企业能把 GPU 资源用到极致，为 AI 创新提供强大、稳定、极具性价比的算力支撑。

火山引擎 mGPU 是研发团队基于字节跳动实践和客户服务经验深度自研的关键云原生技术之一，在客户生产环境中已被验证具有强大的 GPU 资源强隔离能力、GPU 资源细粒度感知和调度能力和性能近乎无损的使用体验。经落地验证，mGPU 可以帮助客户实现资源利用率提升超过 50%。

通过 mGPU，我们也希望能帮助更多客户实现基础设施与云原生的深度融合，构建起智能、高效、协同的异构计算生态，打造 AI 时代科技进步的强大引擎。

欢迎感兴趣的用户扫码咨询、使用！