| """ |
| 显存池化、分片与 Zero-Bubble 张量并行调度器 |
| Author: XiaoZhe (Commercial Contact: janejulius119@gmail.com / WeChat: julius119) |
| """ |
|
|
| import asyncio |
| import logging |
| from dataclasses import dataclass |
| from typing import Dict, List |
|
|
| logger = logging.getLogger("VerseFlow.GPUManager") |
|
|
| @dataclass |
| class GPUAllocation: |
| gpu_ids: List[int] |
| memory_allocated_mb: float |
| model_id: str |
| task_id: str |
| tp_size: int |
|
|
| class ProductionGPUMemoryManager: |
| """分布式 GPU 显存管理器，支持多卡张量并行与防死锁并发申请""" |
| def __init__(self, gpu_count: int = 4, memory_per_gpu_mb: float = 24576.0): |
| self.gpu_count = gpu_count |
| self.memory_per_gpu_mb = memory_per_gpu_mb |
| self.gpu_used_mb = [0.0] * gpu_count |
| self.lock = asyncio.Lock() |
| |
| self.model_vram_map = { |
| "minimax_h3": (38000.0, 2),[cite: 1] |
| "wan2.2_14b": (28000.0, 2),[cite: 1] |
| "hunyuan_video_1_5": (24000.0, 1),[cite: 1] |
| "ltx_2_3": (16000.0, 1),[cite: 1] |
| "wan2.2_5b": (8000.0, 1)[cite: 1] |
| } |
|
|
| async def allocate(self, model_id: str, task_id: str, timeout: float = 60.0) -> GPUAllocation: |
| req_mb, req_tp = self.model_vram_map.get(model_id, (16000.0, 1)) |
| per_gpu_req = req_mb / req_tp |
| start_time = asyncio.get_event_loop().time() |
|
|
| while True: |
| async with self.lock: |
| available_gpus = [] |
| for i in range(self.gpu_count): |
| free_mem = self.memory_per_gpu_mb - self.gpu_used_mb[i] - 2000.0 |
| if free_mem >= per_gpu_req: |
| available_gpus.append(i) |
|
|
| if len(available_gpus) >= req_tp: |
| selected_gpus = available_gpus[:req_tp] |
| for gpu_id in selected_gpus: |
| self.gpu_used_mb[gpu_id] += per_gpu_req |
| |
| logger.info(f"显存分配成功: Task={task_id}, Model={model_id}, GPUs={selected_gpus}") |
| return GPUAllocation(selected_gpus, req_mb, model_id, task_id, req_tp) |
|
|
| if asyncio.get_event_loop().time() - start_time > timeout: |
| raise TimeoutError(f"显存申请超时: Task={task_id}, Model={model_id}") |
| await asyncio.sleep(0.5) |
|
|
| async def release(self, alloc: GPUAllocation): |
| async with self.lock: |
| per_gpu_req = alloc.memory_allocated_mb / alloc.tp_size |
| for gpu_id in alloc.gpu_ids: |
| self.gpu_used_mb[gpu_id] = max(0.0, self.gpu_used_mb[gpu_id] - per_gpu_req) |
| logger.info(f"显存成功释放: Task={alloc.task_id}, GPUs={alloc.gpu_ids}") |
