Necessary Changes for long context llama-3-8b (#516)

* all necessary changes

* tests

Author: sam-scale

model-engine/model_engine_server/domain/use_cases/llm_model_endpoint_use_cases.py

@@ -169,6 +169,7 @@
             "llama-2-70b-chat",
             "llama-3-8b",
             "llama-3-8b-instruct",
+            "llama-3-8b-instruct-262k",
             "llama-3-70b",
             "llama-3-70b-instruct",
             "falcon-7b",
@@ -240,6 +241,7 @@
     # Can also see 13B, 34B there too
     "gemma": {"max_model_len": 8192, "max_num_batched_tokens": 8192},
     "llama-2": {"max_model_len": None, "max_num_batched_tokens": 4096},
+    "llama-3-8b-instruct-262k": {"max_model_len": None, "max_num_batched_tokens": 262144},
     "llama-3": {"max_model_len": None, "max_num_batched_tokens": 8192},
     "mistral": {"max_model_len": 8000, "max_num_batched_tokens": 8000},
     "mixtral-8x7b": {"max_model_len": 32768, "max_num_batched_tokens": 32768},
@@ -2211,6 +2213,12 @@ def infer_hardware_from_model_name(model_name: str) -> CreateDockerImageBatchJob
         memory = "800Gi"
         storage = "460Gi"
         gpu_type = GpuType.NVIDIA_AMPERE_A100E
+    elif "llama-3-8b-instruct-262k" in model_name:
+        cpus = "20"
+        gpus = 2
+        memory = "40Gi"
+        storage = "40Gi"
+        gpu_type = GpuType.NVIDIA_AMPERE_A100E
     else:
         numbers = re.findall(r"\d+", model_name)
         if len(numbers) == 0:

model-engine/model_engine_server/infra/repositories/live_tokenizer_repository.py

@@ -43,6 +43,7 @@ def get_default_supported_models_info() -> Dict[str, ModelInfo]:
         "llama-2-70b-chat": ModelInfo("meta-llama/Llama-2-70b-chat-hf", None),
         "llama-3-8b": ModelInfo("meta-llama/Meta-Llama-3-8B", None),
         "llama-3-8b-instruct": ModelInfo("meta-llama/Meta-Llama-3-8B-Instruct", None),
+        "llama-3-8b-instruct-262k": ModelInfo("gradientai/Llama-3-8B-Instruct-262k", None),
         "llama-3-70b": ModelInfo("meta-llama/Meta-Llama-3-70B", None),
         "llama-3-70b-instruct": ModelInfo("meta-llama/Meta-Llama-3-70B-Instruct", None),
         "falcon-7b": ModelInfo("tiiuae/falcon-7b", None),

model-engine/tests/unit/domain/test_llm_use_cases.py

@@ -1809,6 +1809,13 @@ def test_infer_hardware_from_model_name():
     assert hardware.storage == "160Gi"
     assert hardware.gpu_type == GpuType.NVIDIA_AMPERE_A100E
 
+    hardware = infer_hardware_from_model_name("llama-3-8b-instruct-262k")
+    assert hardware.cpus == "20"
+    assert hardware.gpus == 2
+    assert hardware.memory == "40Gi"
+    assert hardware.storage == "40Gi"
+    assert hardware.gpu_type == GpuType.NVIDIA_AMPERE_A100E
+
     with pytest.raises(ObjectHasInvalidValueException):
         infer_hardware_from_model_name("unsupported_model")