使用 NVIDIA MPS 提高共享 GPU 的性能

如果您在共享 Dataflow GPU 上运行多个 SDK 进程，可以通过启用 NVIDIA 多进程服务 (MPS) 来提高 GPU 效率和利用率。MPS 通过启用进程以共享 CUDA 上下文和调度资源，支持在 GPU 上进行并发处理。MPS 可以减少上下文切换开销、提高并行性并降低存储要求。

目标工作流是指在具有多个 vCPU 的工作器上运行的 Python 流水线。

MPS 是一种 NVIDIA 技术，用于实现 CUDA API，这是一种支持通用 GPU 计算的 NVIDIA 平台。如需了解详情，请参阅 NVIDIA 多进程服务用户指南。

优势

  • 可提高 GPU 流水线的并行处理和总体吞吐量，尤其是对于 GPU 资源使用量较低的工作负载。
  • 可提高 GPU 利用率，这可能会降低您的费用。

支持和限制

  • 只有使用单个 GPU 的 Dataflow 工作器支持 MPS。
  • 流水线无法使用限制并行性的流水线选项。
  • 避免超出可用的 GPU 内存，尤其是在涉及加载大型机器学习模型的应用场景中。平衡 vCPU 和 SDK 进程的数量与这些进程所需的可用 GPU 内存。
  • MPS 不会影响非 GPU 操作的并发性。
  • Dataflow Prime 不支持 MPS。

启用 MPS

使用 GPU 运行流水线时，通过执行以下操作启用 MPS：

  • 在流水线选项 --dataflow_service_options 中，将 use_nvidia_mps 附加到 worker_accelerator 参数。
  • 将 count 设置为 1。
  • 请勿使用流水线选项 --experiments=no_use_multiple_sdk_containers。

流水线选项 --dataflow_service_options 如下所示：

--dataflow_service_options="worker_accelerator=type:GPU_TYPE;count:1;install-nvidia-driver;use_nvidia_mps"

如果您使用 TensorFlow 并启用 MPS，请执行以下操作：

  1. 在 GPU 上启用动态内存分配。使用以下任一 TensorFlow 选项：
    • 通过调用 tf.config.experimental.set_memory_growth(gpu, True) 开启内存增长。
    • 将环境变量 TF_FORCE_GPU_ALLOW_GROWTH 设置为 true。
  2. 使用具有适当内存限制的逻辑设备。
  3. 为获得最佳性能，请尽可能使用软设备配置或手动配置来强制使用 GPU。

后续步骤