Method get_activation_memory_per_layer

llm_analysis/analysis.py:737–843

Get the memory (in bytes) required to store the activations of a transformer layer, given the batch size, sequence length, and whether it is inference or training, the activation recomputation strategy, and the activation data type. Refer to https://arxiv.org/abs/2205.05198 f

(
        self,
        batch_size: int,
        seq_len: int,
        is_inference: bool = True,
        activation_recomputation:
        ActivationRecomputation = ActivationRecomputation.NONE,
        layernorm_dtype_bytes: int = BYTES_FP32,
        flash_attn: bool = True,
        softmax_dropout: bool = False,
        mlp_activation_quant_bits: int = None,
        mlp_1linear_quant_bits: int = None,
        mlp_gelu_input_quant_bits: int = None,
        mlp_2linear_quant_bits: int = None,
        mlp_recompute_gelu: bool = False,
        mlp_gated_linear_units: bool = False,
        return_breakdown: bool = False,
    )

        return self.model_config.vocab_size * batch_size * seq_len * self.dtype_config.activation_bits / BITS_PER_BYTE / self.parallelism_config.tp_size

    def get_activation_memory_per_layer(
        self,
        batch_size: int,
        seq_len: int,
        is_inference: bool = True,
        activation_recomputation:
        ActivationRecomputation = ActivationRecomputation.NONE,
        layernorm_dtype_bytes: int = BYTES_FP32,
        flash_attn: bool = True,
        softmax_dropout: bool = False,
        mlp_activation_quant_bits: int = None,
        mlp_1linear_quant_bits: int = None,
        mlp_gelu_input_quant_bits: int = None,
        mlp_2linear_quant_bits: int = None,
        mlp_recompute_gelu: bool = False,
        mlp_gated_linear_units: bool = False,
        return_breakdown: bool = False,
    ) -> Union[float, tuple]:
        """Get the memory (in bytes) required to store the activations of a
        transformer layer, given the batch size, sequence length, and whether
        it is inference or training, the activation recomputation strategy, and
        the activation data type. Refer to https://arxiv.org/abs/2205.05198 for
        details. For inference, this assumes the maximum tensor buffer reuse.

        Args:
            batch_size (int):
            seq_len (int): sequence length
            is_inference (bool, optional): whether it is inference or not. Return the max memory activation tensor size between layernorm/attn/mlp. Defaults to True.
            activation_recomputation (ActivationRecomputation, optional): \
                activation recomputation strategy. Defaults to ActivationRecomputation.NONE.
            layernorm_dtype_bytes (int, optional): number of bytes in the data type for \
                the layernorm activations. Defaults to BYTES_FP32. Often has to be FP32 in training to maintain model accuracy.
            flash_attn (bool, optional): whether to use Flash Attention. Defaults to True.
            softmax_dropout (bool, optional): whether to apply dropout after softmax. Defaults to False.
            mlp_activation_quant_bits (int, optional): number of bits to quantize MLP activations; if set, override the values for mlp_1linear_quant_bits, mlp_gelu_input_quant_bits and mlp_2linear_quant_bits. Defaults to None.
            mlp_1linear_quant_bits (int, optional): number of bits to quantize the input activations of the first linear layer. Defaults to None.
            mlp_gelu_input_quant_bits (int, optional): number of bits to quantize the GELU input activations. Defaults to None.
            mlp_2linear_quant_bits (int, optional): number of bits to quantize the input activations of the second linear layer. Defaults to None.            mlp_recompute_gelu (bool, optional): whether to recompute the gelu activation in the MLP backward pass. Defaults to False.
            mlp_gated_linear_units (bool, optional): whether to use gated linear units in the MLP. Defaults to False.
        Returns:
            Union[float, tuple]: the memory (in bytes) required to store the activations of a transformer layer or a tuple of its breakdown
        """
        if (not is_inference
            ) and activation_recomputation == ActivationRecomputation.FULL:
            activation_memory_per_layer = (seq_len * batch_size *
                                           self.model_config.hidden_dim *
                                           self.dtype_config.activation_bits /
                                           BITS_PER_BYTE /
                                           self.parallelism_config.tp_size)
            if return_breakdown:
                return activation_memory_per_layer, 0, 0, 0
            else:
                return activation_memory_per_layer

        activation_memory_per_layer_attn = (
            self.get_activation_memory_per_layer_attn(
                batch_size,
                seq_len,