Method inference

llm_analysis/analysis.py:1421–1736

Inference analysis given the configs and inputs. Args: batch_size_per_gpu (int, optional): batch size per gpu. Defaults to 1. seq_len (int, optional): number of input tokens. Defaults to 512. num_tokens_to_generate (int, optional): number of tokens to gen

(
        self,
        batch_size_per_gpu: int = 1,
        seq_len: int = 512,
        num_tokens_to_generate: int = 32,
        use_kv_cache: bool = True,
        ds_zero: DSZeRO = DSZeRO.NONE,
        layernorm_dtype_bytes: int = BYTES_FP16,
        kv_cache_dtype_bytes: int = None,
        cost_per_gpu_hour: float = None,
        output_dir: str = None,
        output_file_suffix: str = "",
    )

            )

    def inference(
        self,
        batch_size_per_gpu: int = 1,
        seq_len: int = 512,
        num_tokens_to_generate: int = 32,
        use_kv_cache: bool = True,
        ds_zero: DSZeRO = DSZeRO.NONE,
        layernorm_dtype_bytes: int = BYTES_FP16,
        kv_cache_dtype_bytes: int = None,
        cost_per_gpu_hour: float = None,
        output_dir: str = None,
        output_file_suffix: str = "",
    ) -> dict:
        """Inference analysis given the configs and inputs.

        Args:
            batch_size_per_gpu (int, optional): batch size per gpu. Defaults to 1.
            seq_len (int, optional): number of input tokens. Defaults to 512.
            num_tokens_to_generate (int, optional): number of tokens to generate for generative models. Defaults to 32.
            use_kv_cache (bool, optional): whether to use kv_cache. Defaults to True.
            ds_zero (DSZeRO, optional): which DeepSpeed ZeRO stage to use. Defaults to DSZeRO.NONE (disabled).
            layernorm_dtype_bytes (int, optional): number of bytes in the data type for the layernorm activations. Defaults to BYTES_FP32. Often has to be at least FP16 in inference to maintain model accuracy.
            kv_cache_dtype_bytes (int, optional): number of bytes in the data type for the kv_cache. Defaults to None. Often has to be at least FP16 in inference to maintain model accuracy.
            cost_per_gpu_hour (float, optional): dollar cost per GPU hour. Defaults to None.
            output_dir (str, optional): if set to a directory path, write the return summary dict out to the directory with the setup. Defaults to None.
            output_dir (str, optional): if set to a directory path, write the return summary dict out to the directory with the setup. Defaults to None.

        Returns:
            dict: a summary dict of the training analysis
        """
        if self.model_config.max_seq_len is not None:
            assert (
                seq_len <= self.model_config.max_seq_len
            ), f"seq_len must be less than model max_seq_len ({self.model_config.max_seq_len})"

        self.print_config("Inference Configs")

        logger.info(f"\n{'Analysis'.center(PRINT_LINE_WIDTH, '-')}")

        if kv_cache_dtype_bytes is None:
            kv_cache_dtype_bytes = (self.dtype_config.activation_bits /
                                    BITS_PER_BYTE)
            logger.info(
                "kv_cache_dtype_bytes not specified, setting to the same as"
                f" the activation data type : {kv_cache_dtype_bytes}")

        num_layers_per_gpu = int(self.model_config.num_layers /
                                 self.parallelism_config.pp_size)
        if self.model_config.num_layers % self.parallelism_config.pp_size:
            logger.info(
                "num_layers not be divisible by pp_size, taking the floor")

        weight_memory_embedding_per_gpu = self.get_memory_embedding(ds_zero)
        weight_memory_layers_per_gpu, weight_memory_attn_per_gpu, weight_memory_mlp_per_gpu, weight_memory_layernorm_per_gpu = [
            x * self.model_config.num_layers
            for x in self.get_weight_memory_per_layer(ds_zero,
                                                      return_breakdown=True)
        ]