This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,350 @@
|
||||
# Copyright (c) ModelScope Contributors. All rights reserved.
|
||||
import asyncio
|
||||
import inspect
|
||||
import lmdeploy
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
from contextlib import contextmanager
|
||||
from copy import deepcopy
|
||||
from lmdeploy import PytorchEngineConfig, TurbomindEngineConfig, VisionConfig, pipeline
|
||||
from lmdeploy.api import autoget_backend_config
|
||||
from lmdeploy.serve import async_engine
|
||||
from packaging import version
|
||||
from PIL import Image
|
||||
from transformers import GenerationConfig
|
||||
from transformers.utils.versions import require_version
|
||||
from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union
|
||||
|
||||
from swift.metrics import Metric
|
||||
from swift.model import get_processor
|
||||
from swift.template import Template
|
||||
from swift.utils import get_logger, get_seed, safe_snapshot_download
|
||||
from .infer_engine import InferEngine
|
||||
from .patch import patch_auto_config, patch_auto_tokenizer
|
||||
from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
|
||||
ChatCompletionStreamResponse, ChatMessage, DeltaMessage, InferRequest, RequestConfig)
|
||||
from .utils import InferStreamer
|
||||
|
||||
try:
|
||||
from lmdeploy import EngineGenerationConfig as LmdeployGenerationConfig
|
||||
except ImportError:
|
||||
# compat lmdeploy >= 0.6.*
|
||||
from lmdeploy import GenerationConfig as LmdeployGenerationConfig
|
||||
|
||||
logger = get_logger()
|
||||
|
||||
|
||||
class LmdeployEngine(InferEngine):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_id_or_path: str,
|
||||
*,
|
||||
template: Optional[Template] = None,
|
||||
torch_dtype: Optional[torch.dtype] = None,
|
||||
model_type: Optional[str] = None,
|
||||
template_type: Optional[str] = None,
|
||||
use_hf: Optional[bool] = None,
|
||||
hub_token: Optional[str] = None,
|
||||
revision: Optional[str] = None,
|
||||
# engine_kwargs
|
||||
tp: int = 1,
|
||||
session_len: Optional[int] = None,
|
||||
cache_max_entry_count: float = 0.8,
|
||||
quant_policy: int = 0, # e.g. 4, 8
|
||||
vision_batch_size: int = 1, # max_batch_size in VisionConfig
|
||||
engine_kwargs: Optional[Dict[str, Any]] = None,
|
||||
devices: Optional[List[int]] = None,
|
||||
) -> None:
|
||||
self.model_id_or_path = model_id_or_path
|
||||
self.torch_dtype = torch_dtype
|
||||
self.model_type = model_type
|
||||
self.use_hf = use_hf
|
||||
self.hub_token = hub_token
|
||||
self.revision = revision
|
||||
self.tp = tp
|
||||
self.session_len = session_len
|
||||
self.cache_max_entry_count = cache_max_entry_count
|
||||
self.quant_policy = quant_policy
|
||||
self.vision_batch_size = vision_batch_size
|
||||
self.devices = devices
|
||||
if template is None:
|
||||
processor = self._get_processor()
|
||||
template = self._get_template(processor, template_type=template_type)
|
||||
else:
|
||||
safe_snapshot_download(
|
||||
model_id_or_path,
|
||||
revision=revision,
|
||||
download_model=True,
|
||||
use_hf=use_hf,
|
||||
ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None),
|
||||
hub_token=hub_token)
|
||||
super().__init__(template)
|
||||
|
||||
if self.max_model_len is not None:
|
||||
self.max_model_len -= 1
|
||||
self._prepare_engine_kwargs(engine_kwargs)
|
||||
self.config.torch_dtype = self.torch_dtype = self.torch_dtype or self.model_info.torch_dtype
|
||||
self._prepare_engine()
|
||||
self._load_generation_config()
|
||||
|
||||
def _get_processor(self):
|
||||
return get_processor(
|
||||
model_id_or_path=self.model_id_or_path,
|
||||
torch_dtype=self.torch_dtype,
|
||||
download_model=True,
|
||||
model_type=self.model_type,
|
||||
use_hf=self.use_hf,
|
||||
hub_token=self.hub_token,
|
||||
revision=self.revision)
|
||||
|
||||
def _prepare_engine_kwargs(self, engine_kwargs):
|
||||
if engine_kwargs is None:
|
||||
engine_kwargs = {}
|
||||
engine_kwargs['tp'] = self.tp
|
||||
engine_kwargs['session_len'] = self.session_len
|
||||
engine_kwargs['cache_max_entry_count'] = self.cache_max_entry_count
|
||||
engine_kwargs['quant_policy'] = self.quant_policy
|
||||
if 'devices' in inspect.signature(TurbomindEngineConfig).parameters:
|
||||
engine_kwargs['devices'] = self.devices
|
||||
backend_config = TurbomindEngineConfig(**engine_kwargs)
|
||||
backend_config = autoget_backend_config(self.model_dir, backend_config)
|
||||
self.backend_config = backend_config
|
||||
logger.info(f'backend_config: {backend_config}')
|
||||
|
||||
pipeline_kwargs = {}
|
||||
is_multimodal = self.model_meta.is_multimodal
|
||||
if is_multimodal:
|
||||
require_version(
|
||||
'lmdeploy<0.9', 'LmdeployEngine will no longer maintain inference for '
|
||||
'multimodal models in lmdeploy>=0.9.')
|
||||
vision_config = VisionConfig(max_batch_size=self.vision_batch_size)
|
||||
pipeline_kwargs['vision_config'] = vision_config
|
||||
logger.info(f'vision_config: {vision_config}')
|
||||
self.pipeline_kwargs = pipeline_kwargs
|
||||
|
||||
@contextmanager
|
||||
def _patch_pipeline(self):
|
||||
_old_best_match_model = async_engine.best_match_model
|
||||
|
||||
def _best_match_model(*args, **kwargs) -> Optional[str]:
|
||||
return self.model_info.model_type
|
||||
|
||||
async_engine.best_match_model = _best_match_model
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
async_engine.best_match_model = _old_best_match_model
|
||||
|
||||
def _prepare_engine(self):
|
||||
with patch_auto_tokenizer(self.tokenizer), patch_auto_config(self.config), self._patch_pipeline():
|
||||
engine = pipeline(self.model_dir, backend_config=self.backend_config, **self.pipeline_kwargs)
|
||||
self.engine = engine
|
||||
|
||||
def _load_generation_config(self):
|
||||
generation_config_path = os.path.join(self.model_dir, 'generation_config.json')
|
||||
if os.path.isfile(generation_config_path):
|
||||
generation_config = GenerationConfig.from_pretrained(self.model_dir)
|
||||
kwargs = generation_config.to_dict()
|
||||
max_new_tokens = kwargs.get('max_new_tokens')
|
||||
if max_new_tokens is None:
|
||||
kwargs.pop('max_new_tokens', None)
|
||||
parameters = inspect.signature(LmdeployGenerationConfig).parameters
|
||||
for k, v in kwargs.copy().items():
|
||||
if k not in parameters or v is None:
|
||||
kwargs.pop(k)
|
||||
self.generation_config = LmdeployGenerationConfig(**kwargs)
|
||||
else:
|
||||
self.generation_config = LmdeployGenerationConfig()
|
||||
|
||||
def _add_stop_words(self, generation_config: LmdeployGenerationConfig, request_config: RequestConfig) -> None:
|
||||
template_meta = self.template.template_meta
|
||||
stop_words = (request_config.stop or []) + (self.generation_config.stop_words or []) + template_meta.stop_words
|
||||
generation_config.stop_words = self._get_stop_token_ids(stop_words)
|
||||
# compat lmdeploy >= 0.6.*
|
||||
generation_config.stop_token_ids = generation_config.stop_words
|
||||
|
||||
def _prepare_generation_config(self, request_config: RequestConfig) -> LmdeployGenerationConfig:
|
||||
kwargs = {'max_new_tokens': request_config.max_tokens}
|
||||
for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']:
|
||||
new_value = getattr(request_config, key)
|
||||
if new_value is None:
|
||||
kwargs[key] = getattr(self.generation_config, key)
|
||||
else:
|
||||
kwargs[key] = new_value
|
||||
if request_config.seed is None:
|
||||
request_config.seed = get_seed()
|
||||
kwargs['random_seed'] = request_config.seed
|
||||
if request_config.temperature == 0:
|
||||
kwargs['temperature'] = 1 # avoid unnecessary process
|
||||
kwargs['top_k'] = 1
|
||||
|
||||
if request_config.logprobs:
|
||||
kwargs['logprobs'] = 1
|
||||
if request_config.top_logprobs is not None:
|
||||
kwargs['logprobs'] = max(1, request_config.top_logprobs)
|
||||
|
||||
res = LmdeployGenerationConfig(**kwargs)
|
||||
return res
|
||||
|
||||
async def _infer_stream_async(
|
||||
self,
|
||||
inputs: Dict[str, Any],
|
||||
generation_config: LmdeployGenerationConfig,
|
||||
request_config: RequestConfig,
|
||||
) -> AsyncIterator[ChatCompletionStreamResponse]:
|
||||
session_id = time.time_ns()
|
||||
kwargs = {'stream_output': True, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True}
|
||||
if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
|
||||
async with self.engine.model_inst(session_id) as inst:
|
||||
context = self.engine.safe_run(inst, session_id, **inputs, **kwargs)
|
||||
else:
|
||||
context = self.engine.safe_run(session_id)
|
||||
|
||||
infer_streamer = InferStreamer(self.template, template_inputs=inputs['template_inputs'])
|
||||
token_idx = 0
|
||||
async with context as gen:
|
||||
if version.parse(lmdeploy.__version__) < version.parse('0.6.5'):
|
||||
generator = await self.engine.get_generator(False, session_id)
|
||||
gen = generator.async_stream_infer(session_id=session_id, **inputs, **kwargs)
|
||||
is_finished = False
|
||||
while not is_finished:
|
||||
try:
|
||||
output = await gen.__anext__()
|
||||
except StopAsyncIteration:
|
||||
is_finished = True
|
||||
delta_text = infer_streamer.get_printable_text(output.token_ids, is_finished)
|
||||
if not delta_text and not is_finished:
|
||||
continue
|
||||
|
||||
logprobs = self._get_logprobs(output.logprobs, output.token_ids[token_idx:],
|
||||
request_config.top_logprobs)
|
||||
token_idx = len(output.token_ids)
|
||||
|
||||
usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token)
|
||||
toolcall = None
|
||||
if is_finished:
|
||||
toolcall = self._get_toolcall(
|
||||
self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs']))
|
||||
finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token,
|
||||
output.status.name == 'FINISH')
|
||||
choices = [
|
||||
ChatCompletionResponseStreamChoice(
|
||||
index=0,
|
||||
delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall),
|
||||
finish_reason=finish_reason,
|
||||
logprobs=logprobs)
|
||||
]
|
||||
yield ChatCompletionStreamResponse(model=self.model_name, choices=choices, usage=usage_info)
|
||||
|
||||
async def _infer_full_async(
|
||||
self,
|
||||
inputs: Dict[str, Any],
|
||||
generation_config: LmdeployGenerationConfig,
|
||||
request_config: RequestConfig,
|
||||
) -> ChatCompletionResponse:
|
||||
session_id = time.time_ns()
|
||||
kwargs = {'stream_output': False, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True}
|
||||
if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
|
||||
async with self.engine.model_inst(session_id) as inst:
|
||||
async with self.engine.safe_run(inst, session_id, **inputs, **kwargs) as gen:
|
||||
async for output in gen:
|
||||
pass
|
||||
if self.engine.backend == 'pytorch':
|
||||
# manually end pytorch session
|
||||
await inst.async_end(session_id)
|
||||
|
||||
else:
|
||||
async with self.engine.safe_run(session_id):
|
||||
generator = await self.engine.get_generator(False, session_id)
|
||||
async for output in generator.async_stream_infer(session_id=session_id, **inputs, **kwargs):
|
||||
pass
|
||||
|
||||
response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs'])
|
||||
logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs)
|
||||
|
||||
usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token)
|
||||
toolcall = self._get_toolcall(response)
|
||||
finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token,
|
||||
output.status.name == 'FINISH')
|
||||
token_ids = output.token_ids if request_config.return_details else None
|
||||
choices = [
|
||||
ChatCompletionResponseChoice(
|
||||
index=0,
|
||||
message=ChatMessage(role='assistant', content=response, tool_calls=toolcall),
|
||||
finish_reason=finish_reason,
|
||||
logprobs=logprobs,
|
||||
token_ids=token_ids)
|
||||
]
|
||||
prompt_token_ids = None
|
||||
images_size = None
|
||||
if request_config.return_details:
|
||||
prompt_token_ids = inputs['input_ids']
|
||||
images = inputs['template_inputs'].images
|
||||
if all(isinstance(image, Image.Image) for image in images):
|
||||
images_size = [image.size for image in images]
|
||||
return ChatCompletionResponse(
|
||||
model=self.model_name,
|
||||
choices=choices,
|
||||
usage=usage_info,
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
images_size=images_size)
|
||||
|
||||
async def infer_async(self,
|
||||
infer_request: InferRequest,
|
||||
request_config: Optional[RequestConfig] = None,
|
||||
*,
|
||||
pre_infer_hook=None,
|
||||
**kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]:
|
||||
request_config = deepcopy(request_config or RequestConfig())
|
||||
self.template.set_mode('lmdeploy')
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
with torch.inference_mode():
|
||||
inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True)
|
||||
images = inputs.pop('images', None)
|
||||
if images:
|
||||
if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'):
|
||||
messages = self.engine._convert_prompts(('', images))
|
||||
messages = await self.engine.async_convert_to_pil_images(messages)
|
||||
results = await self.engine.vl_encoder.preprocess(messages)
|
||||
if self.engine.backend == 'turbomind':
|
||||
results = await self.engine.vl_encoder.async_infer(results)
|
||||
inputs['images'] = [result['content'] for result in results if result['role'] == 'forward'][0]
|
||||
await self.template.prepare_lmdeploy_turbomind_inputs(inputs)
|
||||
else:
|
||||
inputs['images'] = results[1]['content']
|
||||
await self.template.prepare_lmdeploy_pytorch_inputs(inputs)
|
||||
else:
|
||||
inputs['images'] = await self.engine.vl_encoder.async_infer(images)
|
||||
await self.template.prepare_lmdeploy_turbomind_inputs(inputs)
|
||||
|
||||
self.set_default_max_tokens(request_config, inputs)
|
||||
generation_config = self._prepare_generation_config(request_config)
|
||||
self._add_stop_words(generation_config, request_config)
|
||||
kwargs.update({'inputs': inputs, 'generation_config': generation_config, 'request_config': request_config})
|
||||
if pre_infer_hook:
|
||||
kwargs = pre_infer_hook(kwargs)
|
||||
if request_config.stream:
|
||||
return self._infer_stream_async(**kwargs)
|
||||
else:
|
||||
return await self._infer_full_async(**kwargs)
|
||||
|
||||
def _batch_infer_stream(self, *args, **kwargs):
|
||||
if hasattr(self.engine, 'vl_encoder'):
|
||||
self.engine.vl_encoder._loop_task = None
|
||||
if hasattr(self.engine, 'free_insts'):
|
||||
self.engine.free_insts = None
|
||||
return super()._batch_infer_stream(*args, **kwargs)
|
||||
|
||||
def infer(
|
||||
self,
|
||||
infer_requests: List[InferRequest],
|
||||
request_config: Optional[RequestConfig] = None,
|
||||
metrics: Optional[List[Metric]] = None,
|
||||
*,
|
||||
use_tqdm: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]:
|
||||
return super().infer(infer_requests, request_config, metrics, use_tqdm=use_tqdm, **kwargs)
|
||||
Reference in New Issue
Block a user