chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,72 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import deepspeed
|
||||
from deepspeed.ops.op_builder import UtilsBuilder
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from unit.common import DistributedTest
|
||||
|
||||
if not deepspeed.ops.__compatible_ops__[UtilsBuilder.NAME]:
|
||||
pytest.skip(f'Skip tests since {UtilsBuilder.NAME} is not compatible', allow_module_level=True)
|
||||
|
||||
|
||||
def _validate_tensor_cast_properties(typed_tensor, byte_tensor):
|
||||
assert byte_tensor.dtype == torch.uint8
|
||||
assert byte_tensor.numel() == typed_tensor.numel() * typed_tensor.element_size()
|
||||
assert byte_tensor.data_ptr() == typed_tensor.data_ptr()
|
||||
|
||||
|
||||
def _byte_cast_single_tensor(typed_tensor):
|
||||
util_ops = UtilsBuilder().load()
|
||||
byte_tensor = util_ops.cast_to_byte_tensor(typed_tensor)
|
||||
|
||||
_validate_tensor_cast_properties(typed_tensor=typed_tensor, byte_tensor=byte_tensor)
|
||||
|
||||
|
||||
def _byte_cast_multiple_tensors(typed_tensor_list):
|
||||
util_ops = UtilsBuilder().load()
|
||||
byte_tensor_list = util_ops.cast_to_byte_tensor(typed_tensor_list)
|
||||
|
||||
assert len(typed_tensor_list) == len(byte_tensor_list)
|
||||
|
||||
for typed_tensor, byte_tensor in zip(typed_tensor_list, byte_tensor_list):
|
||||
_validate_tensor_cast_properties(typed_tensor=typed_tensor, byte_tensor=byte_tensor)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'dtype',
|
||||
[torch.float32, torch.half, torch.bfloat16, torch.float64, torch.int32, torch.short, torch.int64],
|
||||
)
|
||||
class TestCastSingleTensor(DistributedTest):
|
||||
world_size = 1
|
||||
|
||||
def test_byte_cast_accelerator_tensor(self, dtype):
|
||||
numel = 1024
|
||||
typed_tensor = torch.empty(numel, dtype=dtype).to(get_accelerator().device_name())
|
||||
_byte_cast_single_tensor(typed_tensor)
|
||||
|
||||
@pytest.mark.parametrize("pinned_memory", [True, False])
|
||||
def test_byte_cast_cpu_tensor(self, dtype, pinned_memory):
|
||||
numel = 1024
|
||||
typed_tensor = torch.empty(numel, dtype=dtype, device='cpu')
|
||||
if pinned_memory:
|
||||
typed_tensor = typed_tensor.pin_memory()
|
||||
|
||||
_byte_cast_single_tensor(typed_tensor)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('tensor_count', [1, 8, 15])
|
||||
class TestCastTensorList(DistributedTest):
|
||||
world_size = 1
|
||||
|
||||
def test_byte_cast_accelerator_tensor_list(self, tensor_count):
|
||||
typed_tensor_list = [torch.empty(1024, dtype=torch.half).to(get_accelerator().device_name())] * tensor_count
|
||||
_byte_cast_multiple_tensors(typed_tensor_list)
|
||||
|
||||
def test_byte_cast_cpu_tensor_list(self, tensor_count):
|
||||
typed_tensor_list = [torch.empty(1024, dtype=torch.half, device='cpu')] * tensor_count
|
||||
_byte_cast_multiple_tensors(typed_tensor_list)
|
||||
@@ -0,0 +1,22 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import os
|
||||
import pytest
|
||||
from deepspeed.utils.zero_to_fp32 import get_optim_files
|
||||
|
||||
|
||||
@pytest.mark.parametrize('num_checkpoints', [1, 2, 12, 24])
|
||||
def test_get_optim_files(tmpdir, num_checkpoints):
|
||||
saved_files = []
|
||||
for i in range(num_checkpoints):
|
||||
file_name = "zero_" + str(i) + "_optim_states.pt"
|
||||
path_name = os.path.join(tmpdir, file_name)
|
||||
saved_files.append(path_name)
|
||||
with open(path_name, "w") as f:
|
||||
f.write(file_name)
|
||||
loaded_files = get_optim_files(tmpdir)
|
||||
for lf, sf in zip(loaded_files, saved_files):
|
||||
assert lf == sf
|
||||
@@ -0,0 +1,38 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.utils.groups import _get_expert_parallel_ranks
|
||||
|
||||
|
||||
def test_get_expert_parallel_ranks():
|
||||
"""
|
||||
Example - E + M + D parallel
|
||||
world_size = 16
|
||||
model_degree = 2
|
||||
expert_degree = 4 # number of experts in same group
|
||||
mp_group = [0, 1], [2,3], [4,5] ...
|
||||
data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15]
|
||||
expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15]
|
||||
expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15]
|
||||
"""
|
||||
expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks(world_size=16,
|
||||
tensor_parallel_size_=2,
|
||||
expert_parallel_size_=4)
|
||||
assert expert_parallel_groups == [
|
||||
[0, 2, 4, 6],
|
||||
[8, 10, 12, 14],
|
||||
[1, 3, 5, 7],
|
||||
[9, 11, 13, 15],
|
||||
]
|
||||
assert expert_data_parallel_groups == [
|
||||
[0, 8],
|
||||
[2, 10],
|
||||
[4, 12],
|
||||
[6, 14],
|
||||
[1, 9],
|
||||
[3, 11],
|
||||
[5, 13],
|
||||
[7, 15],
|
||||
]
|
||||
@@ -0,0 +1,28 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import torch
|
||||
import pytest
|
||||
from unit.simple_model import SimpleModel
|
||||
from deepspeed import OnDevice
|
||||
from packaging import version as pkg_version
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from unit.common import DistributedTest
|
||||
|
||||
|
||||
@pytest.mark.parametrize('device', ['meta', get_accelerator().device_name(0)])
|
||||
class TestOnDevice(DistributedTest):
|
||||
world_size = 1
|
||||
|
||||
def test_on_device(self, device):
|
||||
if device == "meta" and pkg_version.parse(torch.__version__) < pkg_version.parse("1.10"):
|
||||
pytest.skip("meta tensors only became stable after torch 1.10")
|
||||
|
||||
with OnDevice(dtype=torch.half, device=device):
|
||||
model = SimpleModel(4)
|
||||
|
||||
for p in model.parameters():
|
||||
assert p.device == torch.device(device)
|
||||
assert p.dtype == torch.half
|
||||
@@ -0,0 +1,139 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import deepspeed.utils.nvtx as ds_nvtx
|
||||
import accelerator.cuda_accelerator as cuda_accelerator
|
||||
from accelerator.cuda_accelerator import CUDA_Accelerator
|
||||
|
||||
|
||||
def _sample_nvtx_function():
|
||||
return "ok"
|
||||
|
||||
|
||||
def test_instrument_w_nvtx_uses_deepspeed_domain(monkeypatch, capsys):
|
||||
calls = []
|
||||
|
||||
class FakeAccelerator:
|
||||
supports_nvtx_domain = True
|
||||
|
||||
def range_push(self, msg, domain=None, category=None):
|
||||
calls.append(("push", msg, domain, category))
|
||||
|
||||
def range_pop(self, domain=None):
|
||||
calls.append(("pop", domain))
|
||||
|
||||
monkeypatch.setattr(ds_nvtx, "enable_nvtx", True)
|
||||
monkeypatch.setattr(ds_nvtx, "is_compiling", lambda: False)
|
||||
monkeypatch.setattr(ds_nvtx, "get_accelerator", lambda: FakeAccelerator())
|
||||
|
||||
wrapped_fn = ds_nvtx.instrument_w_nvtx(_sample_nvtx_function)
|
||||
|
||||
assert wrapped_fn() == "ok"
|
||||
|
||||
with capsys.disabled():
|
||||
print(f"\nNVTX instrumentation calls: {calls}")
|
||||
|
||||
assert calls == [
|
||||
("push", "_sample_nvtx_function", ds_nvtx.DEEPSPEED_NVTX_DOMAIN, None),
|
||||
("pop", ds_nvtx.DEEPSPEED_NVTX_DOMAIN),
|
||||
]
|
||||
|
||||
|
||||
def test_instrument_w_nvtx_supports_legacy_accelerator_methods(monkeypatch, capsys):
|
||||
calls = []
|
||||
|
||||
class LegacyAccelerator:
|
||||
|
||||
def range_push(self, msg):
|
||||
calls.append(("push", msg))
|
||||
|
||||
def range_pop(self):
|
||||
calls.append(("pop", ))
|
||||
|
||||
monkeypatch.setattr(ds_nvtx, "enable_nvtx", True)
|
||||
monkeypatch.setattr(ds_nvtx, "is_compiling", lambda: False)
|
||||
monkeypatch.setattr(ds_nvtx, "get_accelerator", lambda: LegacyAccelerator())
|
||||
|
||||
wrapped_fn = ds_nvtx.instrument_w_nvtx(_sample_nvtx_function)
|
||||
|
||||
assert wrapped_fn() == "ok"
|
||||
|
||||
with capsys.disabled():
|
||||
print(f"\nLegacy NVTX instrumentation calls: {calls}")
|
||||
|
||||
assert calls == [
|
||||
("push", "_sample_nvtx_function"),
|
||||
("pop", ),
|
||||
]
|
||||
|
||||
|
||||
def test_cuda_accelerator_uses_nvtx_domain_when_available(monkeypatch, capsys):
|
||||
|
||||
class FakeDomain:
|
||||
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def push_range(self, message=None, category=None):
|
||||
self.calls.append(("push", message, category))
|
||||
return "domain-push"
|
||||
|
||||
def pop_range(self):
|
||||
self.calls.append(("pop", ))
|
||||
return "domain-pop"
|
||||
|
||||
class FakeNvtx:
|
||||
|
||||
def __init__(self):
|
||||
self.domains = {}
|
||||
|
||||
def get_domain(self, name):
|
||||
self.domains.setdefault(name, FakeDomain())
|
||||
return self.domains[name]
|
||||
|
||||
fake_nvtx = FakeNvtx()
|
||||
accelerator = CUDA_Accelerator.__new__(CUDA_Accelerator)
|
||||
accelerator._nvtx_domains = {}
|
||||
monkeypatch.setattr(cuda_accelerator, "nvtx", fake_nvtx)
|
||||
|
||||
assert accelerator.range_push("my_range", domain="DeepSpeed", category="zero") == "domain-push"
|
||||
assert accelerator.range_pop(domain="DeepSpeed") == "domain-pop"
|
||||
|
||||
with capsys.disabled():
|
||||
print(f"\nCUDA NVTX domain calls: {fake_nvtx.domains['DeepSpeed'].calls}")
|
||||
|
||||
assert fake_nvtx.domains["DeepSpeed"].calls == [
|
||||
("push", "my_range", "zero"),
|
||||
("pop", ),
|
||||
]
|
||||
|
||||
|
||||
def test_cuda_accelerator_falls_back_to_torch_nvtx_without_nvtx_package(monkeypatch, capsys):
|
||||
calls = []
|
||||
|
||||
class FakeTorchNvtx:
|
||||
|
||||
def range_push(self, msg):
|
||||
calls.append(("push", msg))
|
||||
return "torch-push"
|
||||
|
||||
def range_pop(self):
|
||||
calls.append(("pop", ))
|
||||
return "torch-pop"
|
||||
|
||||
accelerator = CUDA_Accelerator.__new__(CUDA_Accelerator)
|
||||
monkeypatch.setattr(cuda_accelerator, "nvtx", None)
|
||||
monkeypatch.setattr(cuda_accelerator.torch.cuda, "nvtx", FakeTorchNvtx()) #ignore-cuda
|
||||
|
||||
assert accelerator.range_push("my_range", domain="DeepSpeed", category="zero") == "torch-push"
|
||||
assert accelerator.range_pop(domain="DeepSpeed") == "torch-pop"
|
||||
|
||||
with capsys.disabled():
|
||||
print(f"\nCUDA torch.nvtx fallback calls: {calls}")
|
||||
|
||||
assert calls == [
|
||||
("push", "my_range"),
|
||||
("pop", ),
|
||||
]
|
||||
@@ -0,0 +1,25 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
from deepspeed.runtime import utils as ds_utils
|
||||
|
||||
|
||||
def check_partition(weights, num_parts, target_diff):
|
||||
result = ds_utils.partition_balanced(weights=weights, num_parts=num_parts)
|
||||
|
||||
parts_sum = []
|
||||
for b, e in zip(result[:-1], result[1:]):
|
||||
parts_sum.append(sum(weights[b:e]))
|
||||
|
||||
assert max(parts_sum) - min(
|
||||
parts_sum
|
||||
) == target_diff, f"ds_utils.partition_balanced(weights={weights}, num_parts={num_parts}) return {result}"
|
||||
|
||||
|
||||
def test_partition_balanced():
|
||||
check_partition([1, 2, 1], 4, target_diff=2)
|
||||
check_partition([1, 1, 1, 1], 4, target_diff=0)
|
||||
check_partition([1, 1, 1, 1, 1], 4, target_diff=1)
|
||||
check_partition([1, 1, 1, 1, 0, 1], 4, target_diff=1)
|
||||
Reference in New Issue
Block a user