File size: 4,950 Bytes
8c9ba62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | # -*- coding: utf-8 -*-
"""Custom vLLM Worker."""
import ray
import torch
import torch.distributed
from trinity.common.models.vllm_patch.worker_patch import patch_vllm_prompt_logprobs
from trinity.manager.synchronizer import Synchronizer
from trinity.utils.distributed import init_process_group
from trinity.utils.log import get_logger
class WorkerExtension:
def apply_patches(self):
"""Apply necessary patches to vLLM."""
from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader
patch_vllm_moe_model_weight_loader(self.model_runner.model)
patch_vllm_prompt_logprobs(self.model_runner)
def init_process_group(
self,
master_address: str,
master_port: int,
rank_offset: int,
world_size: int,
group_name: str,
backend: str = "nccl",
timeout: int = 1200,
state_dict_meta: list = None,
explorer_name: str = None,
namespace: str = None,
):
"""Init torch process group for model weights update"""
rank = torch.distributed.get_rank()
self.logger = get_logger(f"vllm_worker_{rank}")
assert torch.distributed.is_initialized(), "default torch process group must be initialized"
assert group_name != "", "group name must not be empty"
self._state_dict_meta = state_dict_meta
self._weight_update_rank = rank + rank_offset
self.logger.info(
f"vLLM starting init_process_group:\n"
f" > address={master_address}:{master_port}\n"
f" > rank={rank}\n"
f" > rank_offset={rank_offset}\n"
f" > world_size={world_size}"
)
self._model_update_group = init_process_group(
host=master_address,
port=master_port,
group_name=group_name,
backend=backend,
timeout=timeout,
world_size=world_size,
rank=self._weight_update_rank,
device_id=self.device,
)
torch.distributed.barrier(group=self._model_update_group)
self.logger.info("vLLM init_process_group finished.")
self._explorer_name = explorer_name
self._namespace = namespace
self.synchronizer = Synchronizer.get_actor(namespace=self._namespace)
self._checkpoint_converter = None
def update_weight(self, source: str = "student"):
"""Broadcast weight to all vllm workers from source rank 0 (actor model)"""
if not hasattr(self, "_weight_update_rank"):
raise RuntimeError(
"init_process_group must be called before update_weight. "
"Ensure setup_weight_sync_group (or setup_model_level_weight_sync_group) "
"is invoked for all models including auxiliary models (e.g. teacher)."
)
if self._weight_update_rank == 0:
state_dict, model_version = ray.get(
self.synchronizer.get_model_state_dict.remote(source)
)
if isinstance(state_dict, tuple):
# currently only megatron return a tuple
method, checkpoint_dir = state_dict
if method == "megatron":
if self._checkpoint_converter is None:
from trinity.common.models.utils import get_megatron_converter
self._checkpoint_converter = get_megatron_converter(checkpoint_dir)
state_dict = self._checkpoint_converter.get_state_dict(checkpoint_dir)
else:
raise NotImplementedError(f"{method} is not supported")
ray.get(self.synchronizer.set_model_state_dict.remote(state_dict, model_version))
# For NCCL student sync, metadata is provided during init_process_group.
# Fall back to synchronizer metadata for other sources or legacy paths.
state_dict_meta = self._state_dict_meta if source == "student" else None
if state_dict_meta is None:
state_dict_meta = ray.get(self.synchronizer.get_state_dict_meta.remote(source))
if state_dict_meta is None:
raise RuntimeError(f"No state dict metadata found for source `{source}`.")
for name, dtype_str, shape in state_dict_meta:
if self._weight_update_rank == 0:
weight = state_dict[name]
weight = weight.to(self.device)
else:
dtype = getattr(torch, dtype_str.split(".")[-1])
weight = torch.empty(shape, dtype=dtype, device=self.device)
torch.distributed.broadcast(weight, 0, group=self._model_update_group)
weight = weight.type(self.model_config.dtype)
self.model_runner.model.load_weights(weights=[(name, weight)])
del weight
torch.distributed.barrier(group=self._model_update_group)
torch.cuda.synchronize()
torch.cuda.empty_cache()
|