diff --git a/ads/aqua/cli.py b/ads/aqua/cli.py index 9fbc83a68..34b028db7 100644 --- a/ads/aqua/cli.py +++ b/ads/aqua/cli.py @@ -15,6 +15,7 @@ from ads.aqua.model import AquaModelApp from ads.aqua.modeldeployment import AquaDeploymentApp from ads.aqua.verify_policies import AquaVerifyPoliciesApp +from ads.aqua.shaperecommend.recommend import AquaRecommendApp from ads.common.utils import LOG_LEVELS @@ -31,6 +32,7 @@ class AquaCommand: deployment = AquaDeploymentApp evaluation = AquaEvaluationApp verify_policies = AquaVerifyPoliciesApp + recommend = AquaRecommendApp def __init__( self, @@ -96,18 +98,20 @@ def _validate_value(flag, value): "If you intend to chain a function call to the result, please separate the " "flag and the subsequent function call with separator `-`." ) - + @staticmethod def install(): """Install ADS Aqua Extension from wheel file. Set enviroment variable `AQUA_EXTENSTION_PATH` to change the wheel file path. - Return + Return ------ int: Installatation status. """ import subprocess - wheel_file_path = os.environ.get("AQUA_EXTENSTION_PATH", "/ads/extension/adsjupyterlab_aqua_extension*.whl") - status = subprocess.run(f"pip install {wheel_file_path}",shell=True) - return status.check_returncode \ No newline at end of file + wheel_file_path = os.environ.get( + "AQUA_EXTENSTION_PATH", "/ads/extension/adsjupyterlab_aqua_extension*.whl" + ) + status = subprocess.run(f"pip install {wheel_file_path}", shell=True) + return status.check_returncode diff --git a/ads/aqua/common/entities.py b/ads/aqua/common/entities.py index 13203e34b..27262df26 100644 --- a/ads/aqua/common/entities.py +++ b/ads/aqua/common/entities.py @@ -46,6 +46,17 @@ class Config: arbitrary_types_allowed = True protected_namespaces = () +class ComputeRank(Serializable): + """ + Represents the cost and performance ranking for a compute shape. + """ + cost: int = Field( + None, description="The relative rank of the cost of the shape. Range is [10 (cost-effective), 100 (most-expensive)]" + ) + + performance: int = Field( + None, description="The relative rank of the performance of the shape. Range is [10 (lower performance), 110 (highest performance)]" + ) class GPUSpecs(Serializable): """ @@ -61,6 +72,12 @@ class GPUSpecs(Serializable): gpu_type: Optional[str] = Field( default=None, description="The type of GPU (e.g., 'V100, A100, H100')." ) + quantization: Optional[List[str]] = Field( + default_factory=list, description="The quantization format supported by shape. (ex. bitsandbytes, fp8, etc.)" + ) + ranking: Optional[ComputeRank] = Field( + None, description="The relative rank of the cost and performance of the shape." + ) class GPUShapesIndex(Serializable): diff --git a/ads/aqua/common/errors.py b/ads/aqua/common/errors.py index 15b1018f5..20b578a8f 100644 --- a/ads/aqua/common/errors.py +++ b/ads/aqua/common/errors.py @@ -55,6 +55,11 @@ class AquaValueError(AquaError, ValueError): def __init__(self, reason, status=403, service_payload=None): super().__init__(reason, status, service_payload) +class AquaRecommendationError(AquaError): + """Exception raised for models incompatible with shape recommendation tool.""" + + def __init__(self, reason, status=400, service_payload=None): + super().__init__(reason, status, service_payload) class AquaFileNotFoundError(AquaError, FileNotFoundError): """Exception raised for missing target file.""" diff --git a/ads/aqua/common/utils.py b/ads/aqua/common/utils.py index 2d64fd42f..120ea67a4 100644 --- a/ads/aqua/common/utils.py +++ b/ads/aqua/common/utils.py @@ -1253,24 +1253,24 @@ def load_gpu_shapes_index( file_name = "gpu_shapes_index.json" # Try remote load - remote_data: Dict[str, Any] = {} - if CONDA_BUCKET_NS: - try: - auth = auth or authutil.default_signer() - storage_path = ( - f"oci://{CONDA_BUCKET_NAME}@{CONDA_BUCKET_NS}/service_pack/{file_name}" - ) - logger.debug( - "Loading GPU shapes index from Object Storage: %s", storage_path - ) - with fsspec.open(storage_path, mode="r", **auth) as f: - remote_data = json.load(f) - logger.debug( - "Loaded %d shapes from Object Storage", - len(remote_data.get("shapes", {})), - ) - except Exception as ex: - logger.debug("Remote load failed (%s); falling back to local", ex) + # remote_data: Dict[str, Any] = {} + # if CONDA_BUCKET_NS: + # try: + # auth = auth or authutil.default_signer() + # storage_path = ( + # f"oci://{CONDA_BUCKET_NAME}@{CONDA_BUCKET_NS}/service_pack/{file_name}" + # ) + # logger.debug( + # "Loading GPU shapes index from Object Storage: %s", storage_path + # ) + # with fsspec.open(storage_path, mode="r", **auth) as f: + # remote_data = json.load(f) + # logger.debug( + # "Loaded %d shapes from Object Storage", + # len(remote_data.get("shapes", {})), + # ) + # except Exception as ex: + # logger.debug("Remote load failed (%s); falling back to local", ex) # Load local copy local_data: Dict[str, Any] = {} @@ -1287,6 +1287,7 @@ def load_gpu_shapes_index( # Merge: remote shapes override local local_shapes = local_data.get("shapes", {}) + remote_data = {} remote_shapes = remote_data.get("shapes", {}) merged_shapes = {**local_shapes, **remote_shapes} diff --git a/ads/aqua/extension/__init__.py b/ads/aqua/extension/__init__.py index 4c8d9f3f3..ffd2241c6 100644 --- a/ads/aqua/extension/__init__.py +++ b/ads/aqua/extension/__init__.py @@ -13,6 +13,7 @@ from ads.aqua.extension.evaluation_handler import __handlers__ as __eval_handlers__ from ads.aqua.extension.finetune_handler import __handlers__ as __finetune_handlers__ from ads.aqua.extension.model_handler import __handlers__ as __model_handlers__ +from ads.aqua.extension.recommend_handler import __handlers__ as __gpu_handlers__ from ads.aqua.extension.ui_handler import __handlers__ as __ui_handlers__ from ads.aqua.extension.ui_websocket_handler import __handlers__ as __ws_handlers__ @@ -24,6 +25,7 @@ + __ui_handlers__ + __eval_handlers__ + __ws_handlers__ + + __gpu_handlers__ ) diff --git a/ads/aqua/extension/recommend_handler.py b/ads/aqua/extension/recommend_handler.py new file mode 100644 index 000000000..4105f9320 --- /dev/null +++ b/ads/aqua/extension/recommend_handler.py @@ -0,0 +1,46 @@ +from tornado.web import HTTPError + +from ads.aqua.common.decorator import handle_exceptions +from ads.aqua.extension.base_handler import AquaAPIhandler +from ads.aqua.extension.errors import Errors +from ads.aqua.shaperecommend.recommend import AquaRecommendApp + + +class AquaRecommendHandler(AquaAPIhandler): + """ + Handler for Aqua GPU Recommendation REST APIs. + + Methods + ------- + post(self, *args, **kwargs) + Obtains the eligible compute shapes that would fit the specifed model, context length, model weights, and quantization level. + + Raises + ------ + HTTPError: For various failure scenarios such as invalid input format, missing data, etc. + """ + + @handle_exceptions + def post(self, *args, **kwargs): # noqa: ARG002 + """ + Obtains the eligible compute shapes that would fit the specifed model, context length, model weights, and quantization level. + + Returns + ------- + ShapeRecommendationReport + Report containing shape recommendations and troubleshooting advice, if any. + """ + try: + input_data = self.get_json_body() + except Exception as ex: + raise HTTPError(400, Errors.INVALID_INPUT_DATA_FORMAT) from ex + + if not input_data: + raise HTTPError(400, Errors.NO_INPUT_DATA) + + self.finish(AquaRecommendApp().which_gpu(**input_data)) + + +__handlers__ = [ + ("recommendation/?([^/]*)", AquaRecommendHandler), +] diff --git a/ads/aqua/resources/gpu_shapes_index.json b/ads/aqua/resources/gpu_shapes_index.json index c88155e45..8dd701be6 100644 --- a/ads/aqua/resources/gpu_shapes_index.json +++ b/ads/aqua/resources/gpu_shapes_index.json @@ -1,94 +1,152 @@ { "shapes": { - "BM.GPU.A10.4": { - "gpu_count": 4, - "gpu_memory_in_gbs": 96, - "gpu_type": "A10" + "BM.GPU.H200.8": { + "gpu_count": 8, + "gpu_memory_in_gbs": 1128, + "gpu_type": "H200", + "quantization": ["awq", "gptq", "marlin", "fp8", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking": { + "cost": 100, + "performance": 110 + } }, - "BM.GPU.A100-V2.8": { + "BM.GPU.H100.8": { "gpu_count": 8, "gpu_memory_in_gbs": 640, - "gpu_type": "A100" + "gpu_type": "H100", + "quantization": ["awq", "gptq", "marlin", "fp8", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking": { + "cost": 100, + "performance": 100 + } }, - "BM.GPU.B4.8": { + "BM.GPU.MI300X.8": { "gpu_count": 8, - "gpu_memory_in_gbs": 320, - "gpu_type": "A100" + "gpu_memory_in_gbs": 1536, + "gpu_type": "MI300X", + "quantization": ["fp8", "gguf"], + "ranking": { + "cost": 90, + "performance": 90 + } }, - "BM.GPU.H100.8": { + "BM.GPU.A100-V2.8": { "gpu_count": 8, "gpu_memory_in_gbs": 640, - "gpu_type": "H100" + "gpu_type": "A100", + "quantization": ["awq", "gptq", "marlin", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking": { + "cost": 80, + "performance": 70 + } }, - "BM.GPU.H200.8": { + "BM.GPU.B4.8": { "gpu_count": 8, - "gpu_memory_in_gbs": 1128, - "gpu_type": "H200" + "gpu_memory_in_gbs": 320, + "gpu_type": "A100", + "quantization": ["awq", "gptq", "marlin", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking": { + "cost": 70, + "performance": 60 + } }, "BM.GPU.L40S-NC.4": { "gpu_count": 4, "gpu_memory_in_gbs": 192, - "gpu_type": "L40S" + "gpu_type": "L40S", + "quantization": ["awq", "gptq", "marlin", "fp8", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking": { + "cost": 60, + "performance": 80 + } }, "BM.GPU.L40S.4": { "gpu_count": 4, "gpu_memory_in_gbs": 192, - "gpu_type": "L40S" - }, - "BM.GPU.MI300X.8": { - "gpu_count": 8, - "gpu_memory_in_gbs": 1536, - "gpu_type": "MI300X" - }, - "BM.GPU2.2": { - "gpu_count": 2, - "gpu_memory_in_gbs": 32, - "gpu_type": "P100" - }, - "BM.GPU3.8": { - "gpu_count": 8, - "gpu_memory_in_gbs": 128, - "gpu_type": "V100" - }, - "BM.GPU4.8": { - "gpu_count": 8, - "gpu_memory_in_gbs": 320, - "gpu_type": "A100" + "gpu_type": "L40S", + "quantization": ["awq", "gptq", "marlin", "fp8", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking": { + "cost": 60, + "performance": 80 + } }, "VM.GPU.A10.1": { "gpu_count": 1, "gpu_memory_in_gbs": 24, - "gpu_type": "A10" + "gpu_type": "A10", + "quantization": ["awq", "gptq", "marlin", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking" : { + "cost": 20, + "performance": 30 + } }, "VM.GPU.A10.2": { "gpu_count": 2, "gpu_memory_in_gbs": 48, - "gpu_type": "A10" + "gpu_type": "A10", + "quantization": ["awq", "gptq", "marlin", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking" : { + "cost": 40, + "performance": 40 + } }, - "VM.GPU.A10.4": { + "BM.GPU.A10.4": { "gpu_count": 4, "gpu_memory_in_gbs": 96, - "gpu_type": "A10" + "gpu_type": "A10", + "quantization": ["awq", "gptq", "marlin", "int8", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking" : { + "cost": 50, + "performance": 50 + } + }, + "BM.GPU2.2": { + "gpu_count": 2, + "gpu_memory_in_gbs": 32, + "gpu_type": "P100", + "quantization": ["fp16"], + "ranking": { + "cost": 30, + "performance": 20 + } }, "VM.GPU2.1": { "gpu_count": 1, "gpu_memory_in_gbs": 16, - "gpu_type": "P100" + "gpu_type": "P100", + "quantization": ["fp16"], + "ranking": { + "cost": 10, + "performance": 10 + } }, "VM.GPU3.1": { "gpu_count": 1, "gpu_memory_in_gbs": 16, - "gpu_type": "V100" + "gpu_type": "V100", + "quantization" : ["gptq", "bitblas", "aqlm", "bitsandbytes", "deepspeedfp", "gguf"], + "ranking" : { + "cost": 35, + "performance": 10 + } }, "VM.GPU3.2": { "gpu_count": 2, "gpu_memory_in_gbs": 32, - "gpu_type": "V100" + "gpu_type": "V100", + "ranking" : { + "cost": 45, + "performance": 20 + } }, "VM.GPU3.4": { "gpu_count": 4, "gpu_memory_in_gbs": 64, - "gpu_type": "V100" + "gpu_type": "V100", + "ranking" : { + "cost": 55, + "performance": 45 + } } } -} +} \ No newline at end of file diff --git a/ads/aqua/shaperecommend/__init__.py b/ads/aqua/shaperecommend/__init__.py new file mode 100644 index 000000000..dd30edb85 --- /dev/null +++ b/ads/aqua/shaperecommend/__init__.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python +# Copyright (c) 2025 Oracle and/or its affiliates. +# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ +from ads.aqua.shaperecommend.recommend import AquaRecommendApp + +__all__ = ["AquaRecommendApp"] diff --git a/ads/aqua/shaperecommend/constants.py b/ads/aqua/shaperecommend/constants.py new file mode 100644 index 000000000..3a7bef233 --- /dev/null +++ b/ads/aqua/shaperecommend/constants.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python +# Copyright (c) 2024, 2025 Oracle and/or its affiliates. +# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ + +""" +aqua.shaperecommend.constants +~~~~~~~~~~~~~~ + +This module contains constants used in Aqua GPU Recommendation for Models. + +LLAMA_REQUIRED_FIELDS refer to fields necessary for calculating model memory for GQA Architecture Models + +MOE_REQUIRED_FIELDS refer to fields necessary for Mixture of Experts (MoE) Architecture Models + +NEXT_QUANT suggests the next quantization level based on the current quantization (if applied) or the model weights (if no quantization yet) +""" +LLAMA_REQUIRED_FIELDS = [ + "num_hidden_layers", "hidden_size", "num_attention_heads", + "num_key_value_heads", "head_dim", "intermediate_size", "vocab_size" +] + +MOE_REQUIRED_FIELDS = LLAMA_REQUIRED_FIELDS + [ + "num_local_experts", "intermediate_size" +] + +NEXT_QUANT = { + "float32": ["8bit", "4bit"], # bits and bytes does not support bfloat16, pytorch responsibility + "bfloat16": ["8bit", "4bit"], + "float16": ["8bit", "4bit"], + "int8": ["4bit"], + "fp8": ["4bit"], + "8bit": ["4bit"], + "int4": ["No smaller quantization available"], + "4bit": ["No smaller quantization available"] +} + +TEXT_GENERATION = "text_generation" +SAFETENSORS = "safetensors" + +TROUBLESHOOT_MSG = "The selected model is too large to fit on standard GPU shapes with the current configuration.\nAs troubleshooting, we have suggested the two largest available GPU shapes using the smallest quantization level ('4bit') to maximize chances of fitting the model. " + + +QUANT_MAPPING = { + "float32": 4, + "bfloat16": 2, + "float16": 2, + "fp16": 2, + "half": 2, + "int8": 1, + "fp8": 1, + "8bit": 1, + "4bit": 0.5, + "int4": 0.5, + } + + diff --git a/ads/aqua/shaperecommend/estimator.py b/ads/aqua/shaperecommend/estimator.py new file mode 100644 index 000000000..d1fca5d5c --- /dev/null +++ b/ads/aqua/shaperecommend/estimator.py @@ -0,0 +1,356 @@ +#!/usr/bin/env python +# Copyright (c) 2025 Oracle and/or its affiliates. +# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ +from typing import Optional + +from pydantic import BaseModel, Field + +from ads.aqua.app import logger +from ads.aqua.shaperecommend.constants import ( + LLAMA_REQUIRED_FIELDS, + MOE_REQUIRED_FIELDS, + NEXT_QUANT, + QUANT_MAPPING, +) +from ads.aqua.shaperecommend.llm_config import LLMConfig + + +class MemoryEstimator(BaseModel): + """ + The generic estimator for Transformer Architecture models (OPT/ Bloom) + Used as a fallback estimator if model identified is not a MoE or GQA Architecture Model. + Has properties to estimate the KV Cache size, Model size, and total footprint (KV Cache + Model size) + + KV cache: Use num_attention_heads (all heads, no GQA) + Parameter estimation: Standard decoder-only, untied embeddings possible + """ + + llm_config: LLMConfig = Field( + ..., + description="The model's config.json file with the necessary parameters for model size and KV cache estimation.", + ) + batch_size: Optional[int] = ( + 1 # we assume that estimation for batch sizes are not supported yet + ) + seq_len: int = Field( + ..., description="The max-seq-len to estimate the size of the KV cache." + ) + + @property + def kv_cache_memory(self) -> float: + """ + Estimates the KV cache size (in GB) using the LLM config.json parameters. + + Uses num_attention_heads (assumes no GQA, each attention head has its own query, key, value) for estimation. + """ + seq_len = self.seq_len or self.llm_config.max_seq_len + c = self.llm_config + kv_cache_dtype_bytes = QUANT_MAPPING.get( + c.weight_dtype, 2 + ) # vLLM uses model's weight/quantization applied to KV cache + + total_bytes = ( + self.batch_size + * c.num_hidden_layers + * 2 + * c.num_attention_heads + * seq_len + * c.head_dim + * kv_cache_dtype_bytes + ) + return total_bytes / 1e9 + + @property + def model_memory(self) -> float: + """ + Estimates the model size (in GB) based on estimating the model parameter size and model weights. + + Model Parameter estimation: Standard decoder-only, untied/tied embeddings possible. + """ + c = self.llm_config + embedding_count = 1 if getattr(c, "tie_word_embeddings", True) else 2 + embedding_params = ( + embedding_count * c.vocab_size * c.hidden_size + ) # input and output untied + layer_params = 12 * c.num_hidden_layers * (c.hidden_size**2) # GPT-style + num_params = layer_params + embedding_params + + return num_params * c.bytes_per_parameter / 1e9 + + @property + def total_memory(self) -> float: + """ + Computes the total memory footprint of the model (KV cache & model size from estimated parameters). + """ + return self.model_memory + self.kv_cache_memory + + def validate_shape(self, allowed_gpu_memory: float) -> bool: + """ + Validates if a given model estimator fits within the allowed GPU memory budget, using a fixed utilization margin. + + Parameters + ---------- + estimator : MemoryEstimator + The estimator with current shape/memory needs. + allowed_gpu_memory : float + The maximum allowed GPU memory. + + Returns + ------- + bool + True if estimator uses less than adjusted GPU memory, else False. + """ + gpu_utilization = 0.9 + return (allowed_gpu_memory * gpu_utilization) > self.total_memory + + def suggest_param_advice(self, allowed: float) -> str: + """ + Suggests parameter modifications to help a model fit within GPU memory limits. + + Parameters + ---------- + estimator : MemoryEstimator + The memory estimator object. + allowed : float + Allowed GPU memory in GB. + + Returns + ------- + str + Advice message with suggestions. + """ + kv_gb = self.kv_cache_memory + wt_gb = self.model_memory + batch_size = self.batch_size + seq_len = self.seq_len + weight_size = getattr(self.llm_config, "weight_dtype", "unknown") + config = self.llm_config + + suggested_quant_msg = None + quant_advice = ", ".join(getattr(config, "suggested_quantizations", [])) + quantization = getattr(config, "quantization", None) + + if getattr(config, "suggested_quantizations", []): + to_do = f", which is smaller than the current quantization/weight size: {quantization if quantization in NEXT_QUANT else weight_size}." + if "No" in quant_advice: + suggested_quant_msg = "No smaller quantized version exists. Use a model with fewer parameters." + elif not quant_advice: + suggested_quant_msg = ( + "Use a quantized version of the same model (e.g., INT8 or other)" + + to_do + ) + else: + suggested_quant_msg = ( + f"Use the same model with {quant_advice} quantization" + to_do + ) + + kv_advice = ( + f"To reduce KV cache memory usage:\n" + f"1. Reduce maximum context length (set --max-model-len < {seq_len})\n" + f"2. Reduce batch size to less than {batch_size}." + if batch_size != 1 + else "" + ) + + wt_advice = ( + "To reduce model size:\n" + "1. Use a model with fewer parameters.\n" + f"2. {suggested_quant_msg}" + if suggested_quant_msg + else "" + ) + + if kv_gb > wt_gb and kv_gb > allowed * 0.5: + main = "KV cache memory usage is the main limiting factor." + advice = kv_advice + elif wt_gb > kv_gb and wt_gb > allowed * 0.5: + main = "Model weights are the main limiting factor." + advice = wt_advice + else: + main = "Both model weights and KV cache are significant contributors to memory use." + advice = f"{kv_advice}\n{wt_advice}" + return f"{main} (KV cache: {kv_gb:.1f}GB, Weights: {wt_gb:.1f}GB).\n{advice}" + + def limiting_factor( + self, allowed_gpu_memory: float, warn_delta: float = 0.85 + ) -> str: + """ + Determines the memory limiting factor for a model deployment and returns advice. + + Parameters + ---------- + estimator : MemoryEstimator + The memory estimator object with current model configuration. + allowed_gpu_memory : float + The maximum allowed GPU memory (in GBs). + warn_delta : float, optional + The threshold (fraction) of allowed GPU memory to trigger a warning (default=0.85). + + Returns + ------- + str + Advice message about model fit and limiting factors. + """ + required = self.total_memory + batch_size = self.batch_size + seq_len = self.seq_len + weight_size = getattr(self.llm_config, "weight_dtype", "unknown") + quantization = getattr(self.llm_config, "quantization", "None") + + # Warn if required is close to but under allowed + if allowed_gpu_memory > required > allowed_gpu_memory * warn_delta: + model_params = self.suggest_param_advice(allowed_gpu_memory) + advice = ( + f"While the selected compute shape is estimated to work " + f"({required:.1f}GB used / {allowed_gpu_memory:.1f}GB allowed), " + f"the model configuration is close to the GPU memory limit. " + "This estimation is theoretical; actual memory usage may vary at runtime.\n\n" + "If you encounter issues with this shape, consider the following options to reduce memory usage:\n\n" + f"{model_params.lstrip()}" + ) + elif required > allowed_gpu_memory: + model_params = self.suggest_param_advice(allowed_gpu_memory) + advice = ( + f"Model does not fit within GPU memory budget. " + "Consider the following options to reduce memory usage:\n\n" + f"{model_params.lstrip()}" + ) + else: + advice = ( + f"Model fits well within the allowed compute shape " + f"({required:.1f}GB used / {allowed_gpu_memory:.1f}GB allowed).\n" + f"(Batch size: {batch_size}, seq len: {seq_len}, " + f"quantization/weight size: {quantization or weight_size})." + ) + return advice + + +# Specialized estimators: +class LlamaMemoryEstimator(MemoryEstimator): + """ + Estimator for GQA-type architectures. Handles tied (memory savings) and untied embeddings, + and uses grouped attention (GQA) for more efficient KV cache memory estimation. + + KV cache: Use num_attention_heads (assumes GQA) + Model Parameter estimation: Standard decoder-only, untied/tied embeddings possible + """ + + @property + def model_memory(self) -> float: + """ + Returns estimated model parameter memory (in GB), accurately accounting + for Llama-style attention and MLP, and tied or untied embeddings. + """ + c = self.llm_config + + embedding_params, attn_params = self._calc_attn_embed_params() + + # MLP params + gate_proj = c.hidden_size * c.intermediate_size + up_proj = c.hidden_size * c.intermediate_size + down_proj = c.intermediate_size * c.hidden_size + mlp_params = gate_proj + up_proj + down_proj + + # Total per-layer + layer_params = attn_params + mlp_params + # Total params + num_params = c.num_hidden_layers * layer_params + embedding_params + return num_params * c.bytes_per_parameter / 1e9 + + @property + def kv_cache_memory(self) -> float: + """ + Returns estimated KV cache memory in GB for GQA models. + + Grouped Query Attention uses num_key_value_heads, which groups of Q heads share a K and V projection. + num_key_value_heads < num_attention_heads, which reduces the KV Cache size. + """ + c = self.llm_config + seq_len = self.seq_len or getattr(c, "max_seq_len", 2048) + kv_cache_dtype_bytes = QUANT_MAPPING.get(c.weight_dtype, 2) + kv_heads = c.num_key_value_heads + + total_bytes = ( + self.batch_size + * c.num_hidden_layers + * 2 + * kv_heads + * seq_len + * c.head_dim + * kv_cache_dtype_bytes + ) + return total_bytes / 1e9 + + def _calc_attn_embed_params(self) -> tuple: + """ + Returns the embedding parameter count and attention parameter count for Llama-family (GQA) models. + """ + c = self.llm_config + + # Embedding parameters + # assume tied embeddings unless tie_word_embeddings = False + embedding_count = 1 if getattr(c, "tie_word_embeddings", True) else 2 + embedding_params = embedding_count * c.vocab_size * c.hidden_size + + q_proj = c.hidden_size * c.hidden_size + k_proj = c.hidden_size * (c.num_key_value_heads * c.head_dim) + v_proj = c.hidden_size * (c.num_key_value_heads * c.head_dim) + o_proj = c.hidden_size * c.hidden_size + attn_params = q_proj + k_proj + v_proj + o_proj + + return embedding_params, attn_params + + +class MixtureMemoryEstimator(LlamaMemoryEstimator): + """ + Estimator for Mixture-of-Experts (MoE) architectures (e.g., Mixtral, MoE Llama). + Adds extra expert parallelism block parameter count to LlamaMemoryEstimator logic. + """ + + @property + def model_memory(self) -> float: + """ + Accounts for the increase in model parameters due to additional expert MLP blocks in MoE Models. + + Returns the estimated memory size of the MoE Model (in GB). + """ + c = self.llm_config + # Attention parameter count (Llama-style) + embedding_params, attn_params = self._calc_attn_embed_params() + + # MoE MLP params per layer + moe_params_per_layer = ( + c.num_local_experts * 3 * c.hidden_size * c.intermediate_size + ) + total_params = ( + c.num_hidden_layers * (attn_params + moe_params_per_layer) + + embedding_params + ) + + # Convert to GB + return total_params * c.bytes_per_parameter / 1e9 + + +def get_estimator(llm_config, **kwargs) -> MemoryEstimator: + """ + Extracts the correct estimator based on the defined parameters in the config.json + See constants.py for LLMConfig parameters necessary for specific estimators. + Uses MemoryEstimator as a fallback if parameters needed for GQA and MoE Architectures are missing. + + Returns the appropriate MemoryEstimator based on the fields defined by the model's config.json (as represented by LLMConfig). + """ + if all( + hasattr(llm_config, f) and getattr(llm_config, f) is not None + for f in MOE_REQUIRED_FIELDS + ): + return MixtureMemoryEstimator(llm_config=llm_config, **kwargs) + elif all( + hasattr(llm_config, f) and getattr(llm_config, f) is not None + for f in LLAMA_REQUIRED_FIELDS + ): + return LlamaMemoryEstimator(llm_config=llm_config, **kwargs) + else: + logger.warning( + "Falling back to generic GPT estimator: required fields missing from config.json file in model." + ) + return MemoryEstimator(llm_config=llm_config, **kwargs) diff --git a/ads/aqua/shaperecommend/llm_config.py b/ads/aqua/shaperecommend/llm_config.py new file mode 100644 index 000000000..aef2f4f50 --- /dev/null +++ b/ads/aqua/shaperecommend/llm_config.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python +# Copyright (c) 2025 Oracle and/or its affiliates. +# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ + +import re +from typing import Optional + +from pydantic import BaseModel, Field + +from ads.aqua.common.errors import AquaRecommendationError, AquaValueError +from ads.aqua.shaperecommend.constants import NEXT_QUANT, QUANT_MAPPING + + +class LLMConfig(BaseModel): + """ + Standardized configuration object for evaluating the size of Large Language Models (LLMs) + based on their architecture and quantization. + """ + + num_hidden_layers: int = Field( + ..., + description="Number of transformer blocks (layers) in the model’s neural network stack.", + ) + hidden_size: int = Field( + ..., description="Embedding dimension or hidden size of each layer." + ) + vocab_size: int = Field(..., description="Vocabulary size for input/output tokens.") + num_attention_heads: int = Field( + ..., + description="Number of attention heads (used for queries and to determine head_dim).", + ) + + head_dim: int = Field( + ..., + description="Dimension of each attention head. Typically hidden_size // num_attention_heads.", + ) + max_seq_len: Optional[int] = Field( + 8192, description="Maximum input sequence length (context window)." + ) + weight_dtype: Optional[str] = Field( + "float32", description="Parameter data type: 'float32', 'float16', etc." + ) + quantization: Optional[str] = Field( + None, + description="Quantization weight (e.g., '8bit', '4bit') or None if unquantized.", + ) + quantization_type: Optional[str] = Field( + None, + description="Quantization method (e.g., '8bit', '4bit', 'gptq', 'awq') or None if unquantized.", + ) + + num_key_value_heads: Optional[int] = Field( + None, + description="Number of key/value heads (for GQA architectures: Llama, Mistral, Falcon, Qwen, etc.). Used to determine KV cache size", + ) + + num_local_experts: Optional[int] = Field( + None, description="For MoE architectures, the number of experts per MoE layer" + ) + intermediate_size: Optional[int] = Field( + None, description="For MoE architectures, size of the MLP activation layer." + ) + + tie_word_embeddings: Optional[bool] = Field(None) + + @property + def bytes_per_parameter(self) -> float: + """ + Returns the number of bytes used to store a model parameter, + accounting for quantization or weight storage type. + """ + # Quantization takes precedence + q = (self.quantization or "").lower() + + # Direct match in mapping + if q in QUANT_MAPPING: + return QUANT_MAPPING[q] + + # Dynamic bit-width detection + m = re.match(r"(\d+)\s*bit", q) + if m: + bits = int(m[1]) + return bits / 8 # bytes per parameter + + # Fallback to dtype mapping + dtype = (self.weight_dtype or "float32").lower() + return QUANT_MAPPING.get(dtype, QUANT_MAPPING["float32"]) + + @classmethod + def detect_quantization_type(cls, raw: dict) -> Optional[str]: + """ + Detects quantization type (e.g., 'gptq', 'bitsandbytes', 'awq', etc.) from Hugging Face config dict. + """ + qcfg = raw.get("quantization_config", {}) + if raw.get("load_in_8bit") or raw.get("load_in_4bit"): + return "bitsandbytes" + for key in [ + "gptq", + "awq", + "marlin", + "bitblas", + "aqlm", + "deepspeedfp", + "gguf", + "fp8", + ]: + if key in str(qcfg).lower() or key in str(raw).lower(): + return key + return None + + @classmethod + def detect_quantization_bits(cls, raw: dict) -> Optional[str]: + """ + Detects quantization bit-width as a string (e.g., '4bit', '8bit') from Hugging Face config dict. + """ + if raw.get("load_in_8bit"): + return "8bit" + if raw.get("load_in_4bit"): + return "4bit" + if "quantization_config" in raw: + qcfg = raw["quantization_config"] + bits = qcfg.get("bits") or qcfg.get("wbits") + if bits: + return f"{bits}bit" + return None + + @property + def suggested_quantizations(self): + """ + Suggests the next lower quantization options based on the current quantization level/ weight size. + + If model is un-quantized, uses the weight size. + If model is pre-quantized, uses the quantization level. + """ + key = (self.quantization or self.weight_dtype or "float32").lower() + return NEXT_QUANT.get(key, []) + + def calculate_possible_seq_len(self, min_len=2048): + """ + Calculates a list of possible sequence lengths (in tokens). + [2048, ... max-length] (max-length found in model's config.json file) + """ + vals = [] + curr = min_len + max_seq_len = 16384 if not self.max_seq_len else self.max_seq_len + while curr <= max_seq_len: + vals.append(curr) + curr *= 2 + if vals and vals[-1] != max_seq_len: + vals.append(max_seq_len) + return vals + + def optimal_config(self): + """ + Builds a list of optimal configuration parameters (sorted descending). Combination of: + - Quantization / weight sizes: bfloat16 weight size -> 8bit -> 4bit + - max-model-len: power-of-two model lengths from max length (config.json of model) to 2048 tokens. + + Example: + [('bfloat16', max_model_len supported by model) ('bfloat16', 1/2 of max_model_len) ... ('int8', 2048), ('int4', 4096), ('int4', 2048)] + + """ + # Create a copy of the suggested_quantizations list + quantizations = self.suggested_quantizations[:] + quantizations.append("bfloat16") + + lengths = self.calculate_possible_seq_len() + + configs = [] + for quantization in quantizations: + for length in lengths: + configs.append((quantization, length)) + + configs.sort( + key=lambda x: (-QUANT_MAPPING.get(x[0], 0), -x[1]) + ) # (-quant_priority, -max_seq_len) + return configs + + @classmethod + def validate_model_support(cls, raw: dict) -> ValueError: + """ + Validates if model is decoder-only. Check for text-generation model occurs at DataScienceModel level. + """ + excluded_models = {"t5", "gemma", "bart", "bert", "roberta", "albert"} + if ( + raw.get("is_encoder_decoder", False) # exclude encoder-decoder models + or (raw.get("is_decoder") is False) # exclude explicit encoder-only models (altho no text-generation task ones, just dbl check) + or raw.get("model_type", "").lower() # exclude by known model types + in excluded_models + ): + raise AquaRecommendationError( + "Please provide a decoder-only text-generation model (ex. Llama, Falcon, etc). " + "Encoder-decoder models (ex. T5, Gemma) and encoder-only (BERT) are not supported in this tool at this time." + ) + + @classmethod + def from_raw_config(cls, raw: dict) -> "LLMConfig": + """ + Instantiates an LLMConfig from a raw Hugging Face config.json file, + using robust key detection and fallback for architecture. + """ + cls.validate_model_support(raw) + + # Field mappings with fallback + num_hidden_layers = ( + raw.get("num_hidden_layers") or raw.get("n_layer") or raw.get("num_layers") + ) + hidden_size = raw.get("hidden_size") or raw.get("n_embd") or raw.get("d_model") + vocab_size = raw.get("vocab_size") + weight_dtype = str(raw.get("torch_dtype", "float32")) + quantization = cls.detect_quantization_bits(raw) + quantization_type = cls.detect_quantization_type(raw) + + if not quantization and quantization_type in QUANT_MAPPING: + quantization = quantization_type + + num_key_value_heads = ( + raw.get("num_key_value_heads") # GQA models (ex. Llama-type) + ) + + num_attention_heads = ( + raw.get("num_attention_heads") or raw.get("n_head") or raw.get("num_heads") + ) + + head_dim = raw.get("head_dim") or ( + int(hidden_size) // int(num_attention_heads) + if hidden_size and num_attention_heads + else None + ) + max_seq_len = ( + raw.get("max_position_embeddings") + or raw.get("n_positions") + or raw.get("max_seq_len") + or 2048 + ) + + num_local_experts = ( + raw.get("num_local_experts") + or raw.get("n_routed_experts") + or raw.get("num_experts") + ) + intermediate_size = raw.get("moe_intermediate_size") or raw.get( + "intermediate_size" + ) + + # Type safety: minimal assertion + if None in [ + num_hidden_layers, + hidden_size, + vocab_size, + num_attention_heads, + head_dim, + ]: + raise ValueError("Missing required value in model config.") + + return cls( + num_hidden_layers=int(num_hidden_layers), + hidden_size=int(hidden_size), + num_attention_heads=int(num_attention_heads), + num_key_value_heads=num_key_value_heads, + head_dim=int(head_dim), + vocab_size=int(vocab_size), + weight_dtype=weight_dtype, + quantization=quantization, + quantization_type=quantization_type, + max_seq_len=int(max_seq_len), + num_local_experts=num_local_experts, + intermediate_size=intermediate_size, + ) diff --git a/ads/aqua/shaperecommend/recommend.py b/ads/aqua/shaperecommend/recommend.py new file mode 100644 index 000000000..1f1b2134e --- /dev/null +++ b/ads/aqua/shaperecommend/recommend.py @@ -0,0 +1,337 @@ +import json +from typing import List + +from pydantic import ValidationError + +from ads.aqua.app import AquaApp, logger +from ads.aqua.common.entities import ComputeShapeSummary +from ads.aqua.common.errors import ( + AquaFileNotFoundError, + AquaRecommendationError, + AquaValueError, +) +from ads.aqua.common.utils import ( + build_pydantic_error_message, + get_resource_type, + load_config, + load_gpu_shapes_index, +) +from ads.aqua.shaperecommend.constants import ( + SAFETENSORS, + TEXT_GENERATION, + TROUBLESHOOT_MSG, +) +from ads.aqua.shaperecommend.estimator import get_estimator +from ads.aqua.shaperecommend.llm_config import LLMConfig +from ads.aqua.shaperecommend.shape_report import ( + ModelConfig, + RequestRecommend, + ShapeRecommendationReport, + ShapeReport, +) +from ads.model.datascience_model import DataScienceModel + + +class AquaRecommendApp(AquaApp): + """ + Interface for recommending GPU shapes for machine learning model deployments + on Oracle Cloud Infrastructure Data Science service. + + This class provides methods to recommend deployment shapes based on a model's requirements, + handle recommendation details and troubleshooting, and retrieve specific OCI Machine Learning shapes. + Must be used within a properly configured and authenticated OCI environment. + + Methods + ------- + which_gpu(self, **kwargs) -> List[Dict]: + Lists the valid GPU deployment shapes that fit the given model and user-provided settings. + + Note: + Use `ads aqua recommend which_gpu --help` to get more details on available parameters. + """ + + def which_gpu(self, **kwargs) -> ShapeRecommendationReport: + """ + Lists valid GPU deployment shapes for the provided model and configuration. + + Validates input, retrieves the model configuration, checks the requested sequence length, + identifies available and valid compute shapes, and summarizes which shapes are compatible + with the current model settings. + + Parameters + ---------- + model_ocid : str + OCID of the model to recommend feasible compute shapes. + + Returns + ------- + ShapeRecommendationReport + A recommendation report with compatible deployment shapes, or troubleshooting info + citing the largest shapes if no shape is suitable. + + Raises + ------ + AquaValueError + If parameters are missing or invalid, or if no valid sequence length is requested. + """ + try: + request = RequestRecommend(**kwargs) + data = self.get_model_config(request.model_ocid) + llm_config = LLMConfig.from_raw_config(data) + + available_shapes = self.valid_compute_shapes() + recommendations = self.summarize_shapes_for_seq_lens( + llm_config, available_shapes + ) + + # custom error to catch model incompatibility issues + except AquaRecommendationError as error: + return ShapeRecommendationReport( + recommendations=[], troubleshoot=str(error) + ) + + except ValidationError as ex: + custom_errors = build_pydantic_error_message(ex) + raise AquaValueError( + f"Invalid parameters to read config.json of LLM Artifact. Error details: {custom_errors}." + ) from ex + except AquaValueError as ex: + logger.error(f"Error with LLM config: {ex}") + raise + + return recommendations + + @staticmethod + def get_model_config(ocid: str): + """ + Loads the configuration for a given Oracle Cloud Data Science model. + + Validates the resource type associated with the provided OCID, ensures the model + is for text-generation with a supported decoder-only architecture, and loads the model's + configuration JSON from the artifact path. + + Parameters + ---------- + ocid : str + The OCID of the Data Science model. + + Returns + ------- + dict + The parsed configuration dictionary from config.json. + + Raises + ------ + AquaValueError + If the OCID is not for a Data Science model, or if the model type is not supported, + or if required files/tags are not present. + + AquaRecommendationError + If the model OCID provided is not supported (only text-generation decoder models in safetensor format supported). + """ + resource_type = get_resource_type(ocid) + + if resource_type != "datasciencemodel": + raise AquaValueError( + f"The provided OCID '{ocid}' is not a valid Oracle Cloud Data Science Model OCID. " + "Please provide an OCID corresponding to a Data Science model resource. " + "Tip: Data Science model OCIDs typically start with 'ocid1.datasciencemodel...'." + ) + + model = DataScienceModel.from_id(ocid) + + model_task = model.freeform_tags.get("task", "").lower() + model_format = model.freeform_tags.get("model_format", "").lower() + + logger.info(f"Current model task type: {model_task}") + logger.info(f"Current model format: {model_format}") + + if TEXT_GENERATION not in model_task: + raise AquaRecommendationError( + "Please provide a decoder-only text-generation model (ex. Llama, Falcon, etc.). " + f"Only text-generation models are supported in this tool at this time. Current model task type: {model_task}" + ) + if SAFETENSORS not in model_format: + msg = "Please provide a model in Safetensor format." + if model_format: + msg += f"The current model format ({model_format}) is not supported by this tool at this time." + + raise AquaRecommendationError(msg) + + if not model.artifact: + raise AquaValueError( + "Unable to retrieve model artifact. Ensure model is registered and active." + ) + + try: + data = load_config(model.artifact, "config.json") + + except AquaFileNotFoundError as e: + logger.error( + f"config.json not found in model artifact at {model.artifact}: {e}" + ) + raise AquaRecommendationError( + "The configuration file 'config.json' was not found in the specified model directory. " + "Please ensure your model follows the Hugging Face format and includes a 'config.json' with the necessary architecture parameters." + ) from e + + return data + + @staticmethod + def valid_compute_shapes() -> List["ComputeShapeSummary"]: + """ + Returns a filtered list of GPU-only ComputeShapeSummary objects by reading and parsing a JSON file. + + Parameters + ---------- + file : str + Path to the JSON file containing shape data. + + Returns + ------- + List[ComputeShapeSummary] + List of ComputeShapeSummary objects passing the checks. + + Raises + ------ + ValueError + If the file cannot be opened, parsed, or the 'shapes' key is missing. + """ + gpu_shapes_metadata = load_gpu_shapes_index().shapes + + valid_shapes = [] + for name, spec in gpu_shapes_metadata.items(): + valid_shapes.append( + ComputeShapeSummary(name=name, shape_series="GPU", gpu_specs=spec) + ) + valid_shapes.sort( + key=lambda shape: shape.gpu_specs.gpu_memory_in_gbs, reverse=True + ) + return valid_shapes + + @staticmethod + def summarize_shapes_for_seq_lens( + config: LLMConfig, + shapes: List[ComputeShapeSummary], + batch_size: int = 1, + ) -> ShapeRecommendationReport: + """ + Generate a recommendation report for eligible deployment shapes by evaluating + model memory consumption and maximum model length for given configurations. + + Parameters + ---------- + config : LLMConfig + The loaded model configuration. + shapes : List[ComputeShapeSummary] + All candidate deployment shapes. + batch_size : int, optional + Batch size to evaluate (default is 1). + + Returns + ------- + ShapeRecommendationReport + Report containing shape recommendations and troubleshooting advice, if any. + + Raises + ------ + ValueError + If no GPU shapes are available. + + Notes + ----- + - Considers quantization if defined in config, otherwise cycles through optimal configs. + - Applies pareto optimality if too many recommendations. + - Provides troubleshooting options if nothing fits. + """ + recommendations = [] + + if not shapes: + raise ValueError( + "No GPU shapes were passed for recommendation. Ensure shape parsing succeeded." + ) + + # Pre-quantized: only consider different max-seq-len + if config.quantization_type: + deployment_config = config.calculate_possible_seq_len() + for shape in shapes: + if config.quantization_type in shape.gpu_specs.quantization: + allowed_gpu_memory = shape.gpu_specs.gpu_memory_in_gbs + for max_seq_len in deployment_config: + estimator = get_estimator( + llm_config=config, + seq_len=max_seq_len, + batch_size=batch_size, + ) + if estimator.validate_shape(estimator, allowed_gpu_memory): + best_config = [ + ModelConfig.constuct_model_config( + estimator, allowed_gpu_memory + ) + ] + recommendations.append( + ShapeReport( + shape_details=shape, configurations=best_config + ) + ) + break + + # unquantized: consider inflight quantization (4bit and 8bit) + else: + deployment_config = config.optimal_config() + prev_quant = None + for shape in shapes: + allowed_gpu_memory = shape.gpu_specs.gpu_memory_in_gbs + for quantization, max_seq_len in deployment_config: + if quantization != prev_quant: + updated_config = config.model_copy( + update={"quantization": quantization} + ) + prev_quant = quantization + estimator = get_estimator( + llm_config=updated_config, + seq_len=max_seq_len, + batch_size=batch_size, + ) + if estimator.validate_shape(allowed_gpu_memory): + best_config = [ + ModelConfig.constuct_model_config( + estimator, allowed_gpu_memory + ) + ] + recommendations.append( + ShapeReport(shape_details=shape, configurations=best_config) + ) + break + + troubleshoot_msg = "" + + if len(recommendations) > 2: + recommendations = ShapeReport.pareto_front(recommendations) + + if not recommendations: + # Troubleshooting advice if nothing fits + # Assumes shapes is sorted largest to smallest and quantizations 'fp8'/'4bit' exist + troubleshoot_msg += TROUBLESHOOT_MSG + + largest_shapes = ( + [(shapes[0], "fp8"), (shapes[1], "4bit")] if len(shapes) > 1 else [] + ) + for shape, quantization in largest_shapes: + updated_config = config.model_copy( + update={"quantization": quantization} + ) + estimator = get_estimator( + llm_config=updated_config, seq_len=2048, batch_size=batch_size + ) + allowed_gpu_memory = shape.gpu_specs.gpu_memory_in_gbs * 0.9 + best_config = [ + ModelConfig.constuct_model_config(estimator, allowed_gpu_memory) + ] + recommendations.append( + ShapeReport(shape_details=shape, configurations=best_config) + ) + + return ShapeRecommendationReport( + recommendations=recommendations, troubleshoot=troubleshoot_msg + ) diff --git a/ads/aqua/shaperecommend/shape_report.py b/ads/aqua/shaperecommend/shape_report.py new file mode 100644 index 000000000..4968fe289 --- /dev/null +++ b/ads/aqua/shaperecommend/shape_report.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python +# Copyright (c) 2025 Oracle and/or its affiliates. +# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ + +from typing import List, Optional + +from pydantic import BaseModel, Field + +from ads.aqua.common.entities import ComputeShapeSummary +from ads.aqua.shaperecommend.constants import QUANT_MAPPING +from ads.aqua.shaperecommend.estimator import MemoryEstimator + + +class RequestRecommend(BaseModel): + """ + A request to recommend compute shapes and parameters for a given model. + """ + + model_ocid: str = Field(..., description="The OCID of the model to recommend feasible compute shapes.") + + +class DeploymentParams(BaseModel): # noqa: N801 + """ + Recommended parameters for deployment and model inferencing (specific to compute shape & model). + """ + + quantization: Optional[str] = Field( + None, description="Type of quantization (e.g. int8)." + ) + max_model_len: int = Field(..., description="Maximum length of input sequence.") + batch_size: Optional[int] = Field(1, description="Batch size for training.") + + +class ModelDetail(BaseModel): + """ + The estimated memory footprint of a model, KV cache, and its total (model + KV cache). + """ + + model_size_gb: float = Field(..., description="Size of the model in GB.") + kv_cache_size_gb: float = Field(..., description="Size of KV cache in GB.") + total_model_gb: float = Field( + ..., description="Total size of model and cache in GB." + ) + + +class ModelConfig(BaseModel): + """ + The configuration for a model based on specific set of deployment parameters and memory capacity of shape. + """ + + model_details: ModelDetail = Field(..., description="Details about the model.") + deployment_params: DeploymentParams = Field( + ..., description="Parameters for deployment." + ) + recommendation: str = Field(..., description="GPU recommendation for the model.") + + @classmethod + def constuct_model_config(cls, estimator: MemoryEstimator, allowed_gpu_memory: float) -> "ModelConfig": + """ + Assembles a complete ModelConfig, including model details, deployment parameters (vLLM), and recommendations. + + Parameters + ---------- + estimator : MemoryEstimator + Estimator with model details and processed config. + allowed_gpu_memory : float + Maximum allowed GPU memory (in GBs) for this configuration. + + Returns + ------- + ModelConfig + Contains round-tripped model size, kv cache, total, vLLM parameters, and recommendations. + + Notes + ----- + - Rounds all sizes to 3 decimal digits. + - Computes a recommendation string using `limiting_factor`. + """ + deployment_params = DeploymentParams( + quantization=getattr(estimator.llm_config, "quantization", None), + max_model_len=getattr(estimator, "seq_len", None) + ) + model_detail = ModelDetail( + model_size_gb=round(getattr(estimator, "model_memory", 0.0), 3), + kv_cache_size_gb=round(getattr(estimator, "kv_cache_memory", 0.0), 3), + total_model_gb=round(getattr(estimator, "total_memory", 0.0), 3) + ) + return ModelConfig( + model_details=model_detail, + deployment_params=deployment_params, + recommendation= estimator.limiting_factor(allowed_gpu_memory) + ) + + +class ShapeReport(BaseModel): + """ + The feasible deployment configurations for the model per shape. + """ + shape_details: 'ComputeShapeSummary' = Field( + ..., description="Details about the compute shape (ex. VM.GPU.A10.2)." + ) + configurations: List['ModelConfig'] = Field( + default_factory=list, description="List of model configurations." + ) + + def is_dominated(self, others: List['ShapeReport']) -> bool: + """ + Determines whether this shape is dominated by any other shape in a Pareto sense. + + Parameters + ---------- + others : list of ShapeReport + List of other shape/deployment configurations to compare against. + + Returns + ------- + bool + True if this shape is dominated by at least one other, False otherwise. + + Notes + ----- + A shape is dominated if there exists another configuration that is + at least as good in all criteria and strictly better in at least one. + Criteria: + - Cost (to be minimized) + - Performance, quantization level, max sequence length (to be maximized) + """ + + cand_cost = self.shape_details.gpu_specs.ranking.cost + cand_perf = self.shape_details.gpu_specs.ranking.performance + cand_quant = QUANT_MAPPING.get(self.configurations[0].deployment_params.quantization, 0) + cand_maxlen = self.configurations[0].deployment_params.max_model_len + + for other in others: + other_cost = other.shape_details.gpu_specs.ranking.cost + other_perf = other.shape_details.gpu_specs.ranking.performance + other_quant = QUANT_MAPPING.get(other.configurations[0].deployment_params.quantization, 0) + other_maxlen = other.configurations[0].deployment_params.max_model_len + if ( + other_cost <= cand_cost and + other_perf >= cand_perf and + other_quant >= cand_quant and + other_maxlen >= cand_maxlen and + ( + other_cost < cand_cost or + other_perf > cand_perf or + other_quant > cand_quant or + other_maxlen > cand_maxlen + ) + ): + return True + return False + + @classmethod + def pareto_front(cls, shapes: List['ShapeReport']) -> List['ShapeReport']: + """ + Filters a list of shapes/configurations to those on the Pareto frontier. + + Parameters + ---------- + shapes : list of ShapeReport + List of candidate shape/configuration reports to evaluate. + + Returns + ------- + list of ShapeReport + Subset of input shapes that are not dominated by any other (the Pareto front). + + Notes + ----- + The returned set contains non-dominated deployments for maximizing + performance, quantization, and model length, while minimizing cost. + """ + return [shape for shape in shapes if not shape.is_dominated([s for s in shapes if s != shape])] + + +class ShapeRecommendationReport(BaseModel): + """ + Full report of shape fit recommendations and troubleshooting, if applicable. + + Attributes: + recommendations (List[DeploymentShapeSummary]): Recommended deployment shapes + for each tested batch size and max sequence length combination. + troubleshoot (Optional[TroubleshootShapeSummary]): Troubleshooting information + if no valid deployment shapes are available. + """ + + recommendations: List[ShapeReport] = Field( + default_factory=list, description="List of shape fit recommendations." + ) + troubleshoot: Optional[str] = Field( + None, + description="details for troubleshooting if no shapes fit the current model.", + )