wehub-resource-sync
124 lines
4.7 KiB
Python
124 lines
4.7 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import ray
|
|
import time
|
|
import torch
|
|
from typing import Any, Generator, List, Optional, Tuple
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _is_ipc_supported() -> bool:
|
|
"""Check if CUDA IPC is supported (GPU=True, NPU=fallback to SHM)."""
|
|
from transformers.utils import is_torch_npu_available
|
|
return not is_torch_npu_available()
|
|
|
|
|
|
class RolloutAdapter:
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
replica_rank: int = 0,
|
|
rollout_rank: int = 0,
|
|
bucket_size_mb: int = 2048,
|
|
):
|
|
self.replica_rank = replica_rank
|
|
self.rollout_rank = rollout_rank
|
|
self.bucket_size_mb = bucket_size_mb
|
|
self.is_primary = (rollout_rank == 0)
|
|
self.use_shm = not _is_ipc_supported()
|
|
self.zmq_handle = (f'ipc:///tmp/swift-rollout-zmq-replica-{replica_rank}-rank-{rollout_rank}.sock')
|
|
self._server_handle = None
|
|
# Persistent CUDA IPC buffer reused across all update_weights() syncs so the IPC
|
|
# handle stays stable (the vLLM worker's mapping cache hits) and no IPC mapping
|
|
# leaks per step.
|
|
self._ipc_buffer: Optional[torch.Tensor] = None
|
|
|
|
@property
|
|
def server_handle(self) -> Any:
|
|
if self._server_handle is None:
|
|
self._server_handle = ray.get_actor(f'swift_rollout_server_{self.replica_rank}_0')
|
|
return self._server_handle
|
|
|
|
def update_weights(
|
|
self,
|
|
weight_iter: Generator[Tuple[str, torch.Tensor], None, None],
|
|
*,
|
|
vllm_lora_param_names: Optional[set] = None,
|
|
peft_config: Optional[dict] = None,
|
|
base_sync_done: bool = False,
|
|
) -> None:
|
|
"""Push training weights to vLLM via ZMQ IPC.
|
|
|
|
Only the primary rank (rollout_rank == 0) sends weights;
|
|
other ranks are no-op.
|
|
|
|
Args:
|
|
weight_iter: Iterator of (name, tensor) from bridge.export_weights.
|
|
vllm_lora_param_names: When set, remap dense param names to
|
|
LoRA-wrapped names (*.base_layer.weight) for full-weight sync
|
|
with vllm_enable_lora.
|
|
peft_config: When provided with base_sync_done=True, vLLM loads
|
|
weights as a LoRA adapter via TensorLoRARequest.
|
|
base_sync_done: Indicates this is an adapter-only sync after
|
|
the initial full weight sync has completed.
|
|
"""
|
|
if vllm_lora_param_names:
|
|
from swift.rlhf_trainers.utils import add_base_layer_suffix_by_param_names
|
|
weight_iter = add_base_layer_suffix_by_param_names(weight_iter, vllm_lora_param_names)
|
|
|
|
if not self.is_primary:
|
|
for _ in weight_iter:
|
|
pass
|
|
return
|
|
|
|
from ..rollout.weight_transfer import BucketedWeightSender
|
|
|
|
start_time = time.time()
|
|
|
|
# Lazily (re)allocate the persistent IPC buffer; reused across syncs so the
|
|
# handle signature stays stable and the worker-side IPC cache hits.
|
|
external_buffer = None
|
|
if not self.use_shm:
|
|
from swift.utils import get_current_device
|
|
bucket_size = self.bucket_size_mb << 20
|
|
if self._ipc_buffer is None or self._ipc_buffer.numel() < bucket_size:
|
|
self._ipc_buffer = torch.empty(bucket_size, dtype=torch.uint8, device=get_current_device())
|
|
external_buffer = self._ipc_buffer
|
|
|
|
async def _do_ipc_sync():
|
|
sender = BucketedWeightSender(
|
|
zmq_handle=self.zmq_handle,
|
|
bucket_size_mb=self.bucket_size_mb,
|
|
use_shm=self.use_shm,
|
|
external_buffer=external_buffer,
|
|
)
|
|
try:
|
|
async with sender:
|
|
rpc_ref = self.server_handle.update_weights_ipc.remote(self.zmq_handle, self.use_shm, 600,
|
|
peft_config, base_sync_done)
|
|
await sender.handshake()
|
|
await sender.send_weights(weight_iter)
|
|
await asyncio.get_running_loop().run_in_executor(None, ray.get, rpc_ref)
|
|
finally:
|
|
sender.cleanup()
|
|
|
|
asyncio.run(_do_ipc_sync())
|
|
logger.debug('RolloutAdapter: update_weights done (replica=%d, adapter_only=%s, %.2fs)', self.replica_rank,
|
|
base_sync_done,
|
|
time.time() - start_time)
|
|
|
|
def reset_prefix_cache(self) -> None:
|
|
if not self.is_primary:
|
|
return
|
|
ray.get(self.server_handle.reset_prefix_cache.remote())
|
|
|
|
def get_model_param_names(self) -> List[str]:
|
|
if not self.is_primary:
|
|
return []
|
|
return ray.get(self.server_handle.get_model_param_names.remote()) or []
|