Add SwiGLU support

[gordicaleksa]

Implemented SwiGLU - swish GLU activation function from the "GLU Variants Improve Transformer" paper.

Note: there is an increase in memory footprint as a consequence of adding an additional FC layer (params/grads/optimizer states + 2 new activation buffers (L+1)*B*T*4*C in total for the activations). I'm sure we can optimize this, can be done in the subsequent PR as well.

Tests:
I ran an A/B experiment: trained a 124M GPT-2 on 10B tokens (FineWeb subset) with:
a) GELU
b) SwiGLU (note: currently SwiGLU model actually has 152M params)

Results:

Conclusion: SwiGLU converges to a lower loss, so I'm confident the implementation is correct, but unclear whether the perf comes from SwiGLU or more params (152M vs 124M).

After normalizing for the number of params (multiplying the inner FFN module dimension by 2/3) I get:

Conclusion: SwiGLU does start converging faster than GELU but they end up at a pretty much same loss. Actually the trend seems to go in favor of GELU looking at the gradient of the curves.

Given that SwiGLU complicates the code, it's unclear to me that it offers a clear advantage over GELU/ReLU/etc. In general i doubt any of these activation functions are more powerful than ReLU. You just want more params / compute. :)

Next steps (can be done in follow-up PRs):

• Reduce memory footprint
• Speed up kernels / logic

Appendix:
In case it helps here is a diagram I drew to make it easier for me to implement the backward pass:

[gordicaleksa]

PR is ready and will be merged into the LLaMA 3 fork.