diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index f05a53867c93..7e1143ba424b 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -64,14 +64,15 @@ def input(msg): def split_half_float_double(tensors): - device_type = get_accelerator().device_name() - dtypes = [ - "torch.{}.HalfTensor".format(device_type), "torch.{}.FloatTensor".format(device_type), - "torch.{}.DoubleTensor".format(device_type), "torch.{}.BFloat16Tensor".format(device_type) - ] + # Legacy type strings omit the device prefix on CPU ("torch.FloatTensor"), so + # building them from the accelerator device name matches nothing on CPU and + # silently drops every gradient bucket, skipping the all-reduce entirely. + # Compare dtypes directly so the buckets are device-independent. Sparse layouts + # never matched the legacy type strings on any device, so they stay excluded. + dtypes = [torch.half, torch.float, torch.double, torch.bfloat16] buckets = [] for i, dtype in enumerate(dtypes): - bucket = [t for t in tensors if t.type() == dtype] + bucket = [t for t in tensors if t.dtype == dtype and not t.is_sparse] if bucket: buckets.append(bucket) return buckets diff --git a/tests/unit/common.py b/tests/unit/common.py index 7d69365ed78b..6705c2b40181 100644 --- a/tests/unit/common.py +++ b/tests/unit/common.py @@ -575,8 +575,8 @@ def wrapper(*args: Any, **kwargs: Any): def reduce_boolean_flags(flag: bool, op=all) -> bool: if not dist.is_initialized(): return flag - device = get_accelerator().current_device() - tensor_flag = torch.tensor(1 if flag else 0, dtype=torch.int, device=device) + device = get_accelerator().current_device_name() + tensor_flag = torch.tensor([1 if flag else 0], dtype=torch.int, device=device) world_size = dist.get_world_size() tensor_flag_buf = torch.zeros(world_size, dtype=torch.int, device=device) dist.all_gather_into_tensor(tensor_flag_buf, tensor_flag) diff --git a/tests/unit/v1/zero/test_zero.py b/tests/unit/v1/zero/test_zero.py index a56d873d8408..0fff79d2156a 100644 --- a/tests/unit/v1/zero/test_zero.py +++ b/tests/unit/v1/zero/test_zero.py @@ -24,12 +24,33 @@ from deepspeed.runtime.engine import DeepSpeedEngine from deepspeed.runtime.bf16_optimizer import BF16_Optimizer from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus +from deepspeed.runtime.zero.stage_1_and_2 import split_half_float_double from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint from deepspeed.runtime.zero.utils import ZeRORuntimeException from deepspeed.accelerator import get_accelerator from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad +class TestSplitHalfFloatDouble: + + def test_device_independent_buckets_exclude_sparse(self): + # Pins two fixed membership bugs: the legacy accelerator-prefixed type + # strings matched nothing on CPU, silently dropping every bucket, and + # dtype-only matching would admit sparse layouts, which cannot be + # flattened into a dense all-reduce buffer. + dense_grads = [ + torch.zeros(2, dtype=dtype) for dtype in (torch.half, torch.float, torch.double, torch.bfloat16) + ] + sparse_grad = torch.sparse_coo_tensor(torch.tensor([[0]]), torch.tensor([1.0]), (1, )) + + buckets = split_half_float_double(dense_grads + [sparse_grad]) + + assert len(buckets) == 4 + for bucket, grad in zip(buckets, dense_grads): + assert len(bucket) == 1 + assert bucket[0] is grad + + @pytest.mark.parametrize("zero_stage", [0, 1, 2]) @pytest.mark.parametrize("gradient_allreduce_op,expected_scale", [("mean", 1.0), ("sum", 2.0)]) @pytest.mark.parametrize(