Add support for GPT-NeoX (Pythia) (#50)

2023-04-28 00:32:10 -07:00
parent aa50b17ca7
commit a96d63c21d
9 changed files with 436 additions and 71 deletions
--- a/cacheflow/models/opt.py
+++ b/cacheflow/models/opt.py
@@ -327,4 +327,4 @@ class OPTForCausalLM(nn.Module):

    def initialize_dummy_weights(self) -> None:
        for param in self.state_dict().values():
-            param.data.uniform_(-0.1, 0.1)
+            param.data.uniform_(-1e-3, 1e-3)