llm_config.py Source File

llm_config.py Source File#

Mobilint SDK qb Compiler: llm_config.py Source File
Mobilint SDK qb Compiler v0.11.0.1
MCS002-EN
import inspect
from functools import wraps
 
from typeguard import typechecked
from . import ConfigABC
 
"""
max_data_len : maximum data length, sum of maximum sequence lengths on all NPU cores. Defaults to 4096.
max_seq_len : maximum sequence length. Defaults to 4096.
max_cache_len : maximum cache length. Defaults to 4096.
random_seq_len_calib : sequence length for random calibration. Defaults to 80.
use_full_seq_len_calib : whether to use full sequence length for calibration for the last attention layer. Defaults to False.
use_custom_mask_input : whether to use custom mask input. Defaults to False.
start_dram_offset : start dram offset. Defaults to 0.
use_global_core : whether to use global core. Defaults to False.
llm_batch_size : batch size for LLM. Defaults to 1.
npu_core_ids : NPU core IDs to use(batched LLM). Defaults to [0]. If the llm_batch_size is greater than 1, then defaults to [0, 1, 2, 3, 4, 5, 6, 7].
"""
 
 
class LlmConfig(ConfigABC):
    py_to_cpp_name = {
        "max_data_len": "maxDataLen",
        "max_seq_len": "maxSeqLen",
        "max_cache_len": "maxCacheLen",
        "random_seq_len_calib": "randomSeqLenCalib",
        "use_full_seq_len_calib": "useFullSeqLenCalib",
        "use_custom_mask_input": "useCustomMaskInput",
        "start_dram_offset": "startDramOffset",
        "use_global_core": "useGlobalCore",
        "llm_batch_size": "llmBatchSize",
        "npu_core_ids": "npuCoreIds",
    }
 
    min_max_check = {
        "max_data_len": [0, None],
        "max_seq_len": [0, None],
        "max_cache_len": [0, None],
        "random_seq_len_calib": [0, None],
        "use_full_seq_len_calib": [0, 1],
        "use_custom_mask_input": [0, 1],
        "start_dram_offset": [0, None],
        "use_global_core": [0, 1],
        "llm_batch_size": [0, None],
    }
 
    @typechecked
    def __init__(
        self,
        max_data_len: int,
        max_seq_len: int,
        max_cache_len: int,
        random_seq_len_calib: int,
        use_full_seq_len_calib: bool,
        use_custom_mask_input: bool,
        start_dram_offset: int,
        use_global_core: bool,
        llm_batch_size: int,
        npu_core_ids: list[int],
        optional: bool = True,
    ):
        super().__init__(optional)
        self.max_data_len = max_data_len
        self.max_seq_len = max_seq_len
        self.max_cache_len = max_cache_len
        self.random_seq_len_calib = random_seq_len_calib
        self.use_full_seq_len_calib = use_full_seq_len_calib
        self.use_custom_mask_input = use_custom_mask_input
        self.start_dram_offset = start_dram_offset
        self.use_global_core = use_global_core
        self.llm_batch_size = llm_batch_size
        self.npu_core_ids = npu_core_ids
        self.check_valid()
 
    @classmethod
    def default_config(cls):
        max_data_len = 4096
        max_seq_len = 4096
        max_cache_len = 4096
        random_seq_len_calib = 80
        use_full_seq_len_calib = False
        use_custom_mask_input = False
        start_dram_offset = 0
        use_global_core = False
        llm_batch_size = 1
        npu_core_ids = [0]
 
        return LlmConfig(
            optional=True,
            max_data_len=max_data_len,
            max_seq_len=max_seq_len,
            max_cache_len=max_cache_len,
            random_seq_len_calib=random_seq_len_calib,
            use_full_seq_len_calib=use_full_seq_len_calib,
            use_custom_mask_input=use_custom_mask_input,
            start_dram_offset=start_dram_offset,
            use_global_core=use_global_core,
            llm_batch_size=llm_batch_size,
            npu_core_ids=npu_core_ids,
        )
 
    def to_dict(self):
        llm_config_dict = dict()
        llm_params = dict()
        for py_name, cpp_name in LlmConfig.py_to_cpp_name.items():
            value = getattr(self, py_name)
            llm_params[cpp_name] = value
        llm_config_dict["LLM"] = llm_params
        return llm_config_dict
 
 
def track_passed_args(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        sig = inspect.signature(fn)
        bound_args = sig.bind_partial(*args, **kwargs)
 
        # If llm_batch_size is larger than 1, set npu_core_ids to [0, 1, 2, 3, 4, 5, 6, 7]
        npu_core_ids = []
        if "llm_batch_size" in bound_args.arguments:
            llm_batch_size = bound_args.arguments["llm_batch_size"]
            if llm_batch_size > 1 and "npu_core_ids" not in bound_args.arguments:
                npu_core_ids = [0, 1, 2, 3, 4, 5, 6, 7]
 
        # Update default npu_core_ids if necessary
        bound_args.apply_defaults()
        if npu_core_ids:
            bound_args.arguments["npu_core_ids"] = npu_core_ids
 
        return fn(**bound_args.arguments)
 
    return wrapper
 
 
@track_passed_args
def get_llm_config(
    max_data_len: int = 4096,
    max_seq_len: int = 4096,
    max_cache_len: int = 4096,
    random_seq_len_calib: int = 80,
    use_full_seq_len_calib: bool = False,
    use_custom_mask_input: bool = False,
    start_dram_offset: int = 0,
    use_global_core: bool = False,
    llm_batch_size: int = 1,
    npu_core_ids: list[int] = [0],
) -> LlmConfig:
    return LlmConfig(
        max_data_len=max_data_len,
        max_seq_len=max_seq_len,
        max_cache_len=max_cache_len,
        random_seq_len_calib=random_seq_len_calib,
        use_full_seq_len_calib=use_full_seq_len_calib,
        use_custom_mask_input=use_custom_mask_input,
        start_dram_offset=start_dram_offset,
        use_global_core=use_global_core,
        llm_batch_size=llm_batch_size,
        npu_core_ids=npu_core_ids,
    )