oracle-ads 2.13.18rc0__py3-none-any.whl → 2.13.19__py3-none-any.whl

ads/aqua/shaperecommend/estimator.py

@@ -0,0 +1,384 @@
+#!/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 (
+    IN_FLIGHT_QUANTIZATION,
+    LLAMA_REQUIRED_FIELDS,
+    MOE_REQUIRED_FIELDS,
+    NEXT_QUANT,
+    QUANT_MAPPING,
+    VLLM_PARAMS,
+)
+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 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, gpu_utilization: float = 0.9
+    ) -> 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.
+        """
+        return (allowed_gpu_memory * gpu_utilization) > self.total_memory
+
+    def construct_deployment_params(self) -> str:
+        """
+        Constructs a deployment parameter string for the model.
+
+        This method assembles runtime configuration parameters to be passed
+        during model deployment. It:
+        - Overrides the max sequence length if a shorter length is provided.
+        - Suggests in-flight quantization **only if the model is unquantized**
+            and in-flight quantization (such as '4bit') is requested in config.
+
+        Returns
+        -------
+            str: Parameter string for model deployment.
+        """
+        c = self.llm_config
+        params = []
+        if self.seq_len < c.max_seq_len:
+            params.append(VLLM_PARAMS["max_model_len"])
+            params.append(str(self.seq_len))
+
+        # Only suggest in-flight quantization for unquantized models when such quantization is requested
+        if not c.quantization and c.in_flight_quantization in IN_FLIGHT_QUANTIZATION:
+            # vLLM only supports 4bit in-flight quantization
+            params.append(VLLM_PARAMS["in_flight_quant"])
+
+        params = " ".join(params) if params else ""
+        return params
+
+    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(config.suggested_quantizations)
+        quantization = getattr(config, "quantization", None)
+
+        advice = []
+
+        if config.suggested_quantizations:
+            to_do = f", which is smaller than the current {quantization if quantization in NEXT_QUANT else weight_size} format."
+            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"Either use a pre-quantized model at {quant_advice}, or apply in-flight {quant_advice} quantization"
+                    + to_do
+                )
+
+        kv_advice = [f"Reduce maximum context length (set --max-model-len < {seq_len})"]
+
+        if batch_size != 1:
+            kv_advice.append(f"Reduce batch size to less than {batch_size}.")
+
+        wt_advice = [
+            "Use a model with fewer parameters.",
+            f"{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 = kv_advice
+            advice.extend(wt_advice)
+
+        advice_str = "\n".join(f"{i}. {item}" for i, item in enumerate(advice, 1))
+
+        return (
+            f"{advice_str}\n\n{main} (KV cache: {kv_gb:.1f}GB, Weights: {wt_gb:.1f}GB)."
+        )
+
+    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
+
+        # 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.\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"No override PARAMS needed. \n\nModel fits well within the allowed compute shape "
+                f"({required:.1f}GB used / {allowed_gpu_memory:.1f}GB allowed)."
+            )
+        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)

ads/aqua/shaperecommend/llm_config.py

@@ -0,0 +1,283 @@
+#!/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
+from ads.aqua.shaperecommend.constants import (
+    BITS_AND_BYTES_4BIT,
+    BITS_AND_BYTES_8BIT,
+    DEFAULT_WEIGHT_SIZE,
+    NEXT_QUANT,
+    QUANT_MAPPING,
+    QUANT_METHODS,
+)
+
+
+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(
+        4096, description="Maximum input sequence length (context window)."
+    )
+    weight_dtype: Optional[str] = Field(
+        DEFAULT_WEIGHT_SIZE,
+        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.",
+    )
+
+    in_flight_quantization: Optional[str] = Field(
+        None,
+        description="By setting this, enables recalculation of model footprint using 4bit in-flight quantization",
+    )
+
+    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
+
+        # consider in-flight quantization
+        if self.in_flight_quantization in QUANT_MAPPING:
+            return QUANT_MAPPING[self.in_flight_quantization]
+
+        # Fallback to dtype mapping
+        dtype = (self.weight_dtype or DEFAULT_WEIGHT_SIZE).lower()
+        return QUANT_MAPPING.get(dtype, QUANT_MAPPING[DEFAULT_WEIGHT_SIZE])
+
+    @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 QUANT_METHODS:
+            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 BITS_AND_BYTES_8BIT
+        if raw.get("load_in_4bit"):
+            return BITS_AND_BYTES_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.in_flight_quantization
+            or self.weight_dtype
+            or DEFAULT_WEIGHT_SIZE
+        ).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
+        while curr <= self.max_seq_len:
+            vals.append(curr)
+            curr *= 2
+        if vals and vals[-1] != self.max_seq_len:
+            vals.append(self.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) ... ('4bit', 4096), ('4bit', 2048)]
+
+        """
+        # use later-Create a copy of the suggested_quantizations list
+        # quantizations = self.suggested_quantizations[:]
+        quantizations = ["bfloat16", "4bit"]
+
+        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 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", DEFAULT_WEIGHT_SIZE))
+        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,
+        )