Upload 9 files

config.json

@@ -0,0 +1,60 @@
+{
+  "alibi_mode": "symetric",
+  "architectures": [
+    "FlashT5ForConditionalGeneration"
+  ],
+  "attention_dropout_rate": 0.0,
+  "attention_scale": 1.0,
+  "attention_type": "ref",
+  "auto_map": {
+    "AutoConfig": "configuration_flash_t5.FlashT5Config",
+    "AutoModel": "modeling_flash_t5.FlashT5ForConditionalGeneration",
+    "AutoModelForQuestionAnswering": "custom_heads_flash_t5.FlashT5ForQuestionAnswering",
+    "AutoModelForSeq2SeqLM": "modeling_flash_t5.FlashT5ForConditionalGeneration",
+    "AutoModelForSequenceClassification": "custom_heads_flash_t5.FlashT5ForSequenceClassification",
+    "AutoModelForTokenClassification": "custom_heads_flash_t5.FlashT5ForTokenClassification"
+  },
+  "classifier_dropout": 0.0,
+  "d_ff": 1024,
+  "d_kv": 64,
+  "d_model": 512,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "relu",
+  "dropout_rate": 0.0,
+  "eos_token_id": 1,
+  "feed_forward_proj": "relu",
+  "fire_mlp_width": 32,
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": false,
+  "is_gated_act": false,
+  "label_smoothing": 0.0,
+  "layer_norm_epsilon": 1e-06,
+  "max_sequence_length": 1024,
+  "model_type": "flash_t5",
+  "num_decoder_layers": 8,
+  "num_heads": 6,
+  "num_layers": 8,
+  "pad_token_id": 3,
+  "position_encoding_type": "t5",
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "rotary_base": 10000,
+  "rotary_emb_fraction": 1.0,
+  "rotary_interleaved": false,
+  "rotary_scale_base": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.42.3",
+  "use_cache": true,
+  "use_flash_attention": "triton",
+  "use_full_bias_size": false,
+  "use_gelu_act": true,
+  "use_glu_mlp": true,
+  "use_masking": false,
+  "use_randomized_position_encoding": false,
+  "use_triton_crossentropy": true,
+  "use_triton_gated_mlp": false,
+  "use_triton_layernorm": true,
+  "vocab_size": 32768,
+  "z_loss": 0.0001
+}

configuration_flash_t5.py

@@ -0,0 +1,84 @@
+import sys
+from collections import OrderedDict
+from typing import Mapping
+import logging
+
+from transformers import T5Config
+
+AUTO_MAP = {
+    "AutoModel": "modeling_flash_t5.FlashT5ForConditionalGeneration",
+    "AutoModelForSeq2SeqLM": "modeling_flash_t5.FlashT5ForConditionalGeneration",
+    "AutoModelForTokenClassification": "custom_heads_flash_t5.FlashT5ForTokenClassification",
+    "AutoModelForQuestionAnswering": "custom_heads_flash_t5.FlashT5ForQuestionAnswering",
+    "AutoModelForSequenceClassification": "custom_heads_flash_t5.FlashT5ForSequenceClassification",
+}
+
+class FlashT5Config(T5Config):
+
+    model_type = "flash_t5"
+
+    def __init__(
+        self,
+        decoder_start_token_id=0,
+        pad_token_id=-100,
+        use_glu_mlp=False,
+        position_encoding_type="t5",
+        use_randomized_position_encoding=False,
+        label_smoothing=0.0,
+        z_loss=None,
+        attention_type="ref",
+        max_sequence_length=1024,
+        attention_dropout_rate=0.0,
+        alibi_mode="symetric",
+        use_triton_layernorm=False,
+        use_triton_crossentropy=False,
+        use_triton_gated_mlp=False,
+        use_gelu_act=True,
+        use_full_bias_size=False,
+        rotary_emb_fraction=1.0,
+        rotary_base=10000,
+        rotary_interleaved=False,
+        rotary_scale_base=None,
+        fire_mlp_width=32,
+        use_masking=False,
+        attention_scale=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.decoder_start_token_id = decoder_start_token_id
+        self.pad_token_id = pad_token_id
+        self.use_glu_mlp = use_glu_mlp
+        self.position_encoding_type = position_encoding_type
+        self.use_randomized_position_encoding = use_randomized_position_encoding
+        self.label_smoothing = label_smoothing
+        self.z_loss = z_loss
+        self.attention_type = attention_type
+        self.max_sequence_length = max_sequence_length
+        self.alibi_mode = alibi_mode
+        self.attention_dropout_rate = attention_dropout_rate
+        self.use_triton_layernorm = use_triton_layernorm
+        self.use_triton_crossentropy = use_triton_crossentropy
+        self.use_triton_gated_mlp = use_triton_gated_mlp
+        self.use_gelu_act = use_gelu_act
+        self.use_full_bias_size = use_full_bias_size
+        self.rotary_base = rotary_base
+        self.rotary_interleaved = rotary_interleaved
+        self.rotary_scale_base = rotary_scale_base
+        self.rotary_emb_fraction = rotary_emb_fraction
+        self.fire_mlp_width = fire_mlp_width
+        self.use_masking = use_masking
+        self.attention_scale = attention_scale
+
+        self.auto_map = AUTO_MAP
+
+def str_to_class(classname):
+    return getattr(sys.modules[__name__], classname)
+
+# Register model in Auto API
+try:
+    FlashT5Config.register_for_auto_class()
+    for key, value in AUTO_MAP.items():
+        str_to_class(value.split(".")[-1]).register_for_auto_class(key)
+except:
+    logging.warn("AutoRegister isn't available.")

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "decoder_start_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 3,
+  "transformers_version": "4.42.3"
+}

optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa38aa9103205dbc3d61430bc913c4a717a85d272fb395bf7f582856369fa572
+size 621095482

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1fc9ba48bdef31e462c54206ad291f9037263b46df389d578dc91b7bc980fc9
+size 310533314

rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f56d72e430dbd41e1d5be36a03fc30b7ed76db0ffce725ffcb770dd624c2c7a
+size 14244

scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1df75838ab9f2b870a5502f483aa0a4f4634e4583f3682dc92e406347f699efa
+size 1256

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16b1124f4756c43674bdaabb10312885654ecc2f3ba68e50e8365d2409fb1a01
+size 5176