Upload GPTOptim

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f2c240204fac1bf66e112ce3be2384a0097a2ea95b57ed2a4896c6cd01ecf5f7
-size 4040701744
+oid sha256:86048048139b2cae7b486e2da9b4abc53112e9290d190f6d23bec864a1fdfa3b
+size 4040722640

modeling_gpt_optimized.py

@@ -1,5 +1,6 @@
 import torch
 import torch.nn as nn
+import bitsandbytes
 from torch.nn import CrossEntropyLoss, functional as F
 from transformers import PreTrainedModel, GPT2PreTrainedModel
 from .configuration_gpt_optimized import GPTOptimConfig
@@ -145,8 +146,8 @@ class GPT(nn.Module):
         self.config = config
 
         self.transformer = nn.ModuleDict(dict(
-            wte = nn.Embedding(config.vocab_size, config.n_embd),
-            wpe = nn.Embedding(config.block_size, config.n_embd),
+            wte = bitsandbytes.nn.StableEmbedding(config.vocab_size, config.n_embd),
+            wpe = bitsandbytes.nn.StableEmbedding(config.block_size, config.n_embd),
             h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
             ln_f = nn.LayerNorm(config.n_embd),
         ))