555

LICENSE

@@ -1,202 +0,0 @@
-
-                                 Apache License
-                           Version 2.0, January 2004
-                        http://www.apache.org/licenses/
-
-   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
-
-   1. Definitions.
-
-      "License" shall mean the terms and conditions for use, reproduction,
-      and distribution as defined by Sections 1 through 9 of this document.
-
-      "Licensor" shall mean the copyright owner or entity authorized by
-      the copyright owner that is granting the License.
-
-      "Legal Entity" shall mean the union of the acting entity and all
-      other entities that control, are controlled by, or are under common
-      control with that entity. For the purposes of this definition,
-      "control" means (i) the power, direct or indirect, to cause the
-      direction or management of such entity, whether by contract or
-      otherwise, or (ii) ownership of fifty percent (50%) or more of the
-      outstanding shares, or (iii) beneficial ownership of such entity.
-
-      "You" (or "Your") shall mean an individual or Legal Entity
-      exercising permissions granted by this License.
-
-      "Source" form shall mean the preferred form for making modifications,
-      including but not limited to software source code, documentation
-      source, and configuration files.
-
-      "Object" form shall mean any form resulting from mechanical
-      transformation or translation of a Source form, including but
-      not limited to compiled object code, generated documentation,
-      and conversions to other media types.
-
-      "Work" shall mean the work of authorship, whether in Source or
-      Object form, made available under the License, as indicated by a
-      copyright notice that is included in or attached to the work
-      (an example is provided in the Appendix below).
-
-      "Derivative Works" shall mean any work, whether in Source or Object
-      form, that is based on (or derived from) the Work and for which the
-      editorial revisions, annotations, elaborations, or other modifications
-      represent, as a whole, an original work of authorship. For the purposes
-      of this License, Derivative Works shall not include works that remain
-      separable from, or merely link (or bind by name) to the interfaces of,
-      the Work and Derivative Works thereof.
-
-      "Contribution" shall mean any work of authorship, including
-      the original version of the Work and any modifications or additions
-      to that Work or Derivative Works thereof, that is intentionally
-      submitted to Licensor for inclusion in the Work by the copyright owner
-      or by an individual or Legal Entity authorized to submit on behalf of
-      the copyright owner. For the purposes of this definition, "submitted"
-      means any form of electronic, verbal, or written communication sent
-      to the Licensor or its representatives, including but not limited to
-      communication on electronic mailing lists, source code control systems,
-      and issue tracking systems that are managed by, or on behalf of, the
-      Licensor for the purpose of discussing and improving the Work, but
-      excluding communication that is conspicuously marked or otherwise
-      designated in writing by the copyright owner as "Not a Contribution."
-
-      "Contributor" shall mean Licensor and any individual or Legal Entity
-      on behalf of whom a Contribution has been received by Licensor and
-      subsequently incorporated within the Work.
-
-   2. Grant of Copyright License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      copyright license to reproduce, prepare Derivative Works of,
-      publicly display, publicly perform, sublicense, and distribute the
-      Work and such Derivative Works in Source or Object form.
-
-   3. Grant of Patent License. Subject to the terms and conditions of
-      this License, each Contributor hereby grants to You a perpetual,
-      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
-      (except as stated in this section) patent license to make, have made,
-      use, offer to sell, sell, import, and otherwise transfer the Work,
-      where such license applies only to those patent claims licensable
-      by such Contributor that are necessarily infringed by their
-      Contribution(s) alone or by combination of their Contribution(s)
-      with the Work to which such Contribution(s) was submitted. If You
-      institute patent litigation against any entity (including a
-      cross-claim or counterclaim in a lawsuit) alleging that the Work
-      or a Contribution incorporated within the Work constitutes direct
-      or contributory patent infringement, then any patent licenses
-      granted to You under this License for that Work shall terminate
-      as of the date such litigation is filed.
-
-   4. Redistribution. You may reproduce and distribute copies of the
-      Work or Derivative Works thereof in any medium, with or without
-      modifications, and in Source or Object form, provided that You
-      meet the following conditions:
-
-      (a) You must give any other recipients of the Work or
-          Derivative Works a copy of this License; and
-
-      (b) You must cause any modified files to carry prominent notices
-          stating that You changed the files; and
-
-      (c) You must retain, in the Source form of any Derivative Works
-          that You distribute, all copyright, patent, trademark, and
-          attribution notices from the Source form of the Work,
-          excluding those notices that do not pertain to any part of
-          the Derivative Works; and
-
-      (d) If the Work includes a "NOTICE" text file as part of its
-          distribution, then any Derivative Works that You distribute must
-          include a readable copy of the attribution notices contained
-          within such NOTICE file, excluding those notices that do not
-          pertain to any part of the Derivative Works, in at least one
-          of the following places: within a NOTICE text file distributed
-          as part of the Derivative Works; within the Source form or
-          documentation, if provided along with the Derivative Works; or,
-          within a display generated by the Derivative Works, if and
-          wherever such third-party notices normally appear. The contents
-          of the NOTICE file are for informational purposes only and
-          do not modify the License. You may add Your own attribution
-          notices within Derivative Works that You distribute, alongside
-          or as an addendum to the NOTICE text from the Work, provided
-          that such additional attribution notices cannot be construed
-          as modifying the License.
-
-      You may add Your own copyright statement to Your modifications and
-      may provide additional or different license terms and conditions
-      for use, reproduction, or distribution of Your modifications, or
-      for any such Derivative Works as a whole, provided Your use,
-      reproduction, and distribution of the Work otherwise complies with
-      the conditions stated in this License.
-
-   5. Submission of Contributions. Unless You explicitly state otherwise,
-      any Contribution intentionally submitted for inclusion in the Work
-      by You to the Licensor shall be under the terms and conditions of
-      this License, without any additional terms or conditions.
-      Notwithstanding the above, nothing herein shall supersede or modify
-      the terms of any separate license agreement you may have executed
-      with Licensor regarding such Contributions.
-
-   6. Trademarks. This License does not grant permission to use the trade
-      names, trademarks, service marks, or product names of the Licensor,
-      except as required for reasonable and customary use in describing the
-      origin of the Work and reproducing the content of the NOTICE file.
-
-   7. Disclaimer of Warranty. Unless required by applicable law or
-      agreed to in writing, Licensor provides the Work (and each
-      Contributor provides its Contributions) on an "AS IS" BASIS,
-      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
-      implied, including, without limitation, any warranties or conditions
-      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
-      PARTICULAR PURPOSE. You are solely responsible for determining the
-      appropriateness of using or redistributing the Work and assume any
-      risks associated with Your exercise of permissions under this License.
-
-   8. Limitation of Liability. In no event and under no legal theory,
-      whether in tort (including negligence), contract, or otherwise,
-      unless required by applicable law (such as deliberate and grossly
-      negligent acts) or agreed to in writing, shall any Contributor be
-      liable to You for damages, including any direct, indirect, special,
-      incidental, or consequential damages of any character arising as a
-      result of this License or out of the use or inability to use the
-      Work (including but not limited to damages for loss of goodwill,
-      work stoppage, computer failure or malfunction, or any and all
-      other commercial damages or losses), even if such Contributor
-      has been advised of the possibility of such damages.
-
-   9. Accepting Warranty or Additional Liability. While redistributing
-      the Work or Derivative Works thereof, You may choose to offer,
-      and charge a fee for, acceptance of support, warranty, indemnity,
-      or other liability obligations and/or rights consistent with this
-      License. However, in accepting such obligations, You may act only
-      on Your own behalf and on Your sole responsibility, not on behalf
-      of any other Contributor, and only if You agree to indemnify,
-      defend, and hold each Contributor harmless for any liability
-      incurred by, or claims asserted against, such Contributor by reason
-      of your accepting any such warranty or additional liability.
-
-   END OF TERMS AND CONDITIONS
-
-   APPENDIX: How to apply the Apache License to your work.
-
-      To apply the Apache License to your work, attach the following
-      boilerplate notice, with the fields enclosed by brackets "[]"
-      replaced with your own identifying information. (Don't include
-      the brackets!)  The text should be enclosed in the appropriate
-      comment syntax for the file format. We also recommend that a
-      file or class name and description of purpose be included on the
-      same "printed page" as the copyright notice for easier
-      identification within third-party archives.
-
-   Copyright 2024 Databricks
-
-   Licensed under the Apache License, Version 2.0 (the "License");
-   you may not use this file except in compliance with the License.
-   You may obtain a copy of the License at
-
-       http://www.apache.org/licenses/LICENSE-2.0
-
-   Unless required by applicable law or agreed to in writing, software
-   distributed under the License is distributed on an "AS IS" BASIS,
-   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-   See the License for the specific language governing permissions and
-   limitations under the License.

README.md

@@ -50,12 +50,14 @@ We demonstrate generations as long as 80k tokens on a single A100-80GB GPU in ou
 
 * [MPT-7B-Instruct](https://huggingface.co/mosaicml/mpt-7b-instruct): a model for short-form instruction following.
 Built by finetuning MPT-7B on a [dataset](https://huggingface.co/datasets/mosaicml/dolly_hhrlhf) we also release, derived from the [Databricks Dolly-15k](https://huggingface.co/datasets/databricks/databricks-dolly-15k) and the [Anthropic Helpful and Harmless (HH-RLHF)](https://huggingface.co/datasets/Anthropic/hh-rlhf) datasets.
-  * License: Apache 2.0
+  * License: _CC-By-SA-3.0_
+  * [Demo on Hugging Face Spaces](https://huggingface.co/spaces/mosaicml/mpt-7b-instruct)
 
 * [MPT-7B-Chat](https://huggingface.co/mosaicml/mpt-7b-chat): a chatbot-like model for dialogue generation.
 Built by finetuning MPT-7B on the [ShareGPT-Vicuna](https://huggingface.co/datasets/jeffwan/sharegpt_vicuna), [HC3](https://huggingface.co/datasets/Hello-SimpleAI/HC3),
  [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca), [HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf), and [Evol-Instruct](https://huggingface.co/datasets/victor123/evol_instruct_70k) datasets.
   * License: _CC-By-NC-SA-4.0_
+  * [Demo on Hugging Face Spaces](https://huggingface.co/spaces/mosaicml/mpt-7b-chat)
 
 ## Model Date
 
@@ -130,22 +132,6 @@ from transformers import AutoTokenizer
 tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')
 ```
 
-The model can then be used, for example, within a text-generation pipeline.  
-Note: when running Torch modules in lower precision, it is best practice to use the [torch.autocast context manager](https://pytorch.org/docs/stable/amp.html).
-
-```python
-from transformers import pipeline
-
-pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, device='cuda:0')
-
-with torch.autocast('cuda', dtype=torch.bfloat16):
-    print(
-        pipe('Here is a recipe for vegan banana bread:\n',
-            max_new_tokens=100,
-            do_sample=True,
-            use_cache=True))
-```
-
 ## Model Description
 
 The architecture is a modification of a standard decoder-only transformer.
@@ -237,10 +223,10 @@ Please cite this model using the following format:
 @online{MosaicML2023Introducing,
     author    = {MosaicML NLP Team},
     title     = {Introducing MPT-7B: A New Standard for Open-Source,
-    Commercially Usable LLMs},
+    ly Usable LLMs},
     year      = {2023},
     url       = {www.mosaicml.com/blog/mpt-7b},
-    note      = {Accessed: 2023-05-05},
-    urldate   = {2023-05-05}
+    note      = {Accessed: 2023-03-28}, % change this date
+    urldate   = {2023-03-28} % change this date
 }
 ```

adapt_tokenizer.py

@@ -1,8 +1,9 @@
-from typing import Any
-from transformers import AutoTokenizer, PreTrainedTokenizerBase
+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: PreTrainedTokenizerBase) -> None:
+def adapt_tokenizer_for_denoising(tokenizer: Tokenizer):
     """Adds sentinel tokens and padding token (if missing).
 
     Expands the tokenizer vocabulary to include sentinel tokens
@@ -33,7 +34,7 @@ class AutoTokenizerForMOD(AutoTokenizer):
     """
 
     @classmethod
-    def from_pretrained(cls, *args: Any, **kwargs: Any) -> PreTrainedTokenizerBase:
+    def from_pretrained(cls, *args, **kwargs):
         """See `AutoTokenizer.from_pretrained` docstring."""
         tokenizer = super().from_pretrained(*args, **kwargs)
         adapt_tokenizer_for_denoising(tokenizer)

attention.py

@@ -1,42 +1,15 @@
 """Attention layers."""
 import math
 import warnings
-from typing import Any, Optional
+from typing import Optional
 import torch
 import torch.nn as nn
-import transformers
 from einops import rearrange
 from packaging import version
 from torch import nn
-from .fc import FC_CLASS_REGISTRY
-from .norm import NORM_CLASS_REGISTRY
+from .norm import LPLayerNorm
 
-def is_flash_v2_installed(v2_version: str='2.0.0'):
-    assert version.parse(v2_version) >= version.parse('2.0.0')
-    try:
-        import flash_attn as flash_attn
-    except:
-        return False
-    return version.parse(flash_attn.__version__) >= version.parse(v2_version)
-
-def is_flash_v1_installed():
-    try:
-        import flash_attn as flash_attn
-    except:
-        return False
-    return version.parse(flash_attn.__version__) < version.parse('2.0.0')
-
-def is_transformers_version_gte(hf_version: str) -> bool:
-    return version.parse(transformers.__version__) >= version.parse(hf_version)
-
-def check_alibi_support(attention_impl: str) -> bool:
-    return attention_impl != 'flash' or is_flash_v2_installed(v2_version='v2.4.2')
-if is_flash_v1_installed():
-    import transformers
-    transformers.utils.is_flash_attn_available = lambda : False
-from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
-
-def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool) -> 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.')
@@ -44,21 +17,9 @@ def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_cau
             return False
     return original_is_causal
 
-def repeat_kv_for_gqa(hidden: torch.Tensor, n_rep: int) -> torch.Tensor:
-    """Perform repeat of kv heads along a particular dimension.
-
-    hidden.shape expected to be: (batch size, seq len, kv_n_heads, head_dim)
-    n_rep: amount of repetitions of kv_n_heads
-    Unlike torch.repeat_interleave, this function avoids allocating new memory.
-    """
-    if n_rep == 1:
-        return hidden
-    (b, s, kv_n_heads, d) = hidden.shape
-    hidden = hidden[:, :, :, None, :].expand(b, s, kv_n_heads, n_rep, d)
-    return hidden.reshape(b, s, kv_n_heads * n_rep, d)
-
-def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: int, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
+def scaled_multihead_dot_product_attention(query, key, value, n_heads, past_key_value=None, 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)
+    kv_n_heads = 1 if multiquery else n_heads
     k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)
     v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)
     if past_key_value is not None:
@@ -68,9 +29,6 @@ def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tenso
         past_key_value = (k, v)
     (b, _, s_q, d) = q.shape
     s_k = k.size(-1)
-    if kv_n_heads > 1 and kv_n_heads < n_heads:
-        k = repeat_kv_for_gqa(k.transpose(1, 2), n_heads // kv_n_heads).transpose(1, 2)
-        v = repeat_kv_for_gqa(v.transpose(1, 2), n_heads // kv_n_heads).transpose(1, 2)
     if softmax_scale is None:
         softmax_scale = 1 / math.sqrt(d)
     attn_weight = q.matmul(k) * softmax_scale
@@ -84,11 +42,11 @@ def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tenso
     min_val = torch.finfo(q.dtype).min
     if key_padding_mask is not None:
         if attn_bias 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('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 and (not q.size(2) == 1):
         s = max(s_q, s_k)
-        causal_mask = attn_weight.new_ones(s, s, dtype=torch.float32)
+        causal_mask = attn_weight.new_ones(s, s, dtype=torch.float16)
         causal_mask = causal_mask.tril()
         causal_mask = causal_mask.to(torch.bool)
         causal_mask = ~causal_mask
@@ -97,79 +55,56 @@ def scaled_multihead_dot_product_attention(query: torch.Tensor, key: torch.Tenso
     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)
-    out = attn_weight.to(v.dtype).matmul(v)
+    out = attn_weight.matmul(v)
     out = rearrange(out, 'b h s d -> b s (h d)')
     if needs_weights:
         return (out, attn_weight, past_key_value)
     return (out, None, past_key_value)
 
-def check_valid_inputs(*tensors: torch.Tensor, valid_dtypes: Optional[list[torch.dtype]]=None):
-    if valid_dtypes is None:
-        valid_dtypes = [torch.float16, torch.bfloat16]
+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}.')
         if not tensor.is_cuda:
             raise TypeError(f'Inputs must be cuda tensors (tensor.is_cuda={tensor.is_cuda!r}).')
 
-def flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: int, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False, multiquery: bool=False, should_repeat_kv_for_gqa: Optional[bool]=True, sliding_window_size: int=-1, alibi_slopes: Optional[torch.Tensor]=None, flash_attn_padding_info: Optional[dict[str, torch.Tensor]]=None) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
-    if key_padding_mask is not None:
-        raise ValueError('key_padding_mask should be None for flash attn.')
-    del key_padding_mask
-    if flash_attn_padding_info is None:
-        raise ValueError('flash_attn_padding_info is required for flash attn.')
+def flash_attn_fn(query, key, value, n_heads, past_key_value=None, 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.9 or flash-attn==2.3.6')
+        raise RuntimeError('Please install flash-attn==1.0.3.post0')
     check_valid_inputs(query, key, value)
     if past_key_value is not None:
         if len(past_key_value) != 0:
             key = torch.cat([past_key_value[0], key], dim=1)
             value = torch.cat([past_key_value[1], value], dim=1)
         past_key_value = (key, value)
+    if attn_bias is not None:
+        _s_q = max(0, attn_bias.size(2) - query.size(1))
+        _s_k = max(0, attn_bias.size(3) - key.size(1))
+        attn_bias = attn_bias[:, :, _s_q:, _s_k:]
     if attn_bias is not None:
         raise NotImplementedError(f'attn_bias not implemented for flash attn.')
     (batch_size, seqlen) = query.shape[:2]
-    indices_q = flash_attn_padding_info['indices_q']
-    indices_k = flash_attn_padding_info['indices_k']
-    indices_v = flash_attn_padding_info['indices_v']
-    cu_seqlens_q = flash_attn_padding_info['cu_seqlens_q']
-    cu_seqlens_k = flash_attn_padding_info['cu_seqlens_k']
-    max_seqlen_q = flash_attn_padding_info['max_seqlen_q']
-    max_seqlen_k = flash_attn_padding_info['max_seqlen_k']
-    query_unpad = bert_padding.index_first_axis(rearrange(query, 'b s ... -> (b s) ...'), indices_q)
+    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 = bert_padding.index_first_axis(rearrange(key, 'b s ... -> (b s) ...'), indices_k)
-    key_unpad = rearrange(key_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
-    value_unpad = bert_padding.index_first_axis(rearrange(value, 'b s ... -> (b s) ...'), indices_v)
-    value_unpad = rearrange(value_unpad, 'nnz (h d) -> nnz h d', h=kv_n_heads)
-    if kv_n_heads < n_heads and (not is_flash_v2_installed()) and (not should_repeat_kv_for_gqa):
-        raise ValueError('For Grouped Query Attention or Multi Query Attention, should_repeat_kv_for_gqa should be set to True if not using Flash Attention v2.')
-    if should_repeat_kv_for_gqa:
-        if kv_n_heads == 1:
-            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))
-        elif kv_n_heads < n_heads:
-            key_unpad = repeat_kv_for_gqa(key_unpad.view(1, key_unpad.size(0), kv_n_heads, -1), n_heads // kv_n_heads).view(key_unpad.size(0), n_heads, -1)
-            value_unpad = repeat_kv_for_gqa(value_unpad.view(1, value_unpad.size(0), kv_n_heads, -1), n_heads // kv_n_heads).view(value_unpad.size(0), n_heads, -1)
+    (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)
+    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))
     dropout_p = dropout_p if training else 0.0
     reset_is_causal = _reset_is_causal(query.size(1), key.size(1), is_causal)
-    if is_flash_v1_installed():
-        output_unpad = flash_attn_interface.flash_attn_unpadded_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights)
-    elif is_flash_v2_installed():
-        alibi_kwargs = {}
-        if check_alibi_support('flash'):
-            alibi_kwargs = {'alibi_slopes': alibi_slopes}
-        elif alibi_slopes is not None:
-            raise ValueError('alibi_slopes is only supported for flash-attn>=2.4.2')
-        output_unpad = flash_attn_interface.flash_attn_varlen_func(q=query_unpad, k=key_unpad, v=value_unpad, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k=cu_seqlens_k, max_seqlen_q=max_seqlen_q, max_seqlen_k=max_seqlen_k, dropout_p=dropout_p, softmax_scale=softmax_scale, causal=reset_is_causal, return_attn_probs=needs_weights, window_size=(sliding_window_size, sliding_window_size), **alibi_kwargs)
-    else:
-        raise RuntimeError('flash-attn==1.0.9 or flash-attn==2.4.2 is required.')
+    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, past_key_value)
 
-def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, n_heads: int, kv_n_heads: int, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, softmax_scale: Optional[float]=None, attn_bias: Optional[torch.Tensor]=None, key_padding_mask: Optional[torch.Tensor]=None, is_causal: bool=False, dropout_p: float=0.0, training: bool=False, needs_weights: bool=False) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
+def triton_flash_attn_fn(query, key, value, n_heads, past_key_value=None, 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_triton import flash_attn_func
     except:
@@ -181,7 +116,7 @@ def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Te
             except:
                 _installed = False
         if not _installed:
-            raise RuntimeError('Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU ' + 'and `pip install .[gpu]` if installing from llm-foundry source or ' + '`pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` ' + 'if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). ' + 'Note: (1) requires you have CMake and PyTorch already installed.')
+            raise RuntimeError('Requirements for `attn_impl: triton` not installed. Either (1) have a CUDA-compatible GPU and `pip install .[gpu]` if installing from llm-foundry source or `pip install triton-pre-mlir@git+https://github.com/vchiley/triton.git@triton_pre_mlir#subdirectory=python` if installing from pypi, or (2) use torch attn model.attn_config.attn_impl=torch (torch attn_impl will be slow). Note: (1) requires you have CMake and PyTorch already installed.')
     check_valid_inputs(query, key, value)
     if past_key_value is not None:
         if len(past_key_value) != 0:
@@ -194,7 +129,6 @@ def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Te
         attn_bias = attn_bias[:, :, _s_q:, _s_k:]
     if dropout_p:
         raise NotImplementedError(f'Dropout not implemented for attn_impl: triton.')
-    dropout_p = dropout_p if training else 0.0
     if needs_weights:
         raise NotImplementedError(f'attn_impl: triton cannot return attn weights.')
     if key_padding_mask is not None:
@@ -204,144 +138,124 @@ def triton_flash_attn_fn(query: torch.Tensor, key: torch.Tensor, value: torch.Te
             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=kv_n_heads)
-    value = rearrange(value, 'b s (h d) -> b s h d', h=kv_n_heads)
-    if kv_n_heads == 1:
-        key = key.repeat(1, 1, n_heads, 1)
-        value = value.repeat(1, 1, n_heads, 1)
-    elif kv_n_heads < n_heads:
-        key = repeat_kv_for_gqa(key, n_heads // kv_n_heads)
-        value = repeat_kv_for_gqa(value, n_heads // kv_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_func(query, key, value, attn_bias, reset_is_causal, softmax_scale)
     output = attn_output.view(*attn_output.shape[:2], -1)
     return (output, None, past_key_value)
 
-class GroupedQueryAttention(nn.Module):
-    """Grouped Query Attention (GQA) is a generalization of Multi-head (MHA).
-
-    and Multi-query attention (MQA).
+class MultiheadAttention(nn.Module):
+    """Multi-head self attention.
 
-    This allows the user to set a variable of number of kv_n_heads, rather than
-    just n_heads or 1, as in MHA and MQA. Using torch or triton attention
-    implementation enables user to also use additive bias.
+    Using torch or triton attention implemetation enables user to also use
+    additive bias.
     """
 
-    def __init__(self, d_model: int, n_heads: int, kv_n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, qk_gn: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True, sliding_window_size: int=-1):
+    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, verbose: int=0, device: Optional[str]=None):
         super().__init__()
         self.attn_impl = attn_impl
         self.clip_qkv = clip_qkv
         self.qk_ln = qk_ln
-        self.qk_gn = qk_gn
         self.d_model = d_model
         self.n_heads = n_heads
-        self.kv_n_heads = kv_n_heads
-        self.sliding_window_size = sliding_window_size
-        self.head_dim = d_model // n_heads
-        if self.kv_n_heads <= 0:
-            raise ValueError('kv_n_heads should be greater than zero.')
-        if self.kv_n_heads > self.n_heads:
-            raise ValueError('The number of KV heads should be less than or equal to Q heads.')
-        if self.n_heads % self.kv_n_heads != 0:
-            raise ValueError('Each Q head should get the same number of KV heads, so n_heads must be divisible by kv_n_heads.')
-        if qk_ln and qk_gn:
-            raise ValueError('Only one of qk_ln and qk_gn can be set to True.')
         self.softmax_scale = softmax_scale
         if self.softmax_scale is None:
             self.softmax_scale = 1 / math.sqrt(self.d_model / self.n_heads)
         self.attn_dropout_p = attn_pdrop
-        fc_kwargs: dict[str, Any] = {'bias': bias}
-        if fc_type != 'te':
-            fc_kwargs['device'] = device
-        self.Wqkv = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model + 2 * self.kv_n_heads * self.head_dim, **fc_kwargs)
-        fuse_splits = [i * self.head_dim for i in range(1, self.n_heads + 2 * self.kv_n_heads)]
+        self.Wqkv = nn.Linear(self.d_model, 3 * self.d_model, device=device)
+        fuse_splits = (d_model, 2 * d_model)
         self.Wqkv._fused = (0, fuse_splits)
-        if self.qk_ln or self.qk_gn:
-            norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
-            norm_size = self.head_dim if qk_gn else d_model
-            self.q_ln = norm_class(norm_size, device=device)
-            if qk_ln:
-                norm_size = self.head_dim * kv_n_heads
-            self.k_ln = norm_class(norm_size, device=device)
+        if self.qk_ln:
+            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':
             self.attn_fn = flash_attn_fn
         elif self.attn_impl == 'triton':
             self.attn_fn = triton_flash_attn_fn
+            if verbose:
+                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() and verbose:
+                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.')
-        self.out_proj = FC_CLASS_REGISTRY[fc_type](self.d_model, self.d_model, **fc_kwargs)
+        self.out_proj = nn.Linear(self.d_model, self.d_model, device=device)
         self.out_proj._is_residual = True
 
-    def forward(self, x: torch.Tensor, past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None, rotary_emb_w_meta_info: Optional[dict]=None, is_causal: bool=True, needs_weights: bool=False, alibi_slopes: Optional[torch.Tensor]=None, flash_attn_padding_info: Optional[dict[str, torch.Tensor]]=None) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor, torch.Tensor]]]:
+    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 = qkv.clamp(min=-self.clip_qkv, max=self.clip_qkv)
-        (query, key, value) = qkv.split([self.d_model, self.kv_n_heads * self.head_dim, self.kv_n_heads * self.head_dim], dim=2)
+            qkv.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
+        (query, key, value) = qkv.chunk(3, dim=2)
         key_padding_mask = attention_mask
-        if self.qk_ln or self.qk_gn:
-            (q_shape, k_shape) = (query.shape, key.shape)
-            if self.qk_gn:
-                (b, s) = query.shape[:2]
-                query = query.view(b, s, self.n_heads, -1)
-                key = key.view(b, s, self.kv_n_heads, -1)
+        if self.qk_ln:
             dtype = query.dtype
-            query = self.q_ln(query).to(dtype).view(q_shape)
-            key = self.k_ln(key).to(dtype).view(k_shape)
-        if rotary_emb_w_meta_info is not None:
-            rotary_emb = rotary_emb_w_meta_info['rotary_emb']
-            seq_len = rotary_emb_w_meta_info['seq_len']
-            offset_info = rotary_emb_w_meta_info['offset_info']
-            (bsz, seqlen) = query.shape[:2]
-            query = query.view(bsz, seqlen, -1, self.head_dim)
-            key = key.view(bsz, seqlen, -1, self.head_dim)
-            if rotary_emb_w_meta_info['impl'] == 'dail':
-                value = value.view(bsz, seqlen, -1, self.head_dim)
-                kv = torch.stack([key, value], dim=2)
-                (query, kv) = rotary_emb(query, kv, seqlen_offset=offset_info, max_seqlen=seq_len)
-                [key, value] = torch.unbind(kv, dim=2)
-                value = value.view(bsz, seqlen, self.kv_n_heads * self.head_dim)
-            elif rotary_emb_w_meta_info['impl'] == 'hf':
-                (cos, sin) = rotary_emb(value, seq_len)
-                if is_transformers_version_gte('4.36'):
-                    (query, key) = apply_rotary_pos_emb(query, key, cos, sin, offset_info, unsqueeze_dim=2)
-                else:
-                    query = query.transpose(1, 2)
-                    key = key.transpose(1, 2)
-                    (query, key) = apply_rotary_pos_emb(query, key, cos, sin, offset_info)
-                    query = query.transpose(1, 2)
-                    key = key.transpose(1, 2)
-            query = query.view(bsz, seqlen, self.d_model)
-            key = key.view(bsz, seqlen, self.kv_n_heads * self.head_dim)
-        extra_attn_kwargs = {}
-        if self.attn_impl == 'flash':
-            key_padding_mask = None
-            extra_attn_kwargs = {'should_repeat_kv_for_gqa': not is_flash_v2_installed(), 'sliding_window_size': self.sliding_window_size, 'alibi_slopes': alibi_slopes, 'flash_attn_padding_info': flash_attn_padding_info}
-        (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, self.kv_n_heads, past_key_value=past_key_value, 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, **extra_attn_kwargs)
+            query = self.q_ln(query).to(dtype)
+            key = self.k_ln(key).to(dtype)
+        (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, 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 MultiheadAttention(GroupedQueryAttention):
-    """Multi-head self attention.
-
-    Using torch or triton attention implementation enables user to also use
-    additive bias.
-    """
-
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, qk_gn: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True, sliding_window_size: int=-1):
-        super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=n_heads, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, qk_gn=qk_gn, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias, sliding_window_size=sliding_window_size)
-
-class MultiQueryAttention(GroupedQueryAttention):
+class MultiQueryAttention(nn.Module):
     """Multi-Query self attention.
 
-    Using torch or triton attention implementation enables user to also use
+    Using torch or triton attention implemetation enables user to also use
     additive bias.
     """
 
-    def __init__(self, d_model: int, n_heads: int, attn_impl: str='triton', clip_qkv: Optional[float]=None, qk_ln: bool=False, qk_gn: bool=False, softmax_scale: Optional[float]=None, attn_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, bias: bool=True, sliding_window_size: int=-1):
-        super().__init__(d_model=d_model, n_heads=n_heads, kv_n_heads=1, attn_impl=attn_impl, clip_qkv=clip_qkv, qk_ln=qk_ln, qk_gn=qk_gn, softmax_scale=softmax_scale, attn_pdrop=attn_pdrop, norm_type=norm_type, fc_type=fc_type, device=device, bias=bias, sliding_window_size=sliding_window_size)
+    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, verbose: int=0, device: Optional[str]=None):
+        super().__init__()
+        self.attn_impl = attn_impl
+        self.clip_qkv = clip_qkv
+        self.qk_ln = qk_ln
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.softmax_scale = softmax_scale
+        if self.softmax_scale is None:
+            self.softmax_scale = 1 / math.sqrt(self.head_dim)
+        self.attn_dropout_p = attn_pdrop
+        self.Wqkv = nn.Linear(d_model, d_model + 2 * self.head_dim, device=device)
+        fuse_splits = (d_model, d_model + self.head_dim)
+        self.Wqkv._fused = (0, fuse_splits)
+        if self.qk_ln:
+            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':
+            self.attn_fn = flash_attn_fn
+        elif self.attn_impl == 'triton':
+            self.attn_fn = triton_flash_attn_fn
+            if verbose:
+                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() and verbose:
+                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.')
+        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):
+        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)
+        key_padding_mask = attention_mask
+        if self.qk_ln:
+            dtype = query.dtype
+            query = self.q_ln(query).to(dtype)
+            key = self.k_ln(key).to(dtype)
+        (context, attn_weights, past_key_value) = self.attn_fn(query, key, value, self.n_heads, past_key_value=past_key_value, 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: str, n_heads: int, seq_len: int, alibi: bool, prefix_lm: bool, causal: bool, use_sequence_id: bool) -> Optional[tuple[int, int, int, int]]:
+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']:
@@ -355,7 +269,7 @@ def attn_bias_shape(attn_impl: str, n_heads: int, seq_len: int, alibi: bool, pre
     else:
         raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
 
-def build_attn_bias(attn_impl: str, attn_bias: torch.Tensor, n_heads: int, seq_len: int, causal: bool=False, alibi: bool=False, alibi_bias_max: int=8) -> Optional[torch.Tensor]:
+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']:
@@ -366,18 +280,16 @@ def build_attn_bias(attn_impl: str, attn_bias: torch.Tensor, n_heads: int, seq_l
     else:
         raise ValueError(f'attn_impl={attn_impl!r} is an invalid setting.')
 
-def gen_slopes(n_heads: int, alibi_bias_max: int=8, device: Optional[torch.device]=None, return_1d: bool=False) -> torch.Tensor:
+def gen_slopes(n_heads, alibi_bias_max=8, device=None):
     _n_heads = 2 ** math.ceil(math.log2(n_heads))
     m = torch.arange(1, _n_heads + 1, dtype=torch.float32, device=device)
     m = m.mul(alibi_bias_max / _n_heads)
     slopes = 1.0 / torch.pow(2, m)
     if _n_heads != n_heads:
         slopes = torch.concat([slopes[1::2], slopes[::2]])[:n_heads]
-    if return_1d:
-        return slopes
     return slopes.view(1, n_heads, 1, 1)
 
-def build_alibi_bias(n_heads: int, seq_len: int, full: bool=False, alibi_bias_max: int=8, device: Optional[torch.device]=None, dtype: Optional[torch.dtype]=None) -> torch.Tensor:
+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)
@@ -385,4 +297,4 @@ def build_alibi_bias(n_heads: int, seq_len: int, full: bool=False, alibi_bias_ma
     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, 'grouped_query_attention': GroupedQueryAttention}
+ATTN_CLASS_REGISTRY = {'multihead_attention': MultiheadAttention, 'multiquery_attention': MultiQueryAttention}

blocks.py

@@ -1,55 +1,41 @@
 """GPT Blocks used for the GPT Model."""
-from typing import Any, Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple
 import torch
 import torch.nn as nn
 from .attention import ATTN_CLASS_REGISTRY
-from .ffn import FFN_CLASS_REGISTRY, build_ffn
 from .norm import NORM_CLASS_REGISTRY
-try:
-    from flash_attn.bert_padding import unpad_input, pad_input
-except:
-    (unpad_input, pad_input) = (None, None)
-attn_config_defaults: Dict = {'attn_type': 'multihead_attention', 'attn_pdrop': 0.0, 'attn_impl': 'triton', 'qk_ln': False, 'qk_gn': False, 'clip_qkv': None, 'softmax_scale': None, 'prefix_lm': False, 'attn_uses_sequence_id': False, 'sliding_window_size': -1, 'alibi': False, 'alibi_bias_max': 8, 'rope': False, 'rope_theta': 10000, 'rope_impl': 'dail', 'rope_dail_config': {'type': 'original', 'pos_idx_in_fp32': True, 'xpos_scale_base': 512}, 'rope_hf_config': {'type': 'no_scaling', 'factor': 1.0}}
+
+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.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: Optional[Dict]=None, ffn_config: Optional[Dict]=None, resid_pdrop: float=0.0, norm_type: str='low_precision_layernorm', fc_type: str='torch', device: Optional[str]=None, no_bias: bool=False, use_pad_tok_in_ffn: bool=True, **kwargs: Any):
-        if attn_config is None:
-            attn_config = attn_config_defaults
-        if ffn_config is None:
-            ffn_config = {'ffn_type': 'mptmlp'}
+    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', verbose: int=0, device: Optional[str]=None, **kwargs):
         del kwargs
         super().__init__()
         norm_class = NORM_CLASS_REGISTRY[norm_type.lower()]
-        assert isinstance(attn_config['attn_type'], str)
         attn_class = ATTN_CLASS_REGISTRY[attn_config['attn_type']]
-        args_to_exclude_in_attn_class = {'attn_type', 'prefix_lm', 'alibi', 'attn_uses_sequence_id', 'alibi_bias_max', 'rope', 'rope_theta', 'rope_impl', 'rope_dail_config', 'rope_hf_config'}
-        attn_config_subset_for_attn_class = {k: v for (k, v) in attn_config.items() if k not in args_to_exclude_in_attn_class}
         self.norm_1 = norm_class(d_model, device=device)
-        self.attn = attn_class(d_model=d_model, n_heads=n_heads, fc_type=fc_type, device=device, **attn_config_subset_for_attn_class, bias=not no_bias)
-        self.norm_2 = None
-        if not getattr(FFN_CLASS_REGISTRY[ffn_config['ffn_type']], '_has_norm', False):
-            self.norm_2 = norm_class(d_model, device=device)
-        self.ffn = build_ffn(d_model=d_model, expansion_ratio=expansion_ratio, device=device, bias=not no_bias, **ffn_config)
+        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, verbose=verbose, 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.resid_attn_dropout = nn.Dropout(resid_pdrop)
         self.resid_ffn_dropout = nn.Dropout(resid_pdrop)
-        self.use_pad_tok_in_ffn = use_pad_tok_in_ffn
 
-    def forward(self, x: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]]=None, attn_bias: Optional[torch.Tensor]=None, rotary_emb_w_meta_info: Optional[Dict]=None, attention_mask: Optional[torch.ByteTensor]=None, is_causal: bool=True, output_attentions: bool=False, alibi_slopes: Optional[torch.Tensor]=None, flash_attn_padding_info: Optional[dict[str, torch.Tensor]]=None) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, 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, attn_weights, past_key_value) = self.attn(a, past_key_value=past_key_value, attn_bias=attn_bias, rotary_emb_w_meta_info=rotary_emb_w_meta_info, attention_mask=attention_mask, is_causal=is_causal, needs_weights=output_attentions, alibi_slopes=alibi_slopes, flash_attn_padding_info=flash_attn_padding_info)
+        (b, attn_weights, 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 = x
-        if self.norm_2 is not None:
-            m = self.norm_2(x)
-        (batch_size, seq_len) = m.size()[:2]
-        indices = None
-        if not self.use_pad_tok_in_ffn:
-            assert unpad_input is not None
-            (m, indices, _, _) = unpad_input(m, attention_mask)
+        m = self.norm_2(x)
         n = self.ffn(m)
-        if not self.use_pad_tok_in_ffn:
-            assert pad_input is not None
-            n = pad_input(n, indices, batch_size, seq_len)
         x = x + self.resid_ffn_dropout(n)
         return (x, attn_weights, past_key_value)

configuration_mpt.py

@@ -1,38 +1,30 @@
 """A HuggingFace-style model configuration."""
-import warnings
-from typing import Any, Dict, Optional, Union
+from typing import Dict, Optional, Union
 from transformers import PretrainedConfig
-from .attention import check_alibi_support, is_flash_v1_installed, is_flash_v2_installed
-from .blocks import attn_config_defaults
-from .fc import FC_CLASS_REGISTRY
-from .norm import LPLayerNorm
-from .ffn import FFN_CLASS_REGISTRY
-from .warnings import VersionedDeprecationWarning
-ffn_config_defaults: Dict = {'ffn_type': 'mptmlp'}
+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', 'init_div_is_residual': True, 'emb_init_std': None, 'emb_init_uniform_lim': None, 'init_std': None, 'init_gain': 0.0}
 
 class MPTConfig(PretrainedConfig):
     model_type = 'mpt'
 
-    def __init__(self, d_model: int=2048, n_heads: int=16, n_layers: int=24, expansion_ratio: Union[int, float]=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, ffn_config: Dict=ffn_config_defaults, init_device: str='cpu', logit_scale: Optional[Union[float, str]]=None, no_bias: bool=False, embedding_fraction: float=1.0, norm_type: str='low_precision_layernorm', use_cache: bool=False, init_config: Dict=init_config_defaults, fc_type: str='torch', tie_word_embeddings: bool=True, use_pad_tok_in_ffn: bool=True, **kwargs: Any):
+    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:
             d_model (int): The size of the embedding dimension of the model.
             n_heads (int): The number of attention heads.
             n_layers (int): The number of layers in the model.
-            expansion_ratio (Union[int, float]): The ratio of the up/down scale in the ffn.
+            expansion_ratio (int): The ratio of the up/down scale in the MLP.
             max_seq_len (int): The maximum sequence length of the model.
             vocab_size (int): The size of the vocabulary.
             resid_pdrop (float): The dropout probability applied to the attention output before combining with residual.
             emb_pdrop (float): The dropout probability for the embedding layer.
             learned_pos_emb (bool): Whether to use learned positional embeddings
-            attn_config (Dict): A dictionary used to configure the model's attention module:
-                attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention, grouped_query_attention
+            attn_config (Dict):  A dictionary used to configure the model's attention module:
+                attn_type (str): type of attention to use. Options: multihead_attention, multiquery_attention
                 attn_pdrop (float): The dropout probability for the attention layers.
                 attn_impl (str): The attention implementation to use. One of 'torch', 'flash', or 'triton'.
                 qk_ln (bool): Whether to apply layer normalization to the queries and keys in the attention layer.
-                qk_gn (bool): Whether to apply group normalization to the queries and keys in the attention layer.
                 clip_qkv (Optional[float]): If not None, clip the queries, keys, and values in the attention layer to
                     this value.
                 softmax_scale (Optional[float]): If not None, scale the softmax in the attention layer by this value. If None,
@@ -44,27 +36,15 @@ class MPTConfig(PretrainedConfig):
                     When the model is in `train` mode, this requires passing an extra `sequence_id` argument which indicates
                     which sub-sequence each token belongs to.
                     Defaults to ``False`` meaning any provided `sequence_id` will be ignored.
-                sliding_window_size (int): Window size for sliding window local attention. Defaults to -1, which means no sliding window. Query at position i will only attend to keys between [i + seqlen_k - seqlen_q - window_size, i + seqlen_k - seqlen_q + window_size] inclusive. Only works for flash attention v2.3.0 or higher.
                 alibi (bool): Whether to use the alibi bias instead of position embeddings.
                 alibi_bias_max (int): The maximum value of the alibi bias.
-                rope (bool): Whether to use rotary positional embeddings.
-                rope_theta (int): The base frequency for rope.
-                rope_impl (str): The implementation of rope to use. One of 'hf' (to use the implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py) or 'dail' (to use the implementation from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py).
-                rope_dail_config (Dict): The configuration for the dail implementation of rope.
-                    type (str): The type of rotary position embedding to use. Options: 'original' (for https://arxiv.org/pdf/2104.09864.pdf), 'xpos' (for https://arxiv.org/pdf/2212.10554.pdf).
-                    pos_idx_in_fp32 (bool): If True, the position indices [0, ..., seqlen - 1] are in fp32, otherwise they might be in lower precision. A consequence could be, for example, that bf16 rounds position 1995 to 2000, which leads to them having the same positional embedding.
-                    xpos_scale_base (float): The scale base for XPos (if using XPos).
-                rope_hf_config (Dict): A dictionary used to configure rope's scaling behavior (when scaling beyond the training length).
-                    type (str): Can be one of 'no_scaling', 'linear', or 'dynamic'. 'no_scaling' uses the default implementation for rotary embeddings, 'linear' uses linear scaling as proposed by the Reddit user /u/kaiokendev, and 'dynamic' uses Dynamic NTK scaling as proposed by the Reddit users /u/bloc97 and /u/emozilla.
-                    factor (float): Scaling factor to use if using 'linear' or 'dynamic' as rope_scaling.type.
-                kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads.
-            ffn_config (Dict): A dictionary used to configure the model's ffn module:
-                ffn_type (str): type of ffn to use. Options: mptmlp, mptglu, te_ln_mlp
             init_device (str): The device to use for parameter initialization.
             logit_scale (Optional[Union[float, str]]): If not None, scale the logits by this value.
             no_bias (bool): Whether to use bias in all layers.
+            verbose (int): The verbosity level. 0 is silent.
             embedding_fraction (float): The fraction to scale the gradients of the embedding layer by.
             norm_type (str): choose type of norm to use
+            multiquery_attention (bool): Whether to use multiquery attention implementation.
             use_cache (bool): Whether or not the model should return the last key/values attentions
             init_config (Dict): A dictionary used to configure the model initialization:
                 init_config.name: The parameter initialization scheme to use. Options: 'default_', 'baseline_',
@@ -81,9 +61,6 @@ class MPTConfig(PretrainedConfig):
                 init_nonlinearity (str): The nonlinearity to use for parameter initialization with kaiming initialization schemes.
                 ---
                 See llmfoundry.models.utils.param_init_fns.py for info on other param init config options
-            fc_type (str): choose fc layer implementation. Options: torch and te. te layers support fp8 when using H100 GPUs.
-            tie_word_embeddings (bool): Whether to tie the input embedding and output layers.
-            use_pad_tok_in_ffn (bool): Whether to forward the pad token in the feedforward networks.
         """
         self.d_model = d_model
         self.n_heads = n_heads
@@ -95,37 +72,29 @@ class MPTConfig(PretrainedConfig):
         self.emb_pdrop = emb_pdrop
         self.learned_pos_emb = learned_pos_emb
         self.attn_config = attn_config
-        self.ffn_config = ffn_config
         self.init_device = init_device
         self.logit_scale = logit_scale
         self.no_bias = no_bias
+        self.verbose = verbose
         self.embedding_fraction = embedding_fraction
         self.norm_type = norm_type
         self.use_cache = use_cache
         self.init_config = init_config
-        self.fc_type = fc_type
-        self.use_pad_tok_in_ffn = use_pad_tok_in_ffn
         if 'name' in kwargs:
             del kwargs['name']
         if 'loss_fn' in kwargs:
             del kwargs['loss_fn']
-        if self.attn_config.get('alibi', False) or self.attn_config.get('rope', False):
-            self.learned_pos_emb = False
-            warnings.warn(f'alibi or rope is turned on, setting `learned_pos_emb` to `False.`')
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        super().__init__(**kwargs)
         self._validate_config()
 
-    def _set_config_defaults(self, config: Dict[str, Any], config_defaults: Dict[str, Any]) -> Dict[str, Any]:
+    def _set_config_defaults(self, config, config_defaults):
         for (k, v) in config_defaults.items():
             if k not in config:
                 config[k] = v
-            elif isinstance(v, dict):
-                config[k] = self._set_config_defaults(config[k] if config[k] is not None else {}, v)
         return config
 
-    def _validate_config(self) -> None:
+    def _validate_config(self):
         self.attn_config = self._set_config_defaults(self.attn_config, attn_config_defaults)
-        self.ffn_config = self._set_config_defaults(self.ffn_config, ffn_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')
@@ -135,49 +104,15 @@ class MPTConfig(PretrainedConfig):
             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['attn_impl'] == 'flash' and is_flash_v1_installed():
-            warnings.warn(VersionedDeprecationWarning('Support for Flash Attention v1 is deprecated. Please upgrade to Flash Attention v2.4.2. To install Flash Attention v2.4.2, please run `pip install -e ".[gpu-flash2]"` from the root directory of the llm-foundry repository.', remove_version='0.6.0'))
-        if self.attn_config['attn_impl'] == 'triton' and (not self.attn_config['prefix_lm']):
-            warnings.warn(UserWarning('If not using a Prefix Language Model, we recommend setting "attn_impl" to "flash" instead of "triton".'))
-        if self.attn_config['alibi'] and (not check_alibi_support(self.attn_config['attn_impl'])):
-            raise NotImplementedError('alibi only implemented with torch, triton, and flash (v2.4.2 or higher) attention.')
-        if self.attn_config['attn_uses_sequence_id'] and (not (self.attn_config['attn_impl'] in ['torch', 'triton'] or (self.attn_config['attn_impl'] == 'flash' and is_flash_v2_installed(v2_version='v2.1.2')))):
-            raise NotImplementedError('attn_uses_sequence_id only implemented with torch, triton, and flash (v2.1.2 or higher) attention.')
-        if self.attn_config['rope'] and self.attn_config['rope_impl'] not in ['dail', 'hf']:
-            raise ValueError('If rope is being used then rope_impl should be either "dail", or "hf".')
-        if self.attn_config['rope'] and self.attn_config['rope_impl'] == 'hf' and (self.attn_config['rope_hf_config']['type'] not in ['no_scaling', 'linear', 'dynamic']):
-            raise ValueError('If using hf implementation of rope, the type should be one of "no_scaling", "linear" or "dynamic".')
-        if self.attn_config['rope'] and self.attn_config['rope_impl'] == 'dail':
-            if self.attn_config['rope_dail_config']['type'] not in ['original', 'xpos']:
-                raise ValueError('If using the dail implementation of rope, the type should be one of "original" or "xpos".')
-            if not is_flash_v2_installed(v2_version='2.0.1'):
-                raise ImportError('If using the dail implementation of rope, the flash_attn library v2.0.1 or higher must be installed. Please check the instructions at https://github.com/mosaicml/llm-foundry/blob/main/TUTORIAL.md#what-kinds-of-positional-embeddings-does-llm-foundry-support')
-        if self.attn_config['sliding_window_size'] != -1 and (not (self.attn_config['attn_impl'] == 'flash' and is_flash_v2_installed(v2_version='v2.3.0'))):
-            raise NotImplementedError('sliding window only implemented with flash attention v2.3.0 or higher.')
+        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 or self.attn_config['alibi'] or self.attn_config['rope']):
-            warnings.warn(f'Positional information not being provided to the model using either learned_pos_emb or alibi or rope.')
-        if self.fc_type == 'te' or self.ffn_config['ffn_type'] == 'te_ln_mlp':
-            try:
-                import transformer_engine.pytorch as te
-                del te
-            except:
-                raise ImportError('TransformerEngine import fail. `fc_type: te` requires TransformerEngine be installed. ' + 'The required version of transformer_engine also requires FlashAttention v1.0.6 is installed:\n' + 'pip install flash-attn==1.0.6 --no-build-isolation \n' + 'pip install git+https://github.com/NVIDIA/TransformerEngine.git@144e4888b2cdd60bd52e706d5b7a79cb9c1a7156')
-        if self.ffn_config['ffn_type'] == 'mptgeglu':
-            raise ValueError('API CHANGE: `ffn_type=="mptgeglu"` changed to `ffn_type=="mptglu"`. ' + 'See [#829](https://github.com/mosaicml/llm-foundry/pull/829) for details.')
-        elif self.ffn_config['ffn_type'] in ['mptmlp', 'mptglu']:
-            self.ffn_config['fc_type'] = self.fc_type
-        elif self.ffn_config['ffn_type'] == 'te_ln_mlp':
-            self.ffn_config['bias'] = not self.no_bias
-            if 'ffn_act_fn' in self.ffn_config.keys():
-                raise ValueError(f'Transformer Engine block does not support custom activation functions.')
-        if not self.use_pad_tok_in_ffn:
-            try:
-                from flash_attn.bert_padding import unpad_input, pad_input
-            except:
-                raise ImportError('In order to set `use_pad_tok_in_ffn=False`, please install flash-attn==1.0.9 or flash-attn==2.3.6')
+        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.')

custom_embedding.py

@@ -1,10 +1,12 @@
+import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 
+
 class SharedEmbedding(nn.Embedding):
 
-    def forward(self, input: Tensor, unembed: bool=False) -> Tensor:
+    def forward(self, input: Tensor, unembed: bool = False) -> Tensor:
         if unembed:
             return F.linear(input, self.weight)
         return super().forward(input)

fc.py

@@ -1,7 +0,0 @@
-from torch import nn
-FC_CLASS_REGISTRY = {'torch': nn.Linear}
-try:
-    import transformer_engine.pytorch as te
-    FC_CLASS_REGISTRY['te'] = te.Linear
-except:
-    pass

ffn.py

@@ -1,97 +0,0 @@
-"""MPT Blocks used for the MPT Model."""
-import logging
-from copy import deepcopy
-from functools import partial
-from typing import Any, Callable, Optional, Union
-import torch
-import torch.nn as nn
-from .fc import FC_CLASS_REGISTRY
-try:
-    import transformer_engine.pytorch as te
-except:
-    te = None
-log = logging.getLogger(__name__)
-_FFN_ACT_FN_DEFAULT = {'name': 'gelu', 'approximate': 'none'}
-
-def resolve_ffn_act_fn(config: Optional[dict]=None) -> Callable[[torch.Tensor], torch.Tensor]:
-    """Resolve the activation function for the feed-forward network.
-
-    Args:
-        config (Optional[dict]): The configuration dictionary for the activation function.
-            The dict config must specify the 'name' of a torch.nn.functional activation
-            function. All of other key values pairs are bound to the function as a partial.
-
-    Returns:
-        Callable[[torch.Tensor], torch.Tensor]: The activation function.
-    """
-    if config is None:
-        config = _FFN_ACT_FN_DEFAULT
-    config = deepcopy(config)
-    name = config.pop('name')
-    if not hasattr(torch.nn.functional, name):
-        raise ValueError(f'Unrecognised activation function name ({name}).')
-    act = getattr(torch.nn.functional, name)
-    return partial(act, **config)
-_DEFAULT_ACT_FN = resolve_ffn_act_fn(_FFN_ACT_FN_DEFAULT)
-
-def resolve_ffn_hidden_size(d_model: int, expansion_ratio: Union[int, float], ffn_hidden_size: Optional[int]=None) -> int:
-    """Resolve the hidden size of the feed-forward network.
-
-    Args:
-        d_model (int): The dimension of the input and output of the feed-forward network.
-        expansion_ratio (Union[int, float]): The expansion ratio of the feed-forward network.
-        ffn_hidden_size (Optional[int]): The hidden size of the feed-forward network.
-
-    Returns:
-        int: The hidden size of the feed-forward network.
-    """
-    if ffn_hidden_size is not None:
-        log.info(f'`expansion_ratio` (={expansion_ratio}) ignored when `ffn_hidden_size` (={ffn_hidden_size}) is specified.')
-    else:
-        ffn_hidden_size = int(d_model * expansion_ratio)
-        if ffn_hidden_size != d_model * expansion_ratio:
-            raise ValueError(f'`d_model * expansion_ratio` must be an integer (d_model={d_model!r}; expansion_ratio={expansion_ratio!r}; d_model * expansion_ratio={d_model * expansion_ratio!r}).')
-    return ffn_hidden_size
-
-class MPTMLP(nn.Module):
-
-    def __init__(self, d_model: int, expansion_ratio: Union[int, float], fc_type: str='torch', ffn_hidden_size: Optional[int]=None, act_fn: Callable[[torch.Tensor], torch.Tensor]=_DEFAULT_ACT_FN, device: Optional[str]=None, bias: bool=True):
-        super().__init__()
-        ffn_hidden_size = resolve_ffn_hidden_size(d_model, expansion_ratio, ffn_hidden_size)
-        self.fc_kwargs: dict[str, Any] = {'bias': bias}
-        if fc_type != 'te':
-            self.fc_kwargs['device'] = device
-        self.up_proj = FC_CLASS_REGISTRY[fc_type](d_model, ffn_hidden_size, **self.fc_kwargs)
-        self.act = act_fn
-        self.down_proj = FC_CLASS_REGISTRY[fc_type](ffn_hidden_size, d_model, **self.fc_kwargs)
-        self.down_proj._is_residual = True
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.down_proj(self.act(self.up_proj(x)))
-
-class MPTGLU(MPTMLP):
-
-    def __init__(self, d_model: int, expansion_ratio: Union[int, float], fc_type: str='torch', ffn_hidden_size: Optional[int]=None, act_fn: Callable[[torch.Tensor], torch.Tensor]=_DEFAULT_ACT_FN, device: Optional[str]=None, bias: bool=True):
-        super().__init__(d_model=d_model, expansion_ratio=expansion_ratio, fc_type=fc_type, ffn_hidden_size=ffn_hidden_size, act_fn=act_fn, device=device, bias=bias)
-        self.gate_proj = FC_CLASS_REGISTRY[fc_type](d_model, self.up_proj.out_features, **self.fc_kwargs)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))
-FFN_CLASS_REGISTRY = {'mptmlp': MPTMLP, 'mptglu': MPTGLU}
-if te is not None:
-    te.LayerNormMLP._has_norm = True
-    FFN_CLASS_REGISTRY['te_ln_mlp'] = te.LayerNormMLP
-
-def build_ffn(d_model: int, expansion_ratio: Union[int, float], fc_type: str='torch', ffn_hidden_size: Optional[int]=None, ffn_act_fn: Optional[dict]=None, device: Optional[str]=None, bias: bool=True, **kwargs: Any) -> nn.Module:
-    ffn_type = kwargs.pop('ffn_type')
-    if ffn_type in ['mptmlp', 'mptglu']:
-        if len(kwargs) > 0:
-            raise ValueError(f'MPTMLP (or MPTGLU) got an unexpected keyword argument: {kwargs}')
-        return FFN_CLASS_REGISTRY[ffn_type](d_model=d_model, expansion_ratio=expansion_ratio, fc_type=fc_type, act_fn=resolve_ffn_act_fn(ffn_act_fn), ffn_hidden_size=ffn_hidden_size, device=device, bias=bias)
-    elif ffn_type == 'te_ln_mlp':
-        assert te is not None
-        ffn_hidden_size = resolve_ffn_hidden_size(d_model, expansion_ratio, ffn_hidden_size)
-        if ffn_act_fn is not None:
-            raise ValueError(f'Transformer Engine block does not support custom activation functions.')
-        return te.LayerNormMLP(hidden_size=d_model, ffn_hidden_size=ffn_hidden_size, bias=bias, **kwargs)
-    raise ValueError(f'ffn_type={ffn_type!r} not recognized.')

hf_prefixlm_converter.py

@@ -6,13 +6,23 @@ Causal LM to convert it to a Prefix LM.
 Prefix LMs accepts a `bidirectional_mask` input in `forward`
 and treat the input prompt as the prefix in `generate`.
 """
+import math
+import warnings
 from types import MethodType
-from typing import Any, List, MutableMapping, Optional, Tuple, Union
+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 _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 logging
 from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 from transformers.models.gpt_neo.modeling_gpt_neo import GPTNeoForCausalLM
 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
+logger = logging.get_logger(__name__)
 _SUPPORTED_GPT_MODELS = (GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM)
 CAUSAL_GPT_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM]
 
@@ -80,14 +90,13 @@ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_T
             bidirectional_mask = torch.cat([bidirectional_mask, pad], dim=1)
         bidirectional = bidirectional_mask.unsqueeze(1).unsqueeze(1)
         for attn_module in attn_modules:
-            assert isinstance(attn_module.bias, torch.Tensor)
             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]
         return output
 
-    def generate(self: CAUSAL_GPT_TYPES, *args: Any, **kwargs: Any):
+    def generate(self: CAUSAL_GPT_TYPES, *args: tuple, **kwargs: Dict[str, Any]):
         """Wraps original generate to enable PrefixLM attention."""
         attn_modules = _get_attn_modules(model)
         for attn_module in attn_modules:
@@ -100,8 +109,228 @@ def _convert_gpt_causal_lm_to_prefix_lm(model: CAUSAL_GPT_TYPES) -> CAUSAL_GPT_T
     setattr(model, 'generate', MethodType(generate, model))
     setattr(model, '_prefix_lm_converted', True)
     return model
-_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS
-CAUSAL_LM_TYPES = Union[GPT2LMHeadModel, GPTJForCausalLM, GPTNeoForCausalLM, GPTNeoXForCausalLM]
+
+def _convert_bloom_causal_lm_to_prefix_lm(model: BloomForCausalLM) -> BloomForCausalLM:
+    """Converts a BLOOM Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `BloomForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    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:
+        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)
+            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_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
+        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:
+        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)
+        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)
+            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)
+        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)
+        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
+        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
+        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')
+        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')
+        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)
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        hidden_states = self.word_embeddings_layernorm(inputs_embeds)
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_hidden_states = () if output_hidden_states else None
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values[0] is not None:
+            tmp = past_key_values[0][0]
+            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)
+        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)):
+            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`...')
+                    use_cache = False
+
+                def create_custom_forward(module):
+
+                    def custom_forward(*inputs):
+                        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])
+            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)
+            hidden_states = outputs[0]
+            if use_cache is True:
+                presents = presents + (outputs[1],)
+            if output_attentions:
+                oa = (outputs[2 if use_cache else 1],)
+                all_self_attentions = all_self_attentions + oa
+        hidden_states = self.ln_f(hidden_states)
+        if output_hidden_states:
+            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))
+    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]:
+        """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 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)
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+        loss = None
+        if labels is not None:
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            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))
+        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:
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            bidirectional_mask = None
+            if past[0][0].shape[0] == input_ids.shape[0]:
+                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 model
+
+def _convert_opt_causal_lm_to_prefix_lm(model: OPTForCausalLM) -> OPTForCausalLM:
+    """Converts an OPT Causal LM to a Prefix LM.
+
+    Supported HuggingFace model classes:
+        - `OPTForCausalLM`
+
+    See `convert_hf_causal_lm_to_prefix_lm` for more details.
+    """
+    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'))
+    model.model.decoder.bidirectional_mask = None
+
+    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':
+                (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)
+            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)
+                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)
+        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
+        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):
+
+        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)
+        if bidirectional_mask is None:
+            return call_og_forward()
+        self.model.decoder.bidirectional_mask = bidirectional_mask
+        try:
+            outputs = call_og_forward()
+        except:
+            self.model.decoder.bidirectional_mask = None
+            raise
+        self.model.decoder.bidirectional_mask = None
+        return outputs
+
+    def generate(self: OPTForCausalLM, *args: tuple, **kwargs: Dict[str, Any]):
+        """Wraps original generate to enable PrefixLM-style attention."""
+        self.model.decoder.bidirectional_mask = 'g'
+        try:
+            output = self._original_generate(*args, **kwargs)
+        except:
+            self.model.decoder.bidirectional_mask = None
+            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)
+    return model
+_SUPPORTED_HF_MODELS = _SUPPORTED_GPT_MODELS + (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.
@@ -111,6 +340,8 @@ def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES
         - `GPTNeoForCausalLM`
         - `GPTNeoXForCausalLM`
         - `GPTJForCausalLM`
+        - `BloomForCausalLM`
+        - `OPTForCausalLM`
 
     Conversion to a Prefix LM is done by modifying the `forward` method, and possibly also the
     `generate` method and/or select underlying methods depending on the model class.
@@ -160,10 +391,14 @@ def convert_hf_causal_lm_to_prefix_lm(model: CAUSAL_LM_TYPES) -> CAUSAL_LM_TYPES
     """
     if isinstance(model, _SUPPORTED_GPT_MODELS):
         return _convert_gpt_causal_lm_to_prefix_lm(model)
+    elif isinstance(model, BloomForCausalLM):
+        return _convert_bloom_causal_lm_to_prefix_lm(model)
+    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}')
 
-def add_bidirectional_mask_if_missing(batch: MutableMapping):
+def add_bidirectional_mask_if_missing(batch: Dict[str, Any]):
     """Attempts to add bidirectional_mask to batch if missing.
 
     Raises:

meta_init_context.py

@@ -1,5 +1,4 @@
 from contextlib import contextmanager
-from typing import Any, Callable, Optional
 import torch
 import torch.nn as nn
 
@@ -58,29 +57,25 @@ def init_on_device(device: torch.device, include_buffers: bool=False):
     if include_buffers:
         old_register_buffer = nn.Module.register_buffer
 
-    def register_empty_parameter(self: torch.nn.Module, name: str, param: Optional[torch.nn.Parameter]):
-        old_register_parameter(self, name, param)
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
         if param is not None:
-            parameter = self._parameters[name]
-            assert parameter is not None
-            param_cls = type(parameter)
-            kwargs = parameter.__dict__
-            self._parameters[name] = param_cls(parameter.to(device), **kwargs)
-
-    def register_empty_buffer(self: torch.nn.Module, name: str, tensor: Optional[torch.Tensor], persistent: bool=True):
-        old_register_buffer(self, name, tensor, persistent=persistent)
-        if tensor is not None:
-            named_buffer = self._buffers[name]
-            assert named_buffer is not None
-            self._buffers[name] = named_buffer.to(device)
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            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']}
     else:
         tensor_constructors_to_patch = {}
 
-    def patch_tensor_constructor(fn: Callable):
+    def patch_tensor_constructor(fn):
 
-        def wrapper(*args: Any, **kwargs: Any):
+        def wrapper(*args, **kwargs):
             kwargs['device'] = device
             return fn(*args, **kwargs)
         return wrapper

modeling_mpt.py

@@ -2,174 +2,34 @@
 
 Inspired by https://github.com/karpathy/minGPT/blob/master/mingpt/model.py
 """
-from __future__ import annotations
 import math
 import warnings
-from typing import Any, Dict, List, Mapping, MutableMapping, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from .attention import is_flash_v1_installed, is_flash_v2_installed
-if is_flash_v2_installed():
-    try:
-        from flash_attn import bert_padding
-        from flash_attn.layers.rotary import RotaryEmbedding as DAILRotaryEmbedding
-    except Exception as e:
-        raise e
-if is_flash_v1_installed():
-    try:
-        from flash_attn import bert_padding
-    except Exception as e:
-        raise e
-from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from transformers.models.llama.modeling_llama import LlamaDynamicNTKScalingRotaryEmbedding as HFDynamicNTKScalingRotaryEmbedding
-from transformers.models.llama.modeling_llama import LlamaLinearScalingRotaryEmbedding as HFLinearScalingRotaryEmbedding
-from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding as HFRotaryEmbedding
-from .attention import ATTN_CLASS_REGISTRY, attn_bias_shape, build_attn_bias, gen_slopes
+from .attention import attn_bias_shape, build_attn_bias
 from .blocks import MPTBlock
 from .custom_embedding import SharedEmbedding
-from .fc import FC_CLASS_REGISTRY as FC_CLASS_REGISTRY
-from .ffn import FFN_CLASS_REGISTRY as FFN_CLASS_REGISTRY
-from .ffn import MPTMLP as MPTMLP
-from .ffn import build_ffn as build_ffn
 from .norm import NORM_CLASS_REGISTRY
 from .configuration_mpt import MPTConfig
 from .adapt_tokenizer import AutoTokenizerForMOD, adapt_tokenizer_for_denoising
 from .hf_prefixlm_converter import add_bidirectional_mask_if_missing, convert_hf_causal_lm_to_prefix_lm
 from .meta_init_context import init_empty_weights
-from .param_init_fns import generic_param_init_fn_, MODEL_INIT_REGISTRY
+from .param_init_fns import MODEL_INIT_REGISTRY, generic_param_init_fn_
 try:
-    from .flash_attn_triton import flash_attn_func as flash_attn_func
+    from .flash_attn_triton import flash_attn_func
 except:
     pass
-import logging
-log = logging.getLogger(__name__)
-
-def gen_rotary_embedding(rope_head_dim: int, rope_impl: str, rope_theta: int, rope_dail_config: dict, rope_hf_config: dict, max_seq_len: int):
-    if rope_impl == 'dail':
-        return DAILRotaryEmbedding(dim=rope_head_dim, base=rope_theta, interleaved=False, scale_base=rope_dail_config['xpos_scale_base'] if rope_dail_config['type'] == 'xpos' else None, pos_idx_in_fp32=rope_dail_config['pos_idx_in_fp32'], device='cpu')
-    elif rope_impl == 'hf':
-        if rope_hf_config['type'] == 'no_scaling':
-            return HFRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, device='cpu')
-        elif rope_hf_config['type'] == 'linear':
-            return HFLinearScalingRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, scaling_factor=rope_hf_config['factor'], device='cpu')
-        elif rope_hf_config['type'] == 'dynamic':
-            return HFDynamicNTKScalingRotaryEmbedding(rope_head_dim, max_position_embeddings=max_seq_len, base=rope_theta, scaling_factor=rope_hf_config['factor'], device='cpu')
-    raise ValueError('rope_impl needs to be either dail or hf')
-
-def gen_attention_mask_in_length(sequence_id: Union[None, torch.Tensor], S: int, attn_uses_sequence_id: bool, attn_impl: str, attention_mask: Union[torch.Tensor, None]):
-    """Generates the attention mask used for sequence masking in FA v2.
-
-    Only supports sequence id based sparse attention for no attention masking or attention masking with right padding.
-    In case of left padding:
-        1. Training with left padding is not supported in MPT (see https://github.com/mosaicml/llm-foundry/blob/1eecd4cb8e734499f77f6a35f657b8b20c0adfcb/llmfoundry/models/mpt/modeling_mpt.py#L407).
-        2. For generation with left padding, we only have a single sequence id per sample, so we don't need sequence id based sparse attention.
-
-    Args:
-        sequence_id (Union[None, torch.Tensor]): Tensor containing the sequence id for each token. Shape (batch_size, seq_len).
-        S (int): Sequence length
-        attn_uses_sequence_id (bool): Whether the attention uses sequence id based masking.
-        attn_impl (str): Attention implementation. This function is only creates attention_mask_in_length for flash attention.
-        attention_mask (Union[torch.Tensor, None]): Attention mask tensor of shape (batch_size, seq_len)
-
-    Returns:
-        attention_mask_in_length: (batch, seqlen), int, a nonzero number (e.g., 1, 2, 3, etc.) means length of concatenated sequence in b-th batch, and 0 means none. For example, if batch = 3 and seqlen = 6, the attention_mask_in_length is:
-            ```
-            [
-            [2, 3, 0, 0, 0, 0],
-            [3, 2, 0, 0, 0, 0],
-            [6, 0, 0, 0, 0, 0]
-            ]
-            ```
-        , which refers to the 3D-attention mask:
-            ```
-            [
-            [
-                [1, 0, 0, 0, 0, 0],
-                [1, 1, 0, 0, 0, 0],
-                [0, 0, 1, 0, 0, 0],
-                [0, 0, 1, 1, 0, 0],
-                [0, 0, 1, 1, 1, 0],
-                [0, 0, 0, 0, 0, 1]
-            ],
-            [
-                [1, 0, 0, 0, 0, 0],
-                [1, 1, 0, 0, 0, 0],
-                [1, 1, 1, 0, 0, 0],
-                [0, 0, 0, 1, 0, 0],
-                [0, 0, 0, 1, 1, 0],
-                [0, 0, 0, 0, 0, 1]
-            ],
-            [
-                [1, 0, 0, 0, 0, 0],
-                [1, 1, 0, 0, 0, 0],
-                [1, 1, 1, 0, 0, 0],
-                [1, 1, 1, 1, 0, 0],
-                [1, 1, 1, 1, 1, 0],
-                [1, 1, 1, 1, 1, 1]
-            ]
-            ]
-            ```.
-            (The description above is taken verbatim from https://github.com/Dao-AILab/flash-attention/blob/9356a1c0389660d7e231ff3163c1ac17d9e3824a/flash_attn/bert_padding.py#L125 .)
-    """
-    attention_mask_in_length = None
-    if sequence_id is not None and attn_uses_sequence_id and (attn_impl == 'flash'):
-        if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0]:
-            raise NotImplementedError('Left padding is not supported with flash attention when attn_uses_sequence_id is set to True.')
-        if S != sequence_id.shape[-1]:
-            raise ValueError(f'Sequence length ({S}) does not match length of sequences in sequence_id ({sequence_id.shape[-1]}).')
-        if attention_mask is not None:
-            sequence_id = sequence_id.masked_fill(~attention_mask, 0)
-        attention_mask_in_length = torch.nn.functional.one_hot(sequence_id)
-        if attention_mask is not None:
-            attention_mask_in_length = attention_mask_in_length.masked_fill(~attention_mask.unsqueeze(-1), 0)
-        attention_mask_in_length = attention_mask_in_length.sum(dim=1)
-        attention_mask_in_length = torch.nn.functional.pad(attention_mask_in_length, (0, S - attention_mask_in_length.shape[-1]), mode='constant', value=0)
-    return attention_mask_in_length
-
-def gen_flash_attn_padding_info(bsz: int, S: int, past_key_len: int, device: torch.device, attention_mask_in_length: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None):
-    flash_attn_padding_info = {}
-    if attention_mask_in_length is None:
-        key_padding_mask = attention_mask
-        if key_padding_mask is None:
-            key_padding_mask = torch.ones((bsz, past_key_len + S), dtype=torch.bool, device=device)
-        query_padding_mask = key_padding_mask[:, -S:]
-        unpadding_function = bert_padding.unpad_input
-    else:
-        key_padding_mask = attention_mask_in_length
-        query_padding_mask = attention_mask_in_length
-        unpadding_function = bert_padding.unpad_input_for_concatenated_sequences
-    (_, indices_q, cu_seqlens_q, max_seqlen_q) = unpadding_function(torch.empty(bsz, S, 1, device=device), query_padding_mask)
-    (_, indices_k, cu_seqlens_k, max_seqlen_k) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
-    (_, indices_v, _, _) = unpadding_function(torch.empty(bsz, past_key_len + S, 1, device=device), key_padding_mask)
-    flash_attn_padding_info['indices_q'] = indices_q
-    flash_attn_padding_info['indices_k'] = indices_k
-    flash_attn_padding_info['indices_v'] = indices_v
-    flash_attn_padding_info['cu_seqlens_q'] = cu_seqlens_q
-    flash_attn_padding_info['cu_seqlens_k'] = cu_seqlens_k
-    flash_attn_padding_info['max_seqlen_q'] = max_seqlen_q
-    flash_attn_padding_info['max_seqlen_k'] = max_seqlen_k
-    return flash_attn_padding_info
-
-def apply_sequence_id(attn_bias: torch.Tensor, sequence_id: torch.LongTensor, max_seq_len: int) -> torch.Tensor:
-    seq_len = sequence_id.shape[-1]
-    if seq_len > max_seq_len:
-        raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={max_seq_len}')
-    attn_bias = attn_bias[..., :seq_len, :seq_len]
-    cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
-    min_val = torch.finfo(attn_bias.dtype).min
-    attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
-    return attn_bias
+Tokenizer = Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
 
 class MPTPreTrainedModel(PreTrainedModel):
     config_class = MPTConfig
     base_model_prefix = 'model'
     _no_split_modules = ['MPTBlock']
 
-def _fsdp_wrap_fn(self: Union[MPTModel, MPTForCausalLM], module: nn.Module) -> bool:
-    return isinstance(module, MPTBlock)
-
 class MPTModel(MPTPreTrainedModel):
 
     def __init__(self, config: MPTConfig):
@@ -180,30 +40,19 @@ class MPTModel(MPTPreTrainedModel):
         self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
         self.alibi = config.attn_config['alibi']
         self.alibi_bias_max = config.attn_config['alibi_bias_max']
-        self.learned_pos_emb = config.learned_pos_emb
-        if config.init_device == 'mixed':
-            if dist.get_local_rank() == 0:
-                config.init_device = 'cpu'
-            else:
-                config.init_device = 'meta'
         if config.norm_type.lower() not in NORM_CLASS_REGISTRY.keys():
             norm_options = ' | '.join(NORM_CLASS_REGISTRY.keys())
             raise NotImplementedError(f'Requested norm type ({config.norm_type}) is not implemented within this repo (Options: {norm_options}).')
         norm_class = NORM_CLASS_REGISTRY[config.norm_type.lower()]
         self.embedding_fraction = config.embedding_fraction
         self.wte = SharedEmbedding(config.vocab_size, config.d_model, device=config.init_device)
-        if self.learned_pos_emb:
-            self.wpe = torch.nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
+        if not self.alibi:
+            self.wpe = nn.Embedding(config.max_seq_len, config.d_model, device=config.init_device)
         self.emb_drop = nn.Dropout(config.emb_pdrop)
         self.blocks = nn.ModuleList([MPTBlock(device=config.init_device, **config.to_dict()) for _ in range(config.n_layers)])
         self.norm_f = norm_class(config.d_model, device=config.init_device)
-        self.rope = config.attn_config['rope']
-        self.rope_impl = None
-        if self.rope:
-            self.rope_impl = config.attn_config['rope_impl']
-            self.rotary_embedding = gen_rotary_embedding(rope_head_dim=config.d_model // config.n_heads, rope_impl=self.rope_impl, rope_theta=config.attn_config['rope_theta'], rope_dail_config=config.attn_config['rope_dail_config'], rope_hf_config=config.attn_config['rope_hf_config'], max_seq_len=self.config.max_seq_len)
         if config.init_device != 'meta':
-            log.info(f'We recommend using config.init_device="meta" with Composer + FSDP for faster initialization.')
+            print(f'You are using config.init_device={config.init_device!r}, but you can also use config.init_device="meta" with Composer + FSDP for fast initialization.')
             self.apply(self.param_init_fn)
         self.is_causal = not self.prefix_lm
         self._attn_bias_initialized = False
@@ -212,22 +61,25 @@ class MPTModel(MPTPreTrainedModel):
         if config.no_bias:
             for module in self.modules():
                 if hasattr(module, 'bias') and isinstance(module.bias, nn.Parameter):
-                    log.info(f'Removing bias from module={module!r}.')
+                    if config.verbose:
+                        warnings.warn(f'Removing bias ({module.bias}) from {module}.')
                     module.register_parameter('bias', None)
-                if hasattr(module, 'use_bias'):
-                    log.info(f'Setting use_bias=False for module={module!r}.')
-                    module.use_bias = False
-        log.debug(self)
-        log.debug(f"Using {self.config.init_config['name']} initialization.")
-
-    def get_input_embeddings(self) -> Union[SharedEmbedding, nn.Embedding]:
+        if config.verbose and config.verbose > 2:
+            print(self)
+        if 'verbose' not in self.config.init_config:
+            self.config.init_config['verbose'] = self.config.verbose
+        if self.config.init_config['verbose'] > 1:
+            init_fn_name = self.config.init_config['name']
+            warnings.warn(f'Using {init_fn_name} initialization.')
+
+    def get_input_embeddings(self):
         return self.wte
 
-    def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
+    def set_input_embeddings(self, value):
         self.wte = value
 
     @torch.no_grad()
-    def _attn_bias(self, device: torch.device, dtype: torch.dtype, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None) -> Tuple[Optional[torch.Tensor], Optional[torch.ByteTensor]]:
+    def _attn_bias(self, device, dtype, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None):
         if not self._attn_bias_initialized:
             if self.attn_bias_shape:
                 self.attn_bias = torch.zeros(self.attn_bias_shape, device=device, dtype=dtype)
@@ -244,7 +96,7 @@ class MPTModel(MPTPreTrainedModel):
             attn_bias = self._apply_prefix_mask(attn_bias, prefix_mask)
         if self.attn_uses_sequence_id and sequence_id is not None:
             assert isinstance(attn_bias, torch.Tensor)
-            attn_bias = apply_sequence_id(attn_bias, sequence_id, self.config.max_seq_len)
+            attn_bias = self._apply_sequence_id(attn_bias, sequence_id)
         if attention_mask is not None:
             s_k = attention_mask.shape[-1]
             if attn_bias is None:
@@ -256,9 +108,9 @@ class MPTModel(MPTPreTrainedModel):
                 raise ValueError(f'attention_mask shape={attention_mask.shape} ' + f'and prefix_mask shape={prefix_mask.shape} are not equal.')
             min_val = torch.finfo(attn_bias.dtype).min
             attn_bias = attn_bias.masked_fill(~attention_mask.view(-1, 1, 1, s_k), min_val)
-        return (attn_bias, attention_mask)
+        return (attn_bias, None)
 
-    def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor) -> torch.Tensor:
+    def _apply_prefix_mask(self, attn_bias: torch.Tensor, prefix_mask: torch.Tensor):
         (s_k, s_q) = attn_bias.shape[-2:]
         if s_k != self.config.max_seq_len or s_q != self.config.max_seq_len:
             raise ValueError('attn_bias does not match the expected shape. ' + f'The last two dimensions should both be {self.config.max_length} ' + f'but are {s_k} and {s_q}.')
@@ -273,7 +125,17 @@ class MPTModel(MPTPreTrainedModel):
         attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
         return attn_bias
 
-    def forward(self, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.Tensor]=None) -> BaseModelOutputWithPast:
+    def _apply_sequence_id(self, attn_bias: torch.Tensor, sequence_id: torch.LongTensor):
+        seq_len = sequence_id.shape[-1]
+        if seq_len > self.config.max_seq_len:
+            raise ValueError(f'sequence_id sequence length cannot exceed max_seq_len={self.config.max_seq_len}')
+        attn_bias = attn_bias[..., :seq_len, :seq_len]
+        cannot_attend = torch.logical_not(torch.eq(sequence_id.view(-1, seq_len, 1), sequence_id.view(-1, 1, seq_len))).unsqueeze(1)
+        min_val = torch.finfo(attn_bias.dtype).min
+        attn_bias = attn_bias.masked_fill(cannot_attend, min_val)
+        return attn_bias
+
+    def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None):
         return_dict = return_dict if return_dict is not None else self.config.return_dict
         use_cache = use_cache if use_cache is not None else self.config.use_cache
         if attention_mask is not None:
@@ -285,7 +147,7 @@ class MPTModel(MPTPreTrainedModel):
         if output_attentions:
             if self.attn_impl != 'torch':
                 raise NotImplementedError('output_attentions is not implemented for MPT when using attn_impl `flash` or `triton`.')
-        if self.training and attention_mask is not None and (attention_mask[:, 0].sum() != attention_mask.shape[0]):
+        if attention_mask is not None and attention_mask[:, 0].sum() != attention_mask.shape[0] and self.training:
             raise NotImplementedError('MPT does not support training with left padding.')
         if self.prefix_lm and prefix_mask is None:
             raise ValueError('prefix_mask is a required argument when MPT is configured with prefix_lm=True.')
@@ -294,42 +156,26 @@ class MPTModel(MPTPreTrainedModel):
                 raise ValueError('sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True ' + 'and the model is in train mode.')
             elif self.attn_uses_sequence_id is False and sequence_id is not None:
                 warnings.warn('MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. ' + 'This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.')
-        if input_ids is not None and inputs_embeds is not None:
-            raise ValueError('You cannot specify both input_ids and inputs_embeds.')
-        elif input_ids is not None:
-            bsz = input_ids.size(0)
-            S = input_ids.size(1)
-            x = self.wte(input_ids)
-            input_device = input_ids.device
-        elif inputs_embeds is not None:
-            bsz = inputs_embeds.size(0)
-            S = inputs_embeds.size(1)
-            x = inputs_embeds
-            input_device = inputs_embeds.device
-        else:
-            raise ValueError('You must specify input_ids or inputs_embeds')
+        S = input_ids.size(1)
         assert S <= self.config.max_seq_len, f'Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}'
-        rotary_emb_w_meta_info = None
-        past_position = 0
-        if past_key_values is not None:
-            if len(past_key_values) != self.config.n_layers:
-                raise ValueError(f'past_key_values must provide a past_key_value for each attention ' + f'layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).')
-            past_position = past_key_values[0][0].size(1)
-            if self.attn_impl == 'torch':
-                past_position = past_key_values[0][0].size(3)
-        if self.learned_pos_emb or self.rope:
-            if self.learned_pos_emb and S + past_position > self.config.max_seq_len:
-                raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length ' + f'{S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.')
-            if self.learned_pos_emb or (self.rope and self.rope_impl == 'hf'):
-                pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_device).unsqueeze(0)
-                if attention_mask is not None:
-                    pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
-                if self.learned_pos_emb:
-                    x = x + self.wpe(pos)
-                elif self.rope and self.rope_impl == 'hf':
-                    rotary_emb_w_meta_info = {'impl': self.rope_impl, 'rotary_emb': self.rotary_embedding, 'offset_info': pos, 'seq_len': S + past_position}
-            elif self.rope and self.rope_impl == 'dail':
-                rotary_emb_w_meta_info = {'impl': self.rope_impl, 'rotary_emb': self.rotary_embedding, 'offset_info': past_position, 'seq_len': S + past_position}
+        tok_emb = self.wte(input_ids)
+        if self.alibi:
+            x = tok_emb
+        else:
+            past_position = 0
+            if past_key_values is not None:
+                if len(past_key_values) != self.config.n_layers:
+                    raise ValueError(f'past_key_values must provide a past_key_value for each attention ' + f'layer in the network (len(past_key_values)={len(past_key_values)!r}; self.config.n_layers={self.config.n_layers!r}).')
+                past_position = past_key_values[0][0].size(1)
+                if self.attn_impl == 'torch':
+                    past_position = past_key_values[0][0].size(3)
+            if S + past_position > self.config.max_seq_len:
+                raise ValueError(f'Cannot forward input with past sequence length {past_position} and current sequence length {S + 1}, this model only supports total sequence length <= {self.config.max_seq_len}.')
+            pos = torch.arange(past_position, S + past_position, dtype=torch.long, device=input_ids.device).unsqueeze(0)
+            if attention_mask is not None:
+                pos = torch.clamp(pos - torch.cumsum((~attention_mask).to(torch.int32), dim=1)[:, past_position:], min=0)
+            pos_emb = self.wpe(pos)
+            x = tok_emb + pos_emb
         if self.embedding_fraction == 1:
             x = self.emb_drop(x)
         else:
@@ -337,26 +183,18 @@ class MPTModel(MPTPreTrainedModel):
             assert isinstance(self.emb_drop, nn.Module)
             x = self.emb_drop(x_shrunk)
         (attn_bias, attention_mask) = self._attn_bias(device=x.device, dtype=torch.float32, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id)
-        attention_mask_in_length = gen_attention_mask_in_length(sequence_id=sequence_id, S=S, attn_uses_sequence_id=self.attn_uses_sequence_id, attn_impl=self.attn_impl, attention_mask=attention_mask)
-        alibi_slopes = None
-        if self.alibi and self.attn_impl == 'flash':
-            alibi_slopes = gen_slopes(n_heads=self.config.n_heads, alibi_bias_max=self.alibi_bias_max, device=x.device, return_1d=True)
-        presents = () if use_cache else None
         if use_cache and past_key_values is None:
             past_key_values = [() for _ in range(self.config.n_layers)]
         all_hidden_states = () if output_hidden_states else None
         all_self_attns = () if output_attentions else None
-        flash_attn_padding_info = {}
-        if self.attn_impl == 'flash':
-            flash_attn_padding_info = gen_flash_attn_padding_info(bsz, S, past_position, x.device, attention_mask_in_length, attention_mask)
         for (b_idx, block) in enumerate(self.blocks):
             if output_hidden_states:
                 assert all_hidden_states is not None
                 all_hidden_states = all_hidden_states + (x,)
             past_key_value = past_key_values[b_idx] if past_key_values is not None else None
-            (x, attn_weights, present) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, rotary_emb_w_meta_info=rotary_emb_w_meta_info, attention_mask=attention_mask, is_causal=self.is_causal, output_attentions=bool(output_attentions), alibi_slopes=alibi_slopes, flash_attn_padding_info=flash_attn_padding_info)
-            if presents is not None:
-                presents += (present,)
+            (x, attn_weights, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
+            if past_key_values is not None:
+                past_key_values[b_idx] = past_key_value
             if output_attentions:
                 assert all_self_attns is not None
                 all_self_attns = all_self_attns + (attn_weights,)
@@ -364,33 +202,25 @@ class MPTModel(MPTPreTrainedModel):
         if output_hidden_states:
             assert all_hidden_states is not None
             all_hidden_states = all_hidden_states + (x,)
-        return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=presents, hidden_states=all_hidden_states, attentions=all_self_attns)
+        return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values, hidden_states=all_hidden_states, attentions=all_self_attns)
 
-    def param_init_fn(self, module: nn.Module) -> None:
+    def param_init_fn(self, module):
         init_fn_name = self.config.init_config['name']
         MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
 
-    def fsdp_wrap_fn(self, module: nn.Module) -> bool:
-        return _fsdp_wrap_fn(self, module)
+    def fsdp_wrap_fn(self, module):
+        return isinstance(module, MPTBlock)
 
-    def activation_checkpointing_fn(self, module: nn.Module) -> bool:
+    def activation_checkpointing_fn(self, module):
         return isinstance(module, MPTBlock)
 
 class MPTForCausalLM(MPTPreTrainedModel):
 
     def __init__(self, config: MPTConfig):
         super().__init__(config)
-        log.info(f'Instantiating an MPTForCausalLM model from {__file__}')
-        self.transformer: MPTModel = MPTModel(config)
-        self.lm_head = None
         if not config.tie_word_embeddings:
-            self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False, device=config.init_device)
-            self.lm_head._fsdp_wrap = True
-        for child in self.transformer.children():
-            if isinstance(child, torch.nn.ModuleList):
-                continue
-            if isinstance(child, torch.nn.Module):
-                child._fsdp_wrap = True
+            raise ValueError('MPTForCausalLM only supports tied word embeddings')
+        self.transformer = MPTModel(config)
         self.logit_scale = None
         if config.logit_scale is not None:
             logit_scale = config.logit_scale
@@ -401,89 +231,53 @@ class MPTForCausalLM(MPTPreTrainedModel):
                     raise ValueError(f"logit_scale={logit_scale!r} is not recognized as an option; use numeric value or 'inv_sqrt_d_model'.")
             self.logit_scale = logit_scale
 
-    def get_input_embeddings(self) -> Union[SharedEmbedding, nn.Embedding]:
-        return self.transformer.get_input_embeddings()
+    def get_input_embeddings(self):
+        return self.transformer.wte
 
-    def set_input_embeddings(self, value: Union[SharedEmbedding, nn.Embedding]) -> None:
-        self.transformer.set_input_embeddings(value)
+    def set_input_embeddings(self, value):
+        self.transformer.wte = value
 
-    def get_output_embeddings(self) -> Union[SharedEmbedding, nn.Embedding, nn.Linear]:
-        if self.lm_head is not None:
-            return self.lm_head
-        return self.transformer.get_input_embeddings()
+    def get_output_embeddings(self):
+        return self.transformer.wte
 
-    def set_output_embeddings(self, new_embeddings: Union[SharedEmbedding, nn.Embedding, nn.Linear]) -> None:
-        if self.lm_head is not None:
-            self.lm_head = new_embeddings
-        else:
-            if not isinstance(new_embeddings, (SharedEmbedding, nn.Embedding)):
-                raise ValueError('new_embeddings must be an instance of SharedEmbedding ' + f'or nn.Embedding, but got {type(new_embeddings)}.')
-            warnings.warn('Using `set_output_embeddings` to set the embedding layer of ' + 'MPTForCausalLM with tied weights. Given weights are tied, ' + 'using `set_input_embeddings` is recommended over using ' + '`set_output_embeddings`.')
-            self.transformer.set_input_embeddings(new_embeddings)
-
-    def tie_weights(self) -> None:
-        self.lm_head = None
+    def set_output_embeddings(self, new_embeddings):
+        self.transformer.wte = new_embeddings
 
-    def set_decoder(self, decoder: MPTModel) -> None:
+    def set_decoder(self, decoder):
         self.transformer = decoder
 
-    def get_decoder(self) -> MPTModel:
+    def get_decoder(self):
         return self.transformer
 
-    def forward(self, input_ids: Optional[torch.LongTensor]=None, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.FloatTensor]=None) -> CausalLMOutputWithPast:
+    def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, labels: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None):
         return_dict = return_dict if return_dict is not None else self.config.return_dict
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache, inputs_embeds=inputs_embeds)
-        if self.lm_head is not None:
-            logits = self.lm_head(outputs.last_hidden_state)
-        else:
-            out = outputs.last_hidden_state
-            out = out.to(self.transformer.wte.weight.device)
-            logits = self.transformer.wte(out, True)
+        outputs = self.transformer(input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, prefix_mask=prefix_mask, sequence_id=sequence_id, return_dict=return_dict, output_attentions=output_attentions, output_hidden_states=output_hidden_states, use_cache=use_cache)
+        logits = self.transformer.wte(outputs.last_hidden_state.to(self.transformer.wte.weight.device), True)
         if self.logit_scale is not None:
             if self.logit_scale == 0:
                 warnings.warn(f'Multiplying logits by self.logit_scale={self.logit_scale!r}. This will produce uniform (uninformative) outputs.')
             logits *= self.logit_scale
         loss = None
         if labels is not None:
-            _labels = torch.roll(labels, shifts=-1)
-            _labels[:, -1] = -100
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), _labels.to(logits.device).view(-1))
+            labels = torch.roll(labels, shifts=-1)
+            labels[:, -1] = -100
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.to(logits.device).view(-1))
         return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
 
-    def param_init_fn(self, module: nn.Module) -> None:
+    def param_init_fn(self, module):
         init_fn_name = self.config.init_config['name']
         MODEL_INIT_REGISTRY[init_fn_name](module=module, n_layers=self.config.n_layers, d_model=self.config.d_model, **self.config.init_config)
 
-    def fsdp_wrap_fn(self, module: nn.Module) -> bool:
-        return _fsdp_wrap_fn(self, module)
+    def fsdp_wrap_fn(self, module):
+        return isinstance(module, MPTBlock)
 
-    def activation_checkpointing_fn(self, module: nn.Module) -> bool:
-        act_ckpt_list = getattr(self.config, 'activation_checkpointing_target', None) or ['MPTBlock']
-        if isinstance(act_ckpt_list, str):
-            act_ckpt_list = [act_ckpt_list]
-        elif not isinstance(act_ckpt_list, list):
-            raise ValueError(f'activation_checkpointing_target must be either a single string or a list, but got {type(act_ckpt_list)}')
-        if 'MPTBlock' in act_ckpt_list or 'mptblock' in act_ckpt_list:
-            if len(act_ckpt_list) > 1:
-                log.info('Activation checkpointing MPTBlock only (ignoring other sub-block modules specified in activation_checkpointing_target).')
-            return isinstance(module, MPTBlock)
-        mod_types = ()
-        for mod_name in act_ckpt_list:
-            if mod_name.lower() == 'mptblock':
-                mod_types += (MPTBlock,)
-            elif mod_name in ATTN_CLASS_REGISTRY:
-                mod_types += (ATTN_CLASS_REGISTRY[mod_name],)
-            elif mod_name in FFN_CLASS_REGISTRY:
-                mod_types += (FFN_CLASS_REGISTRY[mod_name],)
-            elif mod_name in NORM_CLASS_REGISTRY:
-                mod_types += (NORM_CLASS_REGISTRY[mod_name],)
-            else:
-                msg = ', '.join(list(ATTN_CLASS_REGISTRY.keys()) + list(FFN_CLASS_REGISTRY.keys()) + list(NORM_CLASS_REGISTRY.keys()) + ['MPTBlock'])
-                raise ValueError(f'{mod_name} (specified in activation_checkpointing_target) is not a recognized option out of available options {msg}.')
-        return isinstance(module, mod_types)
+    def activation_checkpointing_fn(self, module):
+        return isinstance(module, MPTBlock)
 
-    def prepare_inputs_for_generation(self, input_ids: torch.Tensor, past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]]=None, inputs_embeds: Optional[torch.Tensor]=None, **kwargs: Any) -> Dict[str, Any]:
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        if inputs_embeds is not None:
+            raise NotImplementedError('inputs_embeds is not implemented for MPT yet')
         attention_mask = kwargs['attention_mask'].bool()
         if attention_mask[:, -1].sum() != attention_mask.shape[0]:
             raise NotImplementedError('MPT does not support generation with right padding.')
@@ -499,15 +293,10 @@ class MPTForCausalLM(MPTPreTrainedModel):
                 raise NotImplementedError('MPT with prefix_lm=True does not support use_cache=False.')
         else:
             prefix_mask = None
-        if inputs_embeds is not None and past_key_values is None:
-            model_inputs = {'inputs_embeds': inputs_embeds}
-        else:
-            model_inputs = {'input_ids': input_ids}
-        model_inputs.update({'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)})
-        return model_inputs
+        return {'input_ids': input_ids, 'attention_mask': attention_mask, 'prefix_mask': prefix_mask, 'sequence_id': sequence_id, 'past_key_values': past_key_values, 'use_cache': kwargs.get('use_cache', True)}
 
     @staticmethod
-    def _reorder_cache(past_key_values: List[Tuple[torch.Tensor, torch.Tensor]], beam_idx: torch.LongTensor) -> List[Tuple[torch.Tensor, ...]]:
+    def _reorder_cache(past_key_values, beam_idx):
         """Used by HuggingFace generate when using beam search with kv-caching.
 
         See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133

norm.py

@@ -1,7 +1,6 @@
-from typing import Dict, List, Optional, Type, Union
 import torch
 
-def _cast_if_autocast_enabled(tensor: torch.Tensor) -> torch.Tensor:
+def _cast_if_autocast_enabled(tensor):
     if torch.is_autocast_enabled():
         if tensor.device.type == 'cuda':
             dtype = torch.get_autocast_gpu_dtype()
@@ -14,10 +13,10 @@ def _cast_if_autocast_enabled(tensor: torch.Tensor) -> torch.Tensor:
 
 class LPLayerNorm(torch.nn.LayerNorm):
 
-    def __init__(self, normalized_shape: Union[int, List[int], torch.Size], eps: float=1e-05, elementwise_affine: bool=True, device: Optional[torch.device]=None, dtype: Optional[torch.dtype]=None):
+    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: torch.Tensor) -> torch.Tensor:
+    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
@@ -25,15 +24,15 @@ class LPLayerNorm(torch.nn.LayerNorm):
         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)
 
-def rms_norm(x: torch.Tensor, weight: Optional[torch.Tensor]=None, eps: float=1e-05) -> torch.Tensor:
-    output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
+def rms_norm(x, weight=None, eps=1e-05):
+    output = x / torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)
     if weight is not None:
         return output * weight
     return output
 
 class RMSNorm(torch.nn.Module):
 
-    def __init__(self, normalized_shape: Union[int, List[int], torch.Size], eps: float=1e-05, weight: bool=True, dtype: Optional[torch.dtype]=None, device: Optional[torch.device]=None):
+    def __init__(self, normalized_shape, eps=1e-05, weight=True, dtype=None, device=None):
         super().__init__()
         self.eps = eps
         if weight:
@@ -41,17 +40,17 @@ class RMSNorm(torch.nn.Module):
         else:
             self.register_parameter('weight', None)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    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: Union[int, List[int], torch.Size], eps: float=1e-05, weight: bool=True, dtype: Optional[torch.dtype]=None, device: Optional[torch.device]=None):
+    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: torch.Tensor) -> torch.Tensor:
+    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
         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: Dict[str, Type[torch.nn.Module]] = {'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

@@ -2,26 +2,22 @@ import math
 import warnings
 from collections.abc import Sequence
 from functools import partial
-from typing import Any, Callable, Optional, Tuple, Union
+from typing import Optional, Tuple, Union
 import torch
 from torch import nn
-from .fc import FC_CLASS_REGISTRY
 from .norm import NORM_CLASS_REGISTRY
-try:
-    import transformer_engine.pytorch as te
-except:
-    te = None
 
-def torch_default_param_init_fn_(module: nn.Module, **kwargs: Any) -> None:
+def torch_default_param_init_fn_(module: nn.Module, verbose: int=0, **kwargs):
     del kwargs
-    if hasattr(module, 'reset_parameters') and isinstance(module.reset_parameters, Callable):
+    if verbose > 1:
+        warnings.warn(f"Initializing network using module's reset_parameters attribute")
+    if hasattr(module, 'reset_parameters'):
         module.reset_parameters()
 
-def fused_init_helper_(module: nn.Module, init_fn_: Callable) -> None:
+def fused_init_helper_(module: nn.Module, init_fn_):
     _fused = getattr(module, '_fused', None)
     if _fused is None:
         raise RuntimeError(f'Internal logic error')
-    assert isinstance(module.weight, torch.Tensor)
     (dim, splits) = _fused
     splits = (0, *splits, module.weight.size(dim))
     for (s, e) in zip(splits[:-1], splits[1:]):
@@ -29,8 +25,10 @@ def fused_init_helper_(module: nn.Module, init_fn_: Callable) -> None:
         slice_indices[dim] = slice(s, e)
         init_fn_(module.weight[slice_indices])
 
-def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, 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, **kwargs: Any) -> None:
+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.')
     init_div_is_residual = init_div_is_residual
     if init_div_is_residual is False:
         div_is_residual = 1.0
@@ -38,18 +36,20 @@ def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int,
         div_is_residual = math.sqrt(2 * n_layers)
     elif isinstance(init_div_is_residual, float) or isinstance(init_div_is_residual, int):
         div_is_residual = init_div_is_residual
-    elif init_div_is_residual.isnumeric():
+    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}')
-    if isinstance(module, tuple(set(FC_CLASS_REGISTRY.values()))):
+    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.')
+    if isinstance(module, nn.Linear):
         if hasattr(module, '_fused'):
             fused_init_helper_(module, init_fn_)
         else:
             init_fn_(module.weight)
         if module.bias is not None:
-            assert isinstance(module.bias, torch.Tensor)
             torch.nn.init.zeros_(module.bias)
         if init_div_is_residual is not False and getattr(module, '_is_residual', False):
             with torch.no_grad():
@@ -60,6 +60,8 @@ def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int,
             if std == 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}.')
         elif emb_init_uniform_lim is not None:
             lim = emb_init_uniform_lim
             if isinstance(lim, Sequence):
@@ -73,13 +75,17 @@ def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int,
                 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}.')
         else:
             emb_init_fn_ = init_fn_
         emb_init_fn_(module.weight)
     elif isinstance(module, tuple(set(NORM_CLASS_REGISTRY.values()))):
-        if hasattr(module, 'weight') and isinstance(module.weight, torch.Tensor):
+        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:
             torch.nn.init.ones_(module.weight)
-        if hasattr(module, 'bias') and isinstance(module.bias, torch.Tensor):
+        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:
@@ -108,45 +114,32 @@ def generic_param_init_fn_(module: nn.Module, init_fn_: Callable, n_layers: int,
                 module.out_proj.weight.div_(div_is_residual)
         if module.out_proj.bias is not None:
             torch.nn.init.zeros_(module.out_proj.bias)
-    elif te is not None and isinstance(module, te.LayerNormMLP):
-        if isinstance(module.layer_norm_weight, torch.Tensor):
-            torch.nn.init.ones_(module.layer_norm_weight)
-        if isinstance(module.layer_norm_bias, torch.Tensor):
-            torch.nn.init.zeros_(module.layer_norm_bias)
-        init_fn_(module.fc1_weight)
-        if module.fc1_bias is not None:
-            assert isinstance(module.fc1_bias, torch.Tensor)
-            torch.nn.init.zeros_(module.fc1_bias)
-        init_fn_(module.fc2_weight)
-        if module.fc2_bias is not None:
-            assert isinstance(module.fc2_bias, torch.Tensor)
-            torch.nn.init.zeros_(module.fc2_bias)
-        with torch.no_grad():
-            module.fc2_weight.div_(div_is_residual)
     else:
         for _ in module.parameters(recurse=False):
             raise NotImplementedError(f'{module.__class__.__name__} parameters are not initialized by param_init_fn.')
 
-def _normal_init_(std: float, mean: float=0.0) -> Callable:
+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, **kwargs: Any) -> None:
+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)
-    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)
+    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)
 
-def baseline_param_init_fn_(module: nn.Module, init_std: Optional[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, **kwargs: Any) -> None:
+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)
+    _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, **kwargs: Any) -> None:
+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)
+    _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, **kwargs: Any) -> None:
+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)
@@ -155,25 +148,34 @@ 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)
-    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)
+    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)
 
-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', **kwargs: Any) -> None:
+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)
+    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', **kwargs: Any) -> None:
+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)
+    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, **kwargs: Any) -> None:
+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)
-    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)
+    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)
 
-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, **kwargs: Any) -> None:
-    del 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)
-    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)
+    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.py

@@ -1,22 +0,0 @@
-class VersionedDeprecationWarning(DeprecationWarning):
-    """A custom deprecation warning class that includes version information.
-
-    Attributes:
-        message (str): The deprecation message describing why the feature is deprecated.
-        remove_version (str): The version in which the feature will be removed.
-
-    Example:
-        >>> def deprecated_function():
-        ...     warnings.warn(
-        ...         VersionedDeprecationWarning(
-        ...             "Function XYZ is deprecated.",
-        ...             after_version="2.0.0"
-        ...         )
-        ...     )
-        ...
-        >>> deprecated_function()
-        DeprecationWarning: Function XYZ is deprecated. It will be removed in version 2.0.0.
-    """
-
-    def __init__(self, message: str, remove_version: str) -> None:
-        super().__init__(message + f' It will be removed in version {remove_version}.')