Multi-GPU Inference

from transformers import PreTrainedTokenizer, GenerationConfig
from typing import Optional, List, Tuple
import copy
from transformers import AutoTokenizer
from packaging import version

HistoryType = List[Tuple[str, str]]

class vLLMWrapper:
    def __init__(self,
               model_dir: str,
               trust_remote_code: bool = True,
               tensor_parallel_size: int = 1,
               gpu_memory_utilization: float = 0.98,
               dtype: str = "bfloat16",
               **kwargs):

        if dtype not in ("bfloat16", "float16", "float32"):
            raise ValueError(f"Unsupported dtype: {dtype}")

        # Build generation_config
        self.generation_config = GenerationConfig.from_pretrained(
            model_dir, trust_remote_code=trust_remote_code
        )

        # Build tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_dir, trust_remote_code=True
        )
        self.tokenizer.eos_token_id = self.generation_config.eos_token_id

        from vllm import LLM
        import vllm
        
        if version.parse(vllm.__version__) >= version.parse("0.2.2"):
            self.__vllm_support_repetition_penalty = True
        else:
            self.__vllm_support_repetition_penalty = False

        quantization = getattr(kwargs, 'quantization', None)

        self.model = LLM(
            model=model_dir,
            tokenizer=model_dir,
            tensor_parallel_size=tensor_parallel_size,
            trust_remote_code=trust_remote_code,
            quantization=quantization,
            gpu_memory_utilization=gpu_memory_utilization,
            dtype=dtype
        )

    def chat(self,
        query: str,
        history: Optional[HistoryType],
        tokenizer: PreTrainedTokenizer = None,
        system: str = "You are a helpful assistant.",
        generation_config: Optional[GenerationConfig] = None,
        **kwargs):
        
        generation_config = generation_config if generation_config is not None else self.generation_config
        tokenizer = self.tokenizer if tokenizer is None else tokenizer

        if history is None:
            history = []
        else:
            history = copy.deepcopy(history)

        from vllm.sampling_params import SamplingParams
        
        sampling_kwargs = {
            "stop_token_ids": [self.generation_config.eos_token_id],
            "early_stopping": False,
            "top_p": generation_config.top_p,
            "top_k": -1 if generation_config.top_k == 0 else generation_config.top_k,
            "temperature": generation_config.temperature,
            "max_tokens": generation_config.max_new_tokens,
        }
        
        if self.__vllm_support_repetition_penalty:
            sampling_kwargs["repetition_penalty"] = generation_config.repetition_penalty
        
        sampling_params = SamplingParams(**sampling_kwargs)

        # Format prompt with history
        prompt = self._make_context(query, history, system)
        
        req_outputs = self.model.generate(
            [prompt],
            sampling_params=sampling_params
        )
        
        response = req_outputs[0].outputs[0].text
        history.append((query, response))

        return response, history

    def _make_context(self, query, history, system):
        """Format query with history into prompt."""
        prompt = f"{system}\n\n"
        for q, r in history:
            prompt += f"User: {q}\nAssistant: {r}\n\n"
        prompt += f"User: {query}\nAssistant:"
        return prompt

# Usage
if __name__ == '__main__':
    model_dir = 'Qwen/Qwen-72B-Chat'
    tensor_parallel_size = 2

    model = vLLMWrapper(
        model_dir,
        tensor_parallel_size=tensor_parallel_size,
    )

    response, history = model.chat(
        query="你好",
        history=None
    )
    print(response)
    
    response, history = model.chat(
        query="给我讲一个年轻人奋斗创业最终取得成功的故事。",
        history=history
    )
    print(response)
    
    response, history = model.chat(
        query="给这个故事起一个标题",
        history=history
    )
    print(response)