Skip to content

Commit

Permalink
add swiglu operator
Browse files Browse the repository at this point in the history
  • Loading branch information
sneaxiy committed Feb 5, 2024
1 parent 8d6e813 commit bb157ff
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions paddlenlp/transformers/llama/modeling.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,15 @@
except ImportError:
fused_rotary_position_embedding = None

try:
from paddle.incubate.nn.functional import swiglu
except ImportError:

def swiglu(x, y=None):
if y is None:
x, y = paddle.chunk(x, chunks=2, axis=-1)
return F.silu(x) * y


from paddle.utils import try_import

Expand Down Expand Up @@ -568,10 +577,10 @@ def __init__(self, config):

def forward(self, x):
if self.fuse_attention_ffn:
gate_out, up_out = paddle.chunk(self.gate_up_fused_proj(x), chunks=2, axis=-1)
out = self.down_proj(F.silu(gate_out) * up_out)
x = swiglu(self.gate_up_fused_proj(x))
else:
out = self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
x = swiglu(self.gate_proj(x), self.up_proj(x))
out = self.down_proj(x)
return out


Expand Down

0 comments on commit bb157ff

Please sign in to comment.