Upload FlashSTU

config.json

@@ -17,7 +17,7 @@
   "num_eigh": 24,
   "seq_len": 8192,
   "softcap": 50.0,
-  "torch_dtype": "float32",
+  "torch_dtype": "bfloat16",
   "transformers_version": "4.44.0",
   "use_approx": true,
   "use_flash_fft": true,

model.py

@@ -77,7 +77,7 @@ class FlashSTU(PreTrainedModel):
 
         self.flash_stu = nn.ModuleDict(
             dict(
-                tok_emb=nn.Embedding(self.vocab_size, self.n_embd),
+                tok_emb=nn.Embedding(self.vocab_size, self.n_embd, dtype=config.torch_dtype),
                 dropout=nn.Dropout(self.dropout),
                 hidden=nn.ModuleList(
                     [
@@ -88,7 +88,7 @@ class FlashSTU(PreTrainedModel):
                 rn_f=RMSNorm(config.n_embd, dtype=config.torch_dtype)
             )
         )
-        self.lm_head = nn.Linear(self.n_embd, self.vocab_size, bias=self.bias)
+        self.lm_head = nn.Linear(self.n_embd, self.vocab_size, bias=self.bias, dtype=config.torch_dtype)
 
         self.std = (self.n_embd) ** -0.5
         self.apply(self._init_weights)

model.safetensors

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