Python GPU 任务调度：CUDA 显存管理和多进程隔离

一、模型加载失败——"CUDA out of memory"的锅

多个 AI 服务部署在同一台 GPU 机器上，各自用各自的 Python 进程。看似井水不犯河水，结果第一个服务加载了 7B 模型（占 14GB 显存），第二个服务加载 1.5B 模型（占 3GB），第三个服务一加载就报 CUDA out of memory。

明明 GPU 有 24GB 显存，只是两个服务用了 17GB，为什么第三个加载 3GB 就报 OOM？

答案是 CUDA 的内存碎片化。PyTorch 的显存分配不是"先申请一大块，按需分发"，而是"用多少申请多少"。两个进程虽然总占用只有 17GB，但内存碎片加上 PyTorch 缓存机制，导致剩余空间不足以连续分配 3GB。

二、CUDA 显存管理架构

flowchart TB
    subgraph GPU["GPU 显存（24GB HBM）"]
        direction TB
        Frag1["进程 A: 模型参数\n10GB 连续"]
        Cache1["进程 A: PyTorch 缓存\n4GB 预留"]
        Frag2["进程 B: 模型参数\n3GB 连续"]
        Cache2["进程 B: PyTorch 缓存\n1GB 预留"]
        Free1["碎片空间\n2GB（已碎片化）"]
        Free2["可用空间\n4GB"]
    end

    subgraph Manager["显存管理器"]
        MCP["多进程隔离 (CUDA MPS)"]
        MCP --> Limit1["进程 A 显存限制: 14GB"]
        MCP --> Limit2["进程 B 显存限制: 5GB"]
        
        CacheClean["内存清理策略"]
        CacheClean --> Empty["torch.cuda.empty_cache()"]
        CacheClean --> IPC["CUDA IPC 共享"]
    end

    subgraph Schedule["任务调度"]
        Q["任务队列"] --> Alloc["显存检查"]
        Alloc -->|"够用"| Run["执行任务"]
        Alloc -->|"不够"| Wait["等待 + 清理"]
        Wait --> GC["触发 Python GC\n+ 清理 CUDA 缓存"]
        GC --> Alloc
    end

三、生产级显存管理代码

显存管理器

import torch
import gc
import time
import threading
from typing import Optional, Callable, Dict
from dataclasses import dataclass, field
from collections import deque
import logging

logger = logging.getLogger(__name__)

@dataclass
class GPUMemoryBudget:
    """GPU 显存预算——每个进程/任务的显存限额"""
    device: int = 0
    total_memory_mb: int = 0        # 为该任务预留的显存（MB）
    peak_memory_mb: int = 0         # 历史峰值
    current_allocated_mb: int = 0   # 当前分配

class GPUMemoryManager:
    """GPU 显存管理器——负责任务级别的显存隔离和调度"""
    
    def __init__(self, device_id: int = 0, reserved_memory_mb: int = 1024):
        """
        device_id: GPU 设备编号
        reserved_memory_mb: 为系统和 CUDA 上下文预留的显存
        """
        self.device_id = device_id
        self.reserved_mb = reserved_memory_mb
        
        # 检查 GPU 是否可用
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA 不可用，请检查驱动和 PyTorch 版本")
        
        self.total_memory_mb = torch.cuda.get_device_properties(
            device_id
        ).total_memory // (1024 * 1024)
        
        self.available_mb = self.total_memory_mb - reserved_memory_mb
        
        # 任务显存预算表
        self._budgets: Dict[str, GPUMemoryBudget] = {}
        self._lock = threading.Lock()
        
        logger.info(
            f"GPU {device_id}: 总显存 {self.total_memory_mb}MB, "
            f"可用 {self.available_mb}MB (预留 {reserved_memory_mb}MB)"
        )
    
    def register_task(
        self,
        task_id: str,
        memory_budget_mb: int,
    ) -> GPUMemoryBudget:
        """
        注册一个任务并分配显存预算
        如果剩余显存不足，抛出异常
        """
        with self._lock:
            # 检查总预算是否超标
            allocated_sum = sum(
                b.total_memory_mb for b in self._budgets.values()
            )
            if allocated_sum + memory_budget_mb > self.available_mb:
                raise MemoryError(
                    f"任务 {task_id} 需要 {memory_budget_mb}MB，"
                    f"但仅剩 {self.available_mb - allocated_sum}MB"
                )
            
            budget = GPUMemoryBudget(
                device=self.device_id,
                total_memory_mb=memory_budget_mb,
            )
            self._budgets[task_id] = budget
            logger.info(f"任务 {task_id}: 已分配 {memory_budget_mb}MB 预算")
            return budget
    
    def release_task(self, task_id: str):
        """释放任务的显存预算并清理 CUDA 缓存"""
        with self._lock:
            if task_id in self._budgets:
                del self._budgets[task_id]
        
        # 触发 Python GC 和 CUDA 缓存清理
        gc.collect()
        torch.cuda.empty_cache()
        logger.info(f"任务 {task_id}: 已释放显存预算")
    
    def get_current_usage(self) -> Dict[str, int]:
        """获取当前显存使用情况"""
        allocated = torch.cuda.memory_allocated(self.device_id) // (1024 * 1024)
        cached = torch.cuda.memory_reserved(self.device_id) // (1024 * 1024)
        free = self.total_memory_mb - allocated - cached
        
        return {
            "total_mb": self.total_memory_mb,
            "allocated_mb": allocated,
            "cached_mb": cached,
            "free_mb": max(0, free),
        }

    def safe_execute(
        self,
        task_id: str,
        func: Callable,
        *args,
        max_retries: int = 3,
        **kwargs,
    ):
        """
        在显存安全的环境中执行函数
        如果 OOM，自动清理缓存并重试
        """
        for attempt in range(max_retries):
            try:
                return func(*args, **kwargs)
            except torch.cuda.OutOfMemoryError as e:
                logger.warning(
                    f"任务 {task_id} OOM（第 {attempt+1}/{max_retries} 次尝试）"
                )
                
                # 清理策略：先清 Python GC，再清 CUDA 缓存
                gc.collect()
                torch.cuda.empty_cache()
                
                # 检查显存状态
                usage = self.get_current_usage()
                logger.warning(f"当前显存: {usage}")
                
                if attempt == max_retries - 1:
                    raise RuntimeError(
                        f"任务 {task_id} 重试 {max_retries} 次后仍然 OOM"
                    ) from e
                
                # 递增等待（给其他释放显存的机会）
                time.sleep(2 ** attempt)

多进程 GPU 隔离方案

import multiprocessing as mp
from contextlib import contextmanager

class GPUProcessPool:
    """GPU 多进程池——每个进程独立管理自己的显存"""
    
    def __init__(self, gpu_id: int, num_workers: int = 2):
        self.gpu_id = gpu_id
        self.num_workers = num_workers
        self._pool: Optional[mp.Pool] = None
        
        # 每个 worker 的显存限额
        total_mem = torch.cuda.get_device_properties(gpu_id).total_memory
        self.worker_memory_limit = (
            total_mem // num_workers * 80 // 100  # 每个 worker 用 80% 份额
        )
    
    @staticmethod
    def _worker_init(gpu_id: int, memory_limit: int, worker_id: int):
        """
        每个 worker 进程的初始化函数
        ——在 fork 之后、执行任务之前调用
        """
        import os
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
        
        # 设置 PyTorch 显存限制
        torch.cuda.set_per_process_memory_fraction(
            memory_limit / torch.cuda.get_device_properties(gpu_id).total_memory
        )
        
        # 设置 PyTorch 显存分配策略
        # expandable_segments=True 允许内存段动态扩展（减少碎片）
        os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
        
        logger.info(f"Worker {worker_id} 初始化完成, 显存限制: {memory_limit} 字节")
    
    def start(self):
        """启动多进程池"""
        self._pool = mp.Pool(
            processes=self.num_workers,
            initializer=self._worker_init,
            initargs=(self.gpu_id, self.worker_memory_limit, 0),
        )
    
    def submit(self, func, *args, **kwargs):
        """提交任务到进程池"""
        if self._pool is None:
            raise RuntimeError("请先调用 start()")
        return self._pool.apply_async(func, args, kwargs)

模型加载时的显存优化

def load_model_with_memory_budget(
    model_path: str,
    budget_mb: int,
    device: str = "cuda:0",
) -> torch.nn.Module:
    """
    在显存预算限制下加载模型
    ——如果超出预算，自动尝试量化
    """
    # 初步估算：FP16 模型大约每 1B 参数占用 2GB
    # 如果预算不够，尝试 INT8/INT4 量化
    
    file_size_mb = __import__("os").path.getsize(model_path) // (1024 * 1024)
    
    if file_size_mb > budget_mb * 1.5:
        # 模型文件超过预算的 1.5 倍，必须量化
        logger.info(f"模型大小 {file_size_mb}MB 超过预算 {budget_mb}MB，使用 INT4 量化")
        from transformers import BitsAndBytesConfig
        
        quant_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_use_double_quant=True,
        )
        
        model = load_model(model_path, device_map="auto", quantization_config=quant_config)
    else:
        # 预算充足，直接加载
        torch.cuda.set_per_process_memory_fraction(
            budget_mb / torch.cuda.get_device_properties(0).total_memory * (1024 * 1024)
        )
        model = load_model(model_path)
        model = model.to(device)
    
    return model

四、边界分析与 Trade-offs

多进程 vs MPS（Multi-Process Service）：

  • 多进程：简单易用，隔离性强，但有进程间通信开销
  • CUDA MPS：共享 GPU 上下文，减少碎片，但单进程崩溃可能影响其他进程
  • 建议：推理服务用 MPS，训练任务用多进程隔离

显存碎片化：PyTorch 的 expandable_segments 可以缓解碎片问题，但有轻微性能开销。对于需要频繁分配释放显存的场景（如动态 batch size），建议开启。

任务优先级：简单 FIFO 调度可能不够。建议参考 Kubernetes 的 QoS 模型（Guaranteed / Burstable / BestEffort）设计优先级调度。

GPU 共享的经济性：一张 A100（80GB）按 5 个进程平分，每个 16GB。如果某个进程实际只用 8GB，剩余的 8GB 浪费了。可以使用 MIG（Multi-Instance GPU）或 MPS 提高利用率。

五、总结

Python GPU 任务调度的核心挑战是显存碎片化和多进程隔离：

  1. 显存预算制：每个任务注册时声明显存需求，管理器全局调度
  2. 多进程隔离：每个进程独立管理显存，避免相互影响
  3. OOM 自愈：捕获 OOM → 清理缓存 → 递减重试
  4. 量化降级：模型超过显存预算时，自动尝试量化加载

一个实用的显存监控命令：nvidia-smi --query-gpu=memory.used,memory.free --format=csv -l 1，持续观察显存波动。

# 后端 # golang
Logo
https://edu.csdn.net/learn/39067/627173?utm_source=2019755004

汇聚全球AI编程工具，助力开发者即刻编程。

更多推荐

cover

当 AI 修图师装上“万能插头“：MCP 如何改变修图工作流

avatar AI编程社区
cover

2026年8月开发者实战指南：ChatGPT Plus / Pro + Codex 如何协同完成需求分析、编码、测试、调试与代码审查

avatar AI编程社区
cover

投资一台B300，部署Claude Opus 5多久才能回本？

avatar AI编程社区