Class LLMAnalysis

llm_analysis/analysis.py:78–2248

Given the specified model, GPU, data type, parallelism configuration/implementation, LLMAnalysis estimates the latency and memory usage of LLMs for training or inference. Refer to the `train` and `infer` entry functions for usage details.

class LLMAnalysis:
    """Given the specified model, GPU, data type, parallelism
    configuration/implementation, LLMAnalysis estimates the latency and memory usage of
    LLMs for training or inference.

    Refer to the `train` and `infer` entry functions for usage details.
    """

    def __init__(
        self,
        model_config: ModelConfig,
        gpu_config: GPUConfig,
        dtype_config: DtypeConfig = DtypeConfig(),
        parallelism_config: ParallelismConfig = ParallelismConfig(),
        achieved_tflops: float = None,
        achieved_memory_bandwidth_GBs: float = None,
        flops_efficiency: float = None,
        hbm_memory_efficiency: float = None,
        intra_node_memory_efficiency: float = INTRA_NODE_MEMORY_EFFICIENCY,
        inter_node_memory_efficiency: float = INTER_NODE_MEMORY_EFFICIENCY,
    ) -> None:
        """LLMAnalysis constructor.

        Args:
            model_config (ModelConfig): model configuration
            gpu_config (GPUConfig): GPU configuration
            dtype_config (DtypeConfig, optional): data type configuration. Defaults to DtypeConfig().
            parallelism_config (ParallelismConfig, optional): parallelism configuration. Defaults to ParallelismConfig().
            achieved_tflops (float, optional): achieved TFLOPS per GPU. If specified, will override the flops_efficiency passed in. Defaults to None.
            achieved_memory_bandwidth_GBs (float, optional): achieved GPU memory bandwidth in GB/s. If specified, will override the hbm_memory_efficiency passed in. Defaults to None.
            flops_efficiency (float, optional): flops efficiency, ranging from 0 to 1. Defaults to None.
            hbm_memory_efficiency (float, optional): GPU HBM memory efficiency, ranging from 0 to 1. Defaults to None.
            intra_node_memory_efficiency (float, optional): intra-node memory efficiency, ranging from 0 to 1. Defaults to INTRA_NODE_MEMORY_EFFICIENCY.
            inter_node_memory_efficiency (float, optional): inter-node memory efficiency, ranging from 0 to 1. Defaults to INTER_NODE_MEMORY_EFFICIENCY.
        """
        self.model_config = model_config
        self.gpu_config = gpu_config
        self.parallelism_config = parallelism_config
        self.dtype_config = dtype_config
        self.intra_node_memory_efficiency = intra_node_memory_efficiency
        self.inter_node_memory_efficiency = inter_node_memory_efficiency

        if achieved_memory_bandwidth_GBs and hbm_memory_efficiency:
            logger.info(
                "both achieved_memory_bandwidth_GBs and hbm_memory_efficiency are set, using achieved_memory_bandwidth_GBs({achieved_memory_bandwidth_GBs} GB/s) to calculate hbm_memory_efficiency"
            )
            self.hbm_memory_efficiency = (
                achieved_memory_bandwidth_GBs /
                gpu_config.hbm_bandwidth_in_GB_per_sec)
        elif hbm_memory_efficiency:
            self.hbm_memory_efficiency = hbm_memory_efficiency
        elif achieved_memory_bandwidth_GBs:
            self.hbm_memory_efficiency = (
                achieved_memory_bandwidth_GBs /
                gpu_config.hbm_bandwidth_in_GB_per_sec)
        else:
            self.hbm_memory_efficiency = HBM_MEMORY_EFFICIENCY