wehub-resource-sync
187 lines
7.1 KiB
Python
187 lines
7.1 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
"""Weight transfer utilities for training → rollout weight synchronization.
|
|
|
|
BucketedWeightSender is used by the **training** side (MegatronWorker
|
|
via RolloutAdapter) to ship weights to vLLM's WeightSyncWorkerExtension
|
|
through ZMQ IPC. It was originally in vllm_server.py but belongs here
|
|
because it is a training-side concern, not a rollout engine concern.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import torch
|
|
import uuid
|
|
from typing import Any, Dict, Optional
|
|
|
|
from swift.utils import get_current_device, synchronize
|
|
from swift.utils.logger import get_logger
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class BucketedWeightSender:
|
|
"""Streams model weights to vLLM worker via ZMQ IPC with bucketed transfer."""
|
|
|
|
def __init__(
|
|
self,
|
|
zmq_handle: str,
|
|
bucket_size_mb: int = 512,
|
|
use_shm: bool = False,
|
|
timeout_s: int = 600,
|
|
external_buffer: Optional[torch.Tensor] = None,
|
|
):
|
|
self.zmq_handle = zmq_handle
|
|
self.bucket_size = int(bucket_size_mb) << 20
|
|
self.use_shm = use_shm
|
|
self.timeout_ms = int(timeout_s * 1000)
|
|
self.socket = None
|
|
self.buffer = None
|
|
self.shm = None
|
|
self._pending_handshake = None
|
|
# When provided, reuse this caller-owned persistent GPU buffer instead of
|
|
# allocating a fresh one per sync. Reusing the same storage keeps the CUDA IPC
|
|
# handle signature stable across sync rounds, so the vLLM worker's IPC-mapping
|
|
# cache hits and no new mapping is leaked each step.
|
|
self._external_buffer = external_buffer
|
|
self._owns_buffer = True
|
|
|
|
async def __aenter__(self):
|
|
self._init_socket_and_buffer()
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
self.cleanup()
|
|
|
|
def _init_socket_and_buffer(self):
|
|
"""Bind the REQ socket and allocate the bucket buffer."""
|
|
import zmq
|
|
|
|
ctx = zmq.Context.instance()
|
|
self.socket = ctx.socket(zmq.REQ)
|
|
self.socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms)
|
|
self.socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms)
|
|
self.socket.setsockopt(zmq.LINGER, 0)
|
|
self.socket.bind(self.zmq_handle)
|
|
|
|
from torch.multiprocessing.reductions import reduce_tensor
|
|
|
|
if not self.use_shm:
|
|
if self._external_buffer is not None and self._external_buffer.numel() >= self.bucket_size:
|
|
# Reuse the persistent buffer -> stable IPC handle -> worker cache hit.
|
|
self.buffer = self._external_buffer
|
|
self._owns_buffer = False
|
|
else:
|
|
self.buffer = torch.empty(self.bucket_size, dtype=torch.uint8, device=get_current_device())
|
|
self._owns_buffer = True
|
|
self._pending_handshake = reduce_tensor(self.buffer)
|
|
else:
|
|
from multiprocessing import shared_memory
|
|
shm_name = f'swift_weights_{uuid.uuid4().hex}'
|
|
self.shm = shared_memory.SharedMemory(name=shm_name, create=True, size=self.bucket_size)
|
|
self.buffer = torch.frombuffer(self.shm.buf, dtype=torch.uint8)
|
|
self._pending_handshake = {'name': shm_name, 'size': self.bucket_size}
|
|
|
|
async def handshake(self):
|
|
if self._pending_handshake is None:
|
|
raise RuntimeError('BucketedWeightSender.handshake() called before enter or twice: '
|
|
'the handshake payload is consumed on first call, a second handshake '
|
|
'within the same ``async with`` block would ship stale metadata.')
|
|
loop = asyncio.get_running_loop()
|
|
payload = self._pending_handshake
|
|
self._pending_handshake = None
|
|
|
|
def _send_recv():
|
|
self.socket.send_pyobj(payload)
|
|
return self.socket.recv()
|
|
|
|
await loop.run_in_executor(None, _send_recv)
|
|
|
|
async def _stream_weights_inner(self, items_iter, is_async: bool = False):
|
|
"""Shared bucketing logic for sync and async weight iterators."""
|
|
loop = asyncio.get_running_loop()
|
|
offset = 0
|
|
bucket_meta: Dict[str, Dict[str, Any]] = {}
|
|
n_weights = 0
|
|
|
|
def _zmq_send_recv(payload):
|
|
self.socket.send_pyobj(payload)
|
|
return self.socket.recv()
|
|
|
|
async def _flush(is_last: bool):
|
|
nonlocal offset, bucket_meta
|
|
if not bucket_meta and not is_last:
|
|
return
|
|
if self.buffer.device.type != 'cpu':
|
|
synchronize()
|
|
await loop.run_in_executor(None, _zmq_send_recv, {'bucket_meta': bucket_meta, 'is_last': is_last})
|
|
offset = 0
|
|
bucket_meta = {}
|
|
|
|
async def _process(name, weight):
|
|
nonlocal offset, n_weights
|
|
if self.use_shm and weight.device.type != 'cpu':
|
|
weight = weight.cpu()
|
|
if not weight.is_contiguous():
|
|
weight = weight.contiguous()
|
|
nbytes = int(weight.nbytes)
|
|
if nbytes > self.bucket_size:
|
|
raise RuntimeError(f'Weight {name} ({tuple(weight.shape)}, {weight.dtype}) '
|
|
f'is {nbytes} bytes, exceeding bucket size ({self.bucket_size}).')
|
|
if offset + nbytes > self.bucket_size:
|
|
await _flush(False)
|
|
bucket_meta[name] = {
|
|
'name': name,
|
|
'shape': weight.shape,
|
|
'dtype': weight.dtype,
|
|
'offset': offset,
|
|
}
|
|
self.buffer[offset:offset + nbytes].copy_(weight.view(-1).view(torch.uint8), non_blocking=True)
|
|
offset += nbytes
|
|
n_weights += 1
|
|
|
|
if is_async:
|
|
async for name, weight in items_iter:
|
|
await _process(name, weight)
|
|
else:
|
|
for name, weight in items_iter:
|
|
await _process(name, weight)
|
|
|
|
await _flush(True)
|
|
logger.debug('BucketedWeightSender: sent %d weights', n_weights)
|
|
|
|
async def send_weights(self, weights):
|
|
"""Stream weights into buckets. Accepts ``dict`` or iterator."""
|
|
items = weights.items() if isinstance(weights, dict) else weights
|
|
await self._stream_weights_inner(items, is_async=False)
|
|
|
|
async def send_weights_async(self, async_weights):
|
|
"""Stream weights from an async generator into buckets."""
|
|
await self._stream_weights_inner(async_weights, is_async=True)
|
|
|
|
def cleanup(self):
|
|
if self.socket is not None:
|
|
self.socket.close()
|
|
self.socket = None
|
|
if self.zmq_handle.startswith('ipc://'):
|
|
ipc_path = self.zmq_handle[len('ipc://'):]
|
|
try:
|
|
if os.path.exists(ipc_path):
|
|
os.remove(ipc_path)
|
|
except OSError:
|
|
pass
|
|
if self._owns_buffer:
|
|
del self.buffer
|
|
self.buffer = None
|
|
if self.shm is not None:
|
|
try:
|
|
self.shm.close()
|
|
self.shm.unlink()
|
|
except (FileNotFoundError, BufferError):
|
|
pass
|
|
self.shm = None
|
|
self._pending_handshake = None
|
|
from swift.utils import gc_collect, ipc_collect
|
|
gc_collect()
|
|
ipc_collect()
|