enable deepspeed inference (#47)

Signed-off-by: depenglee1707 <[email protected]>

Author: depenglee1707

llmserve/backend/llm/initializers/hf_transformers/deepspeed.py

@@ -45,7 +45,9 @@ def __init__(
         max_tokens: int = 1024,
         use_kernel: bool = False,
         use_meta_tensor: bool = False,
-        injection_policy=None,
+        test_hybrid_engine: bool = False,
+        save_mp_checkpoint_path: bool = False,
+        # injection_policy=None,
         ds_inference_kwargs: Optional[Dict[str, Any]] = None,
         **from_pretrained_kwargs,
     ):
@@ -60,8 +62,10 @@ def __init__(
         self.max_tokens = max_tokens
         self.use_kernel = use_kernel
         self.use_meta_tensor = use_meta_tensor
+        self.test_hybrid_engine = test_hybrid_engine
+        self.save_mp_checkpoint_path = save_mp_checkpoint_path
         # TODO: Allow conversion from strings (need to do dynamic imports)
-        self.injection_policy = injection_policy
+        # self.injection_policy = injection_policy
         self.ds_inference_kwargs = ds_inference_kwargs
 
         if self.use_kernel:
@@ -114,6 +118,8 @@ def _generate_checkpoint_json(
                     for entry in Path(repo_root).rglob("*.[bp][it][n]")
                     if entry.is_file()
                 ]
+
+                # BOOLM ?!
                 data = {"type": "BLOOM",
                         "checkpoints": file_list, "version": 1.0}
                 json.dump(data, f)
@@ -170,58 +176,78 @@ def load_model(self, model_id: str) -> "PreTrainedModel":
         return model
 
     def postprocess_model(self, model: "PreTrainedModel") -> "PreTrainedModel":
-        from transformers import GPTNeoXForCausalLM, LlamaForCausalLM
-
-        injection_policy = self.injection_policy
-        # TODO: remove those later when deepspeed master is updated
-        if injection_policy is None and not self.use_kernel:
-            if isinstance(model, GPTNeoXForCausalLM):
-                from transformers import GPTNeoXLayer
-
-                injection_policy = {
-                    GPTNeoXLayer: ("attention.dense", "mlp.dense_4h_to_h")
-                }
-            elif isinstance(model, LlamaForCausalLM):
-                from transformers.models.llama.modeling_llama import LlamaDecoderLayer
-
-                injection_policy = {
-                    LlamaDecoderLayer: ("self_attn.o_proj", "mlp.down_proj")
-                }
-
-        if self.use_bettertransformer:
-            from optimum.bettertransformer import BetterTransformer
-
-            logger.info("Transforming the model with BetterTransformer...")
-            model = BetterTransformer.transform(model)
-
-        ds_kwargs = self.ds_inference_kwargs or {}
-        ds_kwargs = ds_kwargs.copy()
-        ds_kwargs.update(
-            dict(
-                dtype=self.dtype,
-                mp_size=self.world_size,
-                replace_with_kernel_inject=self.use_kernel,
-                injection_policy=injection_policy,
-                max_tokens=self.max_tokens,
-            )
-        )
         if self.use_meta_tensor:
-            ds_kwargs.update(
-                dict(base_dir=self._repo_root, checkpoint=self._checkpoints_json)
-            )
-
-        logger.info(f"deepspeed.init_inference kwargs: {ds_kwargs}")
-        model = deepspeed.init_inference(
-            model,
-            **ds_kwargs,
-        )
+            ds_kwargs = dict(base_dir=self._repo_root, checkpoint=self._checkpoints_json)
+        else:
+            ds_kwargs = dict()
+
+        # Use DeepSpeed Hybrid Engine for inference
+        if self.test_hybrid_engine:
+            ds_config = {"train_batch_size": 2, "fp16": {"enabled": True if self.dtype==torch.half else False}, "hybrid_engine": {"enabled": True}}
+            model, *_ = deepspeed.initialize(model=model, config=ds_config)
+            model.eval()
+        # If not trying with the HuggingFace baseline, use DeepSpeed Inference Engine
+        else:
+            model = deepspeed.init_inference(model,
+                                        dtype=self.dtype,
+                                        mp_size=self.world_size,
+                                        replace_with_kernel_inject=self.use_kernel,
+                                        max_tokens=self.max_tokens,
+                                        save_mp_checkpoint_path=self.save_mp_checkpoint_path,
+                                        **ds_kwargs
+                                        )     
+        # from transformers import GPTNeoXForCausalLM, LlamaForCausalLM
+
+        # injection_policy = self.injection_policy
+        # # TODO: remove those later when deepspeed master is updated
+        # if injection_policy is None and not self.use_kernel:
+        #     if isinstance(model, GPTNeoXForCausalLM):
+        #         from transformers import GPTNeoXLayer
+
+        #         injection_policy = {
+        #             GPTNeoXLayer: ("attention.dense", "mlp.dense_4h_to_h")
+        #         }
+        #     elif isinstance(model, LlamaForCausalLM):
+        #         from transformers.models.llama.modeling_llama import LlamaDecoderLayer
+
+        #         injection_policy = {
+        #             LlamaDecoderLayer: ("self_attn.o_proj", "mlp.down_proj")
+        #         }
+
+        # if self.use_bettertransformer:
+        #     from optimum.bettertransformer import BetterTransformer
+
+        #     logger.info("Transforming the model with BetterTransformer...")
+        #     model = BetterTransformer.transform(model)
+
+        # ds_kwargs = self.ds_inference_kwargs or {}
+        # ds_kwargs = ds_kwargs.copy()
+        # ds_kwargs.update(
+        #     dict(
+        #         dtype=self.dtype,
+        #         mp_size=self.world_size,
+        #         replace_with_kernel_inject=self.use_kernel,
+        #         injection_policy=injection_policy,
+        #         max_tokens=self.max_tokens,
+        #     )
+        # )
+        # if self.use_meta_tensor:
+        #     ds_kwargs.update(
+        #         dict(base_dir=self._repo_root, checkpoint=self._checkpoints_json)
+        #     )
+
+        # logger.info(f"deepspeed.init_inference kwargs: {ds_kwargs}")
+        # model = deepspeed.init_inference(
+        #     model,
+        #     **ds_kwargs,
+        # )
 
         if self.torch_compile and self.torch_compile["backend"]:
             logger.info("Compiling the model with torch.compile()...")
             model = torch.compile(model, **self.torch_compile)
 
         # Add attributes for compatibility with the pipeline
-        model.use_kernel = self.use_kernel
-        model.device = self.device
-        model = model.to(self.device)
+        # model.use_kernel = self.use_kernel
+        # model.device = self.device
+        # model = model.to(self.device)
         return model

llmserve/backend/llm/pipelines/_base.py

@@ -165,8 +165,9 @@ def __call__(self, inputs: List[Union[str, Prompt]], **kwargs) -> List[Response]
             postprocess_params,
         ) = self._sanitize_parameters(**kwargs)
         model_inputs = self.preprocess(inputs, **preprocess_params)
+
         model_inputs = self._ensure_tensor_on_device(
-            model_inputs, device=self.model.device)
+            model_inputs, device=(self.model.device if hasattr(self.model, 'device') else self.device))
 
         forward_params = self._add_default_generate_kwargs(
             forward_params, model_inputs)

llmserve/backend/llm/pipelines/default_pipeline.py

@@ -1,5 +1,5 @@
 import time
-from typing import List, Optional, Union
+from typing import List, Optional, Union, TYPE_CHECKING
 
 import torch
 from transformers import PreTrainedModel, PreTrainedTokenizer
@@ -56,7 +56,7 @@ def preprocess(self, prompts: List[str], **generate_kwargs):
 
         inputs = self.tokenizer(
             prompt_text, return_tensors="pt", add_special_tokens = generate_kwargs.get("add_special_tokens", True), padding=True
-        ).to(self.model.device)
+        ).to(self.model.device if hasattr(self.model, 'device') else self.device)
 
         if not generate_kwargs.get("return_token_type_ids", True):
             inputs.pop("token_type_ids", None)

llmserve/backend/llm/predictor.py

@@ -52,8 +52,8 @@ def init_model(
     # Lazy import so that the new cache location is used
     torch.backends.cuda.matmul.allow_tf32 = True
     if torch.cuda.is_available():
-        # device = torch.device(f"cuda:{local_rank}")
-        device = torch.device("cuda")
+        device = torch.device(f"cuda:{local_rank}")
+        # device = torch.device("cuda")
     else:
         device = torch.device("cpu")
 
@@ -383,11 +383,12 @@ async def _create_worker_group(
         await asyncio.gather(
             *[
                 worker.init_model.remote(
-                    local_rank,
+                    local_rank = local_rank,
                     num_cpus_per_worker=scaling_config.num_cpus_per_worker,
                     num_gpus_per_worker=scaling_config.num_gpus_per_worker
                 )
                 for worker, local_rank in zip(worker_group, local_ranks)
+                # for worker in worker_group
             ]
         )
 
@@ -425,20 +426,39 @@ def slice_prompts(worker_num: int, worker_index: int, prompts: list[str]):
         logger.info('LLM Predictor do async predict')
 
         async with self._base_worker_group_lock:
+            # prediction = (
+            #     await asyncio.gather(
+            #         *[
+            #             worker.generate.remote(
+            #                 slice_prompts(len(self.base_worker_group), index, prompts),
+            #                 # prompts,
+            #                 timeout_s=timeout_s,
+            #                 start_timestamp=start_timestamp,
+            #                 **self.args.model_config.generation.all_generate_kwargs if self.args.model_config.generation else {},  # pylint:disable=no-member
+            #             ) if len(slice_prompts(len(self.base_worker_group), index, prompts)) > 0 else ray.put([])
+
+            #             for index, worker in enumerate(self.base_worker_group)
+            #             # for worker in self.base_worker_group
+            #         ]
+            #     )
+            # )
+        # return [response for responses in prediction for response in responses]
             prediction = (
                 await asyncio.gather(
                     *[
                         worker.generate.remote(
-                            slice_prompts(len(self.base_worker_group), index, prompts),
-                            # prompts,
+                            # slice_prompts(len(self.base_worker_group), index, prompts),
+                            prompts,
                             timeout_s=timeout_s,
                             start_timestamp=start_timestamp,
                             **self.args.model_config.generation.all_generate_kwargs if self.args.model_config.generation else {},  # pylint:disable=no-member
-                        ) if len(slice_prompts(len(self.base_worker_group), index, prompts)) > 0 else ray.put([])
+                        )
 
-                        for index, worker in enumerate(self.base_worker_group)
-                        # for worker in self.base_worker_group
+                        # for index, worker in enumerate(self.base_worker_group)
+                        for worker in self.base_worker_group
                     ]
                 )
             )
-        return [response for responses in prediction for response in responses]
+
+            return prediction
+

llmserve/backend/server/models.py

@@ -240,6 +240,8 @@ class DeepSpeed(Transformers):
     use_kernel: bool = False
     max_tokens: int = 1024
     use_meta_tensor: bool = False
+    test_hybrid_engine: bool = False
+    save_mp_checkpoint_path: bool = False
     ds_inference_kwargs: Optional[Dict[str, Any]] = None
 
     @root_validator
@@ -257,7 +259,10 @@ def use_kernel_use_meta_tensor(cls, values):  # pylint:disable=no-self-argument
         if not values.get("use_kernel") and values.get("use_meta_tensor"):
             raise ValueError("'use_meta_tensor=True' needs 'use_kernel=True'.")
         return values
-
+    
+    @property
+    def allowed_pipelines(self) -> Set[str]:
+        return {"default"}
 
 class DeviceMap(Transformers):
     type: Literal["DeviceMap"]

models/text-generation--bigscience--bloom-3b.yaml

@@ -0,0 +1,54 @@
+deployment_config:
+  autoscaling_config:
+    min_replicas: 1
+    initial_replicas: 1
+    max_replicas: 8
+    target_num_ongoing_requests_per_replica: 1.0
+    metrics_interval_s: 10.0
+    look_back_period_s: 30.0
+    smoothing_factor: 1.0
+    downscale_delay_s: 300.0
+    upscale_delay_s: 90.0
+  ray_actor_options:
+    num_cpus: 0.1    # for a model deployment, we have 3 actor created, 1 and 2 will cost 0.1 cpu, and the model infrence will cost 6(see the setting in the end of the file)
+model_config:
+  warmup: True
+  model_task: text-generation
+  model_id: bigscience/bloom-3b
+  max_input_words: 800
+  initialization:
+    # s3_mirror_config:
+    #   bucket_uri: s3://large-dl-models-mirror/models--amazon--LightGPT/main-safetensors/
+    initializer:
+      type: DeepSpeed
+      dtype: float32
+      max_tokens: 512
+      use_kernel: true
+      use_meta_tensor: false
+      test_hybrid_engine: false
+      save_mp_checkpoint_path: false
+      from_pretrained_kwargs:
+        use_cache: true
+        trust_remote_code: true
+    pipeline: default
+  generation:
+    max_batch_size: 2
+    batch_wait_timeout_s: 30
+    generate_kwargs:
+      do_sample: false
+      max_new_tokens: 512
+      min_new_tokens: 16
+      temperature: 0.7
+      repetition_penalty: 1.1
+      top_p: 0.8
+      top_k: 50
+    # prompt_format: "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n### Instruction:\n{instruction}\n### Response:\n"
+    # stopping_sequences: ["### Response:", "### End"]
+scaling_config:
+  num_workers: 1
+  num_gpus_per_worker: 0
+  num_cpus_per_worker: 6   # for infrence
+  # resources_per_worker:
+  #   accelerator_type_cpu: 0.01
+  ray_actor_options:
+    num_cpus: 0.1

setup.py

@@ -49,7 +49,7 @@
             "torchaudio==2.1.2",
             "torchvision==0.16.2",
             "accelerate==0.25.0",
-            "deepspeed==0.12.6",
+            "deepspeed==0.14.0",
             "torchmetrics==1.2.1",
             "llama_cpp_python==0.2.20",
             "transformers==4.33.3",