Spaces:
Running on Zero
Running on Zero
File size: 13,318 Bytes
5ed07ee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 | """
Multi-GPU data-parallel inference pool for VoxCPM.
Uses a work-stealing pattern: each GPU runs in its own thread and pulls
batch tasks from a shared queue. Whichever GPU finishes first takes the
next batch, naturally handling uneven batch durations.
Python threads release the GIL during CUDA kernel launches and
``torch.no_grad()`` blocks, so GPU compute across threads is truly parallel.
Usage::
pool = MultiGPUInferencePool(
load_model_fn=lambda dev: load_svs_model(ckpt_dir, device=dev),
devices=["cuda:0", "cuda:1"],
batch_size=64,
generate_kwargs=dict(cfg_value=2.0, inference_timesteps=10, max_len=2000),
)
# prompts: list of str
results = pool.generate_all(prompts) # list of (global_task_idx, audio_tensor)
pool.shutdown()
"""
from __future__ import annotations
import os
import queue
import sys
import threading
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple
import torch
# ---------------------------------------------------------------------------
# Sentinel & data types
# ---------------------------------------------------------------------------
_SHUTDOWN = object()
@dataclass
class _BatchTask:
"""A batch of prompts to be generated."""
batch_id: int
global_indices: List[int] # position in the original prompt list
prompts: List[str]
prompt_audio_feats: Optional[List[Optional[torch.Tensor]]] = None
@dataclass
class _BatchResult:
"""Results returned by one batch generation call."""
batch_id: int
global_indices: List[int]
audio_tensors: List[Optional[torch.Tensor]]
device: str
elapsed: float
latent_tensors: Optional[List[Optional[torch.Tensor]]] = None
# ---------------------------------------------------------------------------
# GPU Worker
# ---------------------------------------------------------------------------
class _GPUWorker:
"""Runs on a dedicated thread, owns one model on one device."""
def __init__(
self,
device: str,
model: Any,
generate_kwargs: Dict[str, Any],
task_queue: "queue.Queue[_BatchTask | object]",
result_list: List[_BatchResult],
result_lock: threading.Lock,
log_fn: Callable[[str], None],
):
self.device = device
self.model = model
self.generate_kwargs = generate_kwargs
self.task_queue = task_queue
self.result_list = result_list
self.result_lock = result_lock
self.log_fn = log_fn
def run(self) -> None:
"""Main loop: pull tasks until shutdown sentinel."""
while True:
task = self.task_queue.get()
if task is _SHUTDOWN:
self.task_queue.task_done()
break
assert isinstance(task, _BatchTask)
t0 = time.time()
latent_tensors: Optional[List[Optional[torch.Tensor]]] = None
want_latents = bool(self.generate_kwargs.get("return_latents", False))
try:
with torch.no_grad():
out = self.model.generate_batch(
target_texts=task.prompts,
prompt_audio_feats=task.prompt_audio_feats,
**self.generate_kwargs,
)
if want_latents:
audio_tensors, latent_tensors = out
else:
audio_tensors = out
except Exception as exc:
self.log_fn(
f" [{self.device}] Batch {task.batch_id} error: {exc}"
)
import traceback
traceback.print_exc()
audio_tensors = [None] * len(task.prompts)
if want_latents:
latent_tensors = [None] * len(task.prompts)
elapsed = time.time() - t0
result = _BatchResult(
batch_id=task.batch_id,
global_indices=task.global_indices,
audio_tensors=audio_tensors,
device=self.device,
elapsed=elapsed,
latent_tensors=latent_tensors,
)
with self.result_lock:
self.result_list.append(result)
self.log_fn(
f" [{self.device}] Batch {task.batch_id} done "
f"({len(task.prompts)} prompts, {elapsed:.1f}s)"
)
self.task_queue.task_done()
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
class MultiGPUInferencePool:
"""Work-stealing thread pool for multi-GPU batch inference.
Parameters
----------
load_model_fn : callable(device_str) -> model
Factory that loads a VoxCPM model onto the given device string.
devices : list[str]
CUDA device strings, e.g. ``["cuda:0", "cuda:1"]``.
batch_size : int
Per-GPU batch size.
generate_kwargs : dict
Keyword arguments forwarded to ``model.generate_batch()``,
e.g. ``cfg_value``, ``inference_timesteps``, ``max_len``.
log_fn : callable, optional
Logging function. Defaults to ``print(..., file=sys.stderr)``.
"""
def __init__(
self,
load_model_fn: Callable[[str], Any],
devices: List[str],
batch_size: int = 64,
generate_kwargs: Optional[Dict[str, Any]] = None,
log_fn: Optional[Callable[[str], None]] = None,
):
self.devices = devices
self.batch_size = batch_size
self.generate_kwargs = generate_kwargs or {}
self.log_fn = log_fn or (lambda msg: print(msg, file=sys.stderr))
self._shut_down = False
if len(devices) > 1:
# This pool runs one Python worker thread per GPU. Keep PyTorch's
# CPU thread pools single-threaded so we don't multiply host-side
# launcher / BLAS threads across workers and starve CUDA dispatch.
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
try:
if torch.get_num_threads() != 1:
torch.set_num_threads(1)
except Exception:
pass
try:
torch.set_num_interop_threads(1)
except RuntimeError:
# Can only be set once and before inter-op work starts.
pass
# ---- Load one model per GPU ----
self.log_fn(
f"[MultiGPU] Loading model on {len(devices)} device(s): {devices}"
)
self.models: List[Any] = []
for dev in devices:
t0 = time.time()
model = load_model_fn(dev)
elapsed = time.time() - t0
self.models.append(model)
self.log_fn(f"[MultiGPU] Model loaded on {dev} ({elapsed:.1f}s)")
# ---- Shared work queue + result collection ----
self._task_queue: queue.Queue = queue.Queue()
self._result_list: List[_BatchResult] = []
self._result_lock = threading.Lock()
# ---- Spawn worker threads (one per GPU) ----
self._workers: List[threading.Thread] = []
self._gpu_workers: List[_GPUWorker] = []
for dev, model in zip(devices, self.models):
worker = _GPUWorker(
device=dev,
model=model,
generate_kwargs=self.generate_kwargs,
task_queue=self._task_queue,
result_list=self._result_list,
result_lock=self._result_lock,
log_fn=self.log_fn,
)
thread = threading.Thread(
target=worker.run, name=f"gpu-worker-{dev}", daemon=True
)
thread.start()
self._workers.append(thread)
self._gpu_workers.append(worker)
self.log_fn(f"[MultiGPU] {len(self._workers)} worker thread(s) started")
# ------------------------------------------------------------------
# Public: generate_all
# ------------------------------------------------------------------
def generate_all(
self,
prompts: List[str],
prompt_audio_feats_list: Optional[List[Optional[torch.Tensor]]] = None,
sort_by: Optional[Callable[[str], float]] = len,
) -> List[Tuple[int, Optional[torch.Tensor]]]:
"""Generate audio for all *prompts* across GPUs.
Prompts are split into batches of ``self.batch_size`` and enqueued.
Each GPU thread pulls batches from the shared queue as it becomes
idle (work-stealing).
When ``sort_by`` is provided (default: ``len``), prompts are first
ordered by the key so that each batch contains samples of similar
expected output length. This minimises the wall-clock penalty of
``_inference_batch`` waiting for the longest sample in a heterogeneous
batch (the AR loop only breaks when ALL samples have stopped).
Pass ``sort_by=None`` to preserve input order.
Returns
-------
list of (global_index, audio_tensor | None)
Ordered by ``global_index`` (the position in the input list).
"""
if self._shut_down:
raise RuntimeError("Pool has been shut down")
# Clear previous results
with self._result_lock:
self._result_list.clear()
# Build dispatch order. Even if sort_by is None, we still keep a
# uniform `order` list so batch construction stays index-based.
if sort_by is not None and len(prompts) > 1:
order = sorted(range(len(prompts)), key=lambda i: sort_by(prompts[i]))
else:
order = list(range(len(prompts)))
# Enqueue batches over the (possibly reordered) sequence, but keep
# ``global_indices`` pointing at the ORIGINAL position so callers
# always receive results in input order.
num_batches = (len(prompts) + self.batch_size - 1) // self.batch_size
for bi in range(num_batches):
start = bi * self.batch_size
end = min(start + self.batch_size, len(prompts))
batch_order = order[start:end]
batch_pa = None
if prompt_audio_feats_list is not None:
batch_pa = [prompt_audio_feats_list[i] for i in batch_order]
task = _BatchTask(
batch_id=bi,
global_indices=batch_order,
prompts=[prompts[i] for i in batch_order],
prompt_audio_feats=batch_pa,
)
self._task_queue.put(task)
self.log_fn(
f"[MultiGPU] Enqueued {num_batches} batch(es) "
f"({len(prompts)} prompts, batch_size={self.batch_size})"
)
# Wait for all batches to complete
self._task_queue.join()
# Collect results. When ``return_latents`` was set in
# ``generate_kwargs``, each entry is ``(gidx, audio, latent)``;
# otherwise ``(gidx, audio)`` for backwards compatibility.
want_latents = bool(self.generate_kwargs.get("return_latents", False))
with self._result_lock:
if want_latents:
all_results_lat: List[Tuple[int, Optional[torch.Tensor], Optional[torch.Tensor]]] = []
for br in self._result_list:
lat_list = br.latent_tensors or [None] * len(br.global_indices)
for j, gidx in enumerate(br.global_indices):
a = br.audio_tensors[j] if j < len(br.audio_tensors) else None
l = lat_list[j] if j < len(lat_list) else None
all_results_lat.append((gidx, a, l))
all_results_lat.sort(key=lambda x: x[0])
return all_results_lat
all_results: List[Tuple[int, Optional[torch.Tensor]]] = []
for br in self._result_list:
for j, gidx in enumerate(br.global_indices):
tensor = (
br.audio_tensors[j]
if j < len(br.audio_tensors)
else None
)
all_results.append((gidx, tensor))
all_results.sort(key=lambda x: x[0])
return all_results
# ------------------------------------------------------------------
# Public: shutdown
# ------------------------------------------------------------------
def shutdown(self) -> None:
"""Stop all worker threads and release models."""
if self._shut_down:
return
self._shut_down = True
# Send one shutdown sentinel per worker
for _ in self._workers:
self._task_queue.put(_SHUTDOWN)
# Wait for threads to finish
for t in self._workers:
t.join(timeout=30)
# Release model references
self.models.clear()
self._gpu_workers.clear()
self.log_fn("[MultiGPU] Pool shut down")
def __del__(self):
try:
self.shutdown()
except Exception:
pass
|