[JAX] Unifying GeLU and GeGLU in LayerNorm MLP by phu0ngng · Pull Request #765 · NVIDIA/TransformerEngine

phu0ngng · 2024-04-09T21:12:50Z

This PR unifies the GeLU and GEGLU implementations in LayerNormMLP via a generalized fused_layernorm_fp8_mlp. Previously, there were two separate APIs for the two mentioned activations. The new routine takes activation_type: Tuple and use_bias: bool as two additional arguments, compared to old routines.

This is a preparation step for adding more activations (i.e. swiglu) later.

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

test_layer.py, and test_praxis_layer.py Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

zlsh80826 · 2024-04-10T06:55:13Z

/te-ci jax

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

denera

LGTM!

denera · 2024-04-15T21:22:16Z

@phu0ngng do you have CI permissions yet? If not, please check in with @ptrendx to get permissions and then trigger the CI run for this. We can merge once the tests come back clean.

denera · 2024-04-15T21:23:25Z

@phu0ngng Also please fix the linting errors here before the CI. Thanks!

************* Module transformer_engine.jax.flax.module
transformer_engine/jax/flax/module.py:690:26: W1401: Anomalous backslash in string: '\m'. String constant might be missing an r prefix. (anomalous-backslash-in-string)
transformer_engine/jax/flax/module.py:690:42: W1401: Anomalous backslash in string: '\s'. String constant might be missing an r prefix. (anomalous-backslash-in-string)
transformer_engine/jax/flax/module.py:690:48: W1401: Anomalous backslash in string: '\m'. String constant might be missing an r prefix. (anomalous-backslash-in-string)
transformer_engine/jax/flax/module.py:690:66: W1401: Anomalous backslash in string: '\e'. String constant might be missing an r prefix. (anomalous-backslash-in-string)
transformer_engine/jax/flax/module.py:691:17: W1401: Anomalous backslash in string: '\g'. String constant might be missing an r prefix. (anomalous-backslash-in-string)
transformer_engine/jax/flax/module.py:696:51: W1401: Anomalous backslash in string: '\g'. String constant might be missing an r prefix. (anomalous-backslash-in-string)
transformer_engine/jax/flax/module.py:702:64: W1401: Anomalous backslash in string: '\g'. String constant might be missing an r prefix. (anomalous-backslash-in-string)

zlsh80826

LGTM

mingxu1067

Kindly remove the unnessary comment in tests/jax/test_custom_call_compute.py#332

mingxu1067

Kindly remove the unnessary comment in tests/jax/test_custom_call_compute.py#332

phu0ngng · 2024-04-19T00:16:13Z

/te-ci jax

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Co-authored-by: Alp Dener <adener@nvidia.com> Signed-off-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

phu0ngng · 2024-04-22T17:38:07Z

/te-ci jax

phu0ngng · 2024-04-22T21:13:41Z

Hi @denera, @mingxu1067,
I resolved all of your change requests.
Please have a look and let me know if you have any other suggestions.

mingxu1067

LGTM

zlsh80826 · 2024-04-26T03:40:02Z

Congratulations on your first pull request! @phu0ngng
This really help the future maintainance for various activation types!

* combined layernorm_geglu with layernorm_gelu into fused_layernorm Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * fixes to pass all unit tests in test_custom_call_compute.py, test_layer.py, and test_praxis_layer.py Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * cleaning and formatting Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * renaming based on reviewers suggestions Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * implemented partial fused layernorm Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * geglu + bias passed tests Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * added partial fused calculation for dbias_1 Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * clean up Co-authored-by: Alp Dener <adener@nvidia.com> Signed-off-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Signed-off-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com> Co-authored-by: Alp Dener <adener@nvidia.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

denera requested review from denera, mingxu1067 and zlsh80826 April 9, 2024 21:28

denera assigned phu0ngng Apr 9, 2024

denera added enhancement New feature or request jax labels Apr 9, 2024

denera linked an issue Apr 9, 2024 that may be closed by this pull request
[JAX] Support fused SwiGLU MLP #708
Closed

phu0ngng added 2 commits April 9, 2024 22:44

combined layernorm_geglu with layernorm_gelu into fused_layernorm
178dab3
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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 force-pushed the main branch from 01e315d to a2e06a0Compare April 9, 2024 22:51

cleaning and formatting
61564c4
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng force-pushed the main branch from 1e0d0e4 to 61564c4Compare April 9, 2024 23:08

zlsh80826 reviewed Apr 10, 2024
View reviewed changes

Comment threadtransformer_engine/jax/mlp.py Outdated
Comment threadtransformer_engine/jax/mlp.py Outdated
Comment threadtransformer_engine/jax/mlp.py Outdated

ptrendx reviewed Apr 10, 2024
View reviewed changes

Comment threadtests/jax/test_custom_call_compute.py Outdated

denera reviewed Apr 12, 2024
View reviewed changes

Comment threadtransformer_engine/jax/mlp.py Outdated

renaming based on reviewers suggestions
de80d42
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng mentioned this pull request Apr 12, 2024
[JAX] SwiGLU Implementation #773
Merged

phu0ngng removed a link to an issue Apr 15, 2024
[JAX] Support fused SwiGLU MLP #708
Closed

denera approved these changes Apr 15, 2024
View reviewed changes

zlsh80826 approved these changes Apr 17, 2024
View reviewed changes

mingxu1067 approved these changes Apr 17, 2024
View reviewed changes

Comment threadtransformer_engine/jax/flax/module.py Outdated
Comment threadtests/jax/test_custom_call_compute.py Outdated
Comment threadtests/jax/test_custom_call_compute.py Outdated
Comment threadtransformer_engine/jax/mlp.py Outdated

mingxu1067 suggested changes Apr 17, 2024
View reviewed changes

denera requested changes Apr 19, 2024
View reviewed changes

Comment threadtests/jax/test_custom_call_compute.py Outdated

phu0ngng added 3 commits April 22, 2024 16:38

implemented partial fused layernorm
262a07f
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

geglu + bias passed tests
7d74768
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

added partial fused calculation for dbias_1
b497765
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngngand others added 2 commits April 22, 2024 16:50

rs conflicts
79a6b1e
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

clean up
a10f543
Co-authored-by: Alp Dener <adener@nvidia.com> Signed-off-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

mingxu1067 approved these changes Apr 22, 2024
View reviewed changes

denera approved these changes Apr 24, 2024
View reviewed changes

denera merged commit dac0001 into NVIDIA:mainApr 24, 2024

zlsh80826 mentioned this pull request May 2, 2024
[JAX] Generalizing Activation Primitives #810
Merged

Uh oh!

Conversation

phu0ngng commented Apr 9, 2024

Uh oh!

zlsh80826 commented Apr 10, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

denera left a comment

Choose a reason for hiding this comment

Uh oh!

denera commented Apr 15, 2024

Uh oh!

denera commented Apr 15, 2024

Uh oh!

zlsh80826 left a comment

Choose a reason for hiding this comment

Uh oh!

mingxu1067 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mingxu1067 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phu0ngng commented Apr 19, 2024

Uh oh!

Uh oh!

phu0ngng commented Apr 22, 2024

Uh oh!

phu0ngng commented Apr 22, 2024

Uh oh!

mingxu1067 left a comment

Choose a reason for hiding this comment

Uh oh!

zlsh80826 commented Apr 26, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

mingxu1067 left a comment •
edited
Loading

mingxu1067 left a comment •
edited
Loading