Merge pull request #9 from brianhie/enable_gelu

feat: option for gelu act
togethercomputer · Feb 18, 2024 · a73336d · a73336d
2 parents 350b4a4 + ca6401c
commit a73336d
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/layers.py b/src/layers.py
@@ -41,7 +41,14 @@ def __init__(
         super().__init__()
 
         multiple_of = config.get("inner_size_multiple_of", 64)
-        self.act = F.silu
+        self.act_type = config.get("mlp_activation", "silu")
+        if self.act_type == "gelu":
+            self.act = F.gelu
+        elif self.act_type == "silu":
+            self.act = F.silu
+        else:
+            raise NotImplementedError
+
         self.multiple_of = multiple_of * config.model_parallel_size
 
         inner_size = int(2 * config.hidden_size * 4 / 3)