Skip to content

Fix for distributed tests on pytorch>=1.12 #2141

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
mrwyattii merged 20 commits into from
Aug 1, 2022
Merged

Fix for distributed tests on pytorch>=1.12 #2141

mrwyattii merged 20 commits into from
Aug 1, 2022

Conversation

@mrwyattii
Copy link
Contributor

@mrwyattii mrwyattii commented Jul 26, 2022 •
edited
Loading

With the recent release of torch 1.12, we saw all unit tests using the @distirbuted_test decorator break (see this issue).

The problem involves changes in torch.multiprocessing and torch.distributed that prevents using fork to spawn new processes in distributed tests. We must now use spawn or forkserver to spin up the distributed processes. However, this change alone is not sufficient as the decorator causes issues when pickling the target function (e.g., dist_init or test_*):
AttributeError: Can't pickle local object and _pickle.PicklingError: Can't pickle <function test_example at 0x7f6c4e5640d0>: it's not the same object as unit.test_example.test_example errors are unavoidable with the distributed test decorator.

The solution implemented here loosely follows how pytorch unit tests are done and adds a DistributedTest class. This will require refactoring all distributed tests like so:

# OLD TEST:
@pytest.mark.parametrize('foo,bar', [(1,2),(3,4)])
def test_example(foo, bar):
    @distributed_test(world_size=[1,4])
    def _go():
        assert foo < bar
    _go()
# NEW TEST:
@pytest.mark.parametrize('foo,bar', [(1,2),(3,4)])
class TestExample(DistributedTest):
    world_size=[1,4]
    def test(self, foo, bar):
        assert foo < bar

@jeffra, @tjruwase, and I will be refactoring existing tests over time to use this new test class. In doing so, we are also organizing the unit tests files to match the directory structure of the DeepSpeed source code.

mrwyattii added 13 commits July 26, 2022 10:40
@mrwyattii mrwyattii marked this pull request as ready for review July 29, 2022 23:30
@mrwyattii mrwyattii enabled auto-merge (squash) August 1, 2022 17:27
@mrwyattii mrwyattii disabled auto-merge August 1, 2022 23:19
@mrwyattii mrwyattii enabled auto-merge (squash) August 1, 2022 23:20
@mrwyattii mrwyattii merged commit 1a71e77 into deepspeedai:master Aug 1, 2022
This was referenced Aug 3, 2022
Merged
Merged
This was referenced Sep 6, 2022
Merged
Merged
@jahatef jahatef mentioned this pull request Feb 14, 2024
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@jeffra jeffra jeffra approved these changes

@samyam samyam Awaiting requested review from samyam

@tjruwase tjruwase Awaiting requested review from tjruwase tjruwase is a code owner

@ShadenSmith ShadenSmith Awaiting requested review from ShadenSmith

@conglongli conglongli Awaiting requested review from conglongli

@awan-10 awan-10 Awaiting requested review from awan-10

@cli99 cli99 Awaiting requested review from cli99

@eltonzheng eltonzheng Awaiting requested review from eltonzheng

@minjiaz minjiaz Awaiting requested review from minjiaz

@RezaYazdaniAminabadi RezaYazdaniAminabadi Awaiting requested review from RezaYazdaniAminabadi

@duli2012 duli2012 Awaiting requested review from duli2012

@yaozhewei yaozhewei Awaiting requested review from yaozhewei

@arashb arashb Awaiting requested review from arashb

@xiaoxiawu-microsoft xiaoxiawu-microsoft Awaiting requested review from xiaoxiawu-microsoft

@samadejacobs samadejacobs Awaiting requested review from samadejacobs

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@mrwyattii @jeffra @tjruwase