From 23c1413dd7812d5adee798421975537a3c11fa59 Mon Sep 17 00:00:00 2001 From: Shaden Smith Date: Tue, 8 Jun 2021 16:51:11 +0000 Subject: [PATCH 1/3] unit test for bugfix #1135 --- tests/unit/test_zero_context.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/unit/test_zero_context.py b/tests/unit/test_zero_context.py index 8b8d51131503..b0bfcb7eee49 100644 --- a/tests/unit/test_zero_context.py +++ b/tests/unit/test_zero_context.py @@ -281,3 +281,31 @@ def test_stage_3_output_type(output_type): loss = loss['loss'] engine.backward(loss) engine.step() + +class ConvX(torch.nn.Conv1d): + def __init__(self, *args): + super().__init__(*args) + # This would not be partitioned before bugfix 5ca8167 + self.param_in = torch.nn.Parameter(torch.FloatTensor(5).uniform_()) + + def forward(self, x): + return x + +class ConvNet(torch.nn.Module): + + def __init__(self): + super().__init__() + self.conv1 = ConvX(1, 3, 4) + self.param = torch.nn.Parameter(torch.FloatTensor(5).uniform_()) + + def forward(self, x): + return x + + + +def test_subclass_param(): + with deepspeed.zero.Init(config=config): + model = ConvNet() + + assert model.param.ds_status == ZeroParamStatus.NOT_AVAILABLE + assert model.conv1.param_in.ds_status == ZeroParamStatus.NOT_AVAILABLE \ No newline at end of file From f914558919546b05b56b76d47c0c11aff41270b1 Mon Sep 17 00:00:00 2001 From: Shaden Smith Date: Tue, 8 Jun 2021 16:52:52 +0000 Subject: [PATCH 2/3] formatter --- tests/unit/test_zero_context.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_zero_context.py b/tests/unit/test_zero_context.py index b0bfcb7eee49..c5a20b9847a1 100644 --- a/tests/unit/test_zero_context.py +++ b/tests/unit/test_zero_context.py @@ -282,17 +282,18 @@ def test_stage_3_output_type(output_type): engine.backward(loss) engine.step() + class ConvX(torch.nn.Conv1d): def __init__(self, *args): super().__init__(*args) # This would not be partitioned before bugfix 5ca8167 self.param_in = torch.nn.Parameter(torch.FloatTensor(5).uniform_()) - + def forward(self, x): return x - -class ConvNet(torch.nn.Module): + +class ConvNet(torch.nn.Module): def __init__(self): super().__init__() self.conv1 = ConvX(1, 3, 4) @@ -302,10 +303,9 @@ def forward(self, x): return x - def test_subclass_param(): with deepspeed.zero.Init(config=config): model = ConvNet() - + assert model.param.ds_status == ZeroParamStatus.NOT_AVAILABLE - assert model.conv1.param_in.ds_status == ZeroParamStatus.NOT_AVAILABLE \ No newline at end of file + assert model.conv1.param_in.ds_status == ZeroParamStatus.NOT_AVAILABLE From 5f19a191cf24cae22d57292d323fe74aec9946a9 Mon Sep 17 00:00:00 2001 From: Shaden Smith Date: Wed, 9 Jun 2021 08:19:55 +0000 Subject: [PATCH 3/3] fix test in presence of mpi4py --- tests/unit/test_zero_context.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_zero_context.py b/tests/unit/test_zero_context.py index c5a20b9847a1..98ee0c7ad00b 100644 --- a/tests/unit/test_zero_context.py +++ b/tests/unit/test_zero_context.py @@ -304,6 +304,7 @@ def forward(self, x): def test_subclass_param(): + setup_serial_env() with deepspeed.zero.Init(config=config): model = ConvNet()