This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,171 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
from __future__ import annotations
|
||||
|
||||
import ray
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from swift.utils.logger import get_logger
|
||||
from .nccl import NCCLCheckpointEngine
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class CheckpointEngineManager:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
train_actors: List[Any],
|
||||
rollout_replicas: List[Any],
|
||||
*,
|
||||
weight_sync_mode: str = 'nccl',
|
||||
is_colocated: bool = False,
|
||||
sleep_level: int = 1,
|
||||
train_group: Any,
|
||||
):
|
||||
self.train_actors = train_actors
|
||||
self._rollout_replicas = rollout_replicas
|
||||
self.rollout_actors = [r.primary for r in rollout_replicas]
|
||||
self._weight_sync_mode = weight_sync_mode
|
||||
self.is_colocated = is_colocated
|
||||
self._train_group = train_group
|
||||
|
||||
if is_colocated and sleep_level >= 2:
|
||||
logger.warning(
|
||||
'sleep_level=%d capped to 1 in colocate mode '
|
||||
'(out-of-process vLLM cannot safely discard all GPU memory).', sleep_level)
|
||||
sleep_level = 1
|
||||
self.sleep_level = sleep_level
|
||||
|
||||
self.base_sync_done: bool = False
|
||||
self._model_keys: Optional[List[str]] = None
|
||||
self._sleeping_tags: set = set()
|
||||
|
||||
def sync_weights(self, merge_and_sync: bool = True) -> None:
|
||||
"""Synchronize weights from training model to rollout replicas."""
|
||||
if self.is_colocated:
|
||||
self.sleep_rollout()
|
||||
self.wake_up_rollout(tags=['weights'])
|
||||
if self._weight_sync_mode == 'naive':
|
||||
self._sync_weights_naive(merge_and_sync)
|
||||
else:
|
||||
self._sync_weights_nccl(merge_and_sync)
|
||||
|
||||
def sleep_rollout(self) -> None:
|
||||
if self._sleeping_tags:
|
||||
return
|
||||
for replica in self._rollout_replicas:
|
||||
replica.sleep(level=self.sleep_level)
|
||||
self._sleeping_tags = {'weights', 'kv_cache'}
|
||||
logger.debug('CheckpointEngineManager: rollout replicas sleeping (level=%d)', self.sleep_level)
|
||||
|
||||
def wake_up_rollout(self, tags: Optional[List[str]] = None) -> None:
|
||||
if not self._sleeping_tags:
|
||||
return
|
||||
for replica in self._rollout_replicas:
|
||||
replica.wake_up(tags=tags)
|
||||
if tags is None:
|
||||
self._sleeping_tags.clear()
|
||||
else:
|
||||
self._sleeping_tags -= set(tags)
|
||||
logger.debug('CheckpointEngineManager: rollout wake_up tags=%s, still_sleeping=%s', tags, self._sleeping_tags)
|
||||
|
||||
def _sync_weights_naive(self, merge_and_sync: bool) -> None:
|
||||
tg = self._train_group
|
||||
adapter_only = self.base_sync_done and not merge_and_sync
|
||||
need_merge = not adapter_only and merge_and_sync
|
||||
if need_merge:
|
||||
tg.merge_lora()
|
||||
if self.is_colocated:
|
||||
tg.offload_to_cpu()
|
||||
try:
|
||||
tg.update_weights(adapter_only=adapter_only)
|
||||
finally:
|
||||
if self.is_colocated:
|
||||
tg.reload_to_gpu()
|
||||
if need_merge:
|
||||
tg.unmerge_lora()
|
||||
if not self.base_sync_done:
|
||||
self.base_sync_done = True
|
||||
logger.debug('CheckpointEngineManager[naive]: initial weight sync done')
|
||||
|
||||
def _sync_weights_nccl(self, merge_and_sync: bool) -> None:
|
||||
"""NCCL broadcast weight sync path.
|
||||
|
||||
Lifecycle:
|
||||
1. prepare_checkpoint_engine on all actors
|
||||
2. build_topology
|
||||
3. init_process_group on all actors (concurrent — required for TCPStore)
|
||||
4. send_weights (train) + receive_weights (rollout) concurrently
|
||||
5. finalize_checkpoint_engine on all actors
|
||||
"""
|
||||
n_train = len(self.train_actors)
|
||||
n_rollout = len(self.rollout_actors)
|
||||
|
||||
# 1. Prepare — train side: rank 0 is master, others are not
|
||||
is_master_flags = [True] + [False] * (n_train - 1)
|
||||
prepare_refs = [
|
||||
actor.prepare_checkpoint_engine.remote(flag) for actor, flag in zip(self.train_actors, is_master_flags)
|
||||
]
|
||||
prepare_results = ray.get(prepare_refs)
|
||||
model_metadata = prepare_results[0]
|
||||
|
||||
# 1b. Prepare — rollout side: all non-master
|
||||
rollout_prepare_refs = [actor.prepare_checkpoint_engine.remote(False) for actor in self.rollout_actors]
|
||||
ray.get(rollout_prepare_refs)
|
||||
|
||||
# 2. Build topology
|
||||
model_kwargs, rollout_kwargs = NCCLCheckpointEngine.build_topology(n_train, n_rollout, [model_metadata])
|
||||
|
||||
# 3. Init process groups (MUST be concurrent — TCPStore server
|
||||
# blocks until all clients connect)
|
||||
train_init_refs = [
|
||||
actor.init_checkpoint_process_group.remote(
|
||||
rank=model_kwargs['rank'][i],
|
||||
world_size=model_kwargs['world_size'][i],
|
||||
master_metadata=model_kwargs['master_metadata'][i],
|
||||
) for i, actor in enumerate(self.train_actors)
|
||||
]
|
||||
rollout_init_refs = [
|
||||
actor.init_checkpoint_process_group.remote(
|
||||
rank=rollout_kwargs['rank'][i],
|
||||
world_size=rollout_kwargs['world_size'][i],
|
||||
master_metadata=rollout_kwargs['master_metadata'][i],
|
||||
) for i, actor in enumerate(self.rollout_actors)
|
||||
]
|
||||
ray.get(train_init_refs + rollout_init_refs)
|
||||
|
||||
# 4. Send/receive weights (concurrent)
|
||||
adapter_only = self.base_sync_done and not merge_and_sync
|
||||
need_merge = not adapter_only and merge_and_sync
|
||||
peft_config = None
|
||||
if adapter_only:
|
||||
peft_config = ray.get(self.train_actors[0].get_peft_config_dict.remote())
|
||||
|
||||
if need_merge:
|
||||
merge_refs = [actor.merge_lora.remote() for actor in self.train_actors]
|
||||
ray.get(merge_refs)
|
||||
|
||||
train_send_refs = [
|
||||
actor.send_checkpoint_weights.remote(adapter_only=adapter_only) for actor in self.train_actors
|
||||
]
|
||||
rollout_recv_refs = [
|
||||
actor.receive_checkpoint_weights.remote(
|
||||
base_sync_done=self.base_sync_done,
|
||||
peft_config=peft_config,
|
||||
) for actor in self.rollout_actors
|
||||
]
|
||||
ray.get(train_send_refs + rollout_recv_refs)
|
||||
|
||||
if need_merge:
|
||||
unmerge_refs = [actor.unmerge_lora.remote() for actor in self.train_actors]
|
||||
ray.get(unmerge_refs)
|
||||
|
||||
# 5. Finalize
|
||||
train_fin_refs = [actor.finalize_checkpoint_engine.remote() for actor in self.train_actors]
|
||||
rollout_fin_refs = [actor.finalize_checkpoint_engine.remote() for actor in self.rollout_actors]
|
||||
ray.get(train_fin_refs + rollout_fin_refs)
|
||||
|
||||
if not self.base_sync_done:
|
||||
self.base_sync_done = True
|
||||
logger.info('CheckpointEngineManager[nccl]: initial weight sync to %d replica(s) '
|
||||
'(lora_only=%s)', n_rollout, not merge_and_sync)
|
||||
Reference in New Issue
Block a user