chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,252 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import copy
|
||||
import torch.nn as nn
|
||||
import pytest
|
||||
|
||||
import torch
|
||||
|
||||
import deepspeed
|
||||
import deepspeed.comm as dist
|
||||
from deepspeed.runtime.pipe.topology import PipeDataParallelTopology
|
||||
from deepspeed.runtime.pipe.module import PipelineModule
|
||||
from unit.alexnet_model import AlexNetPipe, train_cifar
|
||||
from unit.common import DistributedTest
|
||||
from unit.util import skip_on_arch, no_child_process_in_deepspeed_io
|
||||
|
||||
PipeTopo = PipeDataParallelTopology
|
||||
|
||||
config_dict = {
|
||||
"train_batch_size": 4,
|
||||
"grandient_accumulation_steps": 1,
|
||||
"steps_per_print": 20,
|
||||
"optimizer": {
|
||||
"type": "Adam",
|
||||
"params": {
|
||||
"lr": 0.001,
|
||||
"betas": [0.9, 0.999],
|
||||
"eps": 1e-8,
|
||||
"weight_decay": 3e-7
|
||||
}
|
||||
},
|
||||
"zero_optimization": {
|
||||
"stage": 0
|
||||
},
|
||||
"fp16": {
|
||||
"enabled": False
|
||||
},
|
||||
"pipeline": {
|
||||
"seed_layers": True,
|
||||
"activation_checkpoint_interval": 1
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def rel_diff(A, B):
|
||||
return abs(A - B) / abs(A)
|
||||
|
||||
|
||||
@pytest.mark.parametrize('topo_config', [
|
||||
{
|
||||
"num_pp": 1,
|
||||
"num_dp": 4
|
||||
},
|
||||
{
|
||||
"num_pp": 2,
|
||||
"num_dp": 2
|
||||
},
|
||||
{
|
||||
"num_pp": 4,
|
||||
"num_dp": 1
|
||||
},
|
||||
])
|
||||
class TestPipeCifar10(DistributedTest):
|
||||
world_size = 4
|
||||
|
||||
def test_pipe_base(self, topo_config):
|
||||
skip_on_arch(min_arch=7)
|
||||
topo = PipeTopo(**topo_config)
|
||||
steps = 100 # must be >=100
|
||||
|
||||
# Allocate model for consistent initial weights.
|
||||
init_net = AlexNetPipe()
|
||||
|
||||
base_net = copy.deepcopy(init_net)
|
||||
base_model = PipelineModule(layers=base_net.to_layers(), num_stages=1, loss_fn=nn.CrossEntropyLoss())
|
||||
|
||||
# Train with just data parallelism
|
||||
base_losses = train_cifar(base_model, config=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled'])
|
||||
|
||||
test_net = copy.deepcopy(init_net)
|
||||
test_model = PipelineModule(layers=test_net.to_layers(), topology=topo, loss_fn=nn.CrossEntropyLoss())
|
||||
|
||||
test_losses = train_cifar(test_model, config=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled'])
|
||||
|
||||
abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)]
|
||||
rel_diffs = [rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)]
|
||||
if dist.get_rank() == 0:
|
||||
print(f'abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}')
|
||||
print(f'rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}')
|
||||
print(f'first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}')
|
||||
|
||||
for lastX in [1, 10, 100]:
|
||||
base_avg = sum(base_losses[-lastX:]) / lastX
|
||||
test_avg = sum(test_losses[-lastX:]) / lastX
|
||||
print(
|
||||
f'last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}'
|
||||
)
|
||||
|
||||
lastX = 100
|
||||
base = base_losses[-lastX:]
|
||||
base_avg = sum(base) / len(base)
|
||||
test = test_losses[-lastX:]
|
||||
test_avg = sum(test) / len(test)
|
||||
assert rel_diff(base_avg, test_avg) < 0.05 # Originally 0.03, but seeing instability with AMD results
|
||||
|
||||
# def _check_model_params_equal(self, model1, model2):
|
||||
# for p1, p2 in zip(model1.parameters(), model2.parameters()):
|
||||
# if p1.data.ne(p2.data).sum() > 0:
|
||||
# assert False, f"model params not equal"
|
||||
|
||||
def test_pipe_use_reentrant(self, topo_config):
|
||||
skip_on_arch(min_arch=7)
|
||||
|
||||
topo = PipeTopo(**topo_config)
|
||||
steps = 100 # must be >=100
|
||||
|
||||
# Allocate model for consistent initial weights.
|
||||
init_net = AlexNetPipe()
|
||||
|
||||
# Train with not set use_reentrant, default: True
|
||||
base_net = copy.deepcopy(init_net)
|
||||
base_model = PipelineModule(layers=base_net.to_layers(), topology=topo, loss_fn=nn.CrossEntropyLoss())
|
||||
base_losses = train_cifar(base_model, config=config_dict, num_steps=steps, fp16=config_dict['fp16']['enabled'])
|
||||
|
||||
# Train with set use_reentrant=False, this will use ``non_reentrant_checkpoint``
|
||||
test_config_dict = copy.deepcopy(config_dict)
|
||||
test_config_dict['pipeline']['use_reentrant'] = False
|
||||
test_net = copy.deepcopy(init_net)
|
||||
test_model = PipelineModule(layers=test_net.to_layers(), topology=topo, loss_fn=nn.CrossEntropyLoss())
|
||||
test_losses = train_cifar(test_model,
|
||||
config=test_config_dict,
|
||||
num_steps=steps,
|
||||
fp16=config_dict['fp16']['enabled'])
|
||||
|
||||
abs_diffs = [l0 - l1 for l0, l1 in zip(base_losses, test_losses)]
|
||||
rel_diffs = [rel_diff(l0, l1) for l0, l1 in zip(base_losses, test_losses)]
|
||||
if dist.get_rank() == 0:
|
||||
print(f'abs min={min(abs_diffs)} max={max(abs_diffs)} avg={sum(abs_diffs)/len(abs_diffs)}')
|
||||
print(f'rel min={min(rel_diffs)} max={max(rel_diffs)} avg={sum(rel_diffs)/len(rel_diffs)}')
|
||||
print(f'first: base={base_losses[0]} test={test_losses[0]} abs={abs_diffs[0]} rel={rel_diffs[0]}')
|
||||
|
||||
for lastX in [1, 10, 100]:
|
||||
base_avg = sum(base_losses[-lastX:]) / lastX
|
||||
test_avg = sum(test_losses[-lastX:]) / lastX
|
||||
print(
|
||||
f'last-{lastX}: base={base_avg} test={test_avg} abs={base_avg - test_avg} rel={rel_diff(base_avg, test_avg)}'
|
||||
)
|
||||
lastX = 100
|
||||
base = base_losses[-lastX:]
|
||||
base_avg = sum(base) / len(base)
|
||||
test = test_losses[-lastX:]
|
||||
test_avg = sum(test) / len(test)
|
||||
assert rel_diff(base_avg, test_avg) < 0.05
|
||||
|
||||
# the following check could passed on higher version docker: nvcr.io/nvidia/pytorch:23.07-py3(torch2.1.0 cuda12.1)
|
||||
# Check if models have same weights after training
|
||||
# self._check_model_params_equal(base_model, test_model)
|
||||
|
||||
|
||||
class DynamicShapeTestLayer(nn.Module):
|
||||
|
||||
def __init__(self, hidden_size):
|
||||
super().__init__()
|
||||
self.fc = nn.Linear(hidden_size, hidden_size)
|
||||
self.shapes = set()
|
||||
|
||||
def forward(self, x):
|
||||
self.shapes.add(x.shape)
|
||||
y = self.fc(x)
|
||||
return y
|
||||
|
||||
|
||||
class DynamicShapeTestModel(nn.Module):
|
||||
|
||||
def __init__(self, n_layers, hidden_size):
|
||||
super().__init__()
|
||||
self.layers = nn.ModuleList([DynamicShapeTestLayer(hidden_size) for _ in range(n_layers)])
|
||||
|
||||
|
||||
@pytest.mark.parametrize('topo_config', [
|
||||
{
|
||||
"num_pp": 1,
|
||||
"num_dp": 4
|
||||
},
|
||||
{
|
||||
"num_pp": 2,
|
||||
"num_dp": 2
|
||||
},
|
||||
{
|
||||
"num_pp": 4,
|
||||
"num_dp": 1
|
||||
},
|
||||
])
|
||||
class TestPipeDynamicShape(DistributedTest):
|
||||
world_size = 4
|
||||
|
||||
def test_pipe_base(self, topo_config):
|
||||
"""This test checks if the pipeline engine can handle dynamic shapes correctly.
|
||||
We pass inputs of different shapes to the pipeline engine.
|
||||
"""
|
||||
|
||||
n_iter = 10
|
||||
n_layers = 4
|
||||
n_samples = 1024
|
||||
batch_size = 4
|
||||
channel_dims = [8, 16, 32, 64]
|
||||
hidden_size = 16
|
||||
|
||||
topo = PipeTopo(**topo_config)
|
||||
|
||||
model = DynamicShapeTestModel(n_layers, hidden_size)
|
||||
model = PipelineModule(layers=model.layers, topology=topo, loss_fn=nn.MSELoss(), dynamic_shape=True)
|
||||
|
||||
# Each batch has different channel dim but we use the same channel dim in the same batch
|
||||
xs = [
|
||||
torch.randn(channel_dims[(i // batch_size) % len(channel_dims)], hidden_size, dtype=torch.float32)
|
||||
for i in range(n_samples)
|
||||
]
|
||||
ys = [torch.randn_like(x) for x in xs]
|
||||
|
||||
class CustomDataset(torch.utils.data.Dataset):
|
||||
|
||||
def __init__(self, xs, ys):
|
||||
self.xs = xs
|
||||
self.ys = ys
|
||||
|
||||
def __len__(self):
|
||||
return len(self.xs)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.xs[idx], self.ys[idx]
|
||||
|
||||
dataset = CustomDataset(xs, ys)
|
||||
|
||||
config_dict["train_batch_size"] = batch_size
|
||||
|
||||
with no_child_process_in_deepspeed_io():
|
||||
engine, _, _, _ = deepspeed.initialize(config=config_dict,
|
||||
model=model,
|
||||
model_parameters=[p for p in model.parameters()],
|
||||
training_data=dataset)
|
||||
|
||||
for _ in range(n_iter):
|
||||
_ = engine.train_batch()
|
||||
|
||||
# Check if all layers have seen different shapes
|
||||
for layer in model.modules():
|
||||
if isinstance(layer, DynamicShapeTestLayer):
|
||||
assert len(layer.shapes) > 1
|
||||
@@ -0,0 +1,143 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import pytest
|
||||
import deepspeed.runtime.pipe.schedule as schedule
|
||||
|
||||
|
||||
def _count_type(cmds, classtype):
|
||||
return len(list(filter(lambda c: type(c) == classtype, cmds)))
|
||||
|
||||
|
||||
def test_pipe_inference_schedule_singlestage():
|
||||
sched = schedule.InferenceSchedule(micro_batches=4, stages=1, stage_id=0)
|
||||
assert sched.num_micro_batches == 4
|
||||
full = list(iter(sched))
|
||||
for idx, cmds in enumerate(full):
|
||||
assert len(cmds) == 2
|
||||
assert type(cmds[0]) == schedule.LoadMicroBatch
|
||||
assert type(cmds[1]) == schedule.ForwardPass
|
||||
assert cmds[0].buffer_id == cmds[1].buffer_id
|
||||
assert len(full) == sched.num_micro_batches
|
||||
|
||||
|
||||
def test_pipe_train_schedule_singlestage():
|
||||
sched = schedule.TrainSchedule(micro_batches=4, stages=1, stage_id=0)
|
||||
assert sched.num_micro_batches == 4
|
||||
full = list(iter(sched))
|
||||
for idx, cmds in enumerate(full):
|
||||
if (idx % 2) != 0:
|
||||
assert (len(cmds) == 1) or (len(cmds) == 4)
|
||||
assert type(cmds[0]) == schedule.BackwardPass
|
||||
else:
|
||||
assert len(cmds) == 2
|
||||
assert type(cmds[0]) == schedule.LoadMicroBatch
|
||||
assert type(cmds[1]) == schedule.ForwardPass
|
||||
assert cmds[0].buffer_id == cmds[1].buffer_id
|
||||
assert len(full) == sched.num_micro_batches * 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize('micro_batches', [1, 3, 8, 10])
|
||||
def test_pipe_inference_schedule_firststage(micro_batches, stages=3):
|
||||
sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=0)
|
||||
assert sched.num_micro_batches == micro_batches
|
||||
full = list(iter(sched))
|
||||
for idx, cmds in enumerate(full):
|
||||
# Ensure we don't send an activation the first step
|
||||
if idx == 0:
|
||||
assert len(cmds) == 2
|
||||
assert type(cmds[0]) == schedule.LoadMicroBatch
|
||||
assert type(cmds[1]) == schedule.ForwardPass
|
||||
assert cmds[0].buffer_id == cmds[1].buffer_id
|
||||
continue
|
||||
|
||||
# the last active step is only a send
|
||||
if idx == sched.num_micro_batches:
|
||||
assert len(cmds) == 1
|
||||
assert type(cmds[0]) == schedule.SendActivation
|
||||
continue
|
||||
|
||||
# no work later on
|
||||
if idx > sched.num_micro_batches:
|
||||
assert len(cmds) == 0
|
||||
continue
|
||||
|
||||
# Normally we need to load/forward/send
|
||||
assert len(cmds) == 3
|
||||
assert _count_type(cmds, schedule.LoadMicroBatch) == 1
|
||||
assert _count_type(cmds, schedule.ForwardPass) == 1
|
||||
assert _count_type(cmds, schedule.SendActivation) == 1
|
||||
assert len(full) == micro_batches + stages - 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize('micro_batches', [1, 3, 8, 10])
|
||||
def test_pipe_inference_schedule_midstage(micro_batches, stages=3):
|
||||
sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=1)
|
||||
|
||||
full = list(iter(sched))
|
||||
for idx, cmds in enumerate(full):
|
||||
if idx < sched.stage:
|
||||
assert len(cmds) == 0
|
||||
continue
|
||||
if idx == sched.stage + sched.num_micro_batches:
|
||||
assert len(cmds) == 1
|
||||
assert type(cmds[0]) == schedule.SendActivation
|
||||
continue
|
||||
if idx > sched.stage + sched.num_micro_batches:
|
||||
assert len(cmds) == 0
|
||||
continue
|
||||
assert _count_type(cmds, schedule.LoadMicroBatch) == 0
|
||||
assert _count_type(cmds, schedule.ForwardPass) == 1
|
||||
assert _count_type(cmds, schedule.RecvActivation) == 1
|
||||
if idx > sched.stage:
|
||||
assert _count_type(cmds, schedule.SendActivation) == 1
|
||||
assert len(full) == micro_batches + stages - 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize('micro_batches', [1, 3, 8, 10])
|
||||
def test_pipe_inference_schedule_laststage(micro_batches, stages=3):
|
||||
sched = schedule.InferenceSchedule(micro_batches=micro_batches, stages=stages, stage_id=2)
|
||||
full = list(iter(sched))
|
||||
for idx, cmds in enumerate(full):
|
||||
if idx < sched.stage or idx > sched.stage + sched.num_micro_batches:
|
||||
assert len(cmds) == 0
|
||||
continue
|
||||
assert _count_type(cmds, schedule.LoadMicroBatch) == 1
|
||||
assert _count_type(cmds, schedule.ForwardPass) == 1
|
||||
assert _count_type(cmds, schedule.RecvActivation) == 1
|
||||
assert _count_type(cmds, schedule.SendActivation) == 0
|
||||
assert len(full) == micro_batches + stages - 1
|
||||
|
||||
|
||||
def test_pipe_schedule_firststage():
|
||||
sched = schedule.TrainSchedule(micro_batches=8, stages=3, stage_id=0)
|
||||
for cmds in sched:
|
||||
assert all(instr.__class__ != schedule.SendGrad for instr in cmds)
|
||||
assert all(instr.__class__ != schedule.RecvActivation for instr in cmds)
|
||||
for instr in cmds:
|
||||
if isinstance(instr, schedule.BufferOpInstruction):
|
||||
assert 0 <= instr.buffer_id < sched.num_pipe_buffers()
|
||||
|
||||
|
||||
def test_pipe_schedule_laststage():
|
||||
sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=2)
|
||||
assert len(list(iter(sched))) == 2 * (sched.micro_batches + sched.stages - 1)
|
||||
for cmds in sched:
|
||||
assert all(instr.__class__ != schedule.SendActivation for instr in cmds)
|
||||
assert all(instr.__class__ != schedule.RecvGrad for instr in cmds)
|
||||
|
||||
|
||||
def test_pipe_stagequery():
|
||||
sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=0)
|
||||
assert sched.is_first_stage
|
||||
assert not sched.is_last_stage
|
||||
|
||||
sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=1)
|
||||
assert not sched.is_first_stage
|
||||
assert not sched.is_last_stage
|
||||
|
||||
sched = schedule.TrainSchedule(stages=3, micro_batches=4, stage_id=2)
|
||||
assert not sched.is_first_stage
|
||||
assert sched.is_last_stage
|
||||
@@ -0,0 +1,226 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
# DeepSpeed Team
|
||||
|
||||
import pytest
|
||||
|
||||
import torch
|
||||
import deepspeed.comm as dist
|
||||
|
||||
from deepspeed.runtime.pipe.topology import PipelineParallelGrid as Grid
|
||||
from deepspeed.runtime.pipe.topology import ProcessTopology as Topo
|
||||
from deepspeed.runtime.pipe.topology import _prime_factors
|
||||
|
||||
from deepspeed.accelerator import get_accelerator
|
||||
from unit.common import DistributedTest
|
||||
|
||||
|
||||
def test_topology_2d():
|
||||
topo = Topo(axes=['row', 'col'], dims=[2, 2])
|
||||
|
||||
assert topo.world_size() == 4
|
||||
|
||||
assert topo.get_rank(row=0, col=0) == 0
|
||||
assert topo.get_rank(row=0, col=1) == 1
|
||||
assert topo.get_rank(row=1, col=0) == 2
|
||||
assert topo.get_rank(row=1, col=1) == 3
|
||||
|
||||
assert topo.get_axis_list(axis='row', idx=0) == [0, 1]
|
||||
assert topo.get_axis_list(axis='row', idx=1) == [2, 3]
|
||||
assert topo.get_axis_list(axis='col', idx=0) == [0, 2]
|
||||
assert topo.get_axis_list(axis='col', idx=1) == [1, 3]
|
||||
|
||||
|
||||
def test_topology_dims():
|
||||
topo = Topo(axes=['a', 'b', 'c'], dims=[2, 3, 4])
|
||||
assert topo.world_size() == 24
|
||||
assert topo.get_dim('a') == 2
|
||||
assert topo.get_dim('b') == 3
|
||||
assert topo.get_dim('c') == 4
|
||||
|
||||
|
||||
def test_topology_match():
|
||||
topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
|
||||
print(topo.filter_match(pipe=0, data=1))
|
||||
assert topo.filter_match(pipe=0, data=1) == [2, 3]
|
||||
print([topo.get_coord(r) for r in topo.filter_match(pipe=0, data=1)])
|
||||
|
||||
|
||||
def test_topology_rank_repr():
|
||||
topo = Topo(axes=['a', 'b'], dims=[2, 2])
|
||||
assert topo.get_rank_repr(rank=0) == 'a_00-b_00'
|
||||
assert topo.get_rank_repr(rank=1) == 'a_00-b_01'
|
||||
assert topo.get_rank_repr(rank=2) == 'a_01-b_00'
|
||||
assert topo.get_rank_repr(rank=3) == 'a_01-b_01'
|
||||
|
||||
assert topo.get_rank_repr(rank=3, inner_sep='+') == 'a+01-b+01'
|
||||
assert topo.get_rank_repr(rank=3, inner_sep='🤗', outer_sep='_JEFF_') == 'a🤗01_JEFF_b🤗01'
|
||||
|
||||
topo = Topo(axes=['pipe', 'data'], dims=[2, 2])
|
||||
assert topo.get_rank_repr(rank=0) == ''
|
||||
assert topo.get_rank_repr(rank=1) == ''
|
||||
assert topo.get_rank_repr(rank=2) == ''
|
||||
assert topo.get_rank_repr(rank=3) == ''
|
||||
|
||||
assert topo.get_rank_repr(rank=0, omit_axes=['pipe']) == 'data_00'
|
||||
assert topo.get_rank_repr(rank=1, omit_axes=['pipe']) == 'data_01'
|
||||
assert topo.get_rank_repr(rank=2, omit_axes=['pipe']) == 'data_00'
|
||||
assert topo.get_rank_repr(rank=3, omit_axes=['pipe']) == 'data_01'
|
||||
|
||||
assert topo.get_rank_repr(rank=0, omit_axes=[]) == 'pipe_00-data_00'
|
||||
assert topo.get_rank_repr(rank=1, omit_axes=[]) == 'pipe_00-data_01'
|
||||
assert topo.get_rank_repr(rank=2, omit_axes=[]) == 'pipe_01-data_00'
|
||||
assert topo.get_rank_repr(rank=3, omit_axes=[]) == 'pipe_01-data_01'
|
||||
|
||||
topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
|
||||
assert topo.get_rank_repr(rank=0) == 'model_00'
|
||||
assert topo.get_rank_repr(rank=1) == 'model_01'
|
||||
assert topo.get_rank_repr(rank=2) == 'model_00'
|
||||
assert topo.get_rank_repr(rank=3) == 'model_01'
|
||||
assert topo.get_rank_repr(rank=4) == 'model_00'
|
||||
assert topo.get_rank_repr(rank=5) == 'model_01'
|
||||
assert topo.get_rank_repr(rank=6) == 'model_00'
|
||||
assert topo.get_rank_repr(rank=7) == 'model_01'
|
||||
|
||||
|
||||
def test_topology_3d():
|
||||
topo = Topo(axes=['a', 'b', 'c'], dims=[2, 2, 2])
|
||||
|
||||
assert topo.get_rank(a=0, b=0, c=0) == 0
|
||||
assert topo.get_rank(a=0, b=0, c=1) == 1
|
||||
assert topo.get_rank(a=0, b=1, c=0) == 2
|
||||
assert topo.get_rank(a=0, b=1, c=1) == 3
|
||||
assert topo.get_rank(a=1, b=0, c=0) == 4
|
||||
assert topo.get_rank(a=1, b=0, c=1) == 5
|
||||
assert topo.get_rank(a=1, b=1, c=0) == 6
|
||||
assert topo.get_rank(a=1, b=1, c=1) == 7
|
||||
|
||||
assert topo.get_axis_list('a', 0) == [0, 1, 2, 3]
|
||||
assert topo.get_axis_list('a', 1) == [4, 5, 6, 7]
|
||||
assert topo.get_axis_list('b', 0) == [0, 1, 4, 5]
|
||||
assert topo.get_axis_list('b', 1) == [2, 3, 6, 7]
|
||||
assert topo.get_axis_list('c', 0) == [0, 2, 4, 6]
|
||||
assert topo.get_axis_list('c', 1) == [1, 3, 5, 7]
|
||||
|
||||
assert topo.get_coord(0) == topo.ProcessCoord(0, 0, 0)
|
||||
assert topo.get_coord(1) == topo.ProcessCoord(0, 0, 1)
|
||||
assert topo.get_coord(2) == topo.ProcessCoord(0, 1, 0)
|
||||
assert topo.get_coord(3) == topo.ProcessCoord(0, 1, 1)
|
||||
assert topo.get_coord(4) == topo.ProcessCoord(1, 0, 0)
|
||||
assert topo.get_coord(5) == topo.ProcessCoord(1, 0, 1)
|
||||
assert topo.get_coord(6) == topo.ProcessCoord(1, 1, 0)
|
||||
assert topo.get_coord(7) == topo.ProcessCoord(1, 1, 1)
|
||||
|
||||
assert topo.filter_match(a=0) == [0, 1, 2, 3]
|
||||
assert topo.filter_match(b=1, c=1) == [3, 7]
|
||||
assert topo.filter_match(a=1, b=1, c=1) == [7]
|
||||
|
||||
# Easy access method
|
||||
assert topo.get_coord(0).a == 0
|
||||
|
||||
|
||||
def test_topology_comm_list():
|
||||
topo = Topo(axes=['pipe', 'data', 'model'], dims=[2, 2, 2])
|
||||
|
||||
assert topo.get_rank(pipe=0, data=0, model=0) == 0
|
||||
assert topo.get_rank(pipe=0, data=0, model=1) == 1
|
||||
assert topo.get_rank(pipe=0, data=1, model=0) == 2
|
||||
assert topo.get_rank(pipe=0, data=1, model=1) == 3
|
||||
assert topo.get_rank(pipe=1, data=0, model=0) == 4
|
||||
assert topo.get_rank(pipe=1, data=0, model=1) == 5
|
||||
assert topo.get_rank(pipe=1, data=1, model=0) == 6
|
||||
assert topo.get_rank(pipe=1, data=1, model=1) == 7
|
||||
|
||||
pipe_list = [
|
||||
[0, 4], # data=0, model=0
|
||||
[1, 5], # data=0, model=1
|
||||
[2, 6], # data=1, model=0
|
||||
[3, 7], # data=1, model=1
|
||||
]
|
||||
assert topo.get_axis_comm_lists('pipe') == pipe_list
|
||||
|
||||
data_list = [
|
||||
[0, 2], # pipe=0, model=0
|
||||
[1, 3], # pipe=0, model=1
|
||||
[4, 6], # pipe=1, model=0
|
||||
[5, 7], # pipe=1, model=1
|
||||
]
|
||||
assert topo.get_axis_comm_lists('data') == data_list
|
||||
|
||||
model_list = [
|
||||
[0, 1], # pipe=0, data=0
|
||||
[2, 3], # pipe=0, data=1
|
||||
[4, 5], # pipe=1, data=0
|
||||
[6, 7], # pipe=1, data=1
|
||||
]
|
||||
assert topo.get_axis_comm_lists('model') == model_list
|
||||
|
||||
# Handle nonsense. We don't want to RuntimeError because it allows us to write more
|
||||
# generalized code for data/model/pipe parallelism
|
||||
assert topo.get_axis_comm_lists('jeff') == []
|
||||
|
||||
|
||||
class TestDistributedTopology(DistributedTest):
|
||||
world_size = 4
|
||||
|
||||
def test_grid_pipe_data(self):
|
||||
topo = Topo(axes=['pipe', 'data'], dims=[2, 2])
|
||||
grid = Grid(topology=topo)
|
||||
|
||||
assert grid._is_grid_valid()
|
||||
|
||||
rank = dist.get_rank()
|
||||
|
||||
assert grid.is_first_stage == (grid.get_stage_id() == 0)
|
||||
assert grid.is_last_stage == (grid.get_stage_id() == grid.get_pipe_parallel_world_size() - 1)
|
||||
|
||||
# Test collectives along the pipeline parallel process groups
|
||||
rank_tensor = torch.LongTensor(data=[rank]).to(get_accelerator().device_name())
|
||||
dist.all_reduce(rank_tensor, group=grid.get_pipe_parallel_group())
|
||||
pipe_group = grid.pp_group
|
||||
assert torch.all(rank_tensor == sum(pipe_group))
|
||||
|
||||
# Test collectives along the data parallel process groups
|
||||
rank_tensor = torch.LongTensor(data=[rank]).to(get_accelerator().device_name())
|
||||
dist.all_reduce(rank_tensor, group=grid.get_data_parallel_group())
|
||||
data_group = grid.dp_group
|
||||
assert torch.all(rank_tensor == sum(data_group))
|
||||
|
||||
def test_stage_to_global(self):
|
||||
topo = Topo(axes=['pipe', 'data'], dims=[2, 2])
|
||||
grid = Grid(topology=topo)
|
||||
|
||||
assert grid._is_grid_valid()
|
||||
|
||||
assert grid.stage_to_global(stage_id=0, data=0) == 0
|
||||
assert grid.stage_to_global(stage_id=0, data=1) == 1
|
||||
assert grid.stage_to_global(stage_id=1, data=0) == 2
|
||||
assert grid.stage_to_global(stage_id=1, data=1) == 3
|
||||
|
||||
me = topo.get_coord(rank=dist.get_rank())
|
||||
if me.data == 0:
|
||||
assert grid.stage_to_global(stage_id=0) == 0
|
||||
assert grid.stage_to_global(stage_id=1) == 2
|
||||
else:
|
||||
assert grid.stage_to_global(stage_id=0) == 1
|
||||
assert grid.stage_to_global(stage_id=1) == 3
|
||||
|
||||
|
||||
def test_primes():
|
||||
""" Test prime factorizations. """
|
||||
|
||||
def _product(ps):
|
||||
p = 1
|
||||
for num in ps:
|
||||
p *= num
|
||||
return p
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
_prime_factors(0)
|
||||
|
||||
for x in range(1, 30):
|
||||
primes = _prime_factors(x)
|
||||
assert _product(primes) == x
|
||||
for p in primes:
|
||||
assert _prime_factors(p) == [p]
|
||||
Reference in New Issue
Block a user