[Bugfix] fix initialization #4

szhengac · 2023-01-18T07:48:20Z

Description

Remove redundant dist group creation
Remove group destroy which cause nccl error
Add user defined init func.

Checklist

PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
Changes are complete (i.e. I finished coding on this PR)
All changes have test coverage
Code is well-documented

comaniac

Let @chhzh123 double check.
Also please rebase onto the latest main branch so that the CI should work.

slapo/schedule.py

chhzh123 · 2023-01-18T19:10:38Z

Is the long initialization time mainly because of the dist group creation and destruction?

szhengac · 2023-01-18T19:44:54Z

Is the long initialization time mainly because of the dist group creation and destruction?

Mainly because of group creation. Group destruction causes nccl error for gpt example

comaniac

Mostly LGTM. I'll approve when the comments are addressed.
@chhzh123 PTAL.

examples/gpt/deepspeed_hf.py

slapo/schedule.py

tests/test_shard.py

slapo/schedule.py

examples/gpt/deepspeed_hf.py

slapo/autotune/tune.py

tests/test_shard.py

chhzh123 · 2023-01-20T19:49:12Z

LGTM. I'll approve when the interface is finalized.

comaniac

LGTM. Just nits.

comaniac · 2023-01-21T00:21:57Z

tests/test_pipeline_partition.py

@@ -51,5 +53,50 @@ def forward(self, x):
    assert ("stage2.linear.weight", 2) in tie_weights[0]


+def test_analyze_tie_ranks():


I think you could just reuse the above test, as you also call analyze_tie_weights in this test.

not really, the above test uses 3 stages, but we can only use 2 stages when there are only 2 gpus

Huh I just realized that this test may not pass CI because we don't have DeepSpeed installed in the docker image...can you fake a topology? This also addresses the 2 GPU issue.

slapo/schedule.py

comaniac · 2023-01-21T00:23:04Z

slapo/schedule.py

+        for _, _ in tie_ranks:
+            pass


What's this for?

otherwise lint does not pass

Interesting. What's the lint error?

tie_ranks not used. this tie_ranks is a placeholder for you to add additional functionality

slapo/schedule.py

Co-authored-by: Cody Yu <[email protected]>

comaniac reviewed Jan 18, 2023

View reviewed changes

slapo/schedule.py Show resolved Hide resolved

szhengac force-pushed the consolidation branch 3 times, most recently from a6e8142 to 885f763 Compare January 19, 2023 19:14

comaniac reviewed Jan 20, 2023

View reviewed changes

examples/gpt/deepspeed_hf.py Outdated Show resolved Hide resolved

slapo/schedule.py Outdated Show resolved Hide resolved

slapo/schedule.py Outdated Show resolved Hide resolved

tests/test_shard.py Show resolved Hide resolved

zarzen reviewed Jan 20, 2023

View reviewed changes

slapo/schedule.py Show resolved Hide resolved

chhzh123 reviewed Jan 20, 2023

View reviewed changes

examples/gpt/deepspeed_hf.py Outdated Show resolved Hide resolved

slapo/autotune/tune.py Show resolved Hide resolved

tests/test_shard.py Show resolved Hide resolved

szhengac added 20 commits January 20, 2023 20:51

fix initialization

0e3a1cf

fix init order

4e8c438

remove DEBUG

0ed7d16

add init optional

ecb56b6

reformat

651f707

reformat

941ccc3

add license

46326d8

remove unused import

4d16eae

change import

d053d01

group import

850d54d

fix test

986cf6d

fix shard memory layout for axis > 0

0d7bb9b

black test

d6c64d8

clean lint

3f566b2

add bf16 to gpt script

1bf86f1

fix lint

1549e0e

change init_required position

de8fdcf

fix test conv

93274b0

reformat

657e004

remove gather in test conv

79d439f

szhengac added 5 commits January 20, 2023 20:51

test_conv is flaky and skipped

d259601

add test and single device fallback

20a1aa4

rename init_weights

b1b68dc

add tied rank analysis

fa9be92

clean lint

11e669a

szhengac force-pushed the consolidation branch from 3bde9aa to 11e669a Compare January 20, 2023 23:35

szhengac added 2 commits January 21, 2023 00:13

add test

95001cf

fix lint

9a4f2f2

comaniac approved these changes Jan 21, 2023

View reviewed changes

szhengac and others added 5 commits January 20, 2023 16:28

rewrite

2040da4

Co-authored-by: Cody Yu <[email protected]>

fix lint

c2e04aa

fix tie_rank format

07cb9a8

fix lint

cf4a63c

fix typo

4a359d4

szhengac merged commit e5735e6 into awslabs:main Jan 21, 2023

szhengac deleted the consolidation branch January 21, 2023 02:01

chhzh123 pushed a commit to chhzh123/slapo that referenced this pull request Jan 22, 2023

[Bugfix] fix initialization (awslabs#4)

348084e

Co-authored-by: Cody Yu <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Bugfix] fix initialization #4

[Bugfix] fix initialization #4

szhengac commented Jan 18, 2023

comaniac left a comment

chhzh123 commented Jan 18, 2023

szhengac commented Jan 18, 2023

comaniac left a comment

chhzh123 commented Jan 20, 2023

comaniac left a comment

comaniac Jan 21, 2023

szhengac Jan 21, 2023

comaniac Jan 21, 2023

comaniac Jan 21, 2023

szhengac Jan 21, 2023

comaniac Jan 21, 2023

szhengac Jan 21, 2023

		@@ -51,5 +53,50 @@ def forward(self, x):
		assert ("stage2.linear.weight", 2) in tie_weights[0]


		def test_analyze_tie_ranks():

[Bugfix] fix initialization #4

[Bugfix] fix initialization #4

Conversation

szhengac commented Jan 18, 2023

Description

Checklist

comaniac left a comment

Choose a reason for hiding this comment

chhzh123 commented Jan 18, 2023

szhengac commented Jan 18, 2023

comaniac left a comment

Choose a reason for hiding this comment

chhzh123 commented Jan 20, 2023

comaniac left a comment

Choose a reason for hiding this comment

comaniac Jan 21, 2023

Choose a reason for hiding this comment

szhengac Jan 21, 2023

Choose a reason for hiding this comment

comaniac Jan 21, 2023

Choose a reason for hiding this comment

comaniac Jan 21, 2023

Choose a reason for hiding this comment

szhengac Jan 21, 2023

Choose a reason for hiding this comment

comaniac Jan 21, 2023

Choose a reason for hiding this comment

szhengac Jan 21, 2023

Choose a reason for hiding this comment