在利用 MPI（Message Passing Interface）进行多进程并行计算或分布式深度学习训练时，为了提高 GPU 利用率，我们常常会开启 NVIDIA MPS（Multi-Process Service）。MPS 的初衷是允许多个不相关的 CPU 进程（即 MPI Ranks）并发地向同一个 GPU 发送任务，从而利用空闲的流处理器（SM）并减少上下文切换开销。
然而，许多开发者在启用 MPS 后，会遇到部分 MPI 进程在初始化或运行中途突发 CUDA_ERROR_OUT_OF_MEMORY (CUDA error 2) 的诡异现象。而在不开启 MPS（直接多进程排队或时间片轮转）的情况下，程序反而能正常运行。
本文将深入探讨这一现象背后的底层硬件和软件协同机制，并提供针对性的解决策略。
一、 核心原因剖析
开启 MPS 后，MPI 进程发生 OOM 的本质原因，在于 “显存物理共享下的无序竞争” 与 “多路进程上下文累积效应”。
1. 默认状态下的“公地悲剧”（无显存硬隔离）
在默认情况下，MPS 启动后，所有的客户端进程（MPI Ranks）在逻辑上共享同一个 GPU 虚拟地址空间。
- 非 MPS 模式下：如果多个进程强行挤在单张卡上，由于 CUDA 驱动的时间片轮转机制和虚拟内存保护，进程之间的内存分配是相对孤立的。
- MPS 模式下：所有 MPI 进程共享同一个 MPS Server 维护的物理 GPU 上下文。默认情况下，MPS 并不会限制单个 Client 进程可以使用的最大显存量。这意味着，如果 Rank 0 进程由于代码中存在动态分配（如 PyTorch 的 Caching Allocator 机制），在某一时刻瞬间申请了大量显存，就会直接挤占其他 Rank 的生存空间。后启动或执行较慢的 Rank 在申请显存时，就会直接面临
CUDA_ERROR_OUT_OF_MEMORY。
2. 无法忽视的客户端上下文（Context）与 Runtime 堆叠开销
虽然 MPS 大幅降低了多进程共享 GPU 时的上下文开销，但并不意味着开销降为零。
每个接入 MPS Server 的 MPI 进程（Client），在 CUDA 运行时初始化时，依然需要占用一定的固定显存（用于存储客户端特定的元数据、页表映射、cuBLAS/cuDNN 等库的初始化句柄）。
- 单个进程的初始化开销可能只有 100MB ~ 200MB。
- 如果在一张 GPU 上强行挂载了数十个 MPI 进程（例如，在单台 8 卡机器上，每张卡分配 16 个 MPI 进程），仅这些进程的初始化上下文和底层库句柄，就会吃掉数 GB 甚至十几 GB 的显存。这部分显存属于“硬性占用”，直接压缩了用户实际模型和数据的可用显存空间。
3. 深度学习框架的“抢占式”缓存分配机制（Caching Allocator）
如果你的 MPI 进程运行的是 PyTorch、TensorFlow 等深度学习任务，框架自身的内存管理器通常是贪婪的。
以 PyTorch 为例，它的 Caching Allocator 为了避免频繁调用 cudaMalloc 带来性能抖动，会倾向于一次性向系统申请一大块显存（Block），然后自己在用户态进行二次分配。
- 在 MPI 多进程环境下，多个 Rank 几乎在同一时间启动并执行
model.to(device)或前向传播。 - 这会导致多个进程并发向 MPS Server 申请大块显存。由于没有协同和隔离机制，速度稍慢的进程在框架层执行
cudaMalloc时就会被直接拒绝，抛出 OOM 异常。
二、 解决方案与最佳实践
解决该问题的核心思路是：变“无序共享”为“有序隔离”，同时抑制框架的内存吞噬行为。
1. 实施硬性显存隔离（推荐首选）
从 CUDA 11.x 及更高版本开始，MPS 引入了对单个客户端进程的显存配额限制（Memory Provisioning）。这是解决该问题最优雅、最彻底的方法。
你可以在启动 MPS Daemon 时，通过 nvidia-cuda-mps-control 动态限制每个 MPI 进程能够申请的最大物理显存。
操作示例（限制每个进程最多使用 4GB 显存）：
# 1. 启动 MPS Control 守护进程
nvidia-cuda-mps-control -d
# 2. 设置默认的客户端显存限制（支持百分比或具体字节数）
# 方式 A：限制每个 Client 进程最多只能使用当前 GPU 总显存的 20%
echo "set_default_device_mem_limit 0 20" | nvidia-cuda-mps-control
# 方式 B：限制每个 Client 进程最多只能使用 4096MB (4GB) 显存
echo "set_default_device_mem_limit 0 4294967296" | nvidia-cuda-mps-control
注：这里的 0 代表 GPU 的索引（Device ID），如果有多个 GPU，需要针对每个 Device 独立设置，或使用循环进行配置。
通过这种方式，一旦某个 MPI 进程尝试申请超过配额的显存，它自己会在用户态收到内存受限信号，而不会去“侵占”其他 MPI 进程的显存空间，从而保障了整体系统的健壮性。
2. 限制主动线程百分比（Compute Provisioning）
除了显存限制，还可以配合限制每个进程能够使用的最大 SM（流处理器）比例。这可以间接平抑瞬时的计算和内存申请峰值。
# 限制每个进程最多只能使用 25% 的 GPU 计算资源
export CUDA_MPS_ACTIVE_THREAD_PERCENTAGE=25
将此环境变量在 MPI 启动脚本中导出，确保每个 Rank 在初始化 CUDA 时，其物理硬件占有率被严格锁死在指定区间。
3. 调整深度学习框架的内存分配策略
如果你使用的是 PyTorch，必须限制其“贪婪”的显存碎纸机行为：
设置 PyTorch 显存碎片整理与释放阈值：
在运行 MPI 脚本前，设置环境变量：export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb:128"这会阻止 PyTorch 的 Allocator 分配过大的连续显存块，从而给其他 MPI 进程留出申请空间。
主动释放缓存：
在代码的循环迭代或评估阶段，适时调用torch.cuda.empty_cache()。虽然这会带来微小的性能损失，但在多进程高压环境下，能极大地降低 OOM 的概率。如果使用 TensorFlow，必须开启显存自适应增长：
import tensorflow as tf gpus = tf.config.experimental.list_physical_devices('GPU') for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True)
4. 控制合理的 MPI Rank 密度（Oversubscription Ratio）
MPS 虽然强大，但并非万能。单个物理 GPU 上挂载的 MPI 进程数存在一个临界点。
- 经验法则：单张 GPU 上运行的 MPI 进程（Ranks）数量建议控制在 2 到 8 个之间。
- 如果超过 16 个进程，即使每个进程只占用极少的业务显存，CUDA Runtime 本身的上下文累积开销也会将显存蚕食殆尽。如果计算任务确实需要极高的并发度，应考虑增加物理 GPU 数量，或者采用进程池模式，降低单卡并发密度。
三、 排查与诊断工作流
当你的 MPI+MPS 系统再次抛出 OOM 时，请遵循以下步骤进行快速诊断：
[发生 OOM 报错]
│
▼
1. 检查物理显存是否真的耗尽？
👉 运行 `nvidia-smi` 观察在报错瞬间，GPU 整体显存使用率。
├── 若未耗尽 (有空余) ──> 可能是单个进程触发了硬性配额限制 (检查 `set_default_device_mem_limit`)。
└── 若完全耗尽 ───────> 转到步骤 2。
│
▼
2. 确认是否存在进程“显存越界抢占”？
👉 检查各个 MPI Rank 内部处理的数据尺寸是否不均（Data Imbalance）。
👉 排查是否未设置内存限制。
│
▼
3. 诊断上下文开销：
👉 计算：(MPI 进程数) * 150MB。
👉 对比该数值占总显存的比例。如果比例过高，说明进程密度过大，必须削减单卡 Rank 数量。
通过合理配置 set_default_device_mem_limit 实施物理显存硬隔离，结合对框架层 Caching Allocator 的参数微调，可以彻底解决 MPS 环境下 MPI 进程的 OOM 隐患，在保障计算稳定性的同时，榨干 GPU 的最后一滴算力。
