[JAX] SwiGLU Implementation - #773

Merged

phu0ngng merged 20 commits into

NVIDIA:mainfrom

phu0ngng:swiglu

Apr 24, 2024

Commits on Apr 9, 2024

combined layernorm_geglu with layernorm_gelu into fused_layernorm
phu0ngng
committed
178dab3
fixes to pass all unit tests in test_custom_call_compute.py,
phu0ngng
committed
7abd87f
cleaning and formatting
phu0ngng
committed
61564c4

Commits on Apr 12, 2024

Commits on Apr 22, 2024

implemented partial fused layernorm
phu0ngng
committed
262a07f
geglu + bias passed tests
phu0ngng
committed
7d74768
added partial fused calculation for dbias_1
phu0ngng
committed
b497765
rs conflicts
phu0ngng
committed
79a6b1e
clean up

phu0ngng
and
denera
authored
a10f543
- Renamed nvte-*silu to nvte-*swish + generalized GetDbiasDact function
phu0ngng
committed
1050910
rs conflicts
phu0ngng
committed
72bcdcb
rs conflicts
phu0ngng
committed
da5fd0b

Commits on Apr 23, 2024

cleanup
phu0ngng
committed
1b5d103

Commits on Apr 24, 2024

Merge branch 'main' of github.com:NVIDIA/TransformerEngine
phu0ngng
committed
ea8477a
rs conflicts
phu0ngng
committed
1730dbb
Merge branch 'main' into swiglu
phu0ngng
authored
26db608
Merge branch 'main' into swiglu
phu0ngng
authored
f6fa8a2