From d8077876f4b5fe5812916ff6538a9e1bb7005fe3 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 15 Sep 2021 13:22:31 -0700 Subject: [PATCH 1/7] [zero_to_fp32] adapt to 4-bytes alignment in z2 --- deepspeed/utils/zero_to_fp32.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index db3d61713d76..c0b02eb00835 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -10,6 +10,7 @@ import argparse import torch import glob +import math import os from collections import OrderedDict @@ -217,8 +218,16 @@ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir): if zero_stage == 3: offset *= world_size + mismatch = False + if zero_stage == 2: + # Z2 started to align to 4 to improve nccl performance + offset_aligned_to_4 = 4 * math.ceil(offset / 4) + mismatch = offset_aligned_to_4 != avail_numel + elif zero_stage == 3: + mismatch = offset != avail_numel + # Sanity check - if offset != avail_numel: + if mismatch: raise ValueError( f"consumed {offset} numels out of {avail_numel} - something is wrong") From a00805891e2218cd39ffe9ef9a6fcb4240fc4116 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 15 Sep 2021 13:44:35 -0700 Subject: [PATCH 2/7] align both sides --- deepspeed/utils/zero_to_fp32.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index c0b02eb00835..4da891f388ac 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -218,16 +218,17 @@ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir): if zero_stage == 3: offset *= world_size + def align_to_4(x): + return 4 * math.ceil(x / 4) + mismatch = False if zero_stage == 2: # Z2 started to align to 4 to improve nccl performance - offset_aligned_to_4 = 4 * math.ceil(offset / 4) - mismatch = offset_aligned_to_4 != avail_numel - elif zero_stage == 3: - mismatch = offset != avail_numel + offset = align_to_4(offset) + avail_numel = align_to_4(avail_numel) # Sanity check - if mismatch: + if offset != avail_numel: raise ValueError( f"consumed {offset} numels out of {avail_numel} - something is wrong") From a3bc8d98e414b8230aabb87b5343cb9008d4a775 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 15 Sep 2021 13:45:44 -0700 Subject: [PATCH 3/7] cleanup --- deepspeed/utils/zero_to_fp32.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index 4da891f388ac..fc85e080e80e 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -221,7 +221,6 @@ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir): def align_to_4(x): return 4 * math.ceil(x / 4) - mismatch = False if zero_stage == 2: # Z2 started to align to 4 to improve nccl performance offset = align_to_4(offset) From 08b41a8d93f40083b07545f7c90c780c49f30ce2 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 15 Sep 2021 15:09:08 -0700 Subject: [PATCH 4/7] adjust the existing test to reproduce the bug --- tests/unit/test_zero.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unit/test_zero.py b/tests/unit/test_zero.py index fa0e05e1e15a..7673f18b742e 100755 --- a/tests/unit/test_zero.py +++ b/tests/unit/test_zero.py @@ -133,6 +133,7 @@ def _test_zero3_repeat_forward_loop(args, model, hidden_dim): # testing the fix https://github.com/microsoft/DeepSpeed/pull/1227 +# also reproduces the https://github.com/microsoft/DeepSpeed/pull/1372 @pytest.mark.parametrize('zero_stage', [2, 3]) def test_zero_to_fp32(tmpdir, zero_stage): @@ -168,6 +169,11 @@ def __init__(self, hidden_dim, n_layers): self.ll = torch.nn.ModuleList( torch.nn.Linear(hidden_dim, hidden_dim) for i in range(n_layers)) + # to reproduce https://github.com/microsoft/DeepSpeed/pull/1372 it is important that + # the number of params is uneven - the following adds 4+1 params - the linear + # layers are 6 param each + 5 - so total 17 elements (for 1 gpu) + self.classifier = torch.nn.Linear(4, 1) + self.b = self.register_buffer("buffer", torch.ones(15)) self.cross_entropy_loss = torch.nn.CrossEntropyLoss() def forward(self, x, y): From a38e3ae8533405a428fb540ff6c69500ea90926e Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 15 Sep 2021 15:39:20 -0700 Subject: [PATCH 5/7] test only on one gpu --- tests/unit/test_zero.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_zero.py b/tests/unit/test_zero.py index 7673f18b742e..7fd7a401a4d3 100755 --- a/tests/unit/test_zero.py +++ b/tests/unit/test_zero.py @@ -227,14 +227,15 @@ def dump_state_dict(model): orig_state_dict[name] = param.detach().cpu() print(orig_state_dict) - fp32_model = load_state_dict_from_zero_checkpoint(model.module, tmpdir) - #dump_state_dict(fp32_model) - - fp32_state_dict = fp32_model.state_dict() - for name in orig_state_dict.keys(): - # float() workaround for torch<1.6 - assert torch.allclose(orig_state_dict[name].float(), - fp32_state_dict[name].float()) + if dist.get_rank() == 0: + fp32_model = load_state_dict_from_zero_checkpoint(model.module, tmpdir) + #dump_state_dict(fp32_model) + + fp32_state_dict = fp32_model.state_dict() + for name in orig_state_dict.keys(): + # float() workaround for torch<1.6 + assert torch.allclose(orig_state_dict[name].float(), + fp32_state_dict[name].float()) _test_zero_to_fp32() From 9bee134fc16cb7740bd4f7d5b9f62f896f26da4d Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 15 Sep 2021 15:39:34 -0700 Subject: [PATCH 6/7] handle the edge case of param with 1 element --- deepspeed/utils/zero_to_fp32.py | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/deepspeed/utils/zero_to_fp32.py b/deepspeed/utils/zero_to_fp32.py index fc85e080e80e..894aa02ce5cb 100755 --- a/deepspeed/utils/zero_to_fp32.py +++ b/deepspeed/utils/zero_to_fp32.py @@ -117,6 +117,11 @@ def parse_optim_states(files, ds_checkpoint_dir): def zero3_partitioned_param_info(unpartitioned_numel, world_size): + #print("*** ", unpartitioned_numel, world_size, " ***",) + # handle an edge case where there is only 1 element (e.g. bias in a tiny test model) + if unpartitioned_numel == 1: + return 1, 0 + remainder = unpartitioned_numel % world_size padding_numel = (world_size - remainder) if remainder else 0 partitioned_numel = int(unpartitioned_numel / world_size) @@ -206,13 +211,20 @@ def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir): f"{total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}" ) - # XXX: memory usage doubles here (zero3) - state_dict[name] = torch.cat( - tuple(fp32_flat_groups[i].narrow(0, - offset, - partitioned_numel) - for i in range(world_size)), - 0).view(shape) + if unpartitioned_numel > 1: + # XXX: memory usage doubles here (zero3) + state_dict[name] = torch.cat( + tuple(fp32_flat_groups[i].narrow(0, + offset, + partitioned_numel) + for i in range(world_size)), + 0).view(shape) + else: + # handle an edge case where there is only 1 element (e.g. bias in a tiny test model) + state_dict[name] = fp32_flat_groups[0].narrow( + 0, + offset, + partitioned_numel).view(shape) offset += partitioned_numel + partitioned_padding_numel if zero_stage == 3: From 592a7d794235cc4660af3b6bcaaccba5f268675c Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Wed, 15 Sep 2021 20:41:18 -0700 Subject: [PATCH 7/7] remove the buffer --- tests/unit/test_zero.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_zero.py b/tests/unit/test_zero.py index 7fd7a401a4d3..6cd8b971d6af 100755 --- a/tests/unit/test_zero.py +++ b/tests/unit/test_zero.py @@ -173,7 +173,6 @@ def __init__(self, hidden_dim, n_layers): # the number of params is uneven - the following adds 4+1 params - the linear # layers are 6 param each + 5 - so total 17 elements (for 1 gpu) self.classifier = torch.nn.Linear(4, 1) - self.b = self.register_buffer("buffer", torch.ones(15)) self.cross_entropy_loss = torch.nn.CrossEntropyLoss() def forward(self, x, y):