Skip to content

[JAX] Support SP + RoPE + GeLU - #602

Merged
denera merged 20 commits into
NVIDIA:mainfrom
mingxu1067:mingh/sp_rope_gelu
Feb 2, 2024
Merged

[JAX] Support SP + RoPE + GeLU#602
denera merged 20 commits into
NVIDIA:mainfrom
mingxu1067:mingh/sp_rope_gelu

Conversation

@mingxu1067

@mingxu1067mingxu1067 commented Jan 15, 2024

Copy link
Copy Markdown
Collaborator
  • Support Sqeuence Parallelism
  • Support Rotary Position Embedding (RoPE).
  • Port GeLU relevant custom calls.

@mingxu1067
Adding support of sequence parallelism
3a404ca
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding RoPE
604dcf2
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Fix wrong batch_logical_axes
426050d
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Rnaming FSDP outer env var
6f6a927
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Poring RoPE to Praxis layers.
092205f
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Porting GeLU + [FP8 Cast].
66a9b9c
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
WAR to make XLA successfully match FP8 GEMM on FFN1 with GeLU.
0e5d703
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Allowing arbitrary dimension of NVShape for the workspace allocation
6693bad
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Adding checkpoint_name to fused functions of mlp.py to get better per…
90793a9
…f with nn.scan.
Signed-off-by: Ming Huang <mingh@nvidia.com>
@dathudeptraidathudeptrai mentioned this pull request Jan 15, 2024
Closed
nouiz
nouiz reviewed Jan 15, 2024
View reviewed changes

@nouiznouiz 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.

First round of review.

Comment threadtests/jax/test_praxis_layers.py
Comment threadtransformer_engine/jax/csrc/modules.cpp Outdated
Comment threadtransformer_engine/jax/praxis/transformer.py Outdated
Comment threadtests/jax/test_custom_call_compute.py
Comment threadtests/jax/test_custom_call_compute.py Outdated
Comment threadtests/jax/test_custom_call_compute.py
Comment threadtests/jax/test_custom_call_compute.py
zlsh80826
zlsh80826 reviewed Jan 15, 2024
View reviewed changes
Comment threadtests/jax/test_custom_call_compute.py Outdated
Comment threadtransformer_engine/jax/flax/transformer.py Outdated
Comment threadtransformer_engine/jax/flax/transformer.py
@MoFHekaMoFHeka mentioned this pull request Jan 15, 2024
Closed
2 tasks
@mingxu1067
Modify with review feedback.
e7a447a
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Merge remote-tracking branch 'upstream/main' into mingh/sp_rope_gelu
210a8c0
Signed-off-by: Ming Huang <mingh@nvidia.com>
@zlsh80826
zlsh80826 self-requested a review January 17, 2024 07:08
@zlsh80826

zlsh80826 commented Jan 17, 2024

Copy link
Copy Markdown
Collaborator

LGTM

zlsh80826
zlsh80826 reviewed Jan 18, 2024
View reviewed changes
Comment threadtransformer_engine/jax/flax/transformer.py Outdated
@mingxu1067
Fix bugs
3b2e231
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Fix typo.
780a5f6
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067
Fixed for lint
6bcdaf6
Signed-off-by: Ming Huang <mingh@nvidia.com>
@zlsh80826
zlsh80826 self-requested a review January 19, 2024 09:15
nouiz
nouiz reviewed Jan 19, 2024
View reviewed changes
Comment threadtests/jax/utils.py Outdated
Comment threadtransformer_engine/jax/layernorm.py Outdated
Comment threadtransformer_engine/jax/sharding.py
nouiz
nouiz reviewed Jan 19, 2024
View reviewed changes

@nouiznouiz 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.

Can you update one of the example to show SP?
A simple modification, like just a parameter to the example to use or not SP would be great.

Comment threadtransformer_engine/jax/flax/module.py
Comment threadtransformer_engine/jax/flax/module.py
Comment threadtransformer_engine/jax/flax/module.py
Comment threadtransformer_engine/jax/layernorm.py Outdated
Comment threadtransformer_engine/jax/layernorm.py
Comment threadtransformer_engine/jax/flax/transformer.py
Comment threadtransformer_engine/jax/sharding.py
Comment threadtransformer_engine/jax/sharding.py
Comment threadtests/jax/test_custom_call_compute.py
jnp.asarray(ref_s_grad, np.float32),
dtype=FP8Helper.BWD_DTYPE)

@pytest.mark.skipif(not is_fp8_supported, reason=reason)

@nouiznouizJan 19, 2024

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.

Any test for vmap?

@mingxu1067mingxu1067Jan 22, 2024

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

Not implemented yet, but tests of vmap is orthogonal to this PR. We can schedule time to submit another PR to do so.

@mingxu1067
Follow review feedback to modify code.
6414b29
Signed-off-by: Ming Huang <mingh@nvidia.com>
@mingxu1067

mingxu1067 commented Jan 22, 2024

Copy link
Copy Markdown
CollaboratorAuthor

Can you update one of the example to show SP? A simple modification, like just a parameter to the example to use or not SP would be great.

Done, refer to examples/jax/encoder/test_model_parallel_encoder.py

nouiz
nouiz reviewed Jan 22, 2024
View reviewed changes
Comment threadexamples/jax/encoder/test_model_parallel_encoder.py
nouiz
nouiz reviewed Jan 22, 2024
View reviewed changes
Comment threaddocs/api/jax.rst Outdated
Comment threaddocs/api/jax.rst Outdated
@mingxu1067
Fix typo.
f685393
Signed-off-by: Ming Huang <mingh@nvidia.com>
@denera

denera commented Jan 25, 2024

Copy link
Copy Markdown
Collaborator

/te-ci jax

@nouiznouiz mentioned this pull request Jan 27, 2024
Closed
@mingxu1067

mingxu1067 commented Jan 30, 2024

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci

@nouiznouiz mentioned this pull request Jan 31, 2024
Closed
@mingxu1067
Port SP to Praxis
2ea411a
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
@mingxu1067
Fix an issue when enabling both GQA and RoPE.
45993f6
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
@mingxu1067
Update docs
8e7dcdf
Signed-off-by: Ming-Xu Huang <mingh@nvidia.com>
@mingxu1067

mingxu1067 commented Jan 31, 2024

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci jax

@denera

denera commented Feb 2, 2024

Copy link
Copy Markdown
Collaborator

@ashors1 verified partial convergence and Abhinav Goal verified performance. Merging PR -- thank you for the hard work @mingxu1067!

@denera
denera merged commit ce163f9 into NVIDIA:mainFeb 2, 2024
@irhumirhum mentioned this pull request Mar 10, 2024
Closed
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@deneradeneradenera left review comments
+2 more reviewers
@nouiznouiznouiz left review comments
@zlsh80826zlsh80826zlsh80826 approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@mingxu1067@zlsh80826@denera@nouiz