[JAX] Unifying GeLU and GeGLU in LayerNorm MLP #765

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

transformer_engine/jax/mlp.py

tests/jax/test_custom_call_compute.py

transformer_engine/jax/mlp.py

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

transformer_engine/jax/flax/module.py

tests/jax/test_custom_call_compute.py

transformer_engine/jax/mlp.py

mingxu1067

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

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

/te-ci jax

tests/jax/test_custom_call_compute.py

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 a2e06a0 Compare 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 61564c4 Compare April 9, 2024 23:08

zlsh80826 reviewed Apr 10, 2024

View reviewed changes

transformer_engine/jax/mlp.py Outdated Show resolved Hide resolved

transformer_engine/jax/mlp.py Outdated Show resolved Hide resolved

transformer_engine/jax/mlp.py Outdated Show resolved Hide resolved

ptrendx reviewed Apr 10, 2024

View reviewed changes

tests/jax/test_custom_call_compute.py Outdated Show resolved Hide resolved

denera reviewed Apr 12, 2024

View reviewed changes

transformer_engine/jax/mlp.py Outdated Show resolved Hide resolved

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

transformer_engine/jax/flax/module.py Outdated Show resolved Hide resolved

tests/jax/test_custom_call_compute.py Outdated Show resolved Hide resolved

tests/jax/test_custom_call_compute.py Outdated Show resolved Hide resolved

transformer_engine/jax/mlp.py Outdated Show resolved Hide resolved

mingxu1067 requested changes Apr 17, 2024

View reviewed changes

denera requested changes Apr 19, 2024

View reviewed changes

tests/jax/test_custom_call_compute.py Outdated Show resolved Hide resolved

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>

phu0ngng and 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:main Apr 24, 2024
15 checks passed

zlsh80826 mentioned this pull request May 2, 2024

[JAX] Generalizing Activation Primitives #810

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[JAX] Unifying GeLU and GeGLU in LayerNorm MLP #765

[JAX] Unifying GeLU and GeGLU in LayerNorm MLP #765

phu0ngng commented Apr 9, 2024

zlsh80826 commented Apr 10, 2024

denera left a comment

denera commented Apr 15, 2024

denera commented Apr 15, 2024

zlsh80826 left a comment

mingxu1067 left a comment •

edited

Loading

mingxu1067 left a comment •

edited

Loading

phu0ngng commented Apr 19, 2024

phu0ngng commented Apr 22, 2024

phu0ngng commented Apr 22, 2024

mingxu1067 left a comment

zlsh80826 commented Apr 26, 2024

[JAX] Unifying GeLU and GeGLU in LayerNorm MLP #765

[JAX] Unifying GeLU and GeGLU in LayerNorm MLP #765

Conversation

phu0ngng commented Apr 9, 2024

zlsh80826 commented Apr 10, 2024

denera left a comment

Choose a reason for hiding this comment

denera commented Apr 15, 2024

denera commented Apr 15, 2024

zlsh80826 left a comment

Choose a reason for hiding this comment

mingxu1067 left a comment • edited Loading

Choose a reason for hiding this comment

mingxu1067 left a comment • edited Loading

Choose a reason for hiding this comment

phu0ngng commented Apr 19, 2024

phu0ngng commented Apr 22, 2024

phu0ngng commented Apr 22, 2024

mingxu1067 left a comment

Choose a reason for hiding this comment

zlsh80826 commented Apr 26, 2024

mingxu1067 left a comment •

edited

Loading

mingxu1067 left a comment •

edited

Loading