Nethermind/Mpt-Instruct-DotNet-XS

1# Copyright 2022 MosaicML Examples authors2# SPDX-License-Identifier: Apache-2.03 4"""Attention layers."""5 6import math7import warnings8from typing import Optional9 10import torch11import torch.nn as nn12from einops import rearrange13from torch import nn14 15from .low_precision_layernorm import LPLayerNorm16 17 18def _reset_is_causal(num_query_tokens: int, num_key_tokens: int,19                     original_is_causal: bool):20    if original_is_causal and num_query_tokens != num_key_tokens:21        if num_query_tokens != 1:22            raise NotImplementedError(23                'MosaicGPT does not support query and key with different number of tokens, unless number of query tokens is 1.'24            )25        else:26            return False27    return original_is_causal28 29 30def scaled_multihead_dot_product_attention(31    query,32    key,33    value,34    n_heads,35    softmax_scale=None,36    attn_bias=None,37    key_padding_mask=None,38    is_causal=False,39    dropout_p=0.0,40    training=False,41    needs_weights=False,42):43 44    q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)45    k = rearrange(key, 'b s (h d) -> b h d s', h=n_heads)  # includes key.t()46    v = rearrange(value, 'b s (h d) -> b h s d', h=n_heads)47 48    min_val = torch.finfo(q.dtype).min49 50    b, _, s_q, d = q.shape51    s_k = k.size(-1)52 53    if softmax_scale is None:54        softmax_scale = 1 / math.sqrt(d)55 56    attn_weight = q.matmul(k) * softmax_scale57 58    if attn_bias is not None:59        if (attn_bias.size(-1) != 1 and60                attn_bias.size(-1) != s_k) or (attn_bias.size(-2) != 1 and61                                               attn_bias.size(-2) != s_q):62            raise RuntimeError(63                f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.'64            )65        attn_weight = attn_weight + attn_bias66 67    if key_padding_mask is not None:68        if attn_bias is not None:69            warnings.warn(70                'Propogating key_padding_mask to the attention module ' +\71                'and applying it within the attention module can cause ' +\72                'unneccessary computation/memory usage. Consider integrating ' +\73                'into attn_bias once and passing that to each attention ' +\74                'module instead.'75            )76        attn_weight = attn_weight.masked_fill(77            ~key_padding_mask.view((b, 1, 1, s_k)), min_val)78 79    if is_causal:80        s = max(s_q, s_k)81        causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)82        causal_mask = causal_mask.tril()83        causal_mask = causal_mask.to(torch.bool)84        causal_mask = ~causal_mask85        causal_mask = causal_mask[-s_q:, -s_k:]86        attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k),87                                              min_val)88 89    attn_weight = torch.softmax(attn_weight, dim=-1)90 91    if dropout_p:92        attn_weight = torch.nn.functional.dropout(attn_weight,93                                                  p=dropout_p,94                                                  training=training,95                                                  inplace=True)96 97    out = attn_weight.matmul(v)98    out = rearrange(out, 'b h s d -> b s (h d)')99 100    if needs_weights:101        return out, attn_weight102    return out, None103 104 105def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):106    for tensor in tensors:107        if tensor.dtype not in valid_dtypes:108            raise TypeError(f'{tensor.dtype=} must be in {valid_dtypes=}.')109        if not tensor.is_cuda:110            raise TypeError(f'Inputs must be cuda tensors ({tensor.is_cuda=}).')111 112 113def flash_attn_fn(114    query,115    key,116    value,117    n_heads,118    softmax_scale=None,119    attn_bias=None,120    key_padding_mask=None,121    is_causal=False,122    dropout_p=0.0,123    training=False,124    needs_weights=False,125):126    try:127        from flash_attn import bert_padding, flash_attn_interface128    except:129        raise RuntimeError('Please install flash_attn==0.2.8')130 131    check_valid_inputs(query, key, value)132 133    if attn_bias is not None:134        raise NotImplementedError(f'attn_bias not implemented for flash attn.')135 136    batch_size, seqlen = query.shape[:2]137 138    if key_padding_mask is None:139        key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)140    query_padding_mask = key_padding_mask[:, -query.size(1):]141 142    query_unpad, indices_q, cu_seqlens_q, max_seqlen_q = bert_padding.unpad_input(143        query, query_padding_mask)144    query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)145 146    key_unpad, _, cu_seqlens_k, max_seqlen_k = bert_padding.unpad_input(147        key, key_padding_mask)148    key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=n_heads)149 150    value_unpad, _, _, _ = bert_padding.unpad_input(value, key_padding_mask)151    value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=n_heads)152 153    dropout_p = dropout_p if training else 0.0154 155    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)156 157    output_unpad = flash_attn_interface.flash_attn_unpadded_func(158        query_unpad,159        key_unpad,160        value_unpad,161        cu_seqlens_q,162        cu_seqlens_k,163        max_seqlen_q,164        max_seqlen_k,165        dropout_p,166        softmax_scale=softmax_scale,167        causal=reset_is_causal,168        return_attn_probs=needs_weights)169 170    output = bert_padding.pad_input(171        rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size,172        seqlen)173    return output, None174 175 176def triton_flash_attn_fn(177    query,178    key,179    value,180    n_heads,181    softmax_scale=None,182    attn_bias=None,183    key_padding_mask=None,184    is_causal=False,185    dropout_p=0.0,186    training=False,187    needs_weights=False,188):189    try:190        from flash_attn import flash_attn_triton  # type: ignore191    except:192        raise RuntimeError('Please install flash_attn==0.2.8 and triton==2.0.0.dev20221202.')193 194    check_valid_inputs(query, key, value)195 196    if dropout_p:197        raise NotImplementedError(198            f'Dropout not implemented for attn_impl: triton.')199 200    if needs_weights:201        raise NotImplementedError(202            f'attn_impl: triton cannot return attn weights.')203 204    if key_padding_mask is not None:205        warnings.warn(206            'Propagating key_padding_mask to the attention module ' +\207            'and applying it within the attention module can cause ' +\208            'unnecessary computation/memory usage. Consider integrating ' +\209            'into attn_bias once and passing that to each attention ' +\210            'module instead.'211        )212        b_size, s_k = key_padding_mask.shape[:2]213 214        if attn_bias is None:215            attn_bias = query.new_zeros(b_size, 1, 1, s_k)216 217        attn_bias = attn_bias.masked_fill(218            ~key_padding_mask.view((b_size, 1, 1, s_k)),219            torch.finfo(query.dtype).min)220 221    query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)222    key = rearrange(key, 'b s (h d) -> b s h d', h=n_heads)223    value = rearrange(value, 'b s (h d) -> b s h d', h=n_heads)224 225    reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)226    attn_output = flash_attn_triton.flash_attn_func(query, key, value,227                                                    attn_bias, reset_is_causal,228                                                    softmax_scale)229 230    output = attn_output.view(*attn_output.shape[:2], -1)231 232    return output, None233 234 235class MultiheadAttention(nn.Module):236    """Multi-head self attention.237 238    Using torch or triton attention implemetation enables user to also use239    additive bias.240    """241 242    def __init__(243        self,244        d_model: int,245        n_heads: int,246        attn_impl: str = 'triton',247        attn_clip_qkv: Optional[float] = None,248        attn_qk_ln: bool = False,249        softmax_scale: Optional[float] = None,250        attn_pdrop: float = 0.0,251        low_precision_layernorm: bool = False,252        device: Optional[str] = None,253    ):254        super().__init__()255 256        self.attn_impl = attn_impl257        self.clip_qkv = attn_clip_qkv258        self.attn_qk_ln = attn_qk_ln259 260        self.d_model = d_model261        self.n_heads = n_heads262        self.softmax_scale = softmax_scale263        if self.softmax_scale is None:264            self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)265        self.attn_dropout_p = attn_pdrop266 267        self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)268        # for param init fn; enables shape based init of fused layers269        fuse_splits = (d_model, 2 * d_model)270        self.Wqkv._fused = (0, fuse_splits)  # type: ignore271 272        if self.attn_qk_ln:273            layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm274            self.q_ln = layernorm_class(self.d_model, device=device)275            self.k_ln = layernorm_class(self.d_model, device=device)276 277        if self.attn_impl == 'flash':278            self.attn_fn = flash_attn_fn279        elif self.attn_impl == 'triton':280            self.attn_fn = triton_flash_attn_fn281            warnings.warn(282                'While `attn_impl: triton` can be faster than `attn_impl: flash` ' +\283                'it uses more memory. When training larger models this can trigger '  +\284                'alloc retries which hurts performance. If encountered, we recommend ' +\285                'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')286        elif self.attn_impl == 'torch':287            self.attn_fn = scaled_multihead_dot_product_attention288            if torch.cuda.is_available():289                warnings.warn(290                    'Using `attn_impl: torch`. If your model does not use `alibi` or ' +\291                    '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' +\292                    'we recommend using `attn_impl: triton`.'293                )294        else:295            raise ValueError(f'{attn_impl=} is an invalid setting.')296 297        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)298        self.out_proj._is_residual = True  # type: ignore299 300    def forward(self,301                x,302                past_key_value=None,303                attn_bias=None,304                attention_mask=None,305                is_causal=True,306                needs_weights=False):307        qkv = self.Wqkv(x)308 309        if self.clip_qkv:310            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)311 312        query, key, value = qkv.chunk(3, dim=2)313 314        key_padding_mask = attention_mask315 316        if self.attn_qk_ln:317            # Applying layernorm to qk318            dtype = query.dtype319            query = self.q_ln(query).to(dtype)320            key = self.k_ln(key).to(dtype)321 322        if past_key_value is not None:323            if len(past_key_value) != 0:324                key = torch.cat([past_key_value[0], key], dim=1)325                value = torch.cat([past_key_value[1], value], dim=1)326 327            past_key_value = (key, value)328 329        if attn_bias is not None:330            attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]331 332        context, attn_weights = self.attn_fn(333            query,334            key,335            value,336            self.n_heads,337            softmax_scale=self.softmax_scale,338            attn_bias=attn_bias,339            key_padding_mask=key_padding_mask,340            is_causal=is_causal,341            dropout_p=self.attn_dropout_p,342            training=self.training,343            needs_weights=needs_weights,344        )345 346        return self.out_proj(context), attn_weights, past_key_value347 348 349def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal,350                    use_sequence_id):351    if attn_impl == 'flash':352        return None353    elif attn_impl in ['torch', 'triton']:354        if alibi:355            if (prefix_lm or not causal) or use_sequence_id:356                return (1, n_heads, seq_len, seq_len)357            return (1, n_heads, 1, seq_len)358        elif prefix_lm or use_sequence_id:359            return (1, 1, seq_len, seq_len)360        return None361    else:362        raise ValueError(f'{attn_impl=} is an invalid setting.')363 364 365def attn_bias(attn_impl,366              attn_bias,367              n_heads,368              seq_len,369              causal=False,370              alibi=False,371              alibi_bias_max=8):372    if attn_impl == 'flash':373        return None374    elif attn_impl in ['torch', 'triton']:375        if alibi:376            # in place add alibi to attn bias377            device, dtype = attn_bias.device, attn_bias.dtype378            attn_bias = attn_bias.add(379                alibi_bias(n_heads,380                           seq_len,381                           full=not causal,382                           alibi_bias_max=alibi_bias_max,383                           device=device,384                           dtype=dtype))385        return attn_bias386    else:387        raise ValueError(f'{attn_impl=} is an invalid setting.')388 389 390def alibi_bias(n_heads,391               seq_len,392               full=False,393               alibi_bias_max=8,394               device=None,395               dtype=None):396    alibi_bias = torch.arange(1 - seq_len, 1, dtype=dtype,397                              device=device).view(1, 1, 1, seq_len)398    if full:399        # generate 1 x Heads x SeqLen x SeqLen alibi bias mask400        # otherwise the mask is 1 x Heads x 1 x SeqLen (which is broadcast to the appropriate size)401        alibi_bias = alibi_bias - torch.arange(402            1 - seq_len, 1, dtype=dtype, device=device).view(1, 1, seq_len, 1)403        alibi_bias = alibi_bias.abs().mul(-1)404 405    m = torch.arange(1, n_heads + 1, dtype=dtype, device=device)406    m = m.mul(alibi_bias_max / n_heads)407    alibi_bias = alibi_bias * (1. / (2**m.view(1, n_heads, 1, 1)))408    return alibi_bias409