add support for Falcon LLM

models/falcon/7b_config.json

@@ -0,0 +1,21 @@
+{
+  "emb_size": 4544,
+  "feedforward_size": 18176,
+  "hidden_size": 4544,
+  "hidden_act": "gelu_fast",
+  "heads_num": 71,
+  "layers_num": 32,
+  "dropout": 0.0,
+  "data_processor": "lm",
+  "max_seq_length": 2048,
+  "embedding": ["word"],
+  "remove_transformer_bias": true,
+  "remove_embedding_layernorm": true,
+  "attention": "flash_attention",
+  "encoder": "transformer",
+  "mask": "causal",
+  "layernorm_positioning": "parallel_attn",
+  "layernorm": "normal_torch",
+  "layernorm_eps": 1e-5,
+  "target": ["lm"]
+}

models/falcon_special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>",
+  "cls_token": "<|endoftext|>",
+  "sep_token": "<|endoftext|>",
+  "mask_token": "<mask>"
+}

preprocess.py

@@ -39,6 +39,7 @@ def main():
                         help="Probability of truncating sequence."
                              "The larger value, the higher probability of using short (truncated) sequence.")
     parser.add_argument("--full_sentences", action="store_true", help="Full sentences.")
+    parser.add_argument("--remove_cls_token", action="store_true", help="Preprocess without CLS tokens.")
     parser.add_argument("--seed", type=int, default=7, help="Random seed.")
 
     # Masking options.

scripts/convert_falcon_from_huggingface_to_tencentpretrain.py

@@ -0,0 +1,56 @@
+import argparse
+import collections
+import torch
+import os
+import json
+
+
+parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+parser.add_argument("--input_model_path", type=str, default="models/falcon-7b/",
+                    help=".")
+parser.add_argument("--output_model_path", type=str, default="models/falcon-7b.bin",
+                    help=".")
+parser.add_argument("--layers_num", type=int, default=32,
+                    help=".")
+
+args = parser.parse_args()
+
+files = os.listdir(args.input_model_path)
+model_files = [f for f in files if f[-4:] == ".bin"]
+input_models = {f: torch.load(os.path.join(args.input_model_path, f), map_location="cpu") for f in model_files}
+
+with open(os.path.join(args.input_model_path, "pytorch_model.bin.index.json")) as f:
+    model_index = json.load(f)
+    weight_map = model_index["weight_map"]
+
+
+output_model = collections.OrderedDict()
+
+def get_weight_from_name(layer_name):
+    return input_models[weight_map[layer_name]][layer_name]
+
+
+output_model["embedding.word.embedding.weight"] = get_weight_from_name("transformer.word_embeddings.weight")
+
+for i in range(args.layers_num):
+
+    output_model["encoder.transformer." + str(i) + ".layer_norm_1.weight"] = \
+        get_weight_from_name("transformer.h." + str(i) + ".input_layernorm.weight")
+    output_model["encoder.transformer." + str(i) + ".layer_norm_1.bias"] = \
+        get_weight_from_name("transformer.h." + str(i) + ".input_layernorm.bias")
+
+    output_model["encoder.transformer." + str(i) + ".self_attn.query_key_value.weight"] = \
+        get_weight_from_name("transformer.h." + str(i) + ".self_attention.query_key_value.weight")
+    output_model["encoder.transformer." + str(i) + ".self_attn.dense.weight"] = \
+        get_weight_from_name("transformer.h." + str(i) + ".self_attention.dense.weight")
+
+    output_model["encoder.transformer." + str(i) + ".feed_forward.linear_1.weight"] = \
+        get_weight_from_name("transformer.h." + str(i) + ".mlp.dense_h_to_4h.weight")
+    output_model["encoder.transformer." + str(i) + ".feed_forward.linear_2.weight"] = \
+        get_weight_from_name("transformer.h." + str(i) + ".mlp.dense_4h_to_h.weight")
+
+output_model["encoder.layer_norm.weight"] = get_weight_from_name("transformer.ln_f.weight")
+output_model["encoder.layer_norm.bias"] = get_weight_from_name("transformer.ln_f.bias")
+output_model["target.lm.output_layer.weight"] = get_weight_from_name("lm_head.weight")
+
+torch.save(output_model, args.output_model_path)

scripts/generate_lm.py

@@ -79,11 +79,13 @@ def top_k_top_p_filtering(logits, top_k, top_p):
 
     args = load_hyperparam(args)
 
+    args.tokenizer_type = args.tokenizer
     args.tokenizer = str2tokenizer[args.tokenizer](args)
 
     model = GenerateLm(args)
     model = load_model(model, args.load_model_path)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     model = model.to(device)
 
     model.eval()
@@ -110,6 +112,8 @@ def top_k_top_p_filtering(logits, top_k, top_p):
 
         f.write(line + "\n")
         tokens = [token_id.item() for token_id in src_tensor[0]]
+        if args.tokenizer_type in ["hfpretrained"]:
+            generated_sentence = args.tokenizer.decode(tokens)
         if args.tokenizer.sp_model is not None:
             generated_sentence = args.tokenizer.sp_model.decode(tokens)
         else:

tencentpretrain/encoders/transformer_encoder.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 from tencentpretrain.utils.rope import precompute_freqs_cis
 from tencentpretrain.layers.transformer import TransformerLayer
-from tencentpretrain.layers.layer_norm import *
+from tencentpretrain.layers import *
 from tencentpretrain.layers.relative_position_embedding import RelativePositionEmbedding
 
 class TransformerEncoder(nn.Module):
@@ -36,13 +36,8 @@ def __init__(self, args):
             self.transformer = nn.ModuleList(
                 [TransformerLayer(args) for _ in range(self.layers_num)]
             )
-        if self.layernorm_positioning == "pre":
-            if args.layernorm == "t5":
-                self.layer_norm = T5LayerNorm(args.hidden_size)
-            elif args.layernorm == "rms":
-                self.layer_norm = RMSNorm(args.hidden_size)
-            else:
-                self.layer_norm = LayerNorm(args.hidden_size)
+
+        self.layer_norm = str2layernorm[args.layernorm](args.hidden_size, eps=args.layernorm_eps)
 
         if self.relative_position_embedding:
             self.relative_pos_emb = RelativePositionEmbedding(bidirectional=True, heads_num=args.heads_num,
@@ -143,7 +138,7 @@ def custom_forward(*inputs):
                                                             has_residual_attention=self.has_residual_attention,
                                                             prev_attn=prev_attn, freqs_cis=freqs_cis)
 
-        if self.layernorm_positioning == "pre":
+        if self.layernorm_positioning in ["pre", "parallel_attn"]:
             return self.layer_norm(hidden)
         else:
             return hidden

tencentpretrain/layers/__init__.py

@@ -0,0 +1,15 @@
+from tencentpretrain.layers.layer_norm import *
+from tencentpretrain.layers.multi_headed_attn import *
+from tencentpretrain.layers.position_ffn import *
+import torch.nn as nn
+
+str2layernorm = {"t5": T5LayerNorm, "rms": RMSNorm, "normal_torch": nn.LayerNorm, "normal": LayerNorm}
+
+str2attention = {"multi_head": MultiHeadedAttention, "flash_attention": FlashAttention}
+
+str2feedforward = {"gated": GatedFeedForward, "dense": PositionwiseFeedForward}
+
+__all__ = ["T5LayerNorm", "RMSNorm", "LayerNorm", "MultiHeadedAttention",
+           "FlashAttention", "GatedFeedForward", "PositionwiseFeedForward", "str2layernorm",
+           "str2attention", "str2feedforward"]
+

tencentpretrain/layers/multi_headed_attn.py

@@ -1,8 +1,10 @@
 import math
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
 from tencentpretrain.utils.rope import apply_rotary_emb
 from tencentpretrain.utils.lora import LoraLinear
+from tencentpretrain.utils.rope import RotaryEmbedding
 
 class MultiHeadedAttention(nn.Module):
     """
@@ -88,3 +90,141 @@ def unshape(x):
         output = unshape(torch.matmul(probs, value))
         output = self.final_linear(output)
         return output, prev_attn_out
+
+
+class FlashAttention(nn.Module):
+    """
+    Flash Attention used in Falcon.
+    https://huggingface.co/tiiuae/falcon-7b/blob/main/modelling_RW.py#L154
+    """
+
+    def __init__(self, hidden_size, heads_num, attention_head_size, dropout, has_bias=True, with_scale=True,
+                 lora_params=None):
+        super(FlashAttention, self).__init__()
+        self.heads_num = heads_num
+        self.hidden_size = hidden_size
+        self.per_head_size = attention_head_size
+        self.with_scale = with_scale
+        self.inner_hidden_size = heads_num * attention_head_size
+        self.multi_query = True
+
+
+        self.rotary = RotaryEmbedding(self.per_head_size)
+        # Layer-wise attention scaling
+        self.inv_norm_factor = 1.0 / math.sqrt(self.per_head_size)
+        self.beta = self.inv_norm_factor
+
+        self.query_key_value = nn.Linear(
+            self.hidden_size,
+            3 * self.hidden_size if not self.multi_query else (self.hidden_size + 2 * self.per_head_size),
+            bias=has_bias
+        )
+
+        self.dense = nn.Linear(self.hidden_size, self.hidden_size, bias=has_bias)
+        self.attention_dropout = nn.Dropout(dropout)
+        self.num_kv = self.heads_num if not self.multi_query else 1
+
+    def _split_heads(self, fused_qkv: torch.Tensor):
+        """
+        Split the last dimension into (num_heads, head_dim) without making any copies, results share same memory
+        storage as `fused_qkv`
+        Args:
+            fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim]
+        Returns:
+            query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim]
+            value: [batch_size, seq_length, num_heads, head_dim]
+        """
+        if not self.multi_query:
+            batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
+            fused_qkv = fused_qkv.view(batch_size, seq_length, self.heads_num, 3, self.per_head_size)
+            return fused_qkv[..., 0, :], fused_qkv[..., 1, :], fused_qkv[..., 2, :]
+        else:
+            batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
+            fused_qkv = fused_qkv.view(batch_size, seq_length, self.heads_num + 2, self.per_head_size)
+            return fused_qkv[..., :-2, :], fused_qkv[..., [-2], :], fused_qkv[..., [-1], :]
+
+    def _merge_heads(self, x: torch.Tensor):
+        """
+        Merge heads together over the last dimenstion
+        Args:
+            x: (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim]
+        Returns:
+            torch.tensor: [batch_size, seq_length, num_heads * head_dim]
+        """
+        # What we want to achieve is:
+        # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim
+        batch_size_and_num_heads, seq_length, _ = x.shape
+        batch_size = batch_size_and_num_heads // self.heads_num
+
+        # First view to decompose the batch size
+        # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim
+        x = x.view(batch_size, self.heads_num, seq_length, self.per_head_size)
+
+        # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim
+        x = x.permute(0, 2, 1, 3)
+
+        # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim
+        return x.reshape(batch_size, seq_length, self.heads_num * self.per_head_size)
+
+    def forward(self, key, value, query, mask, position_bias=None, has_residual_attention=False, prev_attn=None,
+                freqs_cis=None):
+        """
+        Args:
+            key: [batch_size x seq_length x hidden_size]
+            value: [batch_size x seq_length x hidden_size]
+            query: [batch_size x seq_length x hidden_size]
+            mask: [batch_size x 1 x seq_length x seq_length]
+            position_bias: [1 x heads_num x seq_length x seq_length]
+        Returns:
+            output: [batch_size x seq_length x hidden_size]
+        """
+
+        fused_qkv = self.query_key_value(query)
+
+        # 3 x [batch_size, seq_length, num_heads, head_dim]
+        (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)
+        batch_size, q_length, _, _ = query_layer.shape
+
+        query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.heads_num, q_length, self.per_head_size)
+
+        key_layer = key_layer.transpose(1, 2).reshape(
+            batch_size * self.num_kv,
+            q_length,
+            self.per_head_size,
+            )
+        value_layer = value_layer.transpose(1, 2).reshape(batch_size * self.num_kv, q_length, self.per_head_size)
+
+        query_layer, key_layer = self.rotary(query_layer, key_layer)
+
+        _, kv_length, _ = key_layer.shape
+
+        query_layer_ = query_layer.reshape(batch_size, self.heads_num, -1, self.per_head_size)
+        key_layer_ = key_layer.reshape(batch_size, self.num_kv, -1, self.per_head_size)
+        value_layer_ = value_layer.reshape(batch_size, self.num_kv, -1, self.per_head_size)
+
+        if torch.__version__ < "2.0.0":
+            scores = torch.matmul(query_layer_, key_layer_.transpose(-2, -1))
+            if self.with_scale:
+                scores = scores / math.sqrt(float(self.per_head_size))
+            scores = scores + mask.type_as(scores)
+            prev_attn_out = None
+            if has_residual_attention:
+                if prev_attn is not None:
+                    scores += prev_attn
+                prev_attn_out = scores
+            probs = nn.Softmax(dim=-1)(scores)
+            attn_output = probs @ value_layer_
+
+        else:
+            prev_attn_out = None
+            attn_output = F.scaled_dot_product_attention(
+                query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True
+            )
+
+        x = attn_output.view(batch_size, self.heads_num, q_length, self.per_head_size)
+        x = x.permute(0, 2, 1, 3)
+        attn_output = x.reshape(batch_size, q_length, self.heads_num * self.per_head_size)
+
+        output_tensor = self.dense(attn_output)
+
+        return output_tensor, prev_attn_out