llm_config.py Source File

llm_config.py Source File#

Mobilint SDK qb Compiler: llm_config.py Source File
Mobilint SDK qb Compiler v0.12.0.0
MCS002-EN
from typing import List, Optional
from typeguard import typechecked
from . import ConfigABC
 
 
class LlmConfig(ConfigABC):
    """
    @brief Configuration for Large Language Model (LLM) compilation.
    """
 
    optional = True
    DEFAULTS = {
        "llm_config_apply": False,
        "max_data_length": 4096,
        "max_sequence_length": 4096,
        "max_cache_length": 4096,
        "max_core_data_length": 128,
        "random_seq_length": 80,
        "use_full_seq_length": False,
        "use_custom_mask_input": False,
        "start_dram_offset": 0,
        "use_global_core": False,
        "batch_size": 1,
        "npu_core_ids": [0],
        "dynamic_rope": False,
    }
 
    py_to_cpp_name = {
        "llm_config_apply": "apply",
        "max_data_length": "maxDataLength",
        "max_sequence_length": "maxSequenceLength",
        "max_cache_length": "maxCacheLength",
        "max_core_data_length": "maxCoreDataLength",
        "random_seq_length": "randomSeqLength",
        "use_full_seq_length": "useFullSeqLength",
        "use_custom_mask_input": "useCustomMaskInput",
        "start_dram_offset": "startDramOffset",
        "use_global_core": "useGlobalCore",
        "batch_size": "batchSize",
        "dynamic_rope": "dynamicRope",
    }
 
    min_max_check = {
        "max_data_length": [0, None],
        "max_sequence_length": [0, None],
        "max_cache_length": [0, None],
        "max_core_data_length": [0, None],
        "random_seq_length": [0, None],
        "start_dram_offset": [0, None],
        "batch_size": [0, None],
    }
 
    def __init__(
        self,
        llm_config_apply: bool,
        max_data_length: int,
        max_sequence_length: int,
        max_cache_length: int,
        max_core_data_length: int,
        random_seq_length: int,
        use_full_seq_length: bool,
        use_custom_mask_input: bool,
        start_dram_offset: int,
        use_global_core: bool,
        batch_size: int,
        npu_core_ids: List[int],
        dynamic_rope: bool,
        optional: bool = optional,
    ):
        """
        @brief Initialize the LlmConfig:
        @param llm_config_apply bool. If True, apply LLM-specific configurations.
        @param max_data_length int. Maximum data length.
        @param max_sequence_length int. Maximum sequence length.
        @param max_cache_length int. Maximum cache length.
        @param max_core_data_length int. Maximum core data length.
        @param random_seq_length int. Random sequence length used for calibration.
        @param use_full_seq_length bool. If True, use the full sequence length for calibration.
        @param use_custom_mask_input bool. If True, use a custom mask input.
        @param start_dram_offset int. Starting DRAM offset.
        @param use_global_core bool. If True, use a global core.
        @param batch_size int. Batch size.
        @param npu_core_ids List[int]. List of NPU core IDs.
        @param dynamic_rope bool. If True, enable dynamic RoPE.
        @param optional bool. Indicates whether this config is optional. Use the default value defined as a class variable.
        """
        super().__init__(optional)
        self.llm_config_apply = llm_config_apply
        self.max_data_length = max_data_length
        self.max_sequence_length = max_sequence_length
        self.max_cache_length = max_cache_length
        self.max_core_data_length = max_core_data_length
        self.random_seq_length = random_seq_length
        self.use_full_seq_length = use_full_seq_length
        self.use_custom_mask_input = use_custom_mask_input
        self.start_dram_offset = start_dram_offset
        self.use_global_core = use_global_core
        self.batch_size = batch_size
        self.npu_core_ids = npu_core_ids
        self.dynamic_rope = dynamic_rope
        self.check_valid(except_list=["optional", "npu_core_ids"])
 
    @classmethod
    def default_config(cls):
        return cls(**cls.DEFAULTS, optional=cls.optional)
 
    @classmethod
    def from_kwargs(cls, **kwargs):
        params = cls.DEFAULTS.copy()
        updated_kwargs = dict()
        for key in cls.DEFAULTS.keys():
            if key in kwargs:
                updated_kwargs[key] = kwargs[key]
                params[key] = kwargs[key]
 
        try:
            return cls(**params, optional=cls.optional)
        except:
            raise ValueError(
                f"Please review the configuration kwargs passed to {cls.__name__}. These kwargs were updated and will be used: {updated_kwargs}"
            )
 
    @classmethod
    def from_dict(cls, config_dict: dict):
        """Create LlmConfig from dictionary (JSON structure)"""
        params = cls.DEFAULTS.copy()
 
        if "llm_config_apply" in config_dict:
            params["llm_config_apply"] = config_dict["apply"]
 
        if "attributes" in config_dict:
            attributes = config_dict["attributes"]
 
            if "maxDataLength" in attributes:
                params["max_data_length"] = attributes["maxDataLength"]
            if "maxSequenceLength" in attributes:
                params["max_sequence_length"] = attributes["maxSequenceLength"]
            if "maxCacheLength" in attributes:
                params["max_cache_length"] = attributes["maxCacheLength"]
            if "maxCoreDataLength" in attributes:
                params["max_core_data_length"] = attributes["maxCoreDataLength"]
 
            if "calibration" in attributes:
                calibration = attributes["calibration"]
                if "randomSeqLength" in calibration:
                    params["random_seq_length"] = calibration["randomSeqLength"]
                if "useFullSeqLength" in calibration:
                    params["use_full_seq_length"] = calibration["useFullSeqLength"]
                if "useCustomMaskInput" in calibration:
                    params["use_custom_mask_input"] = calibration["useCustomMaskInput"]
 
            if "runtime" in attributes:
                runtime = attributes["runtime"]
                if "startDramOffset" in runtime:
                    params["start_dram_offset"] = runtime["startDramOffset"]
                if "useGlobalCore" in runtime:
                    params["use_global_core"] = runtime["useGlobalCore"]
                if "batchSize" in runtime:
                    params["batch_size"] = runtime["batchSize"]
                if "npuCoreIds" in runtime:
                    params["npu_core_ids"] = runtime["npuCoreIds"]
                if "dynamicRope" in runtime:
                    params["dynamic_rope"] = runtime["dynamicRope"]
 
        return cls(**params)
 
    def to_dict(self):
        return {
            "llmConfig": {
                "apply": self.llm_config_apply,
                "attributes": {
                    "maxDataLength": self.max_data_length,
                    "maxSequenceLength": self.max_sequence_length,
                    "maxCacheLength": self.max_cache_length,
                    "maxCoreDataLength": self.max_core_data_length,
                    "calibration": {
                        "randomSeqLength": self.random_seq_length,
                        "useFullSeqLength": self.use_full_seq_length,
                        "useCustomMaskInput": self.use_custom_mask_input,
                    },
                    "runtime": {
                        "startDramOffset": self.start_dram_offset,
                        "useGlobalCore": self.use_global_core,
                        "batchSize": self.batch_size,
                        "npuCoreIds": self.npu_core_ids,
                        "dynamicRope": self.dynamic_rope,
                    },
                    "debug": {
                        "apply": False,
                        "batchDebugBundleSize": 0,
                    },
                },
            }
        }
 
 
def get_llm_config(**kwargs) -> LlmConfig:
    """
    @brief Create LlmConfig with partial parameters merged with defaults.
    """
    params = LlmConfig.DEFAULTS.copy()
    params.update(kwargs)
    return LlmConfig(**params, optional=LlmConfig.optional)
 
 
# @}