83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import logging
|
|
from typing import Generator, Optional, Tuple
|
|
from urllib.parse import urlparse
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from sglang.srt.connector import BaseConnector
|
|
from sglang.srt.utils import init_custom_process_group
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RemoteInstanceConnector(BaseConnector):
|
|
|
|
def __init__(self, url: str, device: torch.device = "cpu"):
|
|
assert (
|
|
device.type == "cuda"
|
|
), "RemoteInstanceConnector only supports cuda device."
|
|
super().__init__(url)
|
|
self.url = url
|
|
self.device = device
|
|
|
|
def build_group(
|
|
self,
|
|
gpu_id: int = -1,
|
|
tp_rank: int = -1,
|
|
instance_ip: str = None,
|
|
group_rank: int = 1,
|
|
world_size: int = 2,
|
|
):
|
|
assert (
|
|
self.device.type == "cuda"
|
|
), "RemoteInstanceConnector only supports cuda device."
|
|
assert (
|
|
gpu_id != -1 and tp_rank != -1
|
|
), "gpu_id and tp_rank must be specified for RemoteInstanceConnector. "
|
|
|
|
self.device_id = torch.device(self.device.type, gpu_id)
|
|
|
|
parsed_url = urlparse(self.url)
|
|
master_address = parsed_url.hostname
|
|
master_port = parsed_url.port
|
|
group_name = f"send_weights_{instance_ip}_{master_port}_{tp_rank}"
|
|
backend = "nccl"
|
|
|
|
logger.info(
|
|
f"init custom process group: master_address={master_address}, master_port={master_port}, "
|
|
f"rank_offset={group_rank}, world_size={world_size}, group_name={group_name}, backend={backend}"
|
|
)
|
|
|
|
try:
|
|
self._model_update_group = init_custom_process_group(
|
|
backend=backend,
|
|
init_method=f"tcp://{master_address}:{master_port}",
|
|
world_size=world_size,
|
|
rank=group_rank,
|
|
group_name=group_name,
|
|
device_id=self.device_id,
|
|
)
|
|
dist.barrier(group=self._model_update_group)
|
|
return True, "Succeeded to initialize custom process group."
|
|
except Exception as e:
|
|
message = f"Failed to initialize custom process group: {e}."
|
|
logger.error(message)
|
|
return False, message
|
|
|
|
# Implemented as a no-op to make BaseConnector interface consistent.
|
|
def pull_files(
|
|
self,
|
|
allow_pattern: Optional[list[str]] = None,
|
|
ignore_pattern: Optional[list[str]] = None,
|
|
) -> None:
|
|
return
|
|
|
|
# Implemented as a no-op to make BaseConnector interface consistent.
|
|
def weight_iterator(
|
|
self, rank: int = 0
|
|
) -> Generator[Tuple[str, torch.Tensor], None, None]:
|
|
return
|