wehub-resource-sync
58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
# Copyright (c) ModelScope Contributors. All rights reserved.
|
|
from __future__ import annotations
|
|
|
|
from typing import Optional
|
|
|
|
from .base import CheckpointEngine, MasterMetadata
|
|
|
|
|
|
class CheckpointEngineMixin:
|
|
"""Mixin providing checkpoint engine lifecycle methods for Ray actors.
|
|
|
|
Add this to Ray actors (MegatronWorker, VllmServer) to give them
|
|
checkpoint engine capabilities. The Manager calls these methods
|
|
via ``ray.remote()`` to coordinate NCCL weight sync.
|
|
|
|
Applied to ``MegatronWorker`` (send side) and ``VllmServer``
|
|
(receive side).
|
|
"""
|
|
|
|
_checkpoint_engine: Optional[CheckpointEngine] = None
|
|
_bucket_size: int = 3072 << 20
|
|
|
|
def _get_or_create_checkpoint_engine(self) -> CheckpointEngine:
|
|
"""Get or create the checkpoint engine instance (lazy singleton)."""
|
|
if self._checkpoint_engine is None:
|
|
from transformers.utils import is_torch_npu_available
|
|
if is_torch_npu_available():
|
|
from .hccl import HCCLCheckpointEngine
|
|
self._checkpoint_engine = HCCLCheckpointEngine(self._bucket_size, rebuild_group=False)
|
|
else:
|
|
from .nccl import NCCLCheckpointEngine
|
|
self._checkpoint_engine = NCCLCheckpointEngine(self._bucket_size)
|
|
return self._checkpoint_engine
|
|
|
|
def prepare_checkpoint_engine(self, is_master: bool) -> Optional[MasterMetadata]:
|
|
"""Prepare checkpoint engine for weight sync."""
|
|
engine = self._get_or_create_checkpoint_engine()
|
|
engine.is_master = is_master
|
|
return engine.prepare()
|
|
|
|
def init_checkpoint_process_group(
|
|
self,
|
|
rank: int,
|
|
world_size: int,
|
|
master_metadata: MasterMetadata,
|
|
) -> None:
|
|
engine = self._get_or_create_checkpoint_engine()
|
|
engine.init_process_group(
|
|
rank=rank,
|
|
world_size=world_size,
|
|
master_metadata=master_metadata,
|
|
)
|
|
|
|
def finalize_checkpoint_engine(self) -> None:
|
|
"""Finalize checkpoint engine: release buffers, optionally destroy group."""
|
|
if self._checkpoint_engine is not None:
|
|
self._checkpoint_engine.finalize()
|