Merge pull request #26 from amazon-science/vllm

[#23] 🚀 Added support for vllm language models.

Author: SibaRajendran

src/fmcore/algorithm/vllm.py

@@ -0,0 +1,175 @@
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+from bears import FileMetadata
+from bears.util import EnvUtil, get_default, optional_dependency, set_param_from_alias
+from pydantic import confloat, conint, model_validator
+
+from fmcore.framework._task.text_generation import (
+    GENERATED_TEXTS_COL,
+    GenerativeLM,
+    NextTokens,
+    Prompts,
+    TextGenerationParams,
+    TextGenerationParamsMapper,
+)
+
+with optional_dependency("vllm"):
+    from vllm import LLM, SamplingParams
+
+    class VLLMGenerativeLM(GenerativeLM):
+        aliases = ["vllm"]
+
+        llm: Optional[LLM] = None
+        cache_dir: Optional[Union[FileMetadata, Dict, str]] = None
+
+        class Hyperparameters(GenerativeLM.Hyperparameters):
+            model_name: str
+            tensor_parallel_size: Optional[conint(ge=1)] = None
+            gpu_memory_utilization: confloat(gt=0.0, le=1.0) = 0.95
+            max_model_len: conint(ge=1)
+            generation_params: Union[TextGenerationParams, Dict, str]
+
+            @model_validator(mode="before")
+            @classmethod
+            def set_params(cls, params: Dict) -> Dict:
+                set_param_from_alias(
+                    params,
+                    param="model_name",
+                    alias=[
+                        "model",
+                        "pretrained_model_name_or_path",
+                        "model_name_or_path",
+                    ],
+                )
+                set_param_from_alias(
+                    params,
+                    param="max_model_len",
+                    alias=[
+                        "max_len",
+                        "max_model_len",
+                        "max_sequence_length",
+                        "max_sequence_len",
+                        "max_input_length",
+                        "max_input_len",
+                    ],
+                )
+                params["generation_params"] = TextGenerationParamsMapper.of(
+                    params["generation_params"]
+                ).initialize()
+                if params.get("cache_dir") is not None:
+                    params["cache_dir"] = FileMetadata.of(params["cache_dir"])
+                return params
+
+        def initialize(self, model_dir: Optional[FileMetadata] = None):
+            """Initialize the VLLM model"""
+            tensor_parallel_size: Optional[conint(ge=1)] = get_default(
+                self.hyperparams.tensor_parallel_size,
+                EnvUtil.num_gpus(),  # Use all GPUs by default
+            )
+
+            kwargs = dict(
+                model=self.hyperparams.model_name,
+                tensor_parallel_size=tensor_parallel_size,
+                gpu_memory_utilization=self.hyperparams.gpu_memory_utilization,
+                max_model_len=self.hyperparams.max_model_len,
+            )
+
+            if self.cache_dir is not None:
+                kwargs["download_dir"] = self.cache_dir.path
+
+            print(f"Initializing vllm with kwargs: {kwargs}")
+            self.llm = LLM(**kwargs)
+
+        def predict_step(self, batch: Prompts, **kwargs) -> Dict:
+            """Run prediction on a batch of prompts"""
+            prompts: List[str] = batch.prompts().to_list()
+
+            sampling_params = SamplingParams(
+                min_tokens=self.hyperparams.generation_params.min_new_tokens,
+                max_tokens=self.hyperparams.generation_params.max_new_tokens,
+                temperature=0.0
+                if not self.hyperparams.generation_params.do_sample
+                else self.hyperparams.generation_params.temperature,
+                top_p=self.hyperparams.generation_params.top_p
+                if hasattr(self.hyperparams.generation_params, "top_p")
+                else 1.0,
+                top_k=self.hyperparams.generation_params.top_k
+                if hasattr(self.hyperparams.generation_params, "top_k")
+                else -1,
+                stop=self.hyperparams.generation_params.stop_sequences,
+                logprobs=self.hyperparams.generation_params.output_scores,
+            )
+            outputs = self.llm.generate(
+                prompts,
+                sampling_params=sampling_params,
+            )
+
+            result = {GENERATED_TEXTS_COL: [output.outputs[0].text for output in outputs]}
+
+            if self.hyperparams.generation_params.output_scores:
+                # Get token IDs and logprobs for each generation
+                token_ids = []
+                tokens = []
+                token_scores = []
+
+                for output in outputs:
+                    # Get the first (and only) generation
+                    generation = output.outputs[0]
+
+                    # Extract token IDs, tokens and logprobs
+                    gen_token_ids = generation.token_ids
+                    gen_tokens = generation.tokens
+                    gen_logprobs = generation.logprobs
+
+                    # Convert scores based on output_scores_format
+                    if self.hyperparams.generation_params.output_scores_format == "probabilities":
+                        # Convert from log probabilities to probabilities
+                        gen_logprobs = np.exp(gen_logprobs)
+                        # Filter based on tolerance
+                        if self.hyperparams.generation_params.output_scores_tolerance is not None:
+                            mask = gen_logprobs >= self.hyperparams.generation_params.output_scores_tolerance
+                            gen_token_ids = [t for t, m in zip(gen_token_ids, mask) if m]
+                            gen_tokens = [t for t, m in zip(gen_tokens, mask) if m]
+                            gen_logprobs = [s for s, m in zip(gen_logprobs, mask) if m]
+
+                    elif self.hyperparams.generation_params.output_scores_format == "log-probabilities":
+                        # Already in log probabilities format
+                        # Filter based on tolerance
+                        if self.hyperparams.generation_params.output_scores_tolerance is not None:
+                            mask = gen_logprobs >= self.hyperparams.generation_params.output_scores_tolerance
+                            gen_token_ids = [t for t, m in zip(gen_token_ids, mask) if m]
+                            gen_tokens = [t for t, m in zip(gen_tokens, mask) if m]
+                            gen_logprobs = [s for s, m in zip(gen_logprobs, mask) if m]
+
+                    elif self.hyperparams.generation_params.output_scores_format == "logits":
+                        # Don't filter or modify scores when using raw logits
+                        pass
+
+                    token_ids.append(gen_token_ids)
+                    tokens.append(gen_tokens)
+                    token_scores.append(gen_logprobs)
+
+                result.update(
+                    {
+                        "generated_token_ids": token_ids,
+                        "generated_tokens": tokens,
+                        "generated_token_scores": token_scores,
+                    }
+                )
+
+            return result
+
+        def _create_predictions(self, batch: Prompts, predictions: Any, **kwargs) -> NextTokens:
+            """Convert raw predictions to NextTokens format"""
+            return NextTokens.from_task_data(data=batch, predictions=predictions, **kwargs)
+
+        @property
+        def max_num_generated_tokens(self) -> int:
+            return self.hyperparams.generation_params.max_new_tokens
+
+        def cleanup(self):
+            """Cleanup the llm"""
+            if self.llm is not None:
+                del self.llm
+                self.llm = None

src/fmcore/framework/_evaluator/Evaluator.py

@@ -369,7 +369,7 @@ def evaluate(
         )
 
         Alias.set_return_predictions(kwargs)
-        return_predictions: bool = kwargs.pop("return_predictions", False)
+        return_predictions: bool = kwargs.pop("return_predictions", True)
 
         Alias.set_predictions_destination(kwargs)
         predictions_destination: Optional[Union[io.IOBase, FileMetadata, Dict, str]] = kwargs.pop(
@@ -419,9 +419,8 @@ def evaluate(
         kwargs["tracker"]: Tracker = Tracker.of(kwargs["tracker"])
 
         try:
-            self._evaluator_is_running: bool = (
-                True  ## Ensures we do not accidentally delete the models while running.
-            )
+            ## Ensures we do not accidentally delete the models while running.
+            self._evaluator_is_running: bool = True
             evaluated_predictions, evaluated_metrics = self._run_evaluation(
                 dataset,
                 metrics=metrics,
@@ -430,6 +429,8 @@ def evaluate(
                 progress_bar=progress_bar,
                 **kwargs,
             )
+        except Exception as e:
+            raise e
         finally:
             if self.cache_timeout is not None:  ## Rests the timeout
                 self.cache_timeout.reset_timeout()

src/fmcore/framework/_evaluator/RayEvaluator.py

@@ -126,11 +126,14 @@ def __init__(
             self.actor = actor
             self.request_counter: RequestCounter = request_counter
 
-        def is_available(self) -> bool:
+        def get_evaluator_status(self) -> str:
             try:
-                return self.evaluator is None
-            except Exception:
-                return False
+                if self.evaluator is None:
+                    return "Evaluator not initialized."
+                assert isinstance(self.evaluator, Evaluator)
+                return (self.evaluator.class_name, self.evaluator.model.class_name)
+            except Exception as e:
+                return String.format_exception_msg(e)
 
         def get_ip_address(self) -> Optional[str]:
             try:
@@ -363,6 +366,10 @@ def ray_evaluator_params(cls, params: Dict) -> Dict:
             return params
 
         def initialize(self, reinit_ray: bool = False, **kwargs):
+            if self.model_num_gpus <= 1 or self.AlgorithmClass.class_name == "VLLMGenerativeLM":
+                self.nested_evaluator_name: str = get_default(self.nested_evaluator_name, "local")
+            else:
+                self.nested_evaluator_name: str = get_default(self.nested_evaluator_name, "accelerate")
             ## Connect to the Ray cluster
             if not ray.is_initialized() or reinit_ray is True:
                 ray.init(
@@ -385,7 +392,7 @@ def _load_model(
             **kwargs,
         ) -> List[RayActorComposite]:
             num_actors: int = get_default(num_actors, self.num_actors)
-            progress_bar: Optional[Dict] = self._run_evaluation_progress_bar(progress_bar)
+            progress_bar: Union[Dict, bool] = self._run_evaluation_progress_bar(progress_bar)
             nested_evaluator_params: Dict = self._create_nested_evaluator_params(**kwargs)
 
             def actor_factory(*, request_counter: Any, actor_i: int, actor_id: str, **kwargs):
@@ -475,10 +482,7 @@ def num_actors(self) -> int:
             return num_actors
 
         def _create_nested_evaluator_params(self, **kwargs) -> Dict:
-            nested_evaluator_name: str = get_default(
-                self.nested_evaluator_name,
-                "accelerate" if self.model_num_gpus > 1 else "local",
-            )
+            nested_evaluator_name: str = self.nested_evaluator_name
             if self.model_dir is not None and not self.model_dir.is_remote_storage():
                 raise ValueError(
                     f"When passing `model_dir` to {self.class_name}.of(...), the model directory "
@@ -563,44 +567,44 @@ def _run_evaluation(
             evaluated_predictions: Optional[Predictions] = None
             evaluated_metrics: Optional[List[Metric]] = None
 
-            try:
-                timer: Timer = Timer(silent=True)
-                timer.start()
-                ## Verbosity >= 1: progress bars
-                progress_bar: Optional[Dict] = self._run_evaluation_progress_bar(progress_bar)
-                ## Verbosity >= 2: basic logging
-                main_logger: Callable = partial(
-                    self.ray_logger,
-                    ## Unless we request silence (verbosity=0), print important information.
-                    should_log=self.verbosity >= 2,
-                    tracker=tracker,
-                )
-                ## Verbosity >= 3: detailed logging
-                debug_logger: Callable = partial(
-                    self.ray_logger,
-                    ## Unless we request silence (verbosity=0), print important information.
-                    should_log=self.verbosity >= 3,
-                    tracker=tracker,
+            timer: Timer = Timer(silent=True)
+            timer.start()
+            ## Verbosity >= 1: progress bars
+            progress_bar: Union[Dict, bool] = self._run_evaluation_progress_bar(progress_bar)
+            ## Verbosity >= 2: basic logging
+            main_logger: Callable = partial(
+                self.ray_logger,
+                ## Unless we request silence (verbosity=0), print important information.
+                should_log=self.verbosity >= 2,
+                tracker=tracker,
+            )
+            ## Verbosity >= 3: detailed logging
+            debug_logger: Callable = partial(
+                self.ray_logger,
+                ## Unless we request silence (verbosity=0), print important information.
+                should_log=self.verbosity >= 3,
+                tracker=tracker,
+            )
+            main_logger(self._evaluate_start_msg(tracker=tracker, **kwargs))
+            if batch_size is None:
+                raise ValueError(
+                    f"Could not find batch_size in model hyperparams; "
+                    f"please pass it explicitly like so: {self.class_name}.evaluate(batch_size=...)"
                 )
-                main_logger(self._evaluate_start_msg(tracker=tracker, **kwargs))
-                if batch_size is None:
+            if predictions_destination is not None:
+                if predictions_destination.storage is not Storage.S3:
                     raise ValueError(
-                        f"Could not find batch_size in model hyperparams; "
-                        f"please pass it explicitly like so: {self.class_name}.evaluate(batch_size=...)"
+                        f"Results can only be saved to {Storage.S3}; "
+                        f"found storage {predictions_destination.storage} having path: {predictions_destination.path}"
                     )
-                if predictions_destination is not None:
-                    if predictions_destination.storage is not Storage.S3:
-                        raise ValueError(
-                            f"Results can only be saved to {Storage.S3}; "
-                            f"found storage {predictions_destination.storage} having path: {predictions_destination.path}"
-                        )
-                    if not predictions_destination.is_path_valid_dir():
-                        raise ValueError(
-                            f"Expected predictions destination to be a valid directory; "
-                            f'found: "{predictions_destination.path}"...did you forget a "/" at the end?'
-                        )
-                    assert predictions_destination.format is not None  ## Checked in .evaluate().
+                if not predictions_destination.is_path_valid_dir():
+                    raise ValueError(
+                        f"Expected predictions destination to be a valid directory; "
+                        f'found: "{predictions_destination.path}"...did you forget a "/" at the end?'
+                    )
+                assert predictions_destination.format is not None  ## Checked in .evaluate().
 
+            try:
                 actors_were_created_in_this_call: bool = self.init_model(progress_bar=progress_bar, **kwargs)
                 num_actors_created: int = len(self.model)
                 if actors_were_created_in_this_call:
@@ -869,15 +873,16 @@ def _run_evaluation(
                     )
                 )
                 return evaluated_predictions, evaluated_metrics
+            except Exception as e:
+                raise e
             except KeyboardInterrupt as e:
                 raise e
             finally:
                 if "row_counter" in locals():
                     accumulate(ray.kill(row_counter))
                     del row_counter
-                if (
-                    self.cache_timeout is None
-                ):  ## If we don't have a timeout, delete actors after every execution.
+                ## If we don't have a timeout, delete actors after every execution.
+                if self.cache_timeout is None:
                     self.cleanup_model()
                 return evaluated_predictions, evaluated_metrics
 
@@ -894,10 +899,10 @@ def _get_actor_usages(self) -> List[Tuple[int, float, str]]:
             )
             return actor_usages
 
-        def _run_evaluation_progress_bar(self, progress_bar: Optional[Dict], **kwargs) -> Optional[Dict]:
+        def _run_evaluation_progress_bar(self, progress_bar: Optional[Dict], **kwargs) -> Union[Dict, bool]:
             if self.verbosity >= 2:
                 return progress_bar
-            return None
+            return False
 
         def _evaluate_start_msg(self, *, tracker: Tracker, **kwargs) -> str:
             if tracker.tracker_name == "noop":