[Feature] Integrated Training and Inference -- Part 1

xtuner/chat/backend/__init__.py

@@ -0,0 +1,4 @@
+from .huggingface import HFBackend
+from .lmdeploy import LMDeployBackend
+
+__all__ = ['HFBackend', 'LMDeployBackend']

xtuner/chat/backend/base.py

@@ -0,0 +1,30 @@
+from abc import abstractmethod
+from typing import List, Optional
+
+from xtuner.chat.streamer import SteamerType
+from xtuner.types import (ChatBackendProtocol, ChatMessages, ChatTemplate,
+                          SampleParams)
+
+
+class BaseBackend(ChatBackendProtocol):
+
+    @property
+    def chat_template(self) -> ChatTemplate:
+        pass
+
+    @abstractmethod
+    def create_streamer(self, iterable: bool = False) -> SteamerType:
+        pass
+
+    @abstractmethod
+    def chat(self,
+             messages: ChatMessages,
+             sample_params: Optional[SampleParams] = None,
+             streamer: Optional[SteamerType] = None):
+        pass
+
+    @abstractmethod
+    def batch_infer(self,
+                    messages: List[ChatMessages],
+                    sample_params: Optional[SampleParams] = None):
+        pass

xtuner/chat/backend/huggingface.py

@@ -0,0 +1,153 @@
+from typing import List, Optional
+
+import torch
+from peft import PeftModel
+from transformers import (AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig)
+from transformers import GenerationConfig as HFGenerationConfig
+from transformers import PreTrainedModel, PreTrainedTokenizer
+
+from xtuner.chat.streamer import HFTextIteratorStreamer, HFTextStreamer
+from xtuner.model.utils import LoadWoInit
+from xtuner.tools.utils import get_stop_criteria
+from xtuner.types import ChatMessages, ChatTemplate, SampleParams
+from .base import BaseBackend
+
+
+class HFBackend(BaseBackend):
+
+    def __init__(
+        self,
+        chat_template: ChatTemplate,
+        llm: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer,
+    ) -> None:
+        super().__init__()
+
+        self.llm = llm
+        self.llm.cuda()
+        self.tokenizer = tokenizer
+
+        self._chat_template = chat_template
+
+    @property
+    def chat_template(self) -> ChatTemplate:
+        return self._chat_template
+
+    @property
+    def eos_token_id(self):
+        if self.tokenizer.pad_token_id:
+            return self.tokenizer.eos_token_id
+        else:
+            return self.tokenizer.eos_token_id
+
+    @property
+    def pad_token_id(self):
+        return self.tokenizer.pad_token_id
+
+    def build_llm_and_tokenizer(self,
+                                model_name_or_path,
+                                adapter=None,
+                                bits=None):
+
+        if bits is None:
+            quantization_config = None
+            load_in_8bit = False
+        elif bits == 4:
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                load_in_8bit=False,
+                llm_int8_threshold=6.0,
+                llm_int8_has_fp16_weight=False,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type='nf4')
+            load_in_8bit = False
+        elif bits == 8:
+            quantization_config = None
+            load_in_8bit = True
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name_or_path,
+            trust_remote_code=True,
+            encode_special_tokens=True)
+
+        with LoadWoInit():
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name_or_path,
+                device_map='auto',
+                load_in_8bit=load_in_8bit,
+                quantization_config=quantization_config,
+                trust_remote_code=True,
+                torch_dtype=torch.float16)
+
+        if adapter is not None:
+            model = PeftModel.from_pretrained(model, adapter)
+
+        model.eval()
+        return model, tokenizer
+
+    def create_streamer(self, iterable=False):
+        if iterable:
+            return HFTextIteratorStreamer(
+                self.tokenizer,
+                skip_prompt=True,
+                chat_template=self.chat_template)
+        else:
+            return HFTextStreamer(
+                self.tokenizer,
+                skip_prompt=True,
+                chat_template=self.chat_template)
+
+    def parse_sample_params(self, params: SampleParams) -> HFGenerationConfig:
+
+        if params is None:
+            params = SampleParams()
+
+        hf_gen_config = HFGenerationConfig(
+            max_new_tokens=params.max_new_tokens,
+            do_sample=params.temperature > 0,
+            temperature=params.temperature,
+            top_k=params.top_k,
+            top_p=params.top_p,
+            repetition_penalty=params.repetition_penalty,
+            seed=params.seed,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.pad_token_id)
+
+        stop_words = params.stop_words
+        stop_words.extend(self.chat_template.stop_words)
+
+        return hf_gen_config, stop_words
+
+    def chat(self,
+             messages: ChatMessages,
+             streamer=None,
+             sample_params: Optional[SampleParams] = None):
+
+        prompt = messages.get_prompt(self.chat_template)
+        ids = self.tokenizer.encode(prompt, return_tensors='pt')
+
+        hf_gen_config, stop_words = self.parse_sample_params(sample_params)
+
+        stop_criteria = get_stop_criteria(
+            tokenizer=self.tokenizer, stop_words=stop_words)
+
+        generate_output = self.llm.generate(
+            inputs=ids.cuda(),
+            streamer=streamer,
+            generation_config=hf_gen_config,
+            stopping_criteria=stop_criteria)
+
+        output = self.tokenizer.decode(
+            generate_output[0][len(ids[0]):], skip_special_tokens=True)
+
+        for word in stop_words:
+            output = output.rstrip(word)
+
+        return output
+
+    def batch_infer(self,
+                    messages: List[ChatMessages],
+                    sample_params: SampleParams | None = None):
+        raise NotImplementedError

xtuner/chat/backend/lmdeploy/__init__.py

@@ -0,0 +1,3 @@
+from .backend import LMDeployBackend
+
+__all__ = ['LMDeployBackend']

xtuner/chat/backend/lmdeploy/_engine.py

@@ -0,0 +1,27 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from lmdeploy.serve.async_engine import AsyncEngine
+
+from xtuner.types import ChatMessages, ChatTemplate
+
+
+class _AsyncEngine(AsyncEngine):
+    """Async inference engine."""
+
+    def __init__(self, chat_template: ChatTemplate, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        assert self.model_name == 'base'
+        self.chat_template = chat_template
+
+    async def _get_prompt_input(self, prompt: ChatMessages,
+                                do_preprocess: bool, sequence_start: bool):
+        """get input_ids, embeddings and offsets."""
+
+        decorated = prompt.get_prompt(self.chat_template)
+
+        results = {}
+
+        input_ids = self.tokenizer.encode(decorated, add_bos=sequence_start)
+
+        results['input_ids'] = input_ids
+        results['prompt'] = decorated
+        return results

xtuner/chat/backend/lmdeploy/backend.py

@@ -0,0 +1,94 @@
+import asyncio
+import os
+from typing import List, Optional, Union
+
+from lmdeploy.utils import get_logger
+
+from xtuner.types import ChatMessages, ChatTemplate, SampleParams
+from ...streamer import LMDeployTextIteratorStreamer, LMDeployTextStreamer
+from ..base import BaseBackend
+from ._engine import _AsyncEngine
+
+os.environ['TM_LOG_LEVEL'] = 'ERROR'
+logger = get_logger('lmdeploy')
+logger.setLevel('ERROR')
+
+_StreamerType = Union[LMDeployTextStreamer, LMDeployTextIteratorStreamer]
+
+
+class LMDeployBackend(BaseBackend):
+
+    def __init__(self, chat_template, llm_name_or_path) -> None:
+        super().__init__()
+
+        self._engine = _AsyncEngine(
+            chat_template, model_path=llm_name_or_path, model_name='base')
+
+        self._chat_template = chat_template
+
+    @property
+    def chat_template(self) -> ChatTemplate:
+        return self._chat_template
+
+    def create_streamer(self, iterable=False):
+
+        if iterable:
+            return LMDeployTextIteratorStreamer()
+        else:
+            return LMDeployTextStreamer()
+
+    def parse_sample_params(self, params: SampleParams):
+
+        if params is None:
+            params = SampleParams()
+
+        stop_words = params.stop_words
+        stop_words.extend(self.chat_template.stop_words)
+
+        from lmdeploy.messages import GenerationConfig as LMDGenerationConfig
+        lmd_gen_config = LMDGenerationConfig(
+            max_new_tokens=params.max_new_tokens,
+            temperature=params.temperature,
+            top_k=params.top_k,
+            top_p=params.top_p,
+            repetition_penalty=params.repetition_penalty,
+            random_seed=params.seed,
+            stop_words=stop_words)
+
+        return lmd_gen_config
+
+    def chat(self,
+             messages: ChatMessages,
+             streamer: Optional[_StreamerType] = None,
+             sample_params: Optional[SampleParams] = None):
+
+        lmd_gen_config = self.parse_sample_params(sample_params)
+        self.session_id += 1
+        import random
+
+        generator = self._engine.generate(
+            messages, random.randint(1, 100000), gen_config=lmd_gen_config)
+
+        async def get_response():
+            out = ''
+            async for res in generator:
+                out += res.response
+                if streamer:
+                    streamer.put(res.response)
+            if streamer:
+                streamer.end()
+            return out
+
+        loop = asyncio.new_event_loop()
+        response = loop.run_until_complete(get_response())
+        return response
+
+    def batch_infer(self,
+                    messages: List[ChatMessages],
+                    sample_params: Optional[SampleParams] = None):
+
+        lmd_gen_config = self.parse_sample_params(sample_params)
+
+        results = self._engine.batch_infer(messages, gen_config=lmd_gen_config)
+
+        return [r.text for r in results]

xtuner/chat/streamer/__init__.py

@@ -0,0 +1,12 @@
+from typing import Union
+
+from .huggingface import HFTextIteratorStreamer, HFTextStreamer
+from .lmdeploy import LMDeployTextIteratorStreamer, LMDeployTextStreamer
+
+SteamerType = Union[HFTextIteratorStreamer, HFTextStreamer,
+                    LMDeployTextIteratorStreamer, LMDeployTextStreamer]
+
+__all__ = [
+    'HFTextIteratorStreamer', 'HFTextStreamer', 'LMDeployTextIteratorStreamer',
+    'LMDeployTextStreamer'
+]

xtuner/chat/streamer/huggingface.py

@@ -0,0 +1,37 @@
+from transformers import TextIteratorStreamer, TextStreamer
+from transformers.models.auto import AutoTokenizer
+
+
+class HFTextIteratorStreamer(TextIteratorStreamer):
+
+    def __init__(self,
+                 tokenizer: AutoTokenizer,
+                 skip_prompt: bool = False,
+                 timeout=None,
+                 chat_template=None,
+                 **decode_kwargs):
+        super().__init__(tokenizer, skip_prompt, timeout, **decode_kwargs)
+        self.chat_template = chat_template
+
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+
+        for word in self.chat_template.stop_words:
+            text = text.rstrip(word)
+        super().on_finalized_text(text, stream_end)
+
+
+class HFTextStreamer(TextStreamer):
+
+    def __init__(self,
+                 tokenizer: AutoTokenizer,
+                 skip_prompt: bool = False,
+                 chat_template=None,
+                 **decode_kwargs):
+        super().__init__(tokenizer, skip_prompt, **decode_kwargs)
+        self.chat_template = chat_template
+
+    def on_finalized_text(self, text: str, stream_end: bool = False):
+
+        for word in self.chat_template.stop_words:
+            text = text.rstrip(word)
+        super().on_finalized_text(text, stream_end)