Merging 'main' to 'amangu-lora'

Author: aman2930

experimental/jax/inference/config/config.py

@@ -19,6 +19,7 @@
 
 class ModelId:
   llama_2_7b_chat_hf = "meta-llama/Llama-2-7b-chat-hf"
+  llama_2_70b_chat_hf = "meta-llama/Llama-2-70b-chat-hf"
 
 
 @dataclasses.dataclass
@@ -43,6 +44,15 @@ class Config:
           page_size=128,
           hbm_utilization=0.875,
       ),
+      ModelId.llama_2_70b_chat_hf: InferenceParams(
+          model_id=ModelId.llama_2_70b_chat_hf,
+          batch_size=100,
+          max_seq_length=2048,
+          max_input_length=1024,
+          prefill_chunk_sizes=[128, 256, 512, 1024],
+          page_size=128,
+          hbm_utilization=0.875,
+      ),
   }
 
   @classmethod

jetstream/core/server_lib.py

@@ -266,6 +266,11 @@ def run(
   Returns:
     JetStreamServer that wraps the grpc server and orchestrator driver.
   """
+  # TODO: Deleting the lora_input_adapters_path for now.
+  # Planning to use it in next big PR. Currently accomodating it
+  # to fix the params mismatch between maxText and JetStream
+  del lora_input_adapters_path
+
   server_start_time = time.time()
   logging.info("Kicking off gRPC server.")
   # Setup Prometheus server