[Paddle] Add parallel support by Tom-Zheng · Pull Request #357 · NVIDIA/TransformerEngine

Tom-Zheng · 2023-08-03T14:27:21Z

Add the following parallel rules:

TP
DP
PP
Group sharding (FSDP)

Tom-Zheng · 2023-08-03T14:27:55Z

@jeng1220 Could you take a look?

jeng1220 · 2023-08-07T13:45:55Z

When using TP, NCCL communication needs to be scheduled before GEMM to guarantee overlap.

jeng1220 · 2023-08-07T13:52:18Z

Although the NCCL is always invoked first, to ensure that the GEMM isn't scheduled first, the environment variable CUDA_DEVICE_MAX_CONNECTIONS needs to be set to 1 to force using a single channel.

Is it a PaddlePaddle limitation?

Same as here. Need to check if such limitation is necessary for Paddle.

I think it is not needed for now. In current implementation, TP collective ops run in the same CUDA stream as compute ops. So there will be no overlapping. We may need this optimization later if we change the implementation to consider overlapping.

Consider replacing CUDA_DEVICE_MAX_CONNECTIONS with cudaEventWait (or cudaStreamWaitEvent) if overlapping setup is needed.

zlsh80826 · 2023-08-07T14:57:48Z

/te-ci

jeng1220 · 2023-08-11T11:11:14Z

This PR needs to wait for #361, and refactoring a bit.

Tom-Zheng · 2023-08-17T09:03:42Z

/te-ci

mingxu1067 · 2023-08-17T09:04:47Z

/te-ci

ksivaman · 2023-08-17T23:08:01Z

@Tom-Zheng #361 has been merged, could you resolve the conflicts and make the necessary refactors?

Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

commit 79e2e5fd774e67dcdda9aae01a9f31a6479c5d70 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:39:16 2023 +0000 Add TP test Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 1d40ad60540490f97ed82ba877cc6eda8902cbf6 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:22:25 2023 +0000 Fix tp_size when disabled Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 6632f735a0c8251862355fc74622af59fae3a509 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 05:52:18 2023 +0000 Add TP for attention and transformer layer Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Tom-Zheng · 2023-08-20T14:46:58Z

@ksivaman @timmoon10 This PR is ready for review. Would you please take a look?

cc: @jeng1220

jeng1220 · 2023-08-21T06:45:54Z

@ksivaman @timmoon10 This PR is ready for review. Would you please take a look?
cc: @jeng1220

Check tensor shapes per GPU carefully if sharding is enabled. To ensure that the every tensor shape per GPU meets the expectation.

Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

timmoon10 · 2023-08-21T17:30:41Z

/te-ci

ksivaman · 2023-08-21T23:31:16Z

/te-ci

Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Tom-Zheng · 2023-08-23T01:53:17Z

@ksivaman @timmoon10 This PR is ready for review. Would you please take a look?
cc: @jeng1220
Check tensor shapes per GPU carefully if sharding is enabled. To ensure that the every tensor shape per GPU meets the expectation.

Done

timmoon10

Overall LGTM, although we might see test failures when we add L40S to the CI.

Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

jeng1220 · 2023-08-25T09:51:46Z

LGTM.
@zlsh80826 please help to trigger CI.

zlsh80826 · 2023-08-25T09:56:05Z

/te-ci

timmoon10

It looks reasonable to me, although we're seeing some test failures related to the MNIST example and the Transformer layer. It's strange since I wouldn't expect this PR to have affected them.

Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Tom-Zheng · 2023-08-26T03:20:00Z

It looks reasonable to me, although we're seeing some test failures related to the MNIST example and the Transformer layer. It's strange since I wouldn't expect this PR to have affected them.

It is due to randomness and is fixed now.

Tom-Zheng · 2023-08-26T03:20:35Z

/te_ci

Tom-Zheng · 2023-08-26T03:21:30Z

@zlsh80826 Would you please help to rerun CI?

zlsh80826 · 2023-08-26T04:08:09Z

/te-ci

ksivaman

Approving for v0.12 with more work to be done

* [Paddle] Add TP, DP, PP, FSDP Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Minor fix Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix CI failure Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Remove set_nccl_overlap_warning_if_tp Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Improve variable naming Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Refactor FP8 Buffer Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Stylic changes Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix FP32 parallel training Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix numel performance issue Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Squashed commit of the following: commit 79e2e5fd774e67dcdda9aae01a9f31a6479c5d70 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:39:16 2023 +0000 Add TP test Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 1d40ad60540490f97ed82ba877cc6eda8902cbf6 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:22:25 2023 +0000 Fix tp_size when disabled Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 6632f735a0c8251862355fc74622af59fae3a509 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 05:52:18 2023 +0000 Add TP for attention and transformer layer Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Add shape check Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Add FSDP check for stage 1,2,3 Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Review changes Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix group_sharding test Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Support NVTE_FUSE_ATTN Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix CI errors Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> --------- Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [Paddle] Add TP, DP, PP, FSDP Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Minor fix Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix CI failure Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Remove set_nccl_overlap_warning_if_tp Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Improve variable naming Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Refactor FP8 Buffer Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Stylic changes Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix FP32 parallel training Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix numel performance issue Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Squashed commit of the following: commit 79e2e5fd774e67dcdda9aae01a9f31a6479c5d70 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:39:16 2023 +0000 Add TP test Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 1d40ad60540490f97ed82ba877cc6eda8902cbf6 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:22:25 2023 +0000 Fix tp_size when disabled Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 6632f735a0c8251862355fc74622af59fae3a509 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 05:52:18 2023 +0000 Add TP for attention and transformer layer Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Add shape check Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Add FSDP check for stage 1,2,3 Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Review changes Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix group_sharding test Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Support NVTE_FUSE_ATTN Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix CI errors Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> --------- Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Jan Bielak <jbielak@nvidia.com>

* [Paddle] Add TP, DP, PP, FSDP Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Minor fix Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix CI failure Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Remove set_nccl_overlap_warning_if_tp Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Improve variable naming Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Refactor FP8 Buffer Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Stylic changes Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix FP32 parallel training Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix numel performance issue Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Squashed commit of the following: commit 79e2e5fd774e67dcdda9aae01a9f31a6479c5d70 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:39:16 2023 +0000 Add TP test Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 1d40ad60540490f97ed82ba877cc6eda8902cbf6 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 14:22:25 2023 +0000 Fix tp_size when disabled Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> commit 6632f735a0c8251862355fc74622af59fae3a509 Author: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Date: Sun Aug 20 05:52:18 2023 +0000 Add TP for attention and transformer layer Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Add shape check Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Add FSDP check for stage 1,2,3 Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Review changes Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix group_sharding test Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Support NVTE_FUSE_ATTN Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> * Fix CI errors Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> --------- Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com> Signed-off-by: Rui Wang <rui@helixon.com>

jeng1220 reviewed Aug 7, 2023
View reviewed changes

Tom-Zheng added 7 commits August 18, 2023 06:04

[Paddle] Add TP, DP, PP, FSDP
da1f48b
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Minor fix
e0f54c0
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Fix CI failure
9b245ca
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Remove set_nccl_overlap_warning_if_tp
5182f04
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Improve variable naming
33a2b4a
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Refactor FP8 Buffer
4bd2696
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Stylic changes
9110267
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Tom-Zheng changed the title ~~[Paddle] Add parallel support~~[WIP] [Paddle] Add parallel supportAug 18, 2023

Tom-Zheng marked this pull request as draft August 18, 2023 06:15

Tom-Zheng force-pushed the gh_add_parallel branch from 61f47ed to 9110267Compare August 18, 2023 06:16

Tom-Zheng added 3 commits August 18, 2023 08:04

Fix FP32 parallel training
b1528ee
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Fix numel performance issue
71a58a3
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Tom-Zheng changed the title ~~[WIP] [Paddle] Add parallel support~~[Paddle] Add parallel supportAug 20, 2023

Tom-Zheng marked this pull request as ready for review August 20, 2023 14:45

Add shape check
494bbf7
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

ksivaman requested review from ksivaman and timmoon10 August 21, 2023 18:39

ksivaman added the 0.12.0 label Aug 21, 2023

Merge branch 'main' into gh_add_parallel
b220ec6

Add FSDP check for stage 1,2,3
62d5ca6
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

timmoon10 reviewed Aug 24, 2023
View reviewed changes

Comment threadtransformer_engine/paddle/distributed.py Outdated
Comment threadtests/paddle/test_parallel.py

timmoon10 self-requested a review August 24, 2023 17:49

Tom-Zheng added 3 commits August 25, 2023 05:37

Review changes
e480c49
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Fix group_sharding test
bf5c01c
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Support NVTE_FUSE_ATTN
e580b89
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

timmoon10 reviewed Aug 25, 2023
View reviewed changes

Tom-Zheng added 2 commits August 26, 2023 03:18

Fix CI errors
e05a80c
Signed-off-by: Tian Zheng (Engrg-Hardware 1) <tizheng@nvidia.com>

Merge branch 'main' into gh_add_parallel
6507cbf

ksivaman approved these changes Aug 26, 2023
View reviewed changes

ksivaman merged commit b8ba734 into NVIDIA:mainAug 26, 2023

hungryGeek16 mentioned this pull request May 31, 2026
fix unfused padding causal sdpa #3063
Open

Uh oh!

Conversation

Tom-Zheng commented Aug 3, 2023

Uh oh!

Tom-Zheng commented Aug 3, 2023

Uh oh!

jeng1220Aug 7, 2023

Choose a reason for hiding this comment

Uh oh!

jeng1220Aug 7, 2023

Choose a reason for hiding this comment

Uh oh!

jeng1220Aug 7, 2023

Choose a reason for hiding this comment

Uh oh!

Tom-ZhengAug 8, 2023• 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!

Tom-ZhengAug 11, 2023• 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!

jeng1220Aug 11, 2023• 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!

zlsh80826 commented Aug 7, 2023

Uh oh!

jeng1220 commented Aug 11, 2023

Uh oh!

Tom-Zheng commented Aug 17, 2023

Uh oh!

mingxu1067 commented Aug 17, 2023

Uh oh!

ksivaman commented Aug 17, 2023

Uh oh!

Tom-Zheng commented Aug 20, 2023

Uh oh!

jeng1220 commented Aug 21, 2023

Uh oh!

timmoon10 commented Aug 21, 2023

Uh oh!

ksivaman commented Aug 21, 2023

Uh oh!

Tom-Zheng commented Aug 23, 2023

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jeng1220 commented Aug 25, 2023

Uh oh!

zlsh80826 commented Aug 25, 2023

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

Tom-Zheng commented Aug 26, 2023

Uh oh!

Tom-Zheng commented Aug 26, 2023

Uh oh!

Tom-Zheng commented Aug 26, 2023• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zlsh80826 commented Aug 26, 2023

Uh oh!

ksivaman left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Tom-ZhengAug 8, 2023•
edited
Loading

Tom-ZhengAug 11, 2023•
edited
Loading

jeng1220Aug 11, 2023•
edited
Loading

Tom-Zheng commented Aug 26, 2023•
edited
Loading