[Perf] Reduce peak memory usage of llama (vllm-project#10339)

Signed-off-by: andoorve <[email protected]> Signed-off-by: Maxime Fournioux <[email protected]>
mfournioux · Nov 20, 2024 · 32005ce · 32005ce
1 parent 7a2b6c6
commit 32005ce
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vllm/model_executor/models/llama.py b/vllm/model_executor/models/llama.py
@@ -90,8 +90,8 @@ def __init__(
         self.act_fn = SiluAndMul()
 
     def forward(self, x):
-        gate_up, _ = self.gate_up_proj(x)
-        x = self.act_fn(gate_up)
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
         x, _ = self.down_proj(x)
         return x