Skip to content

[JAX] SwiGLU Implementation - #773

Merged
phu0ngng merged 20 commits into
NVIDIA:mainfrom
phu0ngng:swiglu
Apr 24, 2024
Merged

[JAX] SwiGLU Implementation#773
phu0ngng merged 20 commits into
NVIDIA:mainfrom
phu0ngng:swiglu

Conversation

@phu0ngng

@phu0ngngphu0ngng commented Apr 12, 2024•
edited
Loading

Copy link
Copy Markdown
Collaborator

This PR adds SwiGLU implementation for JAX.
In addition, all previous activation functions were templated to shorten the codebase.
Tests in test_custom_call_compute.py were adapted to be more generalized and address newly added activation.
This PR should be merged after #765.

@phu0ngng
combined layernorm_geglu with layernorm_gelu into fused_layernorm
178dab3
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
fixes to pass all unit tests in test_custom_call_compute.py,
7abd87f
test_layer.py, and test_praxis_layer.py
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
cleaning and formatting
61564c4
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
renaming based on reviewers suggestions
de80d42
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngngphu0ngng added enhancement New feature or request jax labels Apr 12, 2024
@phu0ngng
fixes to pass all unit tests in test_custom_call_compute.py,
4fbf5e1
test_layer.py, and test_praxis_layer.py
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
templating cast_transpose_gelu related rountines
7500ab8
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
implemented swiglu and silu
40dd8d9
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngngphu0ngng linked an issue Apr 15, 2024 that may be closed by this pull request
[JAX] Support fused SwiGLU MLP #708
Closed
denera
denera requested changes Apr 15, 2024
View reviewed changes
Comment threadtransformer_engine/common/activation/activation_template.h Outdated
Comment threadtransformer_engine/common/include/transformer_engine/activation.h Outdated
Comment threadtransformer_engine/common/include/transformer_engine/activation.h Outdated
Comment threadtransformer_engine/common/include/transformer_engine/transpose.h Outdated
Comment threadtransformer_engine/common/transpose/cast_transpose_fusion.cu Outdated
Comment threadtransformer_engine/jax/csrc/extensions.cpp Outdated
Comment threadtransformer_engine/jax/mlp.py Outdated
phu0ngngand others added 8 commits April 22, 2024 16:38
@phu0ngng
implemented partial fused layernorm
262a07f
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
geglu + bias passed tests
7d74768
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
added partial fused calculation for dbias_1
b497765
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
rs conflicts
79a6b1e
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng@denera
clean up
a10f543
Co-authored-by: Alp Dener <adener@nvidia.com>
Signed-off-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>
@phu0ngng
- Renamed nvte-*silu to nvte-*swish + generalized GetDbiasDact function
1050910
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
rs conflicts
72bcdcb
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng
rs conflicts
da5fd0b
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@phu0ngng

phu0ngng commented Apr 22, 2024

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci jax

@phu0ngng
cleanup
1b5d103
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
@denera

denera commented Apr 24, 2024

Copy link
Copy Markdown
Collaborator

@phu0ngng I just merged #765. Could you resolve the conflicts with this PR and re-run the CI? We can merge when it comes back clean. Thanks!

@phu0ngng

phu0ngng commented Apr 24, 2024

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci jax

@phu0ngng

phu0ngng commented Apr 24, 2024

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci jax

denera
denera approved these changes Apr 24, 2024
View reviewed changes

@deneradenera left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@phu0ngng
phu0ngng merged commit f85553e into NVIDIA:mainApr 24, 2024
@phu0ngng
phu0ngng deleted the swiglu branch April 24, 2024 22:38
pggPL pushed a commit to pggPL/TransformerEngine that referenced this pull request May 23, 2024
@phu0ngng@pggPL
[JAX] SwiGLU Implementation (NVIDIA#773)
a06ab9a
* Implemented swiglu and silu
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* Renamed nvte-*silu to nvte-*swish + generalized GetDBiasDact functions
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@deneradeneradenera approved these changes
@mingxu1067mingxu1067Awaiting requested review from mingxu1067
@zlsh80826zlsh80826Awaiting requested review from zlsh80826
@ptrendxptrendxAwaiting requested review from ptrendxptrendx is a code owner

Assignees

No one assigned

Labels

enhancementNew feature or requestjax

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[JAX] Support fused SwiGLU MLP

2 participants

@phu0ngng@denera