Compile GLU (#1049)

mosaicml · Mar 22, 2024 · 26a5fd4 · 26a5fd4
1 parent 3348b59
commit 26a5fd4
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/llmfoundry/models/layers/ffn.py b/llmfoundry/models/layers/ffn.py
@@ -144,6 +144,7 @@ def __init__(
             **self.fc_kwargs,
         )
 
+    @torch.compile
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.down_proj(self.act(self.gate_proj(x)) * self.up_proj(x))