WIP: removed vllm, llm base class.

richwardle · richwardle · commit 07e6d9c18261 · 2024-11-20T17:45:06.000Z
diff --git a/prompting/llms/base_llm.py b/prompting/llms/base_llm.py
diff --git a/prompting/llms/hf_llm.py b/prompting/llms/hf_llm.py
@@ -1,209 +1,50 @@
-import time
-from typing import Optional, Any
-from prompting.utils.cleaners import CleanerPipeline
-from prompting.llms.base_llm import BaseLLM
-from transformers import AutoModelForCausalLM, AutoTokenizer, AwqConfig, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from loguru import logger
 import random
 import numpy as np
 import torch
 from prompting.utils.timer import Timer
+from prompting.settings import settings
 
-
-class HF_LLM(BaseLLM):
-    def __init__(
-        self,
-        llm: Any,
-        system_prompt,
-        max_new_tokens=256,
-        temperature=0.7,
-        top_p=0.95,
-    ):
-        model_kwargs = {
-            "temperature": temperature,
-            "top_p": top_p,
-            "max_tokens": max_new_tokens,
-        }
-        super().__init__(llm, system_prompt, model_kwargs)
-
-        # Keep track of generation data using messages and times
-        self.system_prompt = system_prompt
-        self.messages = [{"content": self.system_prompt, "role": "system"}] if self.system_prompt else []
-        self.times: list[float] = [0]
-        self._role_template = {
-            "system": "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
-            "user": "<|start_header_id|>user<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
-            "assistant": "<|start_header_id|>assistant<|end_header_id|>\n{{{{ {} }}}}<|eot_id|>",
-            "end": "<|start_header_id|>assistant<|end_header_id|>",
-        }
-
-    def query_conversation(
-        self,
-        messages: list[str],
-        roles: list[str],
-        cleaner: Optional[CleanerPipeline] = None,
-    ):
-        """Query LLM with the given lists of conversation history and roles
-
-        Args:
-            messages (list[str]): List of messages in the conversation.
-            roles (list[str]): List of roles for each message.
-            cleaner (Optional[CleanerPipeline], optional): Cleaner pipeline to use, if any.
-        """
-        assert len(messages) == len(roles), "Length of messages and roles must be the same"
-        inputs: list[dict[str, Any]] = [{"content": self.system_prompt, "role": "system"}]
-        for role, message in zip(roles, messages):
-            inputs.append({"content": message, "role": role})
-
-        t0 = time.perf_counter()
-        response = self.forward(messages=inputs)
-        response = self.clean_response(cleaner, response)
-        self.times.extend((0, time.perf_counter() - t0))
-        return response
-
-    def query(
-        self,
-        message: list[str],
-        role: str = "user",
-        cleaner: CleanerPipeline = CleanerPipeline(),
-    ):
-        # Adds the message to the list of messages for tracking purposes, even though it's not used downstream
-        messages = self.messages + [{"content": message, "role": role}]
-
-        t0 = time.time()
-        response = self._forward(messages=messages)
-        response = self.clean_response(cleaner, response)
-
-        self.messages = messages
-        self.messages.append({"content": response, "role": "assistant"})
-        self.times.extend((0, time.time() - t0))
-
-        return response
-
-    def _make_prompt(self, messages: list[dict[str, str]]) -> str:
-        composed_prompt: list[str] = []
-
-        for message in messages:
-            role = message["role"]
-            if role not in self._role_template:
-                continue
-            content = message["content"]
-            composed_prompt.append(self._role_template[role].format(content))
-
-        # Adds final tag indicating the assistant's turn
-        composed_prompt.append(self._role_template["end"])
-        return "".join(composed_prompt)
-
-    def _forward(self, messages: list[dict[str, str]]):
-        # make composed prompt from messages
-        composed_prompt = self._make_prompt(messages)
-        response = self.llm.generate(
-            composed_prompt,
-            max_length=self.model_kwargs["max_tokens"],
-            temperature=self.model_kwargs["temperature"],
-            top_p=self.model_kwargs["top_p"],
-        )[0]
-
-        try:
-            logger.info(
-                f"{self.__class__.__name__} generated the following output:\n{response['generated_text'].strip()}"
-            )
-        except Exception as e:
-            logger.info(f"Response: {response}")
-            logger.error(f"Error logging the response: {e}")
-
-        return response["generated_text"].strip()
-
-
-def set_random_seeds(seed=42):
-    """
-    Set random seeds for reproducibility across all relevant libraries
-    """
-    if seed is not None:
-        random.seed(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        torch.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
-        torch.backends.cudnn.deterministic = True
-        torch.backends.cudnn.benchmark = False
-
-
-class ReproducibleHF:
-    def __init__(self, model_id="Qwen/Qwen2-0.5B", tensor_parallel_size=0, seed=42, **kwargs):
+class ReproducibleHF():
+    def __init__(self, model_id="hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4", settings=None, **kwargs):
         """
         Initialize Hugging Face model with reproducible settings and optimizations
         """
-        self.set_random_seeds(seed)
-
-        # Load model and tokenizer with optimizations
-        model_kwargs = {
-            "device_map": "auto",
-        }
-
-        # get valid params for generation from model config
-        self.valid_generation_params = set(
-            AutoModelForCausalLM.from_pretrained(model_id).generation_config.to_dict().keys()
-        )
-
-        for k, v in kwargs.items():
-            if k not in ["sampling_params"]:  # exclude sampling_params and any other generation-only args
-                model_kwargs[k] = v
-
-        quantization_config = AwqConfig(
-            bits=4,
-            fuse_max_seq_len=512,
-            do_fuse=True,
-        )
-
+        self.seed = self.set_random_seeds(42)
+        quantization_config = settings.QUANTIZATION_CONFIG.get(model_id, None)
+        
         self.model = AutoModelForCausalLM.from_pretrained(
             model_id,
             torch_dtype=torch.float16,
             low_cpu_mem_usage=True,
-            device_map="auto",
+            device_map="cuda:0",
             quantization_config=quantization_config,
         )
-
+        
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+    
+        self.valid_generation_params = set(
+            AutoModelForCausalLM.from_pretrained(model_id).generation_config.to_dict().keys()
+        )
 
-        # self.model.generation_config.cache_implementation = "static"
-        # self.model.forward = torch.compile(self.model.forward, mode="reduce-overhead", fullgraph=True)
-        # self.valid_generation_params = set(self.model.generation_config.to_dict().keys())
-
-        # Enable model optimizations
-        self.model.eval()
-
-        if tensor_parallel_size > 1:
-            self.model = torch.nn.DataParallel(self.model, device_ids=list(range(tensor_parallel_size)))
-
-        # Create pipeline with optimized settings
         self.llm = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)
 
-        # Default sampling parameters
-        self.sampling_params = {
-            "temperature": 0.7,
-            "top_p": 0.95,
-            "top_k": 50,
-            "max_new_tokens": 256,
-            "presence_penalty": 0,
-            "frequency_penalty": 0,
-            "seed": seed,
-            "do_sample": True,
-            "early_stopping": True,  # Enable early stopping
-            "num_beams": 1,  # Use greedy decoding by default
-        }
+        self.sampling_params = settings.SAMPLING_PARAMS
 
     @torch.inference_mode()
     def generate(self, prompts, sampling_params=None):
         """
         Generate text with optimized performance
         """
-
-        # Convert single prompt to list
-        if isinstance(prompts, str):
-            prompts = [prompts]
-
-        inputs = self.tokenizer(prompts, truncation=True, return_tensors="pt").to(self.model.device)
+        
+        inputs = self.tokenizer.apply_chat_template(
+                prompts,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_tensors="pt",
+                return_dict=True,
+                ).to(settings.NEURON_DEVICE)
 
         params = sampling_params if sampling_params else self.sampling_params
         filtered_params = {k: v for k, v in params.items() if k in self.valid_generation_params}
@@ -215,9 +56,9 @@ def generate(self, prompts, sampling_params=None):
                 **filtered_params,
                 eos_token_id=self.tokenizer.eos_token_id,
             )
-
-            results = self.tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-            results = [text.strip() for text in results]
+            
+            outputs = self.model.generate(**inputs, **filtered_params, eos_token_id=self.tokenizer.eos_token_id,)
+            results = self.tokenizer.batch_decode(outputs[:, inputs['input_ids'].shape[1]:], skip_special_tokens=True, )[0]
 
         logger.debug(
             f"PROMPT: {prompts}\n\nRESPONSES: {results}\n\n"
diff --git a/prompting/llms/model_manager.py b/prompting/llms/model_manager.py
@@ -15,7 +15,6 @@
 # through the tasks based on the currently loaded model.
 open_tasks = []
 
-
 class ModelManager(BaseModel):
     always_active_models: list[ModelConfig] = []
     total_ram: float = settings.LLM_MODEL_RAM
@@ -75,7 +74,6 @@ def load_model(self, model_config: ModelConfig, force: bool = True):
             self.active_models[model_config] = model
             self.used_ram += model_config.min_ram
             logger.info(f"Model {model_config.llm_model_id} loaded. Current used RAM: {self.used_ram} GB")
-
             return model
         except Exception as e:
             logger.exception(f"Failed to load model {model_config.llm_model_id}. Error: {str(e)}")
@@ -148,14 +146,7 @@ def generate(
             model = ModelZoo.get_random(max_ram=self.total_ram)
 
         model_instance: ReproducibleHF = self.get_model(model)
-
-        valid_args = {"max_length", "temperature", "top_p", "min_length", "do_sample", "num_return_sequences"}
-        if sampling_params:
-            sampling_params = {k: v for k, v in sampling_params.items() if k in valid_args}
-        else:
-            sampling_params = {"max_length": settings.NEURON_MAX_TOKENS}
-
-        responses = model_instance.generate(prompts=[composed_prompt], sampling_params=sampling_params)
+        responses = model_instance.generate(prompts=[composed_prompt])
 
         return responses
 
diff --git a/prompting/settings.py b/prompting/settings.py
@@ -8,6 +8,7 @@
 from loguru import logger
 from pydantic import Field, model_validator
 from pydantic_settings import BaseSettings
+from transformers import AwqConfig
 
 from prompting.utils.config import config
 
@@ -81,6 +82,7 @@ class Settings(BaseSettings):
     MAX_ALLOWED_VRAM_GB: int = Field(62, env="MAX_ALLOWED_VRAM_GB")
     LLM_MAX_MODEL_LEN: int = Field(4096, env="LLM_MAX_MODEL_LEN")
     LLM_MODEL: str = Field("hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4", env="LLM_MODEL")
+    SAMPLING_PARAMS: dict[str, Any] = {"temperature": 0.7, "top_p": 0.95, "top_k": 50, "max_new_tokens": 256, "do_sample" : True, "seed": None}
     MINER_LLM_MODEL: Optional[str] = Field(None, env="MINER_LLM_MODEL")
     LLM_MODEL_RAM: float = Field(70, env="LLM_MODEL_RAM")
     OPENAI_API_KEY: str | None = Field(None, env="OPENAI_API_KEY")
@@ -178,6 +180,11 @@ def complete_settings(cls, values: dict[str, Any]) -> dict[str, Any]:
                 "You must provide an OpenAI API key as a backup. It is recommended to also provide an SN19 API key + url to avoid incurring API costs."
             )
         return values
+    
+    @cached_property
+    def QUANTIZATION_CONFIG(self) -> AwqConfig:
+        configs = {"hugging-quants/Meta-Llama-3.1-70B-Instruct-AWQ-INT4" : AwqConfig(bits=4, fuse_max_seq_len=512, do_fuse=True)}
+        return configs
 
     @cached_property
     def WALLET(self) -> bt.wallet:
@@ -206,6 +213,7 @@ def METAGRAPH(self) -> bt.metagraph:
     def DENDRITE(self) -> bt.dendrite:
         logger.info(f"Instantiating dendrite with wallet: {self.WALLET}")
         return bt.dendrite(wallet=self.WALLET)
+    
 
 
 settings: Optional[Settings] = None
diff --git a/prompting/tasks/base_task.py b/prompting/tasks/base_task.py
@@ -3,7 +3,6 @@
 from loguru import logger
 from abc import ABC
 from pydantic import BaseModel, Field, ConfigDict, model_validator
-from prompting.llms.vllm_llm import vLLM_LLM
 from prompting.utils.cleaners import CleanerPipeline
 from typing import ClassVar
 from prompting.datasets.base import DatasetEntry
@@ -78,12 +77,10 @@ def make_reference(self, dataset_entry: DatasetEntry) -> str:
     def generate_reference(self, messages: list[str]) -> str:
         """Generates a reference answer to be used for scoring miner completions"""
         logger.info("🤖 Generating reference...")
-        self.reference = vLLM_LLM(
-            llm=model_manager.get_model(self.llm_model).llm, system_prompt=self.reference_system_prompt or ""
-        ).query(cleaner=self.cleaner, message=messages)
-        # self.reference = model_manager.get_model(self.llm_model).generate(prompts=messages)
+        self.reference = model_manager.get_model(settings.LLM_MODEL).generate(prompts=messages)
         if self.reference is None:
             raise Exception("Reference generation failed")
+                
         return self.reference
 
     def generate_query(
diff --git a/prompting/tasks/inference.py b/prompting/tasks/inference.py
