wehub-resource-sync
90 lines
2.6 KiB
Python
90 lines
2.6 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from __future__ import annotations
|
|
|
|
import socket
|
|
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, TypedDict
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
|
|
|
|
class TensorMeta(TypedDict):
|
|
"""Metadata for a tensor in the weight bucket."""
|
|
name: str
|
|
shape: 'torch.Size'
|
|
dtype: 'torch.dtype'
|
|
offset: int
|
|
|
|
|
|
def _find_free_port() -> int:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
s.bind(('', 0))
|
|
return s.getsockname()[1]
|
|
|
|
|
|
def _is_valid_ipv6_address(addr: str) -> bool:
|
|
try:
|
|
socket.inet_pton(socket.AF_INET6, addr)
|
|
return True
|
|
except (OSError, socket.error):
|
|
return False
|
|
|
|
|
|
@dataclass
|
|
class MasterMetadata:
|
|
"""Metadata from the master (trainer rank 0) for topology building."""
|
|
zmq_ip: str
|
|
zmq_port: int
|
|
nccl_store_host: str = ''
|
|
nccl_store_port: int = 0
|
|
|
|
|
|
class CheckpointEngine(ABC):
|
|
rank: Optional[int] = None
|
|
|
|
@abstractmethod
|
|
def prepare(self) -> Dict[str, Any]:
|
|
"""Prepare the checkpoint engine before weight synchronization.
|
|
|
|
Allocate weight transfer buffers, setup communication channels,
|
|
and return metadata needed for topology building.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def build_topology(
|
|
cls,
|
|
trainer_world_size: int,
|
|
rollout_world_size: int,
|
|
metadata: List[Dict],
|
|
) -> Tuple[Dict[str, List[Any]], Dict[str, List[Any]]]:
|
|
"""Build communication topology between trainer and rollout workers.
|
|
|
|
Returns (trainer_kwargs, rollout_kwargs) for init_process_group().
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def init_process_group(self, **kwargs):
|
|
"""Initialize the process group for weight synchronization."""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def finalize(self):
|
|
"""Finalize: free buffers, optionally destroy the process group."""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
async def send_weights(self, weights: Generator[Tuple[str, 'torch.Tensor'], None, None]):
|
|
"""Send model weights to rollout workers."""
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
async def receive_weights(self) -> AsyncGenerator[Tuple[str, 'torch.Tensor'], None]:
|
|
"""Receive model weights from trainer."""
|
|
raise NotImplementedError
|