This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,633 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
|
||||
import os
|
||||
import re
|
||||
import torch
|
||||
from collections import OrderedDict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from dataclasses import dataclass
|
||||
from itertools import repeat
|
||||
from packaging import version
|
||||
from queue import Queue
|
||||
from transformers import GenerationConfig, LogitsProcessor
|
||||
from transformers.generation.streamers import BaseStreamer
|
||||
from typing import List, Optional, Union
|
||||
|
||||
from swift.model.register import fix_do_sample_warning
|
||||
from swift.utils import get_device, synchronize
|
||||
from .protocol import RequestConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class AdapterRequest:
|
||||
name: str
|
||||
path: str
|
||||
|
||||
|
||||
class InferTools:
|
||||
|
||||
@staticmethod
|
||||
def _is_chinese_char(cp: int) -> bool:
|
||||
"""Checks whether CP is the codepoint of a CJK character."""
|
||||
# copy from transformers.generation.streamers.TextStreamer
|
||||
if ((0x4E00 <= cp <= 0x9FFF) or (0x3400 <= cp <= 0x4DBF) or (0x20000 <= cp <= 0x2A6DF)
|
||||
or (0x2A700 <= cp <= 0x2B73F) or (0x2B740 <= cp <= 0x2B81F) or (0x2B820 <= cp <= 0x2CEAF)
|
||||
or (0xF900 <= cp <= 0xFAFF) or (0x2F800 <= cp <= 0x2FA1F)):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
||||
class InferStreamer(InferTools):
|
||||
|
||||
def __init__(self, template, **decode_kwargs):
|
||||
self.template = template
|
||||
self.tokenizer = template.tokenizer
|
||||
|
||||
self.cache_idx = 0 # token idx
|
||||
self.print_idx = 0
|
||||
self.decode_kwargs = decode_kwargs
|
||||
self.first_num_space = -1 # The number of whitespace characters before the first token.
|
||||
self.first_token = True
|
||||
|
||||
def _align_blank_suffix(self, response: str) -> str:
|
||||
# Avoid the occurrence of repeated words in sentence.
|
||||
cur_num_space = len(response) - len(response.lstrip(' '))
|
||||
if self.first_num_space == -1:
|
||||
self.first_num_space = cur_num_space
|
||||
elif cur_num_space < self.first_num_space:
|
||||
response = ' ' * (self.first_num_space - cur_num_space) + response
|
||||
elif cur_num_space > self.first_num_space:
|
||||
response = response[cur_num_space - self.first_num_space:]
|
||||
return response
|
||||
|
||||
def _get_response(self, response: str, is_finished: bool, token_len: int) -> str:
|
||||
# After the symbol for a new line, we flush the cache.
|
||||
if self.first_token:
|
||||
printable_text = response
|
||||
self.first_token = False
|
||||
elif response.endswith('\n') or is_finished:
|
||||
printable_text = response[self.print_idx:]
|
||||
self.cache_idx += token_len
|
||||
self.first_num_space = -1
|
||||
self.print_idx = 0
|
||||
# If the last token is a CJK character, we print the characters.
|
||||
elif len(response) > 0 and self._is_chinese_char(ord(response[-1])):
|
||||
printable_text = response[self.print_idx:]
|
||||
self.print_idx += len(printable_text)
|
||||
# Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
|
||||
# which may change with the subsequent token -- there are probably smarter ways to do this!)
|
||||
else:
|
||||
printable_text = response[self.print_idx:response.rfind(' ') + 1]
|
||||
self.print_idx += len(printable_text)
|
||||
return printable_text
|
||||
|
||||
def get_printable_text(self, raw_tokens: List[int], is_finished: bool) -> str:
|
||||
raw_tokens = raw_tokens[self.cache_idx:]
|
||||
if self.first_token:
|
||||
raw_tokens = []
|
||||
response = self.template.decode_generate_ids(
|
||||
raw_tokens, is_finished=is_finished, first_token=self.first_token, **self.decode_kwargs)
|
||||
response = self._align_blank_suffix(response)
|
||||
return self._get_response(response, is_finished, len(raw_tokens))
|
||||
|
||||
|
||||
class StreamerMixin:
|
||||
|
||||
def __init__(self):
|
||||
self.queue = Queue()
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self) -> torch.Tensor:
|
||||
value = self.queue.get()
|
||||
if value is None:
|
||||
raise StopIteration()
|
||||
else:
|
||||
return value
|
||||
|
||||
|
||||
class TokensIteratorStreamer(StreamerMixin, BaseStreamer):
|
||||
|
||||
def put(self, value: torch.Tensor) -> None:
|
||||
self.queue.put(value)
|
||||
|
||||
def end(self) -> None:
|
||||
self.queue.put(None)
|
||||
|
||||
|
||||
class LogitsStreamer(LogitsProcessor):
|
||||
|
||||
def __init__(self):
|
||||
self.queue = Queue()
|
||||
|
||||
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
|
||||
self.queue.put(scores)
|
||||
return scores
|
||||
|
||||
|
||||
def _set_generation_config_default_value(model_generation_config: GenerationConfig,
|
||||
generation_config: GenerationConfig) -> GenerationConfig:
|
||||
for k, v in model_generation_config.to_dict().items():
|
||||
new_v = getattr(generation_config, k, None)
|
||||
if k in ['max_length']:
|
||||
continue
|
||||
if k in ['no_repeat_ngram_size'] or v is not None and new_v is None:
|
||||
setattr(generation_config, k, v)
|
||||
return generation_config
|
||||
|
||||
|
||||
def prepare_generation_config(model_generation_config: Optional[GenerationConfig], request_config: RequestConfig,
|
||||
tokenizer) -> Optional[GenerationConfig]:
|
||||
if model_generation_config is None or request_config is None:
|
||||
return model_generation_config
|
||||
kwargs = {'max_new_tokens': request_config.max_tokens}
|
||||
# not use: 'n', 'best_of', 'frequency_penalty', 'presence_penalty'
|
||||
for key in ['length_penalty']:
|
||||
kwargs[key] = getattr(request_config, key)
|
||||
for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty', 'num_beams']:
|
||||
new_value = getattr(request_config, key)
|
||||
if new_value is None:
|
||||
kwargs[key] = getattr(model_generation_config, key, None)
|
||||
else:
|
||||
kwargs[key] = new_value
|
||||
|
||||
if kwargs.get('top_k') is not None and kwargs['top_k'] <= 0:
|
||||
kwargs['top_k'] = None
|
||||
|
||||
if not getattr(model_generation_config, 'do_sample', False) and request_config.temperature in {0, None}:
|
||||
kwargs['temperature'] = 0
|
||||
if kwargs['temperature'] == 0:
|
||||
kwargs['do_sample'] = False
|
||||
kwargs['temperature'] = 1
|
||||
kwargs['top_p'] = 1
|
||||
kwargs['top_k'] = 50
|
||||
else:
|
||||
kwargs['do_sample'] = True
|
||||
generation_config = GenerationConfig(**kwargs)
|
||||
generation_config = _set_generation_config_default_value(model_generation_config, generation_config)
|
||||
fix_do_sample_warning(generation_config)
|
||||
|
||||
if generation_config.eos_token_id is None:
|
||||
generation_config.eos_token_id = tokenizer.eos_token_id
|
||||
generation_config.pad_token_id = tokenizer.pad_token_id
|
||||
return generation_config
|
||||
|
||||
|
||||
def patch_lmdeploy(load_weights=False):
|
||||
"""This patch allows lmdeploy selects device and reload state_dict"""
|
||||
import lmdeploy
|
||||
assert version.parse(lmdeploy.__version__) >= version.parse('0.7.0')
|
||||
from lmdeploy.messages import TurbomindEngineConfig
|
||||
from lmdeploy.turbomind.deploy import loader
|
||||
from lmdeploy.turbomind.deploy.loader import create_loader
|
||||
from lmdeploy.turbomind.deploy.source_model import llama
|
||||
|
||||
def _create_loader(model_path: str, pattern: str):
|
||||
if not isinstance(model_path, (str, os.PathLike)):
|
||||
|
||||
def generate():
|
||||
generator = OrderedDict()
|
||||
model_dict = {}
|
||||
if not isinstance(model_path, dict):
|
||||
for key, value in list(model_path):
|
||||
model_dict[key] = value
|
||||
else:
|
||||
model_dict = model_path
|
||||
for key, value in model_dict.items():
|
||||
match = re.findall(pattern, key)
|
||||
if not match:
|
||||
if -1 not in generator:
|
||||
generator[-1] = {}
|
||||
generator[-1][key] = value
|
||||
else:
|
||||
layer = int(match[0])
|
||||
if layer not in generator:
|
||||
generator[layer] = {}
|
||||
generator[layer][key] = value
|
||||
return generator
|
||||
|
||||
return generate()
|
||||
else:
|
||||
return create_loader(model_path, pattern)
|
||||
|
||||
loader.create_loader = _create_loader
|
||||
llama.create_loader = _create_loader
|
||||
|
||||
TurbomindEngineConfig.devices = [0]
|
||||
|
||||
from lmdeploy.turbomind.turbomind import TurboMind
|
||||
from lmdeploy.turbomind.utils import ModelSource
|
||||
|
||||
@contextmanager
|
||||
def patch_threadpool_map():
|
||||
ThreadPoolExecutor.map_origin = ThreadPoolExecutor.map
|
||||
ThreadPoolExecutor.map = lambda *args, **kwargs: []
|
||||
yield
|
||||
ThreadPoolExecutor.map = ThreadPoolExecutor.map_origin
|
||||
del ThreadPoolExecutor.map_origin
|
||||
|
||||
@contextmanager
|
||||
def tm_model_context(self):
|
||||
|
||||
def _get_tm_model(model_path,
|
||||
model_name,
|
||||
chat_template_name,
|
||||
engine_config: TurbomindEngineConfig,
|
||||
group_size: int = None,
|
||||
out_dir: str = None):
|
||||
from lmdeploy.turbomind.deploy.converter import get_tm_model_origin
|
||||
tm_model = get_tm_model_origin(model_path, model_name, chat_template_name, engine_config, group_size,
|
||||
out_dir)
|
||||
self.tm_model = tm_model
|
||||
return tm_model
|
||||
|
||||
from lmdeploy.turbomind.deploy import converter
|
||||
converter.get_tm_model_origin = converter.get_tm_model
|
||||
converter.get_tm_model = _get_tm_model
|
||||
yield
|
||||
converter.get_tm_model = converter.get_tm_model_origin
|
||||
del converter.get_tm_model_origin
|
||||
|
||||
def __init__(self,
|
||||
model_path: str,
|
||||
tokenizer: object,
|
||||
model_name: str = None,
|
||||
chat_template_name: str = None,
|
||||
engine_config: TurbomindEngineConfig = None,
|
||||
model_source: ModelSource = ModelSource.WORKSPACE,
|
||||
**kwargs):
|
||||
self.gpu_list = engine_config.devices
|
||||
with patch_threadpool_map(), tm_model_context(self):
|
||||
self.__origin_init__(model_path, tokenizer, model_name, chat_template_name, engine_config, model_source,
|
||||
**kwargs)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.gpu_count) as e:
|
||||
ranks = [self.node_id * self.gpu_count + device_id for device_id in range(self.gpu_count)]
|
||||
if not load_weights:
|
||||
for _ in e.map(self.model_comm.process_weight, self.gpu_list, ranks):
|
||||
pass
|
||||
if version.parse(lmdeploy.__version__) < version.parse('0.7.2'):
|
||||
for _ in e.map(self.model_comm.create_engine, self.gpu_list, ranks, repeat(self.nccl_params)):
|
||||
pass
|
||||
else:
|
||||
for _ in e.map(self.model_comm.create_engine, self.gpu_list, ranks):
|
||||
pass
|
||||
|
||||
def _create_weight(self, model_comm):
|
||||
"""Allocate weight buffer, load params if from_workspace."""
|
||||
|
||||
# TODO: support mpi
|
||||
self.node_id = 0
|
||||
self.node_num = 1
|
||||
if version.parse(lmdeploy.__version__) < version.parse('0.7.2'):
|
||||
self.nccl_params = model_comm.create_nccl_params(self.node_id)
|
||||
synchronize()
|
||||
|
||||
# create weight
|
||||
def _create_weight_func(index, device_id):
|
||||
rank = self.node_id * self.gpu_count + index
|
||||
model_comm.create_shared_weights(device_id, rank)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=self.gpu_count) as executor:
|
||||
futures = []
|
||||
for idx, device_id in enumerate(self.gpu_list):
|
||||
futures.append(executor.submit(_create_weight_func, idx, device_id))
|
||||
for future in futures:
|
||||
future.result()
|
||||
|
||||
def _get_model_params(self, model_comm, tm_params):
|
||||
"""Get turbomind model params when loading from hf."""
|
||||
|
||||
def _get_params(idx, device_id, que):
|
||||
rank = self.node_id * self.gpu_count + idx
|
||||
out = model_comm.get_params(device_id, rank)
|
||||
que.put(out)
|
||||
|
||||
que = Queue()
|
||||
with ThreadPoolExecutor(max_workers=self.gpu_count) as executor:
|
||||
futures = []
|
||||
for idx, device_id in enumerate(self.gpu_list):
|
||||
futures.append(executor.submit(_get_params, idx, device_id, que))
|
||||
for future in futures:
|
||||
future.result()
|
||||
|
||||
for _ in range(self.gpu_count):
|
||||
tensor_map = que.get()
|
||||
for k, v in tensor_map.items():
|
||||
if k not in tm_params:
|
||||
tm_params[k] = []
|
||||
tm_params[k].append(v)
|
||||
|
||||
def _load_weights(self, state_dict):
|
||||
tm_params = self.tm_model.tm_params
|
||||
self._get_model_params(self.model_comm, tm_params)
|
||||
input_model = self.tm_model.input_model
|
||||
model_path = input_model.model_path
|
||||
input_model.model_path = state_dict
|
||||
self.tm_model.export()
|
||||
input_model.model_path = model_path
|
||||
|
||||
from lmdeploy.turbomind.turbomind import TurboMindInstance
|
||||
|
||||
def create_instance(self, cuda_stream_id=0):
|
||||
return TurboMindInstance(self, self.config, cuda_stream_id, self.gpu_list)
|
||||
|
||||
TurboMind.__origin_init__ = TurboMind.__init__
|
||||
TurboMind.__init__ = __init__
|
||||
TurboMind._create_weight = _create_weight
|
||||
TurboMind._get_model_params = _get_model_params
|
||||
TurboMind.create_instance = create_instance
|
||||
if load_weights:
|
||||
TurboMind.load_weights = _load_weights
|
||||
|
||||
def __init_ins__(self, tm_model, config, cuda_stream_id=0, gpu_list=None):
|
||||
if gpu_list is None:
|
||||
gpu_list = [0]
|
||||
self.gpu_list = gpu_list
|
||||
self.__origin_init__(tm_model, config, cuda_stream_id)
|
||||
|
||||
def _create_model_instance(self, device_id):
|
||||
model_inst = self.tm_model.model_comm.create_model_instance(self.gpu_list[0])
|
||||
return model_inst
|
||||
|
||||
TurboMindInstance.__origin_init__ = TurboMindInstance.__init__
|
||||
TurboMindInstance.__init__ = __init_ins__
|
||||
TurboMindInstance._create_model_instance = _create_model_instance
|
||||
|
||||
|
||||
def patch_npu_vllm(vllm_device: str, *, colocate: bool = False):
|
||||
if isinstance(vllm_device, int):
|
||||
vllm_device = get_device(vllm_device)
|
||||
device_type = vllm_device.split(':')[0]
|
||||
if device_type == 'npu':
|
||||
from swift.model.npu_patch.vllm_ascend import patch_vllm_ascend_runtime
|
||||
from swift.model.npu_patch.vllm_ascend_memory import vllm_ascend_mem_get_info_context
|
||||
patch_vllm_ascend_runtime(colocate=colocate)
|
||||
return vllm_ascend_mem_get_info_context(vllm_device)
|
||||
|
||||
return nullcontext()
|
||||
|
||||
|
||||
def patch_vllm_triton_device_guard():
|
||||
import functools
|
||||
try:
|
||||
from vllm.v1.worker import gpu_worker as _gw
|
||||
_orig_fn = _gw.init_worker_distributed_environment
|
||||
except (ImportError, AttributeError):
|
||||
return
|
||||
|
||||
if getattr(_gw, '_swift_dist_env_patched', False):
|
||||
return
|
||||
|
||||
@functools.wraps(_orig_fn)
|
||||
def _patched_init_worker_distributed_environment(*args, **kwargs):
|
||||
if not torch.cuda.is_available():
|
||||
return _orig_fn(*args, **kwargs)
|
||||
expected_device = torch.cuda.current_device()
|
||||
result = _orig_fn(*args, **kwargs)
|
||||
actual_device = torch.cuda.current_device()
|
||||
if actual_device != expected_device:
|
||||
torch.cuda.set_device(expected_device)
|
||||
return result
|
||||
|
||||
_gw.init_worker_distributed_environment = _patched_init_worker_distributed_environment
|
||||
_gw._swift_dist_env_patched = True
|
||||
|
||||
|
||||
def patch_vllm_memory_leak():
|
||||
# fix vllm 0.7.3 memory leak
|
||||
# https://github.com/vllm-project/vllm/pull/14326
|
||||
import vllm
|
||||
try:
|
||||
vllm_version = version.parse(vllm.__version__)
|
||||
needs_patch = (vllm_version == version.parse('0.7.3'))
|
||||
except version.InvalidVersion:
|
||||
needs_patch = False
|
||||
|
||||
if not needs_patch:
|
||||
return
|
||||
|
||||
def patch_vllm_abort_seq_group():
|
||||
from typing import Dict, Iterable
|
||||
from vllm.core.scheduler import Scheduler
|
||||
from vllm.sequence import SequenceGroup, SequenceGroupBase, SequenceStatus
|
||||
|
||||
def new_abort_seq_group(
|
||||
self,
|
||||
request_id: Union[str, Iterable[str]],
|
||||
seq_id_to_seq_group: Optional[Dict[str, SequenceGroupBase]] = None,
|
||||
) -> None:
|
||||
if isinstance(request_id, str):
|
||||
request_id = (request_id, )
|
||||
request_ids = set(request_id)
|
||||
seq_id_to_seq_group = seq_id_to_seq_group or {}
|
||||
for state_queue in [self.waiting, self.running, self.swapped]:
|
||||
aborted_groups: List[SequenceGroup] = []
|
||||
for seq_group in state_queue:
|
||||
# When n>1, seq_group.request_id looks like
|
||||
# foo_parallel_sample_0, while request_ids is just foo, and we
|
||||
# should resolve it as real_request_id to match.
|
||||
if seq_group.request_id in seq_id_to_seq_group:
|
||||
real_request_id = seq_id_to_seq_group[seq_group.request_id].group_id
|
||||
else:
|
||||
real_request_id = seq_group.request_id
|
||||
if real_request_id in request_ids:
|
||||
# Appending aborted group into pending list.
|
||||
aborted_groups.append(seq_group)
|
||||
# We can't remove real_request_id in request_ids here,
|
||||
# because there may be other seq groups sharing the same
|
||||
# real_request_id
|
||||
for aborted_group in aborted_groups:
|
||||
# Remove the sequence group from the state queue.
|
||||
state_queue.remove(aborted_group)
|
||||
# Remove the aborted request from the Mamba cache.
|
||||
self._finished_requests_ids.append(aborted_group.request_id)
|
||||
for seq in aborted_group.get_seqs():
|
||||
if seq.is_finished():
|
||||
continue
|
||||
seq.status = SequenceStatus.FINISHED_ABORTED
|
||||
self.free_seq(seq)
|
||||
if aborted_group.request_id in seq_id_to_seq_group:
|
||||
del seq_id_to_seq_group[aborted_group.request_id]
|
||||
|
||||
self._free_seq_group_cross_attn_blocks(aborted_group)
|
||||
|
||||
origin_method = Scheduler.abort_seq_group
|
||||
Scheduler._old_abort_seq_group = origin_method
|
||||
Scheduler.abort_seq_group = new_abort_seq_group
|
||||
|
||||
def patch_vllm_engine():
|
||||
from vllm.engine.llm_engine import LLMEngine, SchedulerOutputState
|
||||
from vllm.outputs import PoolingRequestOutput, RequestOutput
|
||||
from vllm.sequence import ExecuteModelRequest
|
||||
|
||||
def new_abort_request(self, request_id) -> None:
|
||||
for scheduler in self.scheduler:
|
||||
scheduler.abort_seq_group(request_id, seq_id_to_seq_group=self.seq_id_to_seq_group)
|
||||
|
||||
origin_method = LLMEngine.abort_request
|
||||
LLMEngine._old_abort_request = origin_method
|
||||
LLMEngine.abort_request = new_abort_request
|
||||
|
||||
def new_step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:
|
||||
if self.parallel_config.pipeline_parallel_size > 1:
|
||||
raise NotImplementedError('Pipeline parallelism is only supported through AsyncLLMEngine '
|
||||
'as performance will be severely degraded otherwise.')
|
||||
|
||||
# For llm_engine, there is no pipeline parallel support, so the engine
|
||||
# used is always 0.
|
||||
virtual_engine = 0
|
||||
|
||||
# These are cached outputs from previous iterations. None if on first
|
||||
# iteration
|
||||
cached_outputs = self.cached_scheduler_outputs[virtual_engine]
|
||||
seq_group_metadata_list = cached_outputs.seq_group_metadata_list
|
||||
scheduler_outputs = cached_outputs.scheduler_outputs
|
||||
allow_async_output_proc = cached_outputs.allow_async_output_proc
|
||||
|
||||
ctx = self.scheduler_contexts[virtual_engine]
|
||||
|
||||
# Clear outputs for each new scheduler iteration
|
||||
ctx.request_outputs.clear()
|
||||
|
||||
# Skip the scheduler if there are any remaining steps in the seq groups.
|
||||
# This ensures that the scheduler is only called again when the current
|
||||
# batch has completed.
|
||||
# The scheduler is also skipped if a single request caused the last
|
||||
# engine step to fail, and the previous schedule needs to be rerun.
|
||||
if not self._has_remaining_steps(seq_group_metadata_list):
|
||||
# Schedule iteration
|
||||
(seq_group_metadata_list, scheduler_outputs,
|
||||
allow_async_output_proc) = self.scheduler[virtual_engine].schedule()
|
||||
|
||||
ctx.seq_group_metadata_list = seq_group_metadata_list
|
||||
ctx.scheduler_outputs = scheduler_outputs
|
||||
|
||||
finished_requests_ids = self.scheduler[virtual_engine].get_and_reset_finished_requests_ids()
|
||||
# When n>1, elements in self.seq_id_to_seq_group should be deleted
|
||||
# here, otherwise memory leaks.
|
||||
for finished_request_id in finished_requests_ids:
|
||||
if finished_request_id in self.seq_id_to_seq_group:
|
||||
del self.seq_id_to_seq_group[finished_request_id]
|
||||
|
||||
# Maybe switch from async mode to sync mode
|
||||
if not allow_async_output_proc and len(ctx.output_queue) > 0:
|
||||
self._process_model_outputs(ctx=ctx)
|
||||
|
||||
if (self.scheduler_config.is_multi_step and scheduler_outputs.num_lookahead_slots > 0):
|
||||
# cache the scheduler outputs for the next iteration if we have
|
||||
# lookahead slots
|
||||
self._cache_scheduler_outputs_for_multi_step(virtual_engine, seq_group_metadata_list,
|
||||
scheduler_outputs, allow_async_output_proc)
|
||||
else:
|
||||
finished_requests_ids = list()
|
||||
|
||||
assert seq_group_metadata_list is not None
|
||||
assert scheduler_outputs is not None
|
||||
|
||||
if not scheduler_outputs.is_empty():
|
||||
|
||||
# Check if we have a cached last_output from the previous iteration.
|
||||
# For supporting PP this is probably the best way to pass the
|
||||
# sampled_token_ids, as a separate broadcast over all the PP stages
|
||||
# will cause one virtual engine's microbatch to block the pipeline.
|
||||
last_sampled_token_ids = \
|
||||
self._get_last_sampled_token_ids(virtual_engine)
|
||||
|
||||
execute_model_req = ExecuteModelRequest(
|
||||
seq_group_metadata_list=seq_group_metadata_list,
|
||||
blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
|
||||
blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
|
||||
blocks_to_copy=scheduler_outputs.blocks_to_copy,
|
||||
num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
|
||||
running_queue_size=scheduler_outputs.running_queue_size,
|
||||
finished_requests_ids=finished_requests_ids,
|
||||
# We use ExecuteModelRequest to pass the last sampled_token_ids
|
||||
# to each of the non-last PP stages for in-place prepare_input.
|
||||
last_sampled_token_ids=last_sampled_token_ids)
|
||||
|
||||
if allow_async_output_proc:
|
||||
execute_model_req.async_callback = self.async_callbacks[virtual_engine]
|
||||
|
||||
outputs = self.model_executor.execute_model(execute_model_req=execute_model_req)
|
||||
|
||||
# We need to do this here so that last step's sampled_token_ids can
|
||||
# be passed to the next iteration for PP.
|
||||
if self.scheduler_config.is_multi_step:
|
||||
self._update_cached_scheduler_output(virtual_engine, outputs)
|
||||
else:
|
||||
# Nothing scheduled => If there is pending async postprocessor,
|
||||
# then finish it here.
|
||||
if len(ctx.output_queue) > 0:
|
||||
self._process_model_outputs(ctx=ctx)
|
||||
# No outputs in this case
|
||||
outputs = []
|
||||
|
||||
# Finish the current step for all the sequence groups.
|
||||
if self.scheduler_config.is_multi_step:
|
||||
for seq_group in seq_group_metadata_list:
|
||||
seq_group.finish_step()
|
||||
|
||||
if not self._has_remaining_steps(seq_group_metadata_list):
|
||||
# clear the cache if we have finished all the steps.
|
||||
if self.scheduler_config.is_multi_step:
|
||||
self.cached_scheduler_outputs[0] = SchedulerOutputState()
|
||||
|
||||
# is_first_step_output is True only when the num_steps of all
|
||||
# the sequences are 1. When the num_steps > 1,
|
||||
# multi_step_model_runner does the first-step output append.
|
||||
is_first_step_output: bool = False if not seq_group_metadata_list \
|
||||
else seq_group_metadata_list[0].state.num_steps == 1
|
||||
|
||||
# Add results to the output_queue
|
||||
ctx.append_output(
|
||||
outputs=outputs,
|
||||
seq_group_metadata_list=seq_group_metadata_list,
|
||||
scheduler_outputs=scheduler_outputs,
|
||||
is_async=allow_async_output_proc,
|
||||
is_last_step=True,
|
||||
is_first_step_output=is_first_step_output)
|
||||
|
||||
if outputs and allow_async_output_proc:
|
||||
assert len(outputs) == 1, ('Async postprocessor expects only a single output set')
|
||||
|
||||
self._advance_to_next_step(outputs[0], seq_group_metadata_list,
|
||||
scheduler_outputs.scheduled_seq_groups)
|
||||
|
||||
# Check if need to run the usual non-async path
|
||||
if not allow_async_output_proc:
|
||||
self._process_model_outputs(ctx=ctx)
|
||||
|
||||
# Log stats.
|
||||
self.do_log_stats(scheduler_outputs, outputs)
|
||||
|
||||
# Tracing
|
||||
self.do_tracing(scheduler_outputs)
|
||||
else:
|
||||
# Multi-step case
|
||||
return ctx.request_outputs
|
||||
|
||||
if not self.has_unfinished_requests():
|
||||
# Drain async postprocessor (if exists)
|
||||
if len(ctx.output_queue) > 0:
|
||||
self._process_model_outputs(ctx=ctx)
|
||||
assert len(ctx.output_queue) == 0
|
||||
|
||||
# Stop the execute model loop in parallel workers until there are
|
||||
# more requests to process. This avoids waiting indefinitely in
|
||||
# torch.distributed ops which may otherwise timeout, and unblocks
|
||||
# the RPC thread in the workers so that they can process any other
|
||||
# queued control plane messages, such as add/remove lora adapters.
|
||||
self.model_executor.stop_remote_worker_execution_loop()
|
||||
|
||||
return ctx.request_outputs
|
||||
|
||||
origin_method = LLMEngine.step
|
||||
LLMEngine._old_step = origin_method
|
||||
LLMEngine.step = new_step
|
||||
|
||||
patch_vllm_abort_seq_group()
|
||||
patch_vllm_engine()
|
||||
Reference in New Issue
Block a user