wehub-resource-sync
162 lines
7.1 KiB
Python
162 lines
7.1 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""
|
|
unit tests for coalesced collectives
|
|
"""
|
|
|
|
import torch
|
|
import deepspeed
|
|
import deepspeed.comm as dist
|
|
from deepspeed.runtime.comm.coalesced_collectives import reduce_scatter_coalesced, all_to_all_quant_reduce
|
|
from deepspeed.accelerator import get_accelerator
|
|
import pytest
|
|
|
|
from unit.common import DistributedTest
|
|
|
|
|
|
class TestReduceScatterCoalesced(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test_single_input(self):
|
|
input = torch.full((6, ), dist.get_rank(), dtype=torch.half, device=get_accelerator().current_device_name())
|
|
|
|
(output, ) = reduce_scatter_coalesced([input], dist.get_world_group())
|
|
|
|
assert output.shape == (3, )
|
|
assert torch.allclose(output, torch.full_like(output, 0.5))
|
|
|
|
def test_two_inputs(self):
|
|
tensor_kwargs = {"device": get_accelerator().current_device_name(), "dtype": torch.half}
|
|
inputs = [
|
|
dist.get_rank() * torch.arange(0, 6, **tensor_kwargs),
|
|
dist.get_rank() * torch.arange(6, 9, **tensor_kwargs),
|
|
]
|
|
|
|
output1, output2 = reduce_scatter_coalesced(inputs, dist.get_world_group())
|
|
|
|
if dist.get_rank() == 0:
|
|
assert output1.shape == (3, )
|
|
assert torch.allclose(output1, torch.arange(0, 3, **tensor_kwargs) / 2)
|
|
assert output2.shape == (2, )
|
|
assert torch.allclose(output2, torch.arange(6, 8, **tensor_kwargs) / 2)
|
|
elif dist.get_rank() == 1:
|
|
assert output1.shape == (3, )
|
|
assert torch.allclose(output1, torch.arange(3, 6, **tensor_kwargs) / 2)
|
|
assert output2.shape == (1, )
|
|
assert torch.allclose(output2, torch.arange(8, 9, **tensor_kwargs) / 2)
|
|
|
|
|
|
class TestReduceScatterCoalescedTensorSmallerThanWorldSize(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test(self):
|
|
input = torch.zeros((1, ), dtype=torch.half, device=get_accelerator().current_device_name())
|
|
|
|
(output, ) = reduce_scatter_coalesced([input], dist.get_world_group())
|
|
|
|
if dist.get_rank() == 0:
|
|
assert output.shape == (1, )
|
|
assert torch.allclose(output, torch.zeros_like(output))
|
|
elif dist.get_rank() == 1:
|
|
assert output.shape == (0, )
|
|
|
|
|
|
# Currently we cannot test all_to_all_quant_reduce in non-fallback cases because we don't support multinodes tests.
|
|
class TestAllToAllQuantReduceFallback(DistributedTest):
|
|
world_size = 2
|
|
|
|
def test_1d_tensor(self):
|
|
# case 1: 1D tensor
|
|
input = torch.zeros((10, ), dtype=torch.half, device=get_accelerator().current_device_name())
|
|
from deepspeed.ops.op_builder import QuantizerBuilder
|
|
if not deepspeed.ops.__compatible_ops__[QuantizerBuilder.NAME]:
|
|
pytest.skip("QuantizerBuilder is not implemented")
|
|
output = all_to_all_quant_reduce([input], {})[0]
|
|
|
|
if dist.get_rank() == 0:
|
|
assert output.shape == (5, )
|
|
assert torch.allclose(output, torch.zeros_like(output))
|
|
elif dist.get_rank() == 1:
|
|
assert output.shape == (5, )
|
|
assert torch.allclose(output, torch.zeros_like(output))
|
|
|
|
def test_non_divisible(self):
|
|
# case 2: tensor size not divisible by global_world_size
|
|
input = torch.zeros((7, 7), dtype=torch.half, device=get_accelerator().current_device_name())
|
|
from deepspeed.ops.op_builder import QuantizerBuilder
|
|
if not deepspeed.ops.__compatible_ops__[QuantizerBuilder.NAME]:
|
|
pytest.skip("QuantizerBuilder is not implemented")
|
|
output = all_to_all_quant_reduce([input], {})[0]
|
|
|
|
if dist.get_rank() == 0:
|
|
assert output.shape == (25, )
|
|
assert torch.allclose(output, torch.zeros_like(output))
|
|
elif dist.get_rank() == 1:
|
|
assert output.shape == (24, )
|
|
assert torch.allclose(output, torch.zeros_like(output))
|
|
|
|
|
|
class TestLocoQuantized(DistributedTest):
|
|
|
|
world_size = 1
|
|
|
|
@pytest.mark.parametrize("num_bits", [4, 8])
|
|
@pytest.mark.parametrize("tensor_size", [(16, 16), (64, 64)])
|
|
@pytest.mark.parametrize("devices_per_node", [4, 8])
|
|
def test_loco_quantized_reduction(self, num_bits, tensor_size, devices_per_node):
|
|
from deepspeed.ops.op_builder import QuantizerBuilder
|
|
if not deepspeed.ops.__compatible_ops__[QuantizerBuilder.NAME]:
|
|
pytest.skip("QuantizerBuilder is not implemented")
|
|
|
|
quantizer_module = QuantizerBuilder().load()
|
|
|
|
tensor = torch.randn(tensor_size, device='cuda', dtype=torch.half)
|
|
|
|
num_nodes = 2 # Fake world size
|
|
total_elements = tensor.numel()
|
|
total_devices = devices_per_node * num_nodes
|
|
num_groups = max(tensor.shape[0], tensor.shape[1], total_devices)
|
|
|
|
# Initialize error_feedback tensor
|
|
error_feedback = torch.randn(tensor_size, device=tensor.device, dtype=tensor.dtype)
|
|
error_feedback_ori = error_feedback.clone()
|
|
# Swizzle the original tensor
|
|
tensor_reshaped = tensor.reshape(num_nodes, devices_per_node, total_elements // total_devices)
|
|
swizzled_tensor = tensor_reshaped.permute(1, 0, 2).reshape(tensor.size())
|
|
|
|
# Perform loco_swizzle_quant
|
|
output, scales = quantizer_module.loco_swizzle_quant(tensor, error_feedback, 0.0, num_groups, num_bits,
|
|
quantizer_module.Symmetric, 1, num_nodes,
|
|
devices_per_node)
|
|
|
|
# Compare swizzled_tensor with the output of loco_swizzle_quant
|
|
dequantized = quantizer_module.dequantize(output, scales, scales.numel(), num_bits,
|
|
quantizer_module.Symmetric).view(tensor.size())
|
|
|
|
assert torch.allclose(swizzled_tensor + error_feedback_ori, dequantized + error_feedback)
|
|
|
|
# Calculate elements per group and groups per partition
|
|
elements_per_group = total_elements // num_groups
|
|
groups_per_partition = num_groups // devices_per_node
|
|
|
|
# Reshape dequantized data to match the grouping in loco_quantized_reduction
|
|
dequantized_reshaped = dequantized.view(devices_per_node, groups_per_partition, elements_per_group)
|
|
|
|
# Perform reduction across devices_per_node dimension
|
|
reduced_dequantized = dequantized_reshaped.cumsum(dim=0)[-1]
|
|
# Initialize error_feedback tensor
|
|
error_feedback = torch.randn(reduced_dequantized.shape, device=tensor.device, dtype=dequantized.dtype)
|
|
error_feedback_ori = error_feedback.clone()
|
|
|
|
# perform loco_quantized_reduction
|
|
output, scales = quantizer_module.loco_quantized_reduction(output, scales, error_feedback, 0.0, num_groups,
|
|
num_groups // devices_per_node, num_bits,
|
|
quantizer_module.Symmetric, devices_per_node)
|
|
|
|
dequantized_reduced = quantizer_module.dequantize(output, scales, scales.numel(), num_bits,
|
|
quantizer_module.Symmetric).view(error_feedback.size())
|
|
|
|
assert torch.allclose(reduced_dequantized + error_feedback_ori, dequantized_reduced + error_feedback)
|