2022 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license a function that calls the internal softmaxbackwarddata pytorch method and that adjusts the arguments according to the torch version detected prune a linear layer to keep only entries in index used to remove heads args layer torch nn linear the layer to prune index torch longtensor the indices to keep in the layer dim int optional defaults to 0 the dimension on which to keep the indices returns torch nn linear the pruned layer as a new layer with requiresgradtrue 1dconvolutional layer as defined by radford et al for openai gpt and also used in gpt2 basically works like a linear layer but the weights are transposed args nf int the number of output features nx int the number of input features prune a conv1d layer to keep only entries in index a conv1d work as a linear layer see e g bert but the weights are transposed used to remove heads args layer pytorchutils conv1d the layer to prune index torch longtensor the indices to keep in the layer dim int optional defaults to 1 the dimension on which to keep the indices returns pytorchutils conv1d the pruned layer as a new layer with requiresgradtrue prune a conv1d or linear layer to keep only entries in index used to remove heads args layer uniontorch nn linear conv1d the layer to prune index torch longtensor the indices to keep in the layer dim int optional the dimension on which to keep the indices returns torch nn linear or pytorchutils conv1d the pruned layer as a new layer with requiresgradtrue this function chunks the inputtensors into smaller input tensor parts of size chunksize over the dimension chunkdim it then applies a layer forwardfn to each chunk independently to save memory if the forwardfn is independent across the chunkdim this function will yield the same result as directly applying forwardfn to inputtensors args forwardfn callable torch tensor the forward function of the model chunksize int the chunk size of a chunked tensor numchunks leninputtensors0 chunksize chunkdim int the dimension over which the inputtensors should be chunked inputtensors tupletorch tensor the input tensors of forwardfn which will be chunked returns torch tensor a tensor with the same shape as the forwardfn would have given if applied examples python rename the usual forward fn to forwardchunk def forwardchunkself hiddenstates hiddenstates self decoderhiddenstates return hiddenstates implement a chunked forward function def forwardself hiddenstates return applychunkingtoforwardself forwardchunk self chunksizelmhead self seqlendim hiddenstates assert leninputtensors 0 finputtensors has to be a tuplelist of tensors inspect signature exist since python 3 5 and is a python method no problem with backward compatibility numargsinforwardchunkfn leninspect signatureforwardfn parameters if numargsinforwardchunkfn leninputtensors raise valueerror fforwardchunkfn expects numargsinforwardchunkfn arguments but only leninputtensors input tensors are given if chunksize 0 tensorshape inputtensors0 shapechunkdim for inputtensor in inputtensors if inputtensor shapechunkdim tensorshape raise valueerror fall input tenors have to be of the same shape tensorshape ffound shape inputtensor shapechunkdim if inputtensors0 shapechunkdim chunksize 0 raise valueerror fthe dimension to be chunked inputtensors0 shapechunkdim has to be a multiple of the chunk fsize chunksize numchunks inputtensors0 shapechunkdim chunksize chunk input tensor into tuples inputtensorschunks tupleinputtensor chunknumchunks dimchunkdim for inputtensor in inputtensors apply forward fn to every tuple outputchunks tupleforwardfninputtensorschunk for inputtensorschunk in zipinputtensorschunks concatenate output at same dimension return torch catoutputchunks dimchunkdim return forwardfninputtensors def findpruneableheadsandindices heads listint nheads int headsize int alreadyprunedheads setint tuplesetint torch longtensor mask torch onesnheads headsize heads setheads alreadyprunedheads convert to set and remove already pruned heads for head in heads compute how many pruned heads are before the head and move the index accordingly head head sum1 if h head else 0 for h in alreadyprunedheads maskhead 0 mask mask view1 contiguous eq1 index torch longtensor torch arangelenmaskmask long return heads index def meshgrid tensors uniontorch tensor listtorch tensor indexing optionalstr none tupletorch tensor return torch meshgridtensors indexingindexing def idtensorstoragetensor torch tensor tupletorch device int int if tensor device type xla and istorchtpuavailable note xla tensors dont have storage use some other unique id to distinguish this is a xla tensor it must be created using torchxla s device so the following import is safe import torchxla uniqueid torchxla xlac xlagettensoridtensor else uniqueid storageptrtensor return tensor device uniqueid storagesizetensor 2022 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license a function that calls the internal _softmax_backward_data pytorch method and that adjusts the arguments according to the torch version detected prune a linear layer to keep only entries in index used to remove heads args layer torch nn linear the layer to prune index torch longtensor the indices to keep in the layer dim int optional defaults to 0 the dimension on which to keep the indices returns torch nn linear the pruned layer as a new layer with requires_grad true 1d convolutional layer as defined by radford et al for openai gpt and also used in gpt 2 basically works like a linear layer but the weights are transposed args nf int the number of output features nx int the number of input features prune a conv1d layer to keep only entries in index a conv1d work as a linear layer see e g bert but the weights are transposed used to remove heads args layer pytorch_utils conv1d the layer to prune index torch longtensor the indices to keep in the layer dim int optional defaults to 1 the dimension on which to keep the indices returns pytorch_utils conv1d the pruned layer as a new layer with requires_grad true prune a conv1d or linear layer to keep only entries in index used to remove heads args layer union torch nn linear conv1d the layer to prune index torch longtensor the indices to keep in the layer dim int optional the dimension on which to keep the indices returns torch nn linear or pytorch_utils conv1d the pruned layer as a new layer with requires_grad true this function chunks the input_tensors into smaller input tensor parts of size chunk_size over the dimension chunk_dim it then applies a layer forward_fn to each chunk independently to save memory if the forward_fn is independent across the chunk_dim this function will yield the same result as directly applying forward_fn to input_tensors args forward_fn callable torch tensor the forward function of the model chunk_size int the chunk size of a chunked tensor num_chunks len input_tensors 0 chunk_size chunk_dim int the dimension over which the input_tensors should be chunked input_tensors tuple torch tensor the input tensors of forward_fn which will be chunked returns torch tensor a tensor with the same shape as the forward_fn would have given if applied examples python rename the usual forward fn to forward_chunk def forward_chunk self hidden_states hidden_states self decoder hidden_states return hidden_states implement a chunked forward function def forward self hidden_states return apply_chunking_to_forward self forward_chunk self chunk_size_lm_head self seq_len_dim hidden_states inspect signature exist since python 3 5 and is a python method no problem with backward compatibility chunk input tensor into tuples apply forward fn to every tuple concatenate output at same dimension finds the heads and their indices taking already_pruned_heads into account args heads list int list of the indices of heads to prune n_heads int the number of heads in the model head_size int the size of each head already_pruned_heads set int a set of already pruned heads returns tuple set int torch longtensor a tuple with the indices of heads to prune taking already_pruned_heads into account and the indices of rows columns to keep in the layer weight convert to set and remove already pruned heads compute how many pruned heads are before the head and move the index accordingly wrapper around torch meshgrid to avoid warning messages about the introduced indexing argument reference https pytorch org docs 1 13 generated torch meshgrid html unique identifier to a tensor storage multiple different tensors can share the same underlying storage for example meta tensors all share the same storage and thus their identifier will all be equal this identifier is guaranteed to be unique and constant for this tensor s storage during its lifetime two tensor storages with non overlapping lifetimes may have the same id note xla tensors dont have storage use some other unique id to distinguish this is a xla tensor it must be created using torch_xla s device so the following import is safe
import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import is_torch_tpu_available, logging


ALL_LAYERNORM_LAYERS = [nn.LayerNorm]

logger = logging.get_logger(__name__)

parsed_torch_version_base = version.parse(version.parse(torch.__version__).base_version)

is_torch_greater_or_equal_than_2_1 = parsed_torch_version_base >= version.parse("2.1")
is_torch_greater_or_equal_than_2_0 = parsed_torch_version_base >= version.parse("2.0")
is_torch_greater_or_equal_than_1_13 = parsed_torch_version_base >= version.parse("1.13")
is_torch_greater_or_equal_than_1_12 = parsed_torch_version_base >= version.parse("1.12")
is_torch_greater_or_equal_than_1_11 = parsed_torch_version_base >= version.parse("1.11")
is_torch_less_than_1_11 = parsed_torch_version_base < version.parse("1.11")
is_torch_1_8_0 = parsed_torch_version_base == version.parse("1.8.0")


def softmax_backward_data(parent, grad_output, output, dim, self):
    

    from torch import _softmax_backward_data

    if is_torch_less_than_1_11:
        return _softmax_backward_data(grad_output, output, parent.dim, self)
    else:
        return _softmax_backward_data(grad_output, output, parent.dim, self.dtype)


def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int = 0) -> nn.Linear:
    
    index = index.to(layer.weight.device)
    W = layer.weight.index_select(dim, index).clone().detach()
    if layer.bias is not None:
        if dim == 1:
            b = layer.bias.clone().detach()
        else:
            b = layer.bias[index].clone().detach()
    new_size = list(layer.weight.size())
    new_size[dim] = len(index)
    new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
    new_layer.weight.requires_grad = False
    new_layer.weight.copy_(W.contiguous())
    new_layer.weight.requires_grad = True
    if layer.bias is not None:
        new_layer.bias.requires_grad = False
        new_layer.bias.copy_(b.contiguous())
        new_layer.bias.requires_grad = True
    return new_layer


class Conv1D(nn.Module):
    

    def __init__(self, nf, nx):
        super().__init__()
        self.nf = nf
        self.weight = nn.Parameter(torch.empty(nx, nf))
        self.bias = nn.Parameter(torch.zeros(nf))
        nn.init.normal_(self.weight, std=0.02)

    def forward(self, x):
        size_out = x.size()[:-1] + (self.nf,)
        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
        x = x.view(size_out)
        return x


def prune_conv1d_layer(layer: Conv1D, index: torch.LongTensor, dim: int = 1) -> Conv1D:
    
    index = index.to(layer.weight.device)
    W = layer.weight.index_select(dim, index).clone().detach()
    if dim == 0:
        b = layer.bias.clone().detach()
    else:
        b = layer.bias[index].clone().detach()
    new_size = list(layer.weight.size())
    new_size[dim] = len(index)
    new_layer = Conv1D(new_size[1], new_size[0]).to(layer.weight.device)
    new_layer.weight.requires_grad = False
    new_layer.weight.copy_(W.contiguous())
    new_layer.weight.requires_grad = True
    new_layer.bias.requires_grad = False
    new_layer.bias.copy_(b.contiguous())
    new_layer.bias.requires_grad = True
    return new_layer


def prune_layer(
    layer: Union[nn.Linear, Conv1D], index: torch.LongTensor, dim: Optional[int] = None
) -> Union[nn.Linear, Conv1D]:
    
    if isinstance(layer, nn.Linear):
        return prune_linear_layer(layer, index, dim=0 if dim is None else dim)
    elif isinstance(layer, Conv1D):
        return prune_conv1d_layer(layer, index, dim=1 if dim is None else dim)
    else:
        raise ValueError(f"Can't prune layer of class {layer.__class__}")


def apply_chunking_to_forward(
    forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors
) -> torch.Tensor:
    

    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"

    
    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
    if num_args_in_forward_chunk_fn != len(input_tensors):
        raise ValueError(
            f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
            "tensors are given"
        )

    if chunk_size > 0:
        tensor_shape = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    f"All input tenors have to be of the same shape: {tensor_shape}, "
                    f"found shape {input_tensor.shape[chunk_dim]}"
                )

        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
                f"size {chunk_size}"
            )

        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size

        
        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
        
        output_chunks = tuple(forward_fn(*input_tensors_chunk) for input_tensors_chunk in zip(*input_tensors_chunks))
        
        return torch.cat(output_chunks, dim=chunk_dim)

    return forward_fn(*input_tensors)


def find_pruneable_heads_and_indices(
    heads: List[int], n_heads: int, head_size: int, already_pruned_heads: Set[int]
) -> Tuple[Set[int], torch.LongTensor]:
    
    mask = torch.ones(n_heads, head_size)
    heads = set(heads) - already_pruned_heads  
    for head in heads:
        
        head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
        mask[head] = 0
    mask = mask.view(-1).contiguous().eq(1)
    index: torch.LongTensor = torch.arange(len(mask))[mask].long()
    return heads, index


def meshgrid(
    *tensors: Union[torch.Tensor, List[torch.Tensor]], indexing: Optional[str] = None
) -> Tuple[torch.Tensor, ...]:
    
    return torch.meshgrid(*tensors, indexing=indexing)


def id_tensor_storage(tensor: torch.Tensor) -> Tuple[torch.device, int, int]:
    
    if tensor.device.type == "xla" and is_torch_tpu_available():
        
        
        
        
        import torch_xla

        unique_id = torch_xla._XLAC._xla_get_tensor_id(tensor)
    else:
        unique_id = storage_ptr(tensor)

    return tensor.device, unique_id, storage_size(tensor)