🎨 format for readability

Signed-off-by: peter szemraj <[email protected]>

adapt_tokenizer.py

@@ -1,8 +1,10 @@
 from typing import Union
 from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+
 Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
 NUM_SENTINEL_TOKENS: int = 100
 
+
 def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
     """Adds sentinel tokens and padding token (if missing).
 
@@ -12,16 +14,17 @@ def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
     All added tokens are added as special tokens. No tokens are
     added if sentinel tokens and padding token already exist.
     """
-    sentinels_to_add = [f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)]
+    sentinels_to_add = [f"<extra_id_{i}>" for i in range(NUM_SENTINEL_TOKENS)]
     tokenizer.add_tokens(sentinels_to_add, special_tokens=True)
     if tokenizer.pad_token is None:
-        tokenizer.add_tokens('<pad>', special_tokens=True)
-        tokenizer.pad_token = '<pad>'
+        tokenizer.add_tokens("<pad>", special_tokens=True)
+        tokenizer.pad_token = "<pad>"
         assert tokenizer.pad_token_id is not None
-    sentinels = ''.join([f'<extra_id_{i}>' for i in range(NUM_SENTINEL_TOKENS)])
+    sentinels = "".join([f"<extra_id_{i}>" for i in range(NUM_SENTINEL_TOKENS)])
     _sentinel_token_ids = tokenizer(sentinels, add_special_tokens=False).input_ids
     tokenizer.sentinel_token_ids = _sentinel_token_ids
 
+
 class AutoTokenizerForMOD(AutoTokenizer):
     """AutoTokenizer + Adaptation for MOD.
 
@@ -38,4 +41,4 @@ class AutoTokenizerForMOD(AutoTokenizer):
         """See `AutoTokenizer.from_pretrained` docstring."""
         tokenizer = super().from_pretrained(*args, **kwargs)
         adapt_tokenizer_for_denoising(tokenizer)
-        return tokenizer
+        return tokenizer

attention.py

@@ -8,18 +8,37 @@ from einops import rearrange
 from torch import nn
 from .norm import LPLayerNorm
 
-def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool):
+
+def _reset_is_causal(
+    num_query_tokens: int, num_key_tokens: int, original_is_causal: bool
+):
     if original_is_causal and num_query_tokens != num_key_tokens:
         if num_query_tokens != 1:
-            raise NotImplementedError('MPT does not support query and key with different number of tokens, unless number of query tokens is 1.')
+            raise NotImplementedError(
+                "MPT does not support query and key with different number of tokens, unless number of query tokens is 1."
+            )
         else:
             return False
     return original_is_causal
 
-def scaled_multihead_dot_product_attention(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
-    q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)
-    k = rearrange(key, 'b s (h d) -> b h d s', h=1 if multiquery else n_heads)
-    v = rearrange(value, 'b s (h d) -> b h s d', h=1 if multiquery else n_heads)
+
+def scaled_multihead_dot_product_attention(
+    query,
+    key,
+    value,
+    n_heads,
+    softmax_scale=None,
+    attn_bias=None,
+    key_padding_mask=None,
+    is_causal=False,
+    dropout_p=0.0,
+    training=False,
+    needs_weights=False,
+    multiquery=False,
+):
+    q = rearrange(query, "b s (h d) -> b h s d", h=n_heads)
+    k = rearrange(key, "b s (h d) -> b h d s", h=1 if multiquery else n_heads)
+    v = rearrange(value, "b s (h d) -> b h s d", h=1 if multiquery else n_heads)
     min_val = torch.finfo(q.dtype).min
     (b, _, s_q, d) = q.shape
     s_k = k.size(-1)
@@ -27,13 +46,27 @@ def scaled_multihead_dot_product_attention(query, key, value, n_heads, softmax_s
         softmax_scale = 1 / math.sqrt(d)
     attn_weight = q.matmul(k) * softmax_scale
     if attn_bias is not None:
-        if attn_bias.size(-1) != 1 and attn_bias.size(-1) != s_k or (attn_bias.size(-2) != 1 and attn_bias.size(-2) != s_q):
-            raise RuntimeError(f'attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}.')
+        if (
+            attn_bias.size(-1) != 1
+            and attn_bias.size(-1) != s_k
+            or (attn_bias.size(-2) != 1 and attn_bias.size(-2) != s_q)
+        ):
+            raise RuntimeError(
+                f"attn_bias (shape: {attn_bias.shape}) is expected to broadcast to shape: {attn_weight.shape}."
+            )
         attn_weight = attn_weight + attn_bias
     if key_padding_mask is not None:
         if attn_bias is not None:
-            warnings.warn('Propogating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unneccessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
-        attn_weight = attn_weight.masked_fill(~key_padding_mask.view((b, 1, 1, s_k)), min_val)
+            warnings.warn(
+                "Propogating key_padding_mask to the attention module "
+                + "and applying it within the attention module can cause "
+                + "unneccessary computation/memory usage. Consider integrating "
+                + "into attn_bias once and passing that to each attention "
+                + "module instead."
+            )
+        attn_weight = attn_weight.masked_fill(
+            ~key_padding_mask.view((b, 1, 1, s_k)), min_val
+        )
     if is_causal:
         s = max(s_q, s_k)
         causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)
@@ -44,74 +77,146 @@ def scaled_multihead_dot_product_attention(query, key, value, n_heads, softmax_s
         attn_weight = attn_weight.masked_fill(causal_mask.view(1, 1, s_q, s_k), min_val)
     attn_weight = torch.softmax(attn_weight, dim=-1)
     if dropout_p:
-        attn_weight = torch.nn.functional.dropout(attn_weight, p=dropout_p, training=training, inplace=True)
+        attn_weight = torch.nn.functional.dropout(
+            attn_weight, p=dropout_p, training=training, inplace=True
+        )
     out = attn_weight.matmul(v)
-    out = rearrange(out, 'b h s d -> b s (h d)')
+    out = rearrange(out, "b h s d -> b s (h d)")
     if needs_weights:
         return (out, attn_weight)
     return (out, None)
 
+
 def check_valid_inputs(*tensors, valid_dtypes=[torch.float16, torch.bfloat16]):
     for tensor in tensors:
         if tensor.dtype not in valid_dtypes:
-            raise TypeError(f'tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}.')
+            raise TypeError(
+                f"tensor.dtype={tensor.dtype!r} must be in valid_dtypes={valid_dtypes!r}."
+            )
         if not tensor.is_cuda:
-            raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
+            raise TypeError(
+                f"Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r})."
+            )
+
 
-def flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
+def flash_attn_fn(
+    query,
+    key,
+    value,
+    n_heads,
+    softmax_scale=None,
+    attn_bias=None,
+    key_padding_mask=None,
+    is_causal=False,
+    dropout_p=0.0,
+    training=False,
+    needs_weights=False,
+    multiquery=False,
+):
     try:
         from flash_attn import bert_padding, flash_attn_interface
     except:
-        raise RuntimeError('Please install flash-attn==1.0.3.post0')
+        raise RuntimeError("Please install flash-attn==1.0.3.post0")
     check_valid_inputs(query, key, value)
     if attn_bias is not None:
-        raise NotImplementedError(f'attn_bias not implemented for flash attn.')
+        raise NotImplementedError(f"attn_bias not implemented for flash attn.")
     (batch_size, seqlen) = query.shape[:2]
     if key_padding_mask is None:
         key_padding_mask = torch.ones_like(key[:, :, 0], dtype=torch.bool)
-    query_padding_mask = key_padding_mask[:, -query.size(1):]
-    (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(query, query_padding_mask)
-    query_unpad = rearrange(query_unpad, 'nnz (h d) -> nnz h d', h=n_heads)
-    (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(key, key_padding_mask)
-    key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
+    query_padding_mask = key_padding_mask[:, -query.size(1) :]
+    (query_unpad, indices_q, cu_seqlens_q, max_seqlen_q) = bert_padding.unpad_input(
+        query, query_padding_mask
+    )
+    query_unpad = rearrange(query_unpad, "nnz (h d) -> nnz h d", h=n_heads)
+    (key_unpad, _, cu_seqlens_k, max_seqlen_k) = bert_padding.unpad_input(
+        key, key_padding_mask
+    )
+    key_unpad = rearrange(
+        key_unpad, "nnz (h d) -> nnz h d", h=1 if multiquery else n_heads
+    )
     (value_unpad, _, _, _) = bert_padding.unpad_input(value, key_padding_mask)
-    value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=1 if multiquery else n_heads)
+    value_unpad = rearrange(
+        value_unpad, "nnz (h d) -> nnz h d", h=1 if multiquery else n_heads
+    )
     if multiquery:
         key_unpad = key_unpad.expand(key_unpad.size(0), n_heads, key_unpad.size(-1))
-        value_unpad = value_unpad.expand(value_unpad.size(0), n_heads, value_unpad.size(-1))
+        value_unpad = value_unpad.expand(
+            value_unpad.size(0), n_heads, value_unpad.size(-1)
+        )
     dropout_p = dropout_p if training else 0.0
     reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
-    output_unpad = flash_attn_interface.flash_attn_unpadded_func(query_unpad, key_unpad, value_unpad, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
-    output = bert_padding.pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'), indices_q, batch_size, seqlen)
+    output_unpad = flash_attn_interface.flash_attn_unpadded_func(
+        query_unpad,
+        key_unpad,
+        value_unpad,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        dropout_p,
+        softmax_scale=softmax_scale,
+        causal=reset_is_causal,
+        return_attn_probs=needs_weights,
+    )
+    output = bert_padding.pad_input(
+        rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices_q, batch_size, seqlen
+    )
     return (output, None)
 
-def triton_flash_attn_fn(query, key, value, n_heads, softmax_scale=None, attn_bias=None, key_padding_mask=None, is_causal=False, dropout_p=0.0, training=False, needs_weights=False, multiquery=False):
+
+def triton_flash_attn_fn(
+    query,
+    key,
+    value,
+    n_heads,
+    softmax_scale=None,
+    attn_bias=None,
+    key_padding_mask=None,
+    is_causal=False,
+    dropout_p=0.0,
+    training=False,
+    needs_weights=False,
+    multiquery=False,
+):
     try:
         from flash_attn import flash_attn_triton
     except:
-        raise RuntimeError('Please install flash-attn==1.0.3.post0 and triton==2.0.0.dev20221202')
+        raise RuntimeError(
+            "Please install flash-attn==1.0.3.post0 and triton==2.0.0.dev20221202"
+        )
     check_valid_inputs(query, key, value)
     if dropout_p:
-        raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
+        raise NotImplementedError(f"Dropout not implemented for attn_impl: triton.")
     if needs_weights:
-        raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
+        raise NotImplementedError(f"attn_impl: triton cannot return attn weights.")
     if key_padding_mask is not None:
-        warnings.warn('Propagating key_padding_mask to the attention module ' + 'and applying it within the attention module can cause ' + 'unnecessary computation/memory usage. Consider integrating ' + 'into attn_bias once and passing that to each attention ' + 'module instead.')
+        warnings.warn(
+            "Propagating key_padding_mask to the attention module "
+            + "and applying it within the attention module can cause "
+            + "unnecessary computation/memory usage. Consider integrating "
+            + "into attn_bias once and passing that to each attention "
+            + "module instead."
+        )
         (b_size, s_k) = key_padding_mask.shape[:2]
         if attn_bias is None:
             attn_bias = query.new_zeros(b_size, 1, 1, s_k)
-        attn_bias = attn_bias.masked_fill(~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min)
-    query = rearrange(query, 'b s (h d) -> b s h d', h=n_heads)
-    key = rearrange(key, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
-    value = rearrange(value, 'b s (h d) -> b s h d', h=1 if multiquery else n_heads)
+        attn_bias = attn_bias.masked_fill(
+            ~key_padding_mask.view((b_size, 1, 1, s_k)), torch.finfo(query.dtype).min
+        )
+    query = rearrange(query, "b s (h d) -> b s h d", h=n_heads)
+    key = rearrange(key, "b s (h d) -> b s h d", h=1 if multiquery else n_heads)
+    value = rearrange(value, "b s (h d) -> b s h d", h=1 if multiquery else n_heads)
     if multiquery:
         key = key.expand(*key.shape[:2], n_heads, key.size(-1))
         value = value.expand(*value.shape[:2], n_heads, value.size(-1))
     reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
-    attn_output = flash_attn_triton.flash_attn_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
+    attn_output = flash_attn_triton.flash_attn_func(
+        query, key, value, attn_bias, reset_is_causal, softmax_scale
+    )
     output = attn_output.view(*attn_output.shape[:2], -1)
     return (output, None)
 
+
 class MultiheadAttention(nn.Module):
     """Multi-head self attention.
 
@@ -119,7 +224,18 @@ class MultiheadAttention(nn.Module):
     additive bias.
     """
 
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, device: Optional[str]=None):
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        attn_impl: str = "triton",
+        clip_qkv: Optional[float] = None,
+        qk_ln: bool = False,
+        softmax_scale: Optional[float] = None,
+        attn_pdrop: float = 0.0,
+        low_precision_layernorm: bool = False,
+        device: Optional[str] = None,
+    ):
         super().__init__()
         self.attn_impl = attn_impl
         self.clip_qkv = clip_qkv
@@ -137,21 +253,38 @@ class MultiheadAttention(nn.Module):
             layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
             self.q_ln = layernorm_class(self.d_model, device=device)
             self.k_ln = layernorm_class(self.d_model, device=device)
-        if self.attn_impl == 'flash':
+        if self.attn_impl == "flash":
             self.attn_fn = flash_attn_fn
-        elif self.attn_impl == 'triton':
+        elif self.attn_impl == "triton":
             self.attn_fn = triton_flash_attn_fn
-            warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
-        elif self.attn_impl == 'torch':
+            warnings.warn(
+                "While `attn_impl: triton` can be faster than `attn_impl: flash` "
+                + "it uses more memory. When training larger models this can trigger "
+                + "alloc retries which hurts performance. If encountered, we recommend "
+                + "using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`."
+            )
+        elif self.attn_impl == "torch":
             self.attn_fn = scaled_multihead_dot_product_attention
             if torch.cuda.is_available():
-                warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
+                warnings.warn(
+                    "Using `attn_impl: torch`. If your model does not use `alibi` or "
+                    + "`prefix_lm` we recommend using `attn_impl: flash` otherwise "
+                    + "we recommend using `attn_impl: triton`."
+                )
         else:
-            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+            raise ValueError(f"attn_impl={attn_impl!r} is an invalid setting.")
         self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
         self.out_proj._is_residual = True
 
-    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
+    def forward(
+        self,
+        x,
+        past_key_value=None,
+        attn_bias=None,
+        attention_mask=None,
+        is_causal=True,
+        needs_weights=False,
+    ):
         qkv = self.Wqkv(x)
         if self.clip_qkv:
             qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
@@ -167,10 +300,23 @@ class MultiheadAttention(nn.Module):
                 value = torch.cat([past_key_value[1], value], dim=1)
             past_key_value = (key, value)
         if attn_bias is not None:
-            attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]
-        (context, attn_weights) = self.attn_fn(query, key, value, self.n_heads, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights)
+            attn_bias = attn_bias[:, :, -query.size(1) :, -key.size(1) :]
+        (context, attn_weights) = self.attn_fn(
+            query,
+            key,
+            value,
+            self.n_heads,
+            softmax_scale=self.softmax_scale,
+            attn_bias=attn_bias,
+            key_padding_mask=key_padding_mask,
+            is_causal=is_causal,
+            dropout_p=self.attn_dropout_p,
+            training=self.training,
+            needs_weights=needs_weights,
+        )
         return (self.out_proj(context), attn_weights, past_key_value)
 
+
 class MultiQueryAttention(nn.Module):
     """Multi-Query self attention.
 
@@ -178,7 +324,18 @@ class MultiQueryAttention(nn.Module):
     additive bias.
     """
 
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, low_precision_layernorm: bool=False, device: Optional[str]=None):
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        attn_impl: str = "triton",
+        clip_qkv: Optional[float] = None,
+        qk_ln: bool = False,
+        softmax_scale: Optional[float] = None,
+        attn_pdrop: float = 0.0,
+        low_precision_layernorm: bool = False,
+        device: Optional[str] = None,
+    ):
         super().__init__()
         self.attn_impl = attn_impl
         self.clip_qkv = clip_qkv
@@ -197,25 +354,44 @@ class MultiQueryAttention(nn.Module):
             layernorm_class = LPLayerNorm if low_precision_layernorm else nn.LayerNorm
             self.q_ln = layernorm_class(d_model, device=device)
             self.k_ln = layernorm_class(self.head_dim, device=device)
-        if self.attn_impl == 'flash':
+        if self.attn_impl == "flash":
             self.attn_fn = flash_attn_fn
-        elif self.attn_impl == 'triton':
+        elif self.attn_impl == "triton":
             self.attn_fn = triton_flash_attn_fn
-            warnings.warn('While `attn_impl: triton` can be faster than `attn_impl: flash` ' + 'it uses more memory. When training larger models this can trigger ' + 'alloc retries which hurts performance. If encountered, we recommend ' + 'using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`.')
-        elif self.attn_impl == 'torch':
+            warnings.warn(
+                "While `attn_impl: triton` can be faster than `attn_impl: flash` "
+                + "it uses more memory. When training larger models this can trigger "
+                + "alloc retries which hurts performance. If encountered, we recommend "
+                + "using `attn_impl: flash` if your model does not use `alibi` or `prefix_lm`."
+            )
+        elif self.attn_impl == "torch":
             self.attn_fn = scaled_multihead_dot_product_attention
             if torch.cuda.is_available():
-                warnings.warn('Using `attn_impl: torch`. If your model does not use `alibi` or ' + '`prefix_lm` we recommend using `attn_impl: flash` otherwise ' + 'we recommend using `attn_impl: triton`.')
+                warnings.warn(
+                    "Using `attn_impl: torch`. If your model does not use `alibi` or "
+                    + "`prefix_lm` we recommend using `attn_impl: flash` otherwise "
+                    + "we recommend using `attn_impl: triton`."
+                )
         else:
-            raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+            raise ValueError(f"attn_impl={attn_impl!r} is an invalid setting.")
         self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
         self.out_proj._is_residual = True
 
-    def forward(self, x, past_key_value=None, attn_bias=None, attention_mask=None, is_causal=True, needs_weights=False):
+    def forward(
+        self,
+        x,
+        past_key_value=None,
+        attn_bias=None,
+        attention_mask=None,
+        is_causal=True,
+        needs_weights=False,
+    ):
         qkv = self.Wqkv(x)
         if self.clip_qkv:
             qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
-        (query, key, value) = qkv.split([self.d_model, self.head_dim, self.head_dim], dim=2)
+        (query, key, value) = qkv.split(
+            [self.d_model, self.head_dim, self.head_dim], dim=2
+        )
         key_padding_mask = attention_mask
         if self.qk_ln:
             dtype = query.dtype
@@ -227,14 +403,30 @@ class MultiQueryAttention(nn.Module):
                 value = torch.cat([past_key_value[1], value], dim=1)
             past_key_value = (key, value)
         if attn_bias is not None:
-            attn_bias = attn_bias[:, :, -query.size(1):, -key.size(1):]
-        (context, attn_weights) = self.attn_fn(query, key, value, self.n_heads, softmax_scale=self.softmax_scale, attn_bias=attn_bias, key_padding_mask=key_padding_mask, is_causal=is_causal, dropout_p=self.attn_dropout_p, training=self.training, needs_weights=needs_weights, multiquery=True)
+            attn_bias = attn_bias[:, :, -query.size(1) :, -key.size(1) :]
+        (context, attn_weights) = self.attn_fn(
+            query,
+            key,
+            value,
+            self.n_heads,
+            softmax_scale=self.softmax_scale,
+            attn_bias=attn_bias,
+            key_padding_mask=key_padding_mask,
+            is_causal=is_causal,
+            dropout_p=self.attn_dropout_p,
+            training=self.training,
+            needs_weights=needs_weights,
+            multiquery=True,
+        )
         return (self.out_proj(context), attn_weights, past_key_value)
 
-def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id):
-    if attn_impl == 'flash':
+
+def attn_bias_shape(
+    attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_sequence_id
+):
+    if attn_impl == "flash":
         return None
-    elif attn_impl in ['torch', 'triton']:
+    elif attn_impl in ["torch", "triton"]:
         if alibi:
             if (prefix_lm or not causal) or use_sequence_id:
                 return (1, n_heads, seq_len, seq_len)
@@ -243,18 +435,31 @@ def attn_bias_shape(attn_impl, n_heads, seq_len, alibi, prefix_lm, causal, use_s
             return (1, 1, seq_len, seq_len)
         return None
     else:
-        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+        raise ValueError(f"attn_impl={attn_impl!r} is an invalid setting.")
+
 
-def build_attn_bias(attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8):
-    if attn_impl == 'flash':
+def build_attn_bias(
+    attn_impl, attn_bias, n_heads, seq_len, causal=False, alibi=False, alibi_bias_max=8
+):
+    if attn_impl == "flash":
         return None
-    elif attn_impl in ['torch', 'triton']:
+    elif attn_impl in ["torch", "triton"]:
         if alibi:
             (device, dtype) = (attn_bias.device, attn_bias.dtype)
-            attn_bias = attn_bias.add(build_alibi_bias(n_heads, seq_len, full=not causal, alibi_bias_max=alibi_bias_max, device=device, dtype=dtype))
+            attn_bias = attn_bias.add(
+                build_alibi_bias(
+                    n_heads,
+                    seq_len,
+                    full=not causal,
+                    alibi_bias_max=alibi_bias_max,
+                    device=device,
+                    dtype=dtype,
+                )
+            )
         return attn_bias
     else:
-        raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
+        raise ValueError(f"attn_impl={attn_impl!r} is an invalid setting.")
+
 
 def gen_slopes(n_heads, alibi_bias_max=8, device=None):
     _n_heads = 2 ** math.ceil(math.log2(n_heads))
@@ -265,12 +470,24 @@ def gen_slopes(n_heads, alibi_bias_max=8, device=None):
         slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
     return slopes.view(1, n_heads, 1, 1)
 
-def build_alibi_bias(n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None):
-    alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, 1, seq_len)
+
+def build_alibi_bias(
+    n_heads, seq_len, full=False, alibi_bias_max=8, device=None, dtype=None
+):
+    alibi_bias = torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(
+        1, 1, 1, seq_len
+    )
     if full:
-        alibi_bias = alibi_bias - torch.arange(1 - seq_len, 1, dtype=torch.int32, device=device).view(1, 1, seq_len, 1)
+        alibi_bias = alibi_bias - torch.arange(
+            1 - seq_len, 1, dtype=torch.int32, device=device
+        ).view(1, 1, seq_len, 1)
         alibi_bias = alibi_bias.abs().mul(-1)
     slopes = gen_slopes(n_heads, alibi_bias_max, device=device)
     alibi_bias = alibi_bias * slopes
     return alibi_bias.to(dtype=dtype)
-ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}
+
+
+ATTN_CLASS_REGISTRY = {
+    "multihead_attention": MultiheadAttention,
+    "multiquery_attention": MultiQueryAttention,
+}

blocks.py

@@ -5,37 +5,84 @@ import torch.nn as nn
 from .attention import ATTN_CLASS_REGISTRY
 from .norm import NORM_CLASS_REGISTRY
 
-class MPTMLP(nn.Module):
 
-    def __init__(self, d_model: int, expansion_ratio: int, device: Optional[str]=None):
+class MPTMLP(nn.Module):
+    def __init__(
+        self, d_model: int, expansion_ratio: int, device: Optional[str] = None
+    ):
         super().__init__()
         self.up_proj = nn.Linear(d_model, expansion_ratio * d_model, device=device)
-        self.act = nn.GELU(approximate='none')
+        self.act = nn.GELU(approximate="none")
         self.down_proj = nn.Linear(expansion_ratio * d_model, d_model, device=device)
         self.down_proj._is_residual = True
 
     def forward(self, x):
         return self.down_proj(self.act(self.up_proj(x)))
 
-class MPTBlock(nn.Module):
 
-    def __init__(self, d_model: int, n_heads: int, expansion_ratio: int, attn_config: Dict={'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', device: Optional[str]=None, **kwargs):
+class MPTBlock(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        expansion_ratio: int,
+        attn_config: Dict = {
+            "attn_type": "multihead_attention",
+            "attn_pdrop": 0.0,
+            "attn_impl": "triton",
+            "qk_ln": False,
+            "clip_qkv": None,
+            "softmax_scale": None,
+            "prefix_lm": False,
+            "attn_uses_sequence_id": False,
+            "alibi": False,
+            "alibi_bias_max": 8,
+        },
+        resid_pdrop: float = 0.0,
+        norm_type: str = "low_precision_layernorm",
+        device: Optional[str] = None,
+        **kwargs
+    ):
         del kwargs
         super().__init__()
         norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
-        attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
+        attn_class = ATTN_CLASS_REGISTRY[attn_config["attn_type"]]
         self.norm_1 = norm_class(d_model, device=device)
-        self.attn = attn_class(attn_impl=attn_config['attn_impl'], clip_qkv=attn_config['clip_qkv'], qk_ln=attn_config['qk_ln'], softmax_scale=attn_config['softmax_scale'], attn_pdrop=attn_config['attn_pdrop'], d_model=d_model, n_heads=n_heads, device=device)
+        self.attn = attn_class(
+            attn_impl=attn_config["attn_impl"],
+            clip_qkv=attn_config["clip_qkv"],
+            qk_ln=attn_config["qk_ln"],
+            softmax_scale=attn_config["softmax_scale"],
+            attn_pdrop=attn_config["attn_pdrop"],
+            d_model=d_model,
+            n_heads=n_heads,
+            device=device,
+        )
         self.norm_2 = norm_class(d_model, device=device)
-        self.ffn = MPTMLP(d_model=d_model, expansion_ratio=expansion_ratio, device=device)
+        self.ffn = MPTMLP(
+            d_model=d_model, expansion_ratio=expansion_ratio, device=device
+        )
         self.resid_attn_dropout = nn.Dropout(resid_pdrop)
         self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
 
-    def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:
+    def forward(
+        self,
+        x: torch.Tensor,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        attn_bias: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.ByteTensor] = None,
+        is_causal: bool = True,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor]]]:
         a = self.norm_1(x)
-        (b, _, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=is_causal)
+        (b, _, past_key_value) = self.attn(
+            a,
+            past_key_value=past_key_value,
+            attn_bias=attn_bias,
+            attention_mask=attention_mask,
+            is_causal=is_causal,
+        )
         x = x + self.resid_attn_dropout(b)
         m = self.norm_2(x)
         n = self.ffn(m)
         x = x + self.resid_ffn_dropout(n)
-        return (x, past_key_value)
+        return (x, past_key_value)

configuration_mpt.py

@@ -1,13 +1,51 @@
 """A HuggingFace-style model configuration."""
 from typing import Dict, Optional, Union
 from transformers import PretrainedConfig
-attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'alibi': False, 'alibi_bias_max': 8}
-init_config_defaults: Dict = {'name': 'kaiming_normal_', 'fan_mode': 'fan_in', 'init_nonlinearity': 'relu'}
+
+attn_config_defaults: Dict = {
+    "attn_type": "multihead_attention",
+    "attn_pdrop": 0.0,
+    "attn_impl": "triton",
+    "qk_ln": False,
+    "clip_qkv": None,
+    "softmax_scale": None,
+    "prefix_lm": False,
+    "attn_uses_sequence_id": False,
+    "alibi": False,
+    "alibi_bias_max": 8,
+}
+init_config_defaults: Dict = {
+    "name": "kaiming_normal_",
+    "fan_mode": "fan_in",
+    "init_nonlinearity": "relu",
+}
+
 
 class MPTConfig(PretrainedConfig):
-    model_type = 'mpt'
+    model_type = "mpt"
 
-    def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: int=4, max_seq_len: int=2048, vocab_size: int=50368, resid_pdrop: float=0.0, emb_pdrop: float=0.0, learned_pos_emb: bool=True, attn_config: Dict=attn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, verbose: int=0, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, **kwargs):
+    def __init__(
+        self,
+        d_model: int = 2048,
+        n_heads: int = 16,
+        n_layers: int = 24,
+        expansion_ratio: int = 4,
+        max_seq_len: int = 2048,
+        vocab_size: int = 50368,
+        resid_pdrop: float = 0.0,
+        emb_pdrop: float = 0.0,
+        learned_pos_emb: bool = True,
+        attn_config: Dict = attn_config_defaults,
+        init_device: str = "cpu",
+        logit_scale: Optional[Union[float, str]] = None,
+        no_bias: bool = False,
+        verbose: int = 0,
+        embedding_fraction: float = 1.0,
+        norm_type: str = "low_precision_layernorm",
+        use_cache: bool = False,
+        init_config: Dict = init_config_defaults,
+        **kwargs,
+    ):
         """The MPT configuration class.
 
         Args:
@@ -80,39 +118,76 @@ class MPTConfig(PretrainedConfig):
         self.norm_type = norm_type
         self.use_cache = use_cache
         self.init_config = init_config
-        if 'name' in kwargs:
-            del kwargs['name']
-        if 'loss_fn' in kwargs:
-            del kwargs['loss_fn']
+        if "name" in kwargs:
+            del kwargs["name"]
+        if "loss_fn" in kwargs:
+            del kwargs["loss_fn"]
         super().__init__(**kwargs)
         self._validate_config()
 
     def _set_config_defaults(self, config, config_defaults):
-        for (k, v) in config_defaults.items():
+        for k, v in config_defaults.items():
             if k not in config:
                 config[k] = v
         return config
 
     def _validate_config(self):
-        self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
-        self.init_config = self._set_config_defaults(self.init_config, init_config_defaults)
+        self.attn_config = self._set_config_defaults(
+            self.attn_config, attn_config_defaults
+        )
+        self.init_config = self._set_config_defaults(
+            self.init_config, init_config_defaults
+        )
         if self.d_model % self.n_heads != 0:
-            raise ValueError('d_model must be divisible by n_heads')
-        if any((prob < 0 or prob > 1 for prob in [self.attn_config['attn_pdrop'], self.resid_pdrop, self.emb_pdrop])):
-            raise ValueError("self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1")
-        if self.attn_config['attn_impl'] not in ['torch', 'flash', 'triton']:
+            raise ValueError("d_model must be divisible by n_heads")
+        if any(
+            (
+                prob < 0 or prob > 1
+                for prob in [
+                    self.attn_config["attn_pdrop"],
+                    self.resid_pdrop,
+                    self.emb_pdrop,
+                ]
+            )
+        ):
+            raise ValueError(
+                "self.attn_config['attn_pdrop'], resid_pdrop, emb_pdrop are probabilities and must be between 0 and 1"
+            )
+        if self.attn_config["attn_impl"] not in ["torch", "flash", "triton"]:
             raise ValueError(f"Unknown attn_impl={self.attn_config['attn_impl']}")
-        if self.attn_config['prefix_lm'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError('prefix_lm only implemented with torch and triton attention.')
-        if self.attn_config['alibi'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError('alibi only implemented with torch and triton attention.')
-        if self.attn_config['attn_uses_sequence_id'] and self.attn_config['attn_impl'] not in ['torch', 'triton']:
-            raise NotImplementedError('attn_uses_sequence_id only implemented with torch and triton attention.')
+        if self.attn_config["prefix_lm"] and self.attn_config["attn_impl"] not in [
+            "torch",
+            "triton",
+        ]:
+            raise NotImplementedError(
+                "prefix_lm only implemented with torch and triton attention."
+            )
+        if self.attn_config["alibi"] and self.attn_config["attn_impl"] not in [
+            "torch",
+            "triton",
+        ]:
+            raise NotImplementedError(
+                "alibi only implemented with torch and triton attention."
+            )
+        if self.attn_config["attn_uses_sequence_id"] and self.attn_config[
+            "attn_impl"
+        ] not in ["torch", "triton"]:
+            raise NotImplementedError(
+                "attn_uses_sequence_id only implemented with torch and triton attention."
+            )
         if self.embedding_fraction > 1 or self.embedding_fraction <= 0:
-            raise ValueError('model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!')
-        if isinstance(self.logit_scale, str) and self.logit_scale != 'inv_sqrt_d_model':
-            raise ValueError(f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
-        if self.init_config.get('name', None) is None:
-            raise ValueError(f"self.init_config={self.init_config!r} 'name' needs to be set.")
-        if not self.learned_pos_emb and (not self.attn_config['alibi']):
-            raise ValueError(f'Positional information must be provided to the model using either learned_pos_emb or alibi.')
+            raise ValueError(
+                "model.embedding_fraction must be between 0 (exclusive) and 1 (inclusive)!"
+            )
+        if isinstance(self.logit_scale, str) and self.logit_scale != "inv_sqrt_d_model":
+            raise ValueError(
+                f"self.logit_scale={self.logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'."
+            )
+        if self.init_config.get("name", None) is None:
+            raise ValueError(
+                f"self.init_config={self.init_config!r} 'name' needs to be set."
+            )
+        if not self.learned_pos_emb and (not self.attn_config["alibi"]):
+            raise ValueError(
+                f"Positional information must be provided to the model using either learned_pos_emb or alibi."
+            )

hf_prefixlm_converter.py

@@ -11,9 +11,17 @@ import warnings
 from types import MethodType
 from typing import Any, Dict, List, Optional, Tuple, Union
 import torch
-from transformers.models.bloom.modeling_bloom import BaseModelOutputWithPastAndCrossAttentions, BloomForCausalLM, BloomModel, CausalLMOutputWithCrossAttentions, CrossEntropyLoss
+from transformers.models.bloom.modeling_bloom import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BloomForCausalLM,
+    BloomModel,
+    CausalLMOutputWithCrossAttentions,
+    CrossEntropyLoss,
+)
 from transformers.models.bloom.modeling_bloom import _expand_mask as _expand_mask_bloom
-from transformers.models.bloom.modeling_bloom import _make_causal_mask as _make_causal_mask_bloom
+from transformers.models.bloom.modeling_bloom import (
+    _make_causal_mask as _make_causal_mask_bloom,
+)
 from transformers.models.bloom.modeling_bloom import logging
 from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
@@ -21,10 +29,21 @@ from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXForCausalLM
 from transformers.models.gptj.modeling_gptj import GPTJForCausalLM
 from transformers.models.opt.modeling_opt import OPTForCausalLM
 from transformers.models.opt.modeling_opt import _expand_mask as _expand_mask_opt
-from transformers.models.opt.modeling_opt import _make_causal_mask as _make_causal_mask_opt
+from transformers.models.opt.modeling_opt import (
+    _make_causal_mask as _make_causal_mask_opt,
+)
+
 logger = logging.get_logger(__name__)
-_SUPPORTED_GPT_MODELS = (GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)
-CAUSAL_GPT_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM]
+_SUPPORTED_GPT_MODELS = (
+    GPT2LMHeadModel,
+    GPTJForCausalLM,
+    GPTNeoForCausalLM,
+    GPTNeoXForCausalLM,
+)
+CAUSAL_GPT_TYPES = Union[
+    GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM
+]
+
 
 def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_TYPES:
     """Converts a GPT-style Causal LM to a Prefix LM.
@@ -37,10 +56,12 @@ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_T
 
     See `convert_hf_causal_lm_to_prefix_lm` for more details.
     """
-    if hasattr(model, '_prefix_lm_converted'):
+    if hasattr(model, "_prefix_lm_converted"):
         return model
     assert isinstance(model, _SUPPORTED_GPT_MODELS)
-    assert model.config.add_cross_attention == False, 'Only supports GPT-style decoder-only models'
+    assert (
+        model.config.add_cross_attention == False
+    ), "Only supports GPT-style decoder-only models"
 
     def _get_attn_modules(model: CAUSAL_GPT_TYPES) -> List[torch.nn.Module]:
         """Helper that gets a list of the model's attention modules.
@@ -56,7 +77,7 @@ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_T
             blocks = model.transformer.h
         for block in blocks:
             if isinstance(model, GPTNeoForCausalLM):
-                if block.attn.attention_type != 'global':
+                if block.attn.attention_type != "global":
                     continue
                 attn_module = block.attn.attention
             elif isinstance(model, GPTNeoXForCausalLM):
@@ -65,17 +86,58 @@ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_T
                 attn_module = block.attn
             attn_modules.append(attn_module)
         return attn_modules
-    setattr(model, '_original_forward', getattr(model, 'forward'))
-    setattr(model, '_original_generate', getattr(model, 'generate'))
 
-    def forward(self: CAUSAL_GPT_TYPES, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[Tuple[torch.Tensor]]]=None, attention_mask: Optional[torch.FloatTensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, token_type_ids: Optional[torch.LongTensor]=None, position_ids: Optional[torch.LongTensor]=None, head_mask: Optional[torch.FloatTensor]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
+    setattr(model, "_original_forward", getattr(model, "forward"))
+    setattr(model, "_original_generate", getattr(model, "generate"))
+
+    def forward(
+        self: CAUSAL_GPT_TYPES,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        bidirectional_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
         """Wraps original forward to enable PrefixLM attention."""
 
         def call_og_forward():
             if isinstance(self, GPTNeoXForCausalLM):
-                return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+                return self._original_forward(
+                    input_ids=input_ids,
+                    past_key_values=past_key_values,
+                    attention_mask=attention_mask,
+                    head_mask=head_mask,
+                    inputs_embeds=inputs_embeds,
+                    labels=labels,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=return_dict,
+                )
             else:
-                return self._original_forward(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+                return self._original_forward(
+                    input_ids=input_ids,
+                    past_key_values=past_key_values,
+                    attention_mask=attention_mask,
+                    token_type_ids=token_type_ids,
+                    position_ids=position_ids,
+                    head_mask=head_mask,
+                    inputs_embeds=inputs_embeds,
+                    labels=labels,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    return_dict=return_dict,
+                )
+
         if bidirectional_mask is None:
             return call_og_forward()
         assert isinstance(bidirectional_mask, torch.Tensor)
@@ -83,14 +145,23 @@ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_T
         (b, s) = bidirectional_mask.shape
         max_length = attn_modules[0].bias.shape[-1]
         if s > max_length:
-            raise ValueError(f'bidirectional_mask sequence length (={s}) exceeds the ' + f'max length allowed by the model ({max_length}).')
+            raise ValueError(
+                f"bidirectional_mask sequence length (={s}) exceeds the "
+                + f"max length allowed by the model ({max_length})."
+            )
         assert s <= max_length
         if s < max_length:
-            pad = torch.zeros((int(b), int(max_length - s)), dtype=bidirectional_mask.dtype, device=bidirectional_mask.device)
+            pad = torch.zeros(
+                (int(b), int(max_length - s)),
+                dtype=bidirectional_mask.dtype,
+                device=bidirectional_mask.device,
+            )
             bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)
         bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)
         for attn_module in attn_modules:
-            attn_module.bias.data = torch.logical_or(attn_module.bias.data, bidirectional)
+            attn_module.bias.data = torch.logical_or(
+                attn_module.bias.data, bidirectional
+            )
         output = call_og_forward()
         for attn_module in attn_modules:
             attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
@@ -105,11 +176,13 @@ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_T
         for attn_module in attn_modules:
             attn_module.bias.data = torch.tril(attn_module.bias.data[0, 0])[None, None]
         return output
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'generate', MethodType(generate, model))
-    setattr(model, '_prefix_lm_converted', True)
+
+    setattr(model, "forward", MethodType(forward, model))
+    setattr(model, "generate", MethodType(generate, model))
+    setattr(model, "_prefix_lm_converted", True)
     return model
 
+
 def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:
     """Converts a BLOOM Causal LM to a Prefix LM.
 
@@ -118,62 +191,137 @@ def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCa
 
     See `convert_hf_causal_lm_to_prefix_lm` for more details.
     """
-    if hasattr(model, '_prefix_lm_converted'):
+    if hasattr(model, "_prefix_lm_converted"):
         return model
     assert isinstance(model, BloomForCausalLM)
-    assert model.config.add_cross_attention == False, 'Only supports BLOOM decoder-only models'
-
-    def _prepare_attn_mask(self: BloomModel, attention_mask: torch.Tensor, bidirectional_mask: Optional[torch.Tensor], input_shape: Tuple[int, int], past_key_values_length: int) -> torch.BoolTensor:
+    assert (
+        model.config.add_cross_attention == False
+    ), "Only supports BLOOM decoder-only models"
+
+    def _prepare_attn_mask(
+        self: BloomModel,
+        attention_mask: torch.Tensor,
+        bidirectional_mask: Optional[torch.Tensor],
+        input_shape: Tuple[int, int],
+        past_key_values_length: int,
+    ) -> torch.BoolTensor:
         combined_attention_mask = None
         device = attention_mask.device
         (_, src_length) = input_shape
         if src_length > 1:
-            combined_attention_mask = _make_causal_mask_bloom(input_shape, device=device, past_key_values_length=past_key_values_length)
+            combined_attention_mask = _make_causal_mask_bloom(
+                input_shape,
+                device=device,
+                past_key_values_length=past_key_values_length,
+            )
             if bidirectional_mask is not None:
                 assert attention_mask.shape == bidirectional_mask.shape
-                expanded_bidirectional_mask = _expand_mask_bloom(bidirectional_mask, tgt_length=src_length)
-                combined_attention_mask = torch.logical_and(combined_attention_mask, expanded_bidirectional_mask)
+                expanded_bidirectional_mask = _expand_mask_bloom(
+                    bidirectional_mask, tgt_length=src_length
+                )
+                combined_attention_mask = torch.logical_and(
+                    combined_attention_mask, expanded_bidirectional_mask
+                )
         expanded_attn_mask = _expand_mask_bloom(attention_mask, tgt_length=src_length)
-        combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
+        combined_attention_mask = (
+            expanded_attn_mask
+            if combined_attention_mask is None
+            else expanded_attn_mask | combined_attention_mask
+        )
         return combined_attention_mask
 
-    def _build_alibi_tensor(self: BloomModel, batch_size: int, query_length: int, key_length: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
+    def _build_alibi_tensor(
+        self: BloomModel,
+        batch_size: int,
+        query_length: int,
+        key_length: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> torch.Tensor:
         num_heads = self.config.n_head
         closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
-        base = torch.tensor(2 ** (-2 ** (-(math.log2(closest_power_of_2) - 3))), device=device, dtype=torch.float32)
-        powers = torch.arange(1, 1 + closest_power_of_2, device=device, dtype=torch.int32)
+        base = torch.tensor(
+            2 ** (-(2 ** (-(math.log2(closest_power_of_2) - 3)))),
+            device=device,
+            dtype=torch.float32,
+        )
+        powers = torch.arange(
+            1, 1 + closest_power_of_2, device=device, dtype=torch.int32
+        )
         slopes = torch.pow(base, powers)
         if closest_power_of_2 != num_heads:
-            extra_base = torch.tensor(2 ** (-2 ** (-(math.log2(2 * closest_power_of_2) - 3))), device=device, dtype=torch.float32)
-            num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
-            extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32)
+            extra_base = torch.tensor(
+                2 ** (-(2 ** (-(math.log2(2 * closest_power_of_2) - 3)))),
+                device=device,
+                dtype=torch.float32,
+            )
+            num_remaining_heads = min(
+                closest_power_of_2, num_heads - closest_power_of_2
+            )
+            extra_powers = torch.arange(
+                1, 1 + 2 * num_remaining_heads, 2, device=device, dtype=torch.int32
+            )
             slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
         qa = torch.arange(query_length, device=device, dtype=torch.int32).view(-1, 1)
         ka = torch.arange(key_length, device=device, dtype=torch.int32).view(1, -1)
         diffs = qa - ka + key_length - query_length
         diffs = -diffs.abs()
-        alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(1, 1, query_length, key_length)
-        alibi = alibi.expand(batch_size, -1, -1, -1).reshape(-1, query_length, key_length)
+        alibi = slopes.view(1, num_heads, 1, 1) * diffs.view(
+            1, 1, query_length, key_length
+        )
+        alibi = alibi.expand(batch_size, -1, -1, -1).reshape(
+            -1, query_length, key_length
+        )
         return alibi.to(dtype)
+
     KeyValueT = Tuple[torch.Tensor, torch.Tensor]
 
-    def forward(self: BloomModel, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.LongTensor]=None, inputs_embeds: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
-        if deprecated_arguments.pop('position_ids', False) is not False:
-            warnings.warn('`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. ' + 'You can safely ignore passing `position_ids`.', FutureWarning)
+    def forward(
+        self: BloomModel,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[KeyValueT, ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        bidirectional_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
+        if deprecated_arguments.pop("position_ids", False) is not False:
+            warnings.warn(
+                "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. "
+                + "You can safely ignore passing `position_ids`.",
+                FutureWarning,
+            )
         if len(deprecated_arguments) > 0:
-            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
         if input_ids is not None and inputs_embeds is not None:
-            raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
         elif input_ids is not None:
             (batch_size, seq_length) = input_ids.shape
         elif inputs_embeds is not None:
             (batch_size, seq_length, _) = inputs_embeds.shape
         else:
-            raise ValueError('You have to specify either input_ids or inputs_embeds')
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
         if past_key_values is None:
             past_key_values = tuple([None] * len(self.h))
         head_mask = self.get_head_mask(head_mask, self.config.n_layer)
@@ -190,28 +338,62 @@ def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCa
             past_key_values_length = tmp.shape[2]
             seq_length_with_past = seq_length_with_past + past_key_values_length
         if attention_mask is None:
-            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
+            attention_mask = torch.ones(
+                (batch_size, seq_length_with_past), device=hidden_states.device
+            )
         else:
             attention_mask = attention_mask.to(hidden_states.device)
-        alibi = self._build_alibi_tensor(batch_size=batch_size, query_length=seq_length, key_length=seq_length_with_past, dtype=hidden_states.dtype, device=hidden_states.device)
-        causal_mask = self._prepare_attn_mask(attention_mask, bidirectional_mask, input_shape=(batch_size, seq_length), past_key_values_length=past_key_values_length)
-        for (i, (block, layer_past)) in enumerate(zip(self.h, past_key_values)):
+        alibi = self._build_alibi_tensor(
+            batch_size=batch_size,
+            query_length=seq_length,
+            key_length=seq_length_with_past,
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
+        causal_mask = self._prepare_attn_mask(
+            attention_mask,
+            bidirectional_mask,
+            input_shape=(batch_size, seq_length),
+            past_key_values_length=past_key_values_length,
+        )
+        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
             if output_hidden_states:
                 hst = (hidden_states,)
                 all_hidden_states = all_hidden_states + hst
             if self.gradient_checkpointing and self.training:
                 if use_cache:
-                    logger.warning('`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...')
+                    logger.warning(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
                     use_cache = False
 
                 def create_custom_forward(module):
-
                     def custom_forward(*inputs):
-                        return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
+                        return module(
+                            *inputs,
+                            use_cache=use_cache,
+                            output_attentions=output_attentions,
+                        )
+
                     return custom_forward
-                outputs = torch.utils.checkpoint.checkpoint(create_custom_forward(block), hidden_states, alibi, causal_mask, head_mask[i])
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    alibi,
+                    causal_mask,
+                    head_mask[i],
+                )
             else:
-                outputs = block(hidden_states, layer_past=layer_past, attention_mask=causal_mask, head_mask=head_mask[i], use_cache=use_cache, output_attentions=output_attentions, alibi=alibi)
+                outputs = block(
+                    hidden_states,
+                    layer_past=layer_past,
+                    attention_mask=causal_mask,
+                    head_mask=head_mask[i],
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    alibi=alibi,
+                )
             hidden_states = outputs[0]
             if use_cache is True:
                 presents = presents + (outputs[1],)
@@ -223,21 +405,77 @@ def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCa
             hst = (hidden_states,)
             all_hidden_states = all_hidden_states + hst
         if not return_dict:
-            return tuple((v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None))
-        return BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=hidden_states, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attentions)
-    setattr(model.transformer, '_prepare_attn_mask', MethodType(_prepare_attn_mask, model.transformer))
-    setattr(model.transformer, '_build_alibi_tensor', MethodType(_build_alibi_tensor, model.transformer))
-    setattr(model.transformer, 'forward', MethodType(forward, model.transformer))
+            return tuple(
+                (
+                    v
+                    for v in [
+                        hidden_states,
+                        presents,
+                        all_hidden_states,
+                        all_self_attentions,
+                    ]
+                    if v is not None
+                )
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=presents,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+    setattr(
+        model.transformer,
+        "_prepare_attn_mask",
+        MethodType(_prepare_attn_mask, model.transformer),
+    )
+    setattr(
+        model.transformer,
+        "_build_alibi_tensor",
+        MethodType(_build_alibi_tensor, model.transformer),
+    )
+    setattr(model.transformer, "forward", MethodType(forward, model.transformer))
     KeyValueT = Tuple[torch.Tensor, torch.Tensor]
 
-    def forward(self: BloomForCausalLM, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[Tuple[KeyValueT, ...]]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.Tensor]=None, head_mask: Optional[torch.Tensor]=None, inputs_embeds: Optional[torch.Tensor]=None, labels: Optional[torch.Tensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None, **deprecated_arguments) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
+    def forward(
+        self: BloomForCausalLM,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[KeyValueT, ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        bidirectional_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithCrossAttentions]:
         """Replacement forward method for BloomCausalLM."""
-        if deprecated_arguments.pop('position_ids', False) is not False:
-            warnings.warn('`position_ids` have no functionality in BLOOM and will be removed ' + 'in v5.0.0. You can safely ignore passing `position_ids`.', FutureWarning)
+        if deprecated_arguments.pop("position_ids", False) is not False:
+            warnings.warn(
+                "`position_ids` have no functionality in BLOOM and will be removed "
+                + "in v5.0.0. You can safely ignore passing `position_ids`.",
+                FutureWarning,
+            )
         if len(deprecated_arguments) > 0:
-            raise ValueError(f'Got unexpected arguments: {deprecated_arguments}')
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-        transformer_outputs = self.transformer(input_ids, past_key_values=past_key_values, attention_mask=attention_mask, bidirectional_mask=bidirectional_mask, head_mask=head_mask, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        transformer_outputs = self.transformer(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            bidirectional_mask=bidirectional_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
         hidden_states = transformer_outputs[0]
         lm_logits = self.lm_head(hidden_states)
         loss = None
@@ -246,13 +484,28 @@ def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCa
             shift_labels = labels[..., 1:].contiguous()
             (batch_size, seq_length, vocab_size) = shift_logits.shape
             loss_fct = CrossEntropyLoss()
-            loss = loss_fct(shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length))
+            loss = loss_fct(
+                shift_logits.view(batch_size * seq_length, vocab_size),
+                shift_labels.view(batch_size * seq_length),
+            )
         if not return_dict:
             output = (lm_logits,) + transformer_outputs[1:]
             return (loss,) + output if loss is not None else output
-        return CausalLMOutputWithCrossAttentions(loss=loss, logits=lm_logits, past_key_values=transformer_outputs.past_key_values, hidden_states=transformer_outputs.hidden_states, attentions=transformer_outputs.attentions)
-
-    def prepare_inputs_for_generation(self: BloomForCausalLM, input_ids: torch.LongTensor, past: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, **kwargs) -> dict:
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self: BloomForCausalLM,
+        input_ids: torch.LongTensor,
+        past: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> dict:
         if past:
             input_ids = input_ids[:, -1].unsqueeze(-1)
             bidirectional_mask = None
@@ -260,12 +513,24 @@ def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCa
                 past = self._convert_to_bloom_cache(past)
         else:
             bidirectional_mask = torch.ones_like(input_ids)
-        return {'input_ids': input_ids, 'past_key_values': past, 'use_cache': True, 'attention_mask': attention_mask, 'bidirectional_mask': bidirectional_mask}
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'prepare_inputs_for_generation', MethodType(prepare_inputs_for_generation, model))
-    setattr(model, '_prefix_lm_converted', True)
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "use_cache": True,
+            "attention_mask": attention_mask,
+            "bidirectional_mask": bidirectional_mask,
+        }
+
+    setattr(model, "forward", MethodType(forward, model))
+    setattr(
+        model,
+        "prepare_inputs_for_generation",
+        MethodType(prepare_inputs_for_generation, model),
+    )
+    setattr(model, "_prefix_lm_converted", True)
     return model
 
+
 def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
     """Converts an OPT Causal LM to a Prefix LM.
 
@@ -274,36 +539,89 @@ def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM
 
     See `convert_hf_causal_lm_to_prefix_lm` for more details.
     """
-    if hasattr(model, '_prefix_lm_converted'):
+    if hasattr(model, "_prefix_lm_converted"):
         return model
     assert isinstance(model, OPTForCausalLM)
-    assert model.config.add_cross_attention == False, 'Only supports OPT decoder-only models'
-    setattr(model, '_original_forward', getattr(model, 'forward'))
-    setattr(model, '_original_generate', getattr(model, 'generate'))
+    assert (
+        model.config.add_cross_attention == False
+    ), "Only supports OPT decoder-only models"
+    setattr(model, "_original_forward", getattr(model, "forward"))
+    setattr(model, "_original_generate", getattr(model, "generate"))
     model.model.decoder.bidirectional_mask = None
 
-    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+    def _prepare_decoder_attention_mask(
+        self, attention_mask, input_shape, inputs_embeds, past_key_values_length
+    ):
         combined_attention_mask = None
         if input_shape[-1] > 1:
-            if self.bidirectional_mask == 'g':
+            if self.bidirectional_mask == "g":
                 (bsz, src_length) = input_shape
-                combined_attention_mask = torch.zeros((bsz, 1, src_length, src_length + past_key_values_length), dtype=inputs_embeds.dtype, device=inputs_embeds.device)
+                combined_attention_mask = torch.zeros(
+                    (bsz, 1, src_length, src_length + past_key_values_length),
+                    dtype=inputs_embeds.dtype,
+                    device=inputs_embeds.device,
+                )
             else:
-                combined_attention_mask = _make_causal_mask_opt(input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length).to(inputs_embeds.device)
+                combined_attention_mask = _make_causal_mask_opt(
+                    input_shape,
+                    inputs_embeds.dtype,
+                    past_key_values_length=past_key_values_length,
+                ).to(inputs_embeds.device)
                 if self.bidirectional_mask is not None:
                     assert attention_mask.shape == self.bidirectional_mask.shape
-                    expanded_bidirectional_mask = _expand_mask_opt(self.bidirectional_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
-                    combined_attention_mask = torch.maximum(expanded_bidirectional_mask, combined_attention_mask)
+                    expanded_bidirectional_mask = _expand_mask_opt(
+                        self.bidirectional_mask,
+                        inputs_embeds.dtype,
+                        tgt_len=input_shape[-1],
+                    ).to(inputs_embeds.device)
+                    combined_attention_mask = torch.maximum(
+                        expanded_bidirectional_mask, combined_attention_mask
+                    )
         if attention_mask is not None:
-            expanded_attn_mask = _expand_mask_opt(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(inputs_embeds.device)
-            combined_attention_mask = expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+            expanded_attn_mask = _expand_mask_opt(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
+            ).to(inputs_embeds.device)
+            combined_attention_mask = (
+                expanded_attn_mask
+                if combined_attention_mask is None
+                else expanded_attn_mask + combined_attention_mask
+            )
         return combined_attention_mask
-    setattr(model.model.decoder, '_prepare_decoder_attention_mask', MethodType(_prepare_decoder_attention_mask, model.model.decoder))
-
-    def forward(self: OPTForCausalLM, input_ids: Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None, bidirectional_mask: Optional[torch.ByteTensor]=None, head_mask: Optional[torch.Tensor]=None, past_key_values: Optional[List[torch.FloatTensor]]=None, inputs_embeds: Optional[torch.FloatTensor]=None, labels: Optional[torch.LongTensor]=None, use_cache: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, return_dict: Optional[bool]=None):
 
+    setattr(
+        model.model.decoder,
+        "_prepare_decoder_attention_mask",
+        MethodType(_prepare_decoder_attention_mask, model.model.decoder),
+    )
+
+    def forward(
+        self: OPTForCausalLM,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        bidirectional_mask: Optional[torch.ByteTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
         def call_og_forward():
-            return self._original_forward(input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict)
+            return self._original_forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                head_mask=head_mask,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+
         if bidirectional_mask is None:
             return call_og_forward()
         self.model.decoder.bidirectional_mask = bidirectional_mask
@@ -317,7 +635,7 @@ def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM
 
     def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):
         """Wraps original generate to enable PrefixLM-style attention."""
-        self.model.decoder.bidirectional_mask = 'g'
+        self.model.decoder.bidirectional_mask = "g"
         try:
             output = self._original_generate(*args, **kwargs)
         except:
@@ -325,12 +643,23 @@ def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM
             raise
         self.model.decoder.bidirectional_mask = None
         return output
-    setattr(model, 'forward', MethodType(forward, model))
-    setattr(model, 'generate', MethodType(generate, model))
-    setattr(model, '_prefix_lm_converted', True)
+
+    setattr(model, "forward", MethodType(forward, model))
+    setattr(model, "generate", MethodType(generate, model))
+    setattr(model, "_prefix_lm_converted", True)
     return model
+
+
 _SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (BloomForCausalLM, OPTForCausalLM)
-CAUSAL_LM_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM, BloomForCausalLM, OPTForCausalLM]
+CAUSAL_LM_TYPES = Union[
+    GPT2LMHeadModel,
+    GPTJForCausalLM,
+    GPTNeoForCausalLM,
+    GPTNeoXForCausalLM,
+    BloomForCausalLM,
+    OPTForCausalLM,
+]
+
 
 def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES:
     """Converts a HuggingFace Causal LM to a Prefix LM.
@@ -396,7 +725,12 @@ def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES
     elif isinstance(model, OPTForCausalLM):
         return _convert_opt_causal_lm_to_prefix_lm(model)
     else:
-        raise TypeError(f'Cannot convert model to Prefix LM. ' + f'Model does not belong to set of supported HF models:' + f'\n{_SUPPORTED_HF_MODELS}')
+        raise TypeError(
+            f"Cannot convert model to Prefix LM. "
+            + f"Model does not belong to set of supported HF models:"
+            + f"\n{_SUPPORTED_HF_MODELS}"
+        )
+
 
 def add_bidirectional_mask_if_missing(batch: Dict[str, Any]):
     """Attempts to add bidirectional_mask to batch if missing.
@@ -404,12 +738,16 @@ def add_bidirectional_mask_if_missing(batch: Dict[str, Any]):
     Raises:
         KeyError if bidirectional_mask is missing and can't be inferred
     """
-    if 'bidirectional_mask' not in batch:
-        if batch.get('mode', None) == 'icl_task':
-            batch['bidirectional_mask'] = batch['attention_mask'].clone()
-            for (i, continuation_indices) in enumerate(batch['continuation_indices']):
-                batch['bidirectional_mask'][i, continuation_indices] = 0
-        elif 'labels' in batch and 'attention_mask' in batch:
-            batch['bidirectional_mask'] = torch.logical_and(torch.eq(batch['attention_mask'], 1), torch.eq(batch['labels'], -100)).type_as(batch['attention_mask'])
+    if "bidirectional_mask" not in batch:
+        if batch.get("mode", None) == "icl_task":
+            batch["bidirectional_mask"] = batch["attention_mask"].clone()
+            for i, continuation_indices in enumerate(batch["continuation_indices"]):
+                batch["bidirectional_mask"][i, continuation_indices] = 0
+        elif "labels" in batch and "attention_mask" in batch:
+            batch["bidirectional_mask"] = torch.logical_and(
+                torch.eq(batch["attention_mask"], 1), torch.eq(batch["labels"], -100)
+            ).type_as(batch["attention_mask"])
         else:
-            raise KeyError('No bidirectional_mask in batch and not sure how to construct one.')
+            raise KeyError(
+                "No bidirectional_mask in batch and not sure how to construct one."
+            )

meta_init_context.py

@@ -2,8 +2,9 @@ from contextlib import contextmanager
 import torch
 import torch.nn as nn
 
+
 @contextmanager
-def init_empty_weights(include_buffers: bool=False):
+def init_empty_weights(include_buffers: bool = False):
     """Meta initialization context manager.
 
     A context manager under which models are initialized with all parameters
@@ -30,11 +31,12 @@ def init_empty_weights(include_buffers: bool=False):
 
     </Tip>
     """
-    with init_on_device(torch.device('meta'), include_buffers=include_buffers) as f:
+    with init_on_device(torch.device("meta"), include_buffers=include_buffers) as f:
         yield f
 
+
 @contextmanager
-def init_on_device(device: torch.device, include_buffers: bool=False):
+def init_on_device(device: torch.device, include_buffers: bool = False):
     """Device initialization context manager.
 
     A context manager under which models are initialized with all parameters
@@ -62,33 +64,47 @@ def init_on_device(device: torch.device, include_buffers: bool=False):
         if param is not None:
             param_cls = type(module._parameters[name])
             kwargs = module._parameters[name].__dict__
-            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+            module._parameters[name] = param_cls(
+                module._parameters[name].to(device), **kwargs
+            )
 
     def register_empty_buffer(module, name, buffer):
         old_register_buffer(module, name, buffer)
         if buffer is not None:
             module._buffers[name] = module._buffers[name].to(device)
+
     if include_buffers:
-        tensor_constructors_to_patch = {torch_function_name: getattr(torch, torch_function_name) for torch_function_name in ['empty', 'zeros', 'ones', 'full']}
+        tensor_constructors_to_patch = {
+            torch_function_name: getattr(torch, torch_function_name)
+            for torch_function_name in ["empty", "zeros", "ones", "full"]
+        }
     else:
         tensor_constructors_to_patch = {}
 
     def patch_tensor_constructor(fn):
-
         def wrapper(*args, **kwargs):
-            kwargs['device'] = device
+            kwargs["device"] = device
             return fn(*args, **kwargs)
+
         return wrapper
+
     try:
         nn.Module.register_parameter = register_empty_parameter
         if include_buffers:
             nn.Module.register_buffer = register_empty_buffer
         for torch_function_name in tensor_constructors_to_patch.keys():
-            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
+            setattr(
+                torch,
+                torch_function_name,
+                patch_tensor_constructor(getattr(torch, torch_function_name)),
+            )
         yield
     finally:
         nn.Module.register_parameter = old_register_parameter
         if include_buffers:
             nn.Module.register_buffer = old_register_buffer
-        for (torch_function_name, old_torch_function) in tensor_constructors_to_patch.items():
-            setattr(torch, torch_function_name, old_torch_function)
+        for (
+            torch_function_name,
+            old_torch_function,
+        ) in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)

norm.py

@@ -1,28 +1,55 @@
 import torch
 
+
 def _cast_if_autocast_enabled(tensor):
     if torch.is_autocast_enabled():
-        if tensor.device.type == 'cuda':
+        if tensor.device.type == "cuda":
             dtype = torch.get_autocast_gpu_dtype()
-        elif tensor.device.type == 'cpu':
+        elif tensor.device.type == "cpu":
             dtype = torch.get_autocast_cpu_dtype()
         else:
             raise NotImplementedError()
         return tensor.to(dtype=dtype)
     return tensor
 
-class LPLayerNorm(torch.nn.LayerNorm):
 
-    def __init__(self, normalized_shape, eps=1e-05, elementwise_affine=True, device=None, dtype=None):
-        super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=elementwise_affine, device=device, dtype=dtype)
+class LPLayerNorm(torch.nn.LayerNorm):
+    def __init__(
+        self,
+        normalized_shape,
+        eps=1e-05,
+        elementwise_affine=True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__(
+            normalized_shape=normalized_shape,
+            eps=eps,
+            elementwise_affine=elementwise_affine,
+            device=device,
+            dtype=dtype,
+        )
 
     def forward(self, x):
         module_device = x.device
         downcast_x = _cast_if_autocast_enabled(x)
-        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
-        downcast_bias = _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
+        downcast_weight = (
+            _cast_if_autocast_enabled(self.weight)
+            if self.weight is not None
+            else self.weight
+        )
+        downcast_bias = (
+            _cast_if_autocast_enabled(self.bias) if self.bias is not None else self.bias
+        )
         with torch.autocast(enabled=False, device_type=module_device.type):
-            return torch.nn.functional.layer_norm(downcast_x, self.normalized_shape, downcast_weight, downcast_bias, self.eps)
+            return torch.nn.functional.layer_norm(
+                downcast_x,
+                self.normalized_shape,
+                downcast_weight,
+                downcast_bias,
+                self.eps,
+            )
+
 
 def rms_norm(x, weight=None, eps=1e-05):
     output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
@@ -30,27 +57,50 @@ def rms_norm(x, weight=None, eps=1e-05):
         return output * weight
     return output
 
-class RMSNorm(torch.nn.Module):
 
-    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
+class RMSNorm(torch.nn.Module):
+    def __init__(
+        self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None
+    ):
         super().__init__()
         self.eps = eps
         if weight:
-            self.weight = torch.nn.Parameter(torch.ones(normalized_shape, dtype=dtype, device=device))
+            self.weight = torch.nn.Parameter(
+                torch.ones(normalized_shape, dtype=dtype, device=device)
+            )
         else:
-            self.register_parameter('weight', None)
+            self.register_parameter("weight", None)
 
     def forward(self, x):
         return rms_norm(x.float(), self.weight, self.eps).to(dtype=x.dtype)
 
-class LPRMSNorm(RMSNorm):
 
-    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
-        super().__init__(normalized_shape=normalized_shape, eps=eps, weight=weight, dtype=dtype, device=device)
+class LPRMSNorm(RMSNorm):
+    def __init__(
+        self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None
+    ):
+        super().__init__(
+            normalized_shape=normalized_shape,
+            eps=eps,
+            weight=weight,
+            dtype=dtype,
+            device=device,
+        )
 
     def forward(self, x):
         downcast_x = _cast_if_autocast_enabled(x)
-        downcast_weight = _cast_if_autocast_enabled(self.weight) if self.weight is not None else self.weight
+        downcast_weight = (
+            _cast_if_autocast_enabled(self.weight)
+            if self.weight is not None
+            else self.weight
+        )
         with torch.autocast(enabled=False, device_type=x.device.type):
             return rms_norm(downcast_x, downcast_weight, self.eps).to(dtype=x.dtype)
-NORM_CLASS_REGISTRY = {'layernorm': torch.nn.LayerNorm, 'low_precision_layernorm': LPLayerNorm, 'rmsnorm': RMSNorm, 'low_precision_rmsnorm': LPRMSNorm}
+
+
+NORM_CLASS_REGISTRY = {
+    "layernorm": torch.nn.LayerNorm,
+    "low_precision_layernorm": LPLayerNorm,
+    "rmsnorm": RMSNorm,
+    "low_precision_rmsnorm": LPRMSNorm,
+}

param_init_fns.py

@@ -7,97 +7,133 @@ import torch
 from torch import nn
 from .norm import NORM_CLASS_REGISTRY
 
-def torch_default_param_init_fn_(module: nn.Module, verbose: int=0, **kwargs):
+
+def torch_default_param_init_fn_(module: nn.Module, verbose: int = 0, **kwargs):
     del kwargs
     if verbose > 1:
         warnings.warn(f"Initializing network using module's reset_parameters attribute")
-    if hasattr(module, 'reset_parameters'):
+    if hasattr(module, "reset_parameters"):
         module.reset_parameters()
 
+
 def fused_init_helper_(module: nn.Module, init_fn_):
-    _fused = getattr(module, '_fused', None)
+    _fused = getattr(module, "_fused", None)
     if _fused is None:
-        raise RuntimeError(f'Internal logic error')
+        raise RuntimeError(f"Internal logic error")
     (dim, splits) = _fused
     splits = (0, *splits, module.weight.size(dim))
-    for (s, e) in zip(splits[:-1], splits[1:]):
+    for s, e in zip(splits[:-1], splits[1:]):
         slice_indices = [slice(None)] * module.weight.ndim
         slice_indices[dim] = slice(s, e)
         init_fn_(module.weight[slice_indices])
 
-def generic_param_init_fn_(module: nn.Module, init_fn_, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+
+def generic_param_init_fn_(
+    module: nn.Module,
+    init_fn_,
+    n_layers: int,
+    d_model: Optional[int] = None,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    verbose: int = 0,
+    **kwargs,
+):
     del kwargs
     if verbose > 1:
-        warnings.warn(f'If model has bias parameters they are initialized to 0.')
+        warnings.warn(f"If model has bias parameters they are initialized to 0.")
     init_div_is_residual = init_div_is_residual
     if init_div_is_residual is False:
         div_is_residual = 1.0
     elif init_div_is_residual is True:
         div_is_residual = math.sqrt(2 * n_layers)
-    elif isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int):
+    elif isinstance(init_div_is_residual, float) or isinstance(
+        init_div_is_residual, int
+    ):
         div_is_residual = init_div_is_residual
     elif isinstance(init_div_is_residual, str) and init_div_is_residual.isnumeric():
         div_is_residual = float(init_div_is_residual)
     else:
         div_is_residual = 1.0
-        raise ValueError(f'Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}')
+        raise ValueError(
+            f"Expected init_div_is_residual to be boolean or numeric, got {init_div_is_residual}"
+        )
     if init_div_is_residual is not False:
         if verbose > 1:
-            warnings.warn(f'Initializing _is_residual layers then dividing them by {div_is_residual:.3f}. ' + f'Set `init_div_is_residual: false` in init config to disable this.')
+            warnings.warn(
+                f"Initializing _is_residual layers then dividing them by {div_is_residual:.3f}. "
+                + f"Set `init_div_is_residual: false` in init config to disable this."
+            )
     if isinstance(module, nn.Linear):
-        if hasattr(module, '_fused'):
+        if hasattr(module, "_fused"):
             fused_init_helper_(module, init_fn_)
         else:
             init_fn_(module.weight)
         if module.bias is not None:
             torch.nn.init.zeros_(module.bias)
-        if init_div_is_residual is not False and getattr(module, '_is_residual', False):
+        if init_div_is_residual is not False and getattr(module, "_is_residual", False):
             with torch.no_grad():
                 module.weight.div_(div_is_residual)
     elif isinstance(module, nn.Embedding):
         if emb_init_std is not None:
             std = emb_init_std
             if std == 0:
-                warnings.warn(f'Embedding layer initialized to 0.')
+                warnings.warn(f"Embedding layer initialized to 0.")
             emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)
             if verbose > 1:
-                warnings.warn(f'Embedding layer initialized using normal distribution with mean=0 and std={std!r}.')
+                warnings.warn(
+                    f"Embedding layer initialized using normal distribution with mean=0 and std={std!r}."
+                )
         elif emb_init_uniform_lim is not None:
             lim = emb_init_uniform_lim
             if isinstance(lim, Sequence):
                 if len(lim) > 2:
-                    raise ValueError(f'Uniform init requires a min and a max limit. User input: {lim}.')
+                    raise ValueError(
+                        f"Uniform init requires a min and a max limit. User input: {lim}."
+                    )
                 if lim[0] == lim[1]:
-                    warnings.warn(f'Embedding layer initialized to {lim[0]}.')
+                    warnings.warn(f"Embedding layer initialized to {lim[0]}.")
             else:
                 if lim == 0:
-                    warnings.warn(f'Embedding layer initialized to 0.')
+                    warnings.warn(f"Embedding layer initialized to 0.")
                 lim = [-lim, lim]
             (a, b) = lim
             emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)
             if verbose > 1:
-                warnings.warn(f'Embedding layer initialized using uniform distribution in range {lim}.')
+                warnings.warn(
+                    f"Embedding layer initialized using uniform distribution in range {lim}."
+                )
         else:
             emb_init_fn_ = init_fn_
         emb_init_fn_(module.weight)
     elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):
         if verbose > 1:
-            warnings.warn(f'Norm weights are set to 1. If norm layer has a bias it is initialized to 0.')
-        if hasattr(module, 'weight') and module.weight is not None:
+            warnings.warn(
+                f"Norm weights are set to 1. If norm layer has a bias it is initialized to 0."
+            )
+        if hasattr(module, "weight") and module.weight is not None:
             torch.nn.init.ones_(module.weight)
-        if hasattr(module, 'bias') and module.bias is not None:
+        if hasattr(module, "bias") and module.bias is not None:
             torch.nn.init.zeros_(module.bias)
     elif isinstance(module, nn.MultiheadAttention):
         if module._qkv_same_embed_dim:
             assert module.in_proj_weight is not None
-            assert module.q_proj_weight is None and module.k_proj_weight is None and (module.v_proj_weight is None)
+            assert (
+                module.q_proj_weight is None
+                and module.k_proj_weight is None
+                and (module.v_proj_weight is None)
+            )
             assert d_model is not None
             _d = d_model
             splits = (0, _d, 2 * _d, 3 * _d)
-            for (s, e) in zip(splits[:-1], splits[1:]):
+            for s, e in zip(splits[:-1], splits[1:]):
                 init_fn_(module.in_proj_weight[s:e])
         else:
-            assert module.q_proj_weight is not None and module.k_proj_weight is not None and (module.v_proj_weight is not None)
+            assert (
+                module.q_proj_weight is not None
+                and module.k_proj_weight is not None
+                and (module.v_proj_weight is not None)
+            )
             assert module.in_proj_weight is None
             init_fn_(module.q_proj_weight)
             init_fn_(module.k_proj_weight)
@@ -109,37 +145,112 @@ def generic_param_init_fn_(module: nn.Module, init_fn_, n_layers: int, d_model:
         if module.bias_v is not None:
             torch.nn.init.zeros_(module.bias_v)
         init_fn_(module.out_proj.weight)
-        if init_div_is_residual is not False and getattr(module.out_proj, '_is_residual', False):
+        if init_div_is_residual is not False and getattr(
+            module.out_proj, "_is_residual", False
+        ):
             with torch.no_grad():
                 module.out_proj.weight.div_(div_is_residual)
         if module.out_proj.bias is not None:
             torch.nn.init.zeros_(module.out_proj.bias)
     else:
         for _ in module.parameters(recurse=False):
-            raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.')
+            raise NotImplementedError(
+                f"{module.__class__.__name__} parameters are not initialized by param_init_fn."
+            )
+
 
 def _normal_init_(std, mean=0.0):
     return partial(torch.nn.init.normal_, mean=mean, std=std)
 
-def _normal_param_init_fn_(module: nn.Module, std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+
+def _normal_param_init_fn_(
+    module: nn.Module,
+    std: float,
+    n_layers: int,
+    d_model: Optional[int] = None,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    verbose: int = 0,
+    **kwargs,
+):
     del kwargs
     init_fn_ = _normal_init_(std=std)
     if verbose > 1:
-        warnings.warn(f'Using torch.nn.init.normal_ init fn mean=0.0, std={std}')
-    generic_param_init_fn_(module=module, init_fn_=init_fn_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+        warnings.warn(f"Using torch.nn.init.normal_ init fn mean=0.0, std={std}")
+    generic_param_init_fn_(
+        module=module,
+        init_fn_=init_fn_,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=init_div_is_residual,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
+
 
-def baseline_param_init_fn_(module: nn.Module, init_std: float, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+def baseline_param_init_fn_(
+    module: nn.Module,
+    init_std: float,
+    n_layers: int,
+    d_model: Optional[int] = None,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    verbose: int = 0,
+    **kwargs,
+):
     del kwargs
     if init_std is None:
-        raise ValueError("You must set model.init_config['init_std'] to a float value to use the default initialization scheme.")
-    _normal_param_init_fn_(module=module, std=init_std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+        raise ValueError(
+            "You must set model.init_config['init_std'] to a float value to use the default initialization scheme."
+        )
+    _normal_param_init_fn_(
+        module=module,
+        std=init_std,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=init_div_is_residual,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
+
 
-def small_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+def small_param_init_fn_(
+    module: nn.Module,
+    n_layers: int,
+    d_model: int,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    verbose: int = 0,
+    **kwargs,
+):
     del kwargs
     std = math.sqrt(2 / (5 * d_model))
-    _normal_param_init_fn_(module=module, std=std, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+    _normal_param_init_fn_(
+        module=module,
+        std=std,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=init_div_is_residual,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
 
-def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, verbose: int=0, **kwargs):
+
+def neox_param_init_fn_(
+    module: nn.Module,
+    n_layers: int,
+    d_model: int,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    verbose: int = 0,
+    **kwargs,
+):
     """From section 2.3.1 of GPT-NeoX-20B:
 
     An Open-Source AutoregressiveLanguage Model — Black et. al. (2022)
@@ -149,33 +260,158 @@ def neox_param_init_fn_(module: nn.Module, n_layers: int, d_model: int, emb_init
     del kwargs
     residual_div = n_layers / math.sqrt(10)
     if verbose > 1:
-        warnings.warn(f'setting init_div_is_residual to {residual_div}')
-    small_param_init_fn_(module=module, d_model=d_model, n_layers=n_layers, init_div_is_residual=residual_div, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+        warnings.warn(f"setting init_div_is_residual to {residual_div}")
+    small_param_init_fn_(
+        module=module,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=residual_div,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
+
 
-def kaiming_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
+def kaiming_uniform_param_init_fn_(
+    module: nn.Module,
+    n_layers: int,
+    d_model: Optional[int] = None,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    init_gain: float = 0,
+    fan_mode: str = "fan_in",
+    init_nonlinearity: str = "leaky_relu",
+    verbose: int = 0,
+    **kwargs,
+):
     del kwargs
     if verbose > 1:
-        warnings.warn(f'Using nn.init.kaiming_uniform_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
-    kaiming_uniform_ = partial(nn.init.kaiming_uniform_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
-    generic_param_init_fn_(module=module, init_fn_=kaiming_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+        warnings.warn(
+            f"Using nn.init.kaiming_uniform_ init fn with parameters: "
+            + f"a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}"
+        )
+    kaiming_uniform_ = partial(
+        nn.init.kaiming_uniform_,
+        a=init_gain,
+        mode=fan_mode,
+        nonlinearity=init_nonlinearity,
+    )
+    generic_param_init_fn_(
+        module=module,
+        init_fn_=kaiming_uniform_,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=init_div_is_residual,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
 
-def kaiming_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, fan_mode: str='fan_in', init_nonlinearity: str='leaky_relu', verbose: int=0, **kwargs):
+
+def kaiming_normal_param_init_fn_(
+    module: nn.Module,
+    n_layers: int,
+    d_model: Optional[int] = None,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    init_gain: float = 0,
+    fan_mode: str = "fan_in",
+    init_nonlinearity: str = "leaky_relu",
+    verbose: int = 0,
+    **kwargs,
+):
     del kwargs
     if verbose > 1:
-        warnings.warn(f'Using nn.init.kaiming_normal_ init fn with parameters: ' + f'a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}')
-    kaiming_normal_ = partial(torch.nn.init.kaiming_normal_, a=init_gain, mode=fan_mode, nonlinearity=init_nonlinearity)
-    generic_param_init_fn_(module=module, init_fn_=kaiming_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+        warnings.warn(
+            f"Using nn.init.kaiming_normal_ init fn with parameters: "
+            + f"a={init_gain}, mode={fan_mode}, nonlinearity={init_nonlinearity}"
+        )
+    kaiming_normal_ = partial(
+        torch.nn.init.kaiming_normal_,
+        a=init_gain,
+        mode=fan_mode,
+        nonlinearity=init_nonlinearity,
+    )
+    generic_param_init_fn_(
+        module=module,
+        init_fn_=kaiming_normal_,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=init_div_is_residual,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
+
 
-def xavier_uniform_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
+def xavier_uniform_param_init_fn_(
+    module: nn.Module,
+    n_layers: int,
+    d_model: Optional[int] = None,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    init_gain: float = 0,
+    verbose: int = 0,
+    **kwargs,
+):
     del kwargs
     xavier_uniform_ = partial(torch.nn.init.xavier_uniform_, gain=init_gain)
     if verbose > 1:
-        warnings.warn(f'Using torch.nn.init.xavier_uniform_ init fn with parameters: ' + f'gain={init_gain}')
-    generic_param_init_fn_(module=module, init_fn_=xavier_uniform_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
+        warnings.warn(
+            f"Using torch.nn.init.xavier_uniform_ init fn with parameters: "
+            + f"gain={init_gain}"
+        )
+    generic_param_init_fn_(
+        module=module,
+        init_fn_=xavier_uniform_,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=init_div_is_residual,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
 
-def xavier_normal_param_init_fn_(module: nn.Module, n_layers: int, d_model: Optional[int]=None, init_div_is_residual: Union[int, float, str, bool]=True, emb_init_std: Optional[float]=None, emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]]=None, init_gain: float=0, verbose: int=0, **kwargs):
+
+def xavier_normal_param_init_fn_(
+    module: nn.Module,
+    n_layers: int,
+    d_model: Optional[int] = None,
+    init_div_is_residual: Union[int, float, str, bool] = True,
+    emb_init_std: Optional[float] = None,
+    emb_init_uniform_lim: Optional[Union[Tuple[float, float], float]] = None,
+    init_gain: float = 0,
+    verbose: int = 0,
+    **kwargs,
+):
     xavier_normal_ = partial(torch.nn.init.xavier_normal_, gain=init_gain)
     if verbose > 1:
-        warnings.warn(f'Using torch.nn.init.xavier_normal_ init fn with parameters: ' + f'gain={init_gain}')
-    generic_param_init_fn_(module=module, init_fn_=xavier_normal_, d_model=d_model, n_layers=n_layers, init_div_is_residual=init_div_is_residual, emb_init_std=emb_init_std, emb_init_uniform_lim=emb_init_uniform_lim, verbose=verbose)
-MODEL_INIT_REGISTRY = {'default_': torch_default_param_init_fn_, 'baseline_': baseline_param_init_fn_, 'kaiming_uniform_': kaiming_uniform_param_init_fn_, 'kaiming_normal_': kaiming_normal_param_init_fn_, 'neox_init_': neox_param_init_fn_, 'small_init_': small_param_init_fn_, 'xavier_uniform_': xavier_uniform_param_init_fn_, 'xavier_normal_': xavier_normal_param_init_fn_}
+        warnings.warn(
+            f"Using torch.nn.init.xavier_normal_ init fn with parameters: "
+            + f"gain={init_gain}"
+        )
+    generic_param_init_fn_(
+        module=module,
+        init_fn_=xavier_normal_,
+        d_model=d_model,
+        n_layers=n_layers,
+        init_div_is_residual=init_div_is_residual,
+        emb_init_std=emb_init_std,
+        emb_init_uniform_lim=emb_init_uniform_lim,
+        verbose=verbose,
+    )
+
+
+MODEL_INIT_REGISTRY = {
+    "default_": torch_default_param_init_fn_,
+    "baseline_": baseline_param_init_fn_,
+    "kaiming_uniform_": kaiming_uniform_param_init_fn_,
+    "kaiming_normal_": kaiming_normal_param_init_fn_,
+    "neox_init_": neox_param_init_fn_,
+    "small_init_": small_param_init_fn_,
+    "xavier_uniform_": xavier_uniform_param_init_fn_,
+    "xavier_normal_": xavier_normal_param_init_fn_,
+}