Skip to content

fix: forward port PR #2122 from verl to use async vllm with megatron #36

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
desaxce wants to merge 1 commit into
base: numina-main
Choose a base branch
from
Open

fix: forward port PR #2122 from verl to use async vllm with megatron #36

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 9 additions & 4 deletions verl/workers/megatron_workers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -235,14 +235,19 @@ def _build_rollout(self, trust_remote_code=False):
model_hf_config=self.actor_model_config,
)
elif vllm_mode == "spmd":
rollout = vLLMRollout(
from verl.workers.rollout.vllm_rollout import vLLMAsyncRollout

vllm_rollout_cls = vLLMRollout if self.config.rollout.mode == "sync" else vLLMAsyncRollout
rollout = vllm_rollout_cls(
model_path=local_path,
config=self.config.rollout,
tokenizer=self.tokenizer,
model_hf_config=self.actor_model_config,
device_mesh=rollout_device_mesh,
trust_remote_code=trust_remote_code,
)
else:
raise NotImplementedError(f"vllm_mode {vllm_mode} is not supported, must be 'customized' or 'spmd'")
log_gpu_memory_usage("After building vllm rollout", logger=logger)

# perform weight resharding between actor and rollout
Expand All @@ -256,6 +261,8 @@ def _build_rollout(self, trust_remote_code=False):
layer_name_mapping=layer_name_mapping,
actor_module=self.actor.actor_module,
weight_converter=weight_converter,
device_mesh=rollout_device_mesh,
offload_param=self._is_offload_param,
)
log_gpu_memory_usage("After building sharding manager", logger=logger)
elif self.config.rollout.name == "sglang":
Expand Down Expand Up @@ -295,6 +302,7 @@ def _build_rollout(self, trust_remote_code=False):
layer_name_mapping=layer_name_mapping,
weight_converter=weight_converter,
device_mesh=rollout_device_mesh,
offload_param=self._is_offload_param,
)
log_gpu_memory_usage("After building sharding manager", logger=logger)
elif self.config.rollout.name == "sglang_async":
Expand Down Expand Up @@ -474,9 +482,6 @@ def update_actor(self, data: DataProto):
@GPUMemoryLogger(role="generate_sequences", logger=logger)
def generate_sequences(self, prompts: DataProto):
assert self._is_rollout
if self._is_offload_param:
load_megatron_model_to_gpu(self.actor_module)
log_gpu_memory_usage("After load actor params during generate_sequences", logger=logger)
prompts.batch = prompts.batch.cuda()
meta_info = {
"eos_token_id": self.generation_config.eos_token_id if self.generation_config is not None else self.tokenizer.eos_token_id,
Expand Down
64 changes: 54 additions & 10 deletions verl/workers/sharding_manager/megatron_vllm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,12 @@
from verl.protocol import all_gather_data_proto
from verl.third_party.vllm import LLM, vllm_version
from verl.third_party.vllm import parallel_state as vllm_ps
from verl.utils.debug import GPUMemoryLogger
from verl.utils.debug import GPUMemoryLogger, log_gpu_memory_usage
from verl.utils.device import get_torch_device
from verl.utils.megatron_utils import (
get_model,
load_megatron_model_to_gpu,
offload_megatron_model_to_cpu,
per_tensor_generator,
unwrap_model,
)
Expand Down Expand Up @@ -274,11 +277,23 @@ def __init__(
layer_name_mapping,
weight_converter: McoreToHFWeightConverterBase,
module: AllGatherPPModel = None,
device_mesh,
offload_param: bool = True,
):
from megatron.core import parallel_state as mpu

self.actor_module = actor_module
self.inference_engine = inference_engine
self.offload_param = offload_param

# For AsyncLLM, inference_engine and model_runner are defer initialized in vLLMAsyncRollout.load_model
if "vllm_v_0_6_3" in str(type(self.inference_engine)) or "vllm_v_0_5_4" in str(type(self.inference_engine)):
# vLLM <= v0.6.3
self.model_runner = self.inference_engine.llm_engine.model_executor.worker.model_runner if self.inference_engine else None
else:
# vLLM > v0.6.3
self.model_runner = self.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner if self.inference_engine else None

self.model_config = model_config
self.transformer_config = transformer_config
self.layer_name_mapping = layer_name_mapping
Expand All @@ -287,11 +302,9 @@ def __init__(
# initialize groups for vllm inference
self.rank = torch.distributed.get_rank()
self.world_size = torch.distributed.get_world_size()
self.infer_tp_size = vllm_ps.get_tensor_model_parallel_world_size()
self.infer_tp_rank = vllm_ps.get_tensor_model_parallel_rank()
self.infer_tp_group = vllm_ps.get_tensor_model_parallel_group()
if vllm_version not in ("0.5.4", "0.6.3"):
self.infer_tp_group = self.infer_tp_group.device_group
self.device_mesh = device_mesh
self.infer_tp_size = self.device_mesh["infer_tp"].size()
self.infer_tp_rank = self.device_mesh["infer_tp"].get_local_rank()
self.train_tp_size = mpu.get_tensor_model_parallel_world_size()
self.train_tp_rank = mpu.get_tensor_model_parallel_rank()
self.train_tp_group = mpu.get_tensor_model_parallel_group()
Expand All @@ -304,6 +317,15 @@ def __init__(
self.need_tp_reshard = self.train_tp_size != self.infer_tp_size
self.train_tp_larger = self.train_tp_size > self.infer_tp_size

self.torch_random_states = get_torch_device().get_rng_state()
if self.device_mesh is not None:
gen_dp_rank = self.device_mesh["dp"].get_local_rank()
get_torch_device().manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states
self.gen_random_states = get_torch_device().get_rng_state()
get_torch_device().set_rng_state(self.torch_random_states)
else:
self.gen_random_states = None

@GPUMemoryLogger(role="megatron vllm sharding_manager", logger=logger)
def __enter__(self):
if vllm_version in (
Expand All @@ -325,14 +347,23 @@ def __enter__(self):
self.transformer_config,
self.layer_name_mapping,
)
model = self.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model
model = self.model_runner.model
patch_vllm_moe_model_weight_loader(model)
loaded_params = model.load_weights(per_tensor_param)
info = f"vLLM load weights, loaded_params: {len(loaded_params)}"
logger.info(info)

if "tags" in inspect.signature(self.inference_engine.wake_up).parameters:
self.inference_engine.wake_up(tags=["kv_cache"])
if self.offload_param:
offload_megatron_model_to_cpu(self.actor_module)
get_torch_device().empty_cache()

if "tags" in inspect.signature(self.inference_engine.wake_up).parameters:
self.inference_engine.wake_up(tags=["kv_cache"])

# important: need to manually set the random states of each tp to be identical.
if self.device_mesh is not None:
self.torch_random_states = get_torch_device().get_rng_state()
get_torch_device().set_rng_state(self.gen_random_states)

@GPUMemoryLogger(role="megatron vllm sharding_manager", logger=logger)
def __exit__(self, exc_type, exc_value, traceback):
Expand All @@ -347,13 +378,26 @@ def __exit__(self, exc_type, exc_value, traceback):
model.train()

torch.cuda.empty_cache()
# restore random states
if self.device_mesh is not None:
self.gen_random_states = get_torch_device().get_rng_state()
get_torch_device().set_rng_state(self.torch_random_states)

@GPUMemoryLogger(role="megatron vllm sharding_manager", logger=logger)
def preprocess_data(self, data: DataProto) -> DataProto:
# DP_COMPUTE_PROTO: all training ranks are dp, the same as fsdp
if self.infer_tp_size == 1:
return data
all_gather_data_proto(data, self.infer_tp_group)
# TODO: Current impl doesn't consider FSDP with torch micro-dp
if vllm_version in (
"0.5.4",
"0.6.3",
):
group = vllm_ps.get_tensor_model_parallel_group()
else:
group = vllm_ps.get_tensor_model_parallel_group().device_group

all_gather_data_proto(data=data, process_group=group)
return data

@GPUMemoryLogger(role="megatron vllm sharding_manager", logger=logger)
Expand Down