def test_win_mutex_full(self):
size = bf.size()
rank = bf.rank()
if size <= 2:
fname = inspect.currentframe().f_code.co_name
warnings.warn(
"Skip {} because it only supports test over at least 3 nodes".
format(fname))
return
bf.set_topology(topology_util.FullyConnectedGraph(size))
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
for dtype in dtypes:
tensor = torch.FloatTensor([DIM_SIZE]).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_mutex_full_{}".format(dtype)
bf.win_create(tensor, window_name)
if rank == 0:
with bf.win_mutex(window_name, for_self=True):
bf.barrier()
time.sleep(1.01)
else:
bf.barrier()
t_start = time.time()
with bf.win_mutex(window_name):
time.sleep(0.001)
t_end = time.time()
assert (t_end - t_start) > 1, \
"The mutex acquire time should be longer than 1 second"
assert (t_end - t_start) < 2, \
"The mutex acquire time should be shorter than 2 second"
def test_win_get(self):
"""Test that the window get operation."""
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
# By default, we use exponential two ring topology.
indegree = int(np.ceil(np.log2(size)))
neighbor_ranks = [(rank - 2**i) % size
for i in range(indegree)] # in-neighbor
avg_value = (rank + np.sum(neighbor_ranks)) / float(indegree + 1)
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_get_{}_{}".format(dim, dtype)
bf.win_create(tensor, window_name)
bf.win_get(window_name)
bf.barrier()
recv_tensor = bf.win_update(window_name, clone=True)
assert (list(recv_tensor.shape) == [DIM_SIZE] *
dim), ("bf.win_get produce wrong shape tensor.")
assert (recv_tensor.data - avg_value).abs().max() < EPSILON, (
"bf.win_get produce wrong tensor value " +
"[{}-{}]!={} at rank {}.".format(
recv_tensor.min(), recv_tensor.max(), avg_value, rank))
def test_win_update_then_collect(self):
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
indegree = int(np.ceil(np.log2(size)))
expected_result = rank * (indegree + 1)
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_update_collect_{}_{}".format(dim, dtype)
bf.win_create(tensor, window_name)
# After the collect ops, the neighbro tensor will become zero.
# So second win_update_then_collect should produce the same value.
for _ in range(2):
collect_tensor = bf.win_update_then_collect(window_name)
assert (list(collect_tensor.shape) == [DIM_SIZE] * dim), (
"bf.win_update_then_collect produces wrong shape tensor.")
assert (collect_tensor.data - expected_result).abs().max(
) < EPSILON, (
"bf.win_update_then_collect produces wrong tensor value " +
"[{0}-{1}]!={2} at rank {2}.".format(
collect_tensor.min(), collect_tensor.max(), rank))
def test_timeline_push_sum(self):
# Use win_accumulate to simulate the push-sum algorithm (sync).
outdegree = len(bf.out_neighbor_ranks())
indegree = len(bf.in_neighbor_ranks())
# we append the p at the last of data.
x = torch.Tensor(
[bf.rank() / (indegree + 1), 1.0 / bf.size() / (indegree + 1)])
# Remember we do not create buffer with 0.
bf.win_create(x, name="x_buff")
x = bf.win_update_then_collect(name="x_buff")
for _ in range(10):
bf.win_accumulate(x,
name="x_buff",
dst_weights={
rank: 1.0 / (outdegree + 1)
for rank in bf.out_neighbor_ranks()
},
require_mutex=True)
x.div_(1 + outdegree)
x = bf.win_update_then_collect(name="x_buff")
bf.barrier()
# Do not forget to sync at last!
x = bf.win_update_then_collect(name="x_buff")
file_name = f"{self.temp_file}{bf.rank()}.json"
with open(file_name, 'r') as tf:
timeline_text = tf.read()
assert 'MPI_WIN_ACCUMULATE' in timeline_text, timeline_text
assert 'ENQUEUE_WIN_ACCUMULATE' in timeline_text, timeline_text
bf.win_free()
def test_get_win_version_with_win_get(self):
"""Test version window is initialized, updated and cleared correctly with win get."""
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
# By default, we use exponential two ring topology.
indegree = int(np.ceil(np.log2(size)))
neighbor_ranks = [(rank - 2**i) % size
for i in range(indegree)] # in-neighbor
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([23] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_version_get_{}_{}".format(dim, dtype)
bf.win_create(tensor, window_name)
original_versions = list(bf.get_win_version(window_name).values())
bf.barrier()
bf.win_get(window_name)
bf.barrier()
versions_after_win_get = list(
bf.get_win_version(window_name).values())
bf.win_update(window_name, clone=True)
versions_after_win_update = list(
bf.get_win_version(window_name).values())
neighbor_ranks_number = len(neighbor_ranks)
zero_number_in_original_versions = len(
original_versions) - np.count_nonzero(original_versions)
assert ((zero_number_in_original_versions) == neighbor_ranks_number
), ("version initialization is wrong.")
zero_number_after_win_update = len(
versions_after_win_update) - np.count_nonzero(
versions_after_win_update)
assert ((zero_number_after_win_update) == neighbor_ranks_number), (
"version clear up is wrong.")
expected_versions_after_win_get = [1] * neighbor_ranks_number
assert (versions_after_win_get == expected_versions_after_win_get
), ("version after win get is wrong.")
for dtype, dim in itertools.product(dtypes, dims):
window_name = "win_version_get_{}_{}".format(dim, dtype)
is_freed = bf.win_free(window_name)
assert is_freed, "bf.win_free do not free window object successfully."
def test_asscoicated_with_p(self):
size = bf.size()
rank = bf.rank()
if size <= 3:
fname = inspect.currentframe().f_code.co_name
warnings.warn(
"Skip {} because it only supports test over at least 3 nodes".
format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU and not bf.nccl_built():
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
bf.set_topology(topology_util.RingGraph(size))
bf.turn_on_win_ops_with_associated_p()
for dtype, send_rank in itertools.product(dtypes, range(size)):
tensor = torch.FloatTensor([23]).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_asscoicate_with_p_{}_{}".format(
dtype, send_rank)
bf.win_create(tensor, window_name)
left_neighbor_rank = (send_rank - 1) % size
right_neighbor_rank = (send_rank + 1) % size
if rank == send_rank:
bf.win_accumulate(tensor,
name=window_name,
self_weight=0.5,
dst_weights={
left_neighbor_rank: 0.5,
right_neighbor_rank: 0.5
})
bf.barrier()
bf.win_update_then_collect(name=window_name)
associated_p = bf.win_associated_p(name=window_name)
if rank == send_rank:
assert associated_p == 0.5, (
"associated_p for sender {} is wrong. Get {}".format(
rank, associated_p))
elif (rank == left_neighbor_rank) or (rank == right_neighbor_rank):
assert (associated_p - 1.5) < EPSILON, (
"associated_p for received neighbor {} is wrong. Get {}".
format(rank, associated_p))
else:
assert associated_p == 1.0, (
"associated_p for untouched node {} is wrong. Get {}".
format(rank, associated_p))
bf.turn_off_win_ops_with_associated_p()
def test_win_update_with_given_weights(self):
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_create_{}_{}".format(dim, dtype)
is_created = bf.win_create(tensor, window_name)
assert is_created, "bf.win_create do not create window object successfully."
# Test simple average rule.
weight = 1.0 / (len(bf.in_neighbor_ranks()) + 1)
sync_result = bf.win_update(
window_name,
self_weight=weight,
neighbor_weights={x: weight
for x in bf.in_neighbor_ranks()})
assert (list(sync_result.shape) == [DIM_SIZE] * dim), (
"bf.win_update (weighted) produces wrong shape tensor.")
assert (sync_result.data - rank).abs().max() < EPSILON, (
"bf.win_update (weighted) produces wrong tensor value " +
"[{0}-{1}]!={2} at rank {2}.".format(sync_result.min(),
sync_result.max(), rank))
def test_set_topology_fail_with_win_create(self):
bf.init()
size = bf.size()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
tensor = torch.FloatTensor([1])
window_name = "win_create_test"
is_created = bf.win_create(tensor, window_name)
assert is_created, "bf.win_create do not create window object successfully."
if size == 1:
expected_topology = nx.from_numpy_array(np.array([[0.5]]),
create_using=nx.DiGraph)
elif size == 2:
expected_topology = nx.from_numpy_array(np.array([[0, 0.2],
[0.2, 0]]),
create_using=nx.DiGraph)
else:
expected_topology = RingGraph(size)
is_set = bf.set_topology(expected_topology)
assert not is_set, "bf.set_topology do not fail due to win_create."
topology = bf.load_topology()
assert isinstance(topology, nx.DiGraph)
assert IsTopologyEquivalent(topology, ExponentialGraph(size))
is_freed = bf.win_free()
assert is_freed, "bf.win_free do not free window object successfully."
def test_win_put_with_varied_tensor_elements(self):
"""Test that the window put operation."""
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
# By default, we use exponential two ring topology.
indegree = int(np.ceil(np.log2(size)))
neighbor_ranks = [(rank - 2**i) % size
for i in range(indegree)] # in-neighbor
avg_value = (rank + np.sum(neighbor_ranks)) / float(indegree + 1)
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
base_tensor = torch.arange(
DIM_SIZE**dim, dtype=torch.float32).view_as(tensor).div(1000)
tensor = self.cast_and_place(tensor, dtype)
base_tensor = self.cast_and_place(base_tensor, dtype)
tensor = tensor + base_tensor
window_name = "win_put_{}_{}".format(dim, dtype)
bf.win_create(tensor, window_name)
bf.win_put(tensor, window_name)
bf.barrier()
sync_result = bf.win_update(window_name)
assert (list(sync_result.shape) == [DIM_SIZE] * dim), (
"bf.win_update after win_put produces wrong shape tensor.")
assert (
(sync_result - base_tensor).data -
avg_value).abs().max() < EPSILON, (
"bf.win_update after win_put produces wrong tensor value "
+ "[{}-{}]!={} at rank {}.".format(
(sync_result - base_tensor).min(),
(sync_result - base_tensor).max(), avg_value, rank))
time.sleep(0.5)
for dtype, dim in itertools.product(dtypes, dims):
window_name = "win_put_{}_{}".format(dim, dtype)
is_freed = bf.win_free(window_name)
assert is_freed, "bf.win_free do not free window object successfully."
def test_win_get_names(self):
tensor_1 = torch.FloatTensor([1])
tensor_2 = torch.FloatTensor([2])
win_names = bf.get_current_created_window_names()
assert not win_names
assert bf.win_create(tensor_1, "1")
win_names = bf.get_current_created_window_names()
assert win_names == ["1"]
assert bf.win_create(tensor_2, "2")
win_names = bf.get_current_created_window_names()
assert win_names == ["1", "2"]
assert bf.win_free("1")
win_names = bf.get_current_created_window_names()
assert win_names == ["2"]
assert bf.win_free()
win_names = bf.get_current_created_window_names()
assert not win_names
def test_asscoicated_with_p_random_test(self):
size = bf.size()
rank = bf.rank()
dtypes = [torch.FloatTensor, torch.DoubleTensor]
# Current, nccl version hasn't supported the associated with p yet.
if TEST_ON_GPU and not bf.nccl_built():
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
dims = [1]
bf.turn_on_win_ops_with_associated_p()
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([23] * dim)).fill_(1)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_asscoicate_with_p_random_{}_{}".format(
dim, dtype)
bf.win_create(tensor, window_name, zero_init=True)
for _ in range(10):
random_weights = np.random.rand(
len(bf.out_neighbor_ranks()) + 1)
random_weights /= random_weights.sum()
self_weight = random_weights[-1]
dst_weights = {
r: random_weights[i]
for i, r in enumerate(bf.out_neighbor_ranks())
}
bf.win_put(tensor,
self_weight=self_weight,
dst_weights=dst_weights,
name=window_name,
require_mutex=True)
bf.win_update(name=window_name, require_mutex=True)
bf.win_accumulate(tensor,
name=window_name,
require_mutex=True,
self_weight=self_weight,
dst_weights=dst_weights)
bf.win_update_then_collect(name=window_name)
bf.barrier()
bf.win_update_then_collect(name=window_name)
associated_p = bf.win_associated_p(name=window_name)
# Because the associated p should operate the same as tensor always
# the following assert should be true no matter what order is excuted.
assert abs(associated_p - tensor.data[0]) < EPSILON
bf.turn_off_win_ops_with_associated_p()
def test_win_put_with_given_destination(self):
"""Test that the window put operation with given destination."""
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
# By default, we use exponential two ring topology.
indegree = int(np.ceil(np.log2(size)))
# We use given destination to form a (right-)ring.
avg_value = (rank * indegree + 1.23 *
((rank - 1) % size)) / float(indegree + 1)
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_put_given_{}_{}".format(dim, dtype)
bf.win_create(tensor, window_name)
bf.win_put(tensor,
window_name,
dst_weights={(rank + 1) % size: 1.23})
bf.barrier()
sync_result = bf.win_update(window_name)
assert (list(sync_result.shape) == [DIM_SIZE] * dim), (
"bf.win_update after win_put given destination produces wrong shape tensor."
)
assert (sync_result.data - avg_value).abs().max() < EPSILON, (
"bf.win_update after win_put given destination produces wrong tensor value "
+ "[{}-{}]!={} at rank {}.".format(
sync_result.min(), sync_result.max(), avg_value, rank))
time.sleep(0.5)
for dtype, dim in itertools.product(dtypes, dims):
window_name = "win_put_given_{}_{}".format(dim, dtype)
is_freed = bf.win_free(window_name)
assert is_freed, "bf.win_free do not free window object successfully."
def _register_window(self):
for param_group in self.param_groups:
for p in param_group["params"]:
name = self._parameter_names.get(p)
if name is None:
raise KeyError(
"Cannot find parameter {} in the _parameter_names dictionary"
.format(name))
if not bf.win_create(p.data, name):
raise ValueError(
"Cannot allocate MPI window for the parameter {}".
format(name))
def test_win_mutex_given_ranks(self):
size = bf.size()
rank = bf.rank()
if size < 4:
fname = inspect.currentframe().f_code.co_name
warnings.warn(
"Skip {} because it only supports test above 4 nodes".format(
fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
for dtype in dtypes:
tensor = torch.FloatTensor([DIM_SIZE]).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_mutex_given_ranks_{}".format(dtype)
bf.win_create(tensor, window_name)
if rank == 0:
with bf.win_mutex(window_name, for_self=True, ranks=[1]):
bf.barrier()
time.sleep(1.01)
elif rank == 1:
bf.barrier()
t_start = time.time()
with bf.win_mutex(window_name, ranks=[0]):
time.sleep(0.001)
t_end = time.time()
assert (t_end - t_start) > 1
elif rank == 2:
bf.barrier()
t_start = time.time()
with bf.win_mutex(window_name, ranks=[0]):
time.sleep(0.001)
t_end = time.time()
assert (t_end - t_start) < 0.1
else:
bf.barrier()
def test_win_accumulate_with_given_destination(self):
"""Test that the window accumulate operation with given destination."""
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
avg_value = rank + ((rank - 1) % size) * 1.23 / 2.0
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_accumulate_{}_{}".format(dim, dtype)
bf.win_create(tensor, window_name)
bf.win_accumulate(tensor,
window_name,
dst_weights={(rank + 1) % size: 1.23})
bf.barrier()
sync_result = bf.win_update(window_name,
self_weight=0.5,
neighbor_weights={
(rank - 1) % size: 0.5
})
assert (list(sync_result.shape) == [DIM_SIZE] * dim), (
"bf.win_update after win_accmulate given destination produces wrong shape tensor."
)
assert (sync_result.data - avg_value).abs().max() < EPSILON, (
"bf.win_update after win_accmulate given destination produces wrong tensor value "
+ "[{}-{}]!={} at rank {}.".format(
sync_result.min(), sync_result.max(), avg_value, rank))
def test_win_get_with_given_sources(self):
"""Test that the window get operation with given sources."""
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
# We use given destination to form a (right-)ring.
avg_value = (rank + 1.23 * ((rank - 1) % size)) / float(2)
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_get_given_{}_{}".format(dim, dtype)
bf.win_create(tensor, window_name)
bf.win_get(window_name, src_weights={(rank - 1) % size: 1.23})
bf.barrier()
recv_tensor = bf.win_update(window_name,
self_weight=0.5,
neighbor_weights={
(rank - 1) % size: 0.5
},
clone=True)
assert (list(recv_tensor.shape) == [DIM_SIZE] * dim), (
"bf.win_get with given sources produces wrong shape tensor.")
assert (recv_tensor.data - avg_value).abs().max() < EPSILON, (
"bf.win_get with given sources produces wrong tensor value " +
"[{}-{}]!={} at rank {}.".format(
recv_tensor.min(), recv_tensor.max(), avg_value, rank))
def test_win_free_all(self):
size = bf.size()
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_create_{}_{}".format(dim, dtype)
is_created = bf.win_create(tensor, window_name)
assert is_created, "bf.win_create do not create window object successfully."
is_freed = bf.win_free()
assert is_freed, "bf.win_free do not free window object successfully."
def _register_window(self):
for param_group in self.param_groups:
for p in param_group["params"]:
name = self._parameter_names.get(p)
if name is None:
raise KeyError(
"Cannot find parameter {} in the _parameter_names dictionary"
.format(name))
ps_weights = torch.Tensor([1.0]).to(p.data.dtype).to(
p.data.device)
self._named_ps_weights[name] = ps_weights
# If do not modify in the C level, it is inevitable to copy
# the parameter once in the cat ops.
extended_parameter = torch.cat((p.data.view(-1), ps_weights),
0)
self._named_extension_parameters[name] = extended_parameter
if not bf.win_create(extended_parameter, name, zero_init=True):
raise ValueError(
"Cannot allocate MPI window for the parameter {}".
format(name))
def test_win_update_with_default_weights(self):
size = bf.size()
rank = bf.rank()
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor]
bf.set_topology(topology_util.StarGraph(size), is_weighted=True)
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_create_{}_{}".format(dim, dtype)
is_created = bf.win_create(tensor, window_name)
assert is_created, "bf.win_create do not create window object successfully."
# Note the buffers store the copy of original value so they will not change.
tensor.mul_(2)
if rank == 0:
expected_result = rank * 2 / size + rank * (size - 1) / size
else:
expected_result = rank / size + rank * 2 * (1 - 1 / size)
sync_result = bf.win_update(window_name)
assert (list(sync_result.shape) == [DIM_SIZE] * dim), (
"bf.win_update (weighted) produces wrong shape tensor.")
assert (
sync_result.data - expected_result).abs().max() < EPSILON, (
"bf.win_update (weighted) produces wrong tensor value " +
"[{0}-{1}]!={2} at rank {2}.".format(
sync_result.min(), sync_result.max(), rank))
assert bf.win_free()
def test_win_create_and_sync_and_free(self):
"""Test that the window create and free objects correctly."""
size = bf.size()
rank = bf.rank()
# OpenMPI implementation seems won't allow win_create on size 1.
if size <= 1:
fname = inspect.currentframe().f_code.co_name
warnings.warn("Skip {} due to size 1".format(fname))
return
dtypes = [torch.FloatTensor, torch.DoubleTensor]
if TEST_ON_GPU:
dtypes += [torch.cuda.FloatTensor, torch.cuda.DoubleTensor]
# By default, we use exponential two ring topology.
dims = [1, 2, 3]
for dtype, dim in itertools.product(dtypes, dims):
tensor = torch.FloatTensor(*([DIM_SIZE] * dim)).fill_(1).mul_(rank)
tensor = self.cast_and_place(tensor, dtype)
window_name = "win_create_{}_{}".format(dim, dtype)
is_created = bf.win_create(tensor, window_name)
assert is_created, "bf.win_create do not create window object successfully."
sync_result = bf.win_update(window_name)
assert (list(sync_result.shape) == [DIM_SIZE] *
dim), ("bf.win_update produce wrong shape tensor.")
assert (sync_result.data.min() == rank), (
"bf.win_update produces wrong tensor value " +
"{0}!={1} at rank {1}.".format(sync_result.data.min(), rank))
assert (sync_result.data.max() == rank), (
"bf.win_update produces wrong tensor value " +
"{0}!={1} at rank {1}.".format(sync_result.data.max(), rank))
for dtype, dim in itertools.product(dtypes, dims):
window_name = "win_create_{}_{}".format(dim, dtype)
is_freed = bf.win_free(window_name)
assert is_freed, "bf.win_free do not free window object successfully."
self_weight = 1 / (len(recv_neighbors) + 1)
x = bf.neighbor_allreduce(x,
name='x',
self_weight=self_weight,
neighbor_weights=neighbor_weights,
send_neighbors=send_neighbors,
enable_topo_check=False)
mse.append(torch.norm(x - x_bar, p=2) / torch.norm(x_bar, p=2))
else:
outdegree = len(bf.out_neighbor_ranks())
indegree = len(bf.in_neighbor_ranks())
if not bf.nccl_built(): # NCCL do not support associated P yet.
bf.turn_on_win_ops_with_associated_p()
bf.win_create(x, name="x", zero_init=True)
for i in range(args.max_iters):
if args.enable_dynamic_topology:
num_out_neighbors = len(bf.out_neighbor_ranks())
sent_neighbor = bf.out_neighbor_ranks()[i % num_out_neighbors]
dst_weights = {sent_neighbor: 0.5}
self_weight = 0.5
else:
dst_weights = {
rank: 1.0 / (outdegree + 1)
for rank in bf.out_neighbor_ranks()
}
self_weight = 1 / (1 + outdegree)
bf.win_accumulate(x,
name="x",
def push_diging(X, y, w_opt, loss, maxite=2000, alpha=1e-1, **kwargs):
if loss == 'logistic_regression':
rho = kwargs.get('rho', 1e-1)
elif loss == 'linear_regression':
rho = 0
else:
raise NotImplementedError(
'Task not supported. This example only supports' +
' linear_regression and logistic_regression')
outdegree = len(bf.out_neighbor_ranks())
indegree = len(bf.in_neighbor_ranks())
# We let w = col{u, y, v}, i.e., u, y, v = w[:n], w[n:2*n], w[2n]
# Insteady of three directed_neighbor_allreduce operations for u, y,
# and v respectively, we exploit one directed_neighbor_allreduce for
# the combo vector w. This guarantees u, y, and v to be transmitted
# simultanesly and avoids the mismatch between them. Experiments
# show directed_neighbor_allreduce(w) is crutial for convergence of
# push_diging.
w = torch.zeros(2 * n + 1, 1).to(torch.double)
x = torch.zeros(n, 1, dtype=torch.double, requires_grad=True)
loss_step(X, y, x, tensor_name='w_buff', loss=loss, rho=rho)
grad = x.grad.data.clone()
w[n:2 * n] = grad
x.grad.data.zero_()
w[-1] = 1.0
grad_prev = w[n:2 * n].clone()
bf.win_create(w, name="w_buff", zero_init=True)
mse = []
for _ in range(maxite):
bf.barrier()
w[:n] = w[:n] - alpha * w[n:2 * n]
bf.win_accumulate(w,
name="w_buff",
dst_weights={
rank: 1.0 / (outdegree * 2)
for rank in bf.out_neighbor_ranks()
},
require_mutex=True)
w.div_(2)
bf.barrier()
w = bf.win_update_then_collect(name="w_buff")
x.data = w[:n] / w[-1]
loss_step(X, y, x, tensor_name='w_buff', loss=loss, rho=rho)
grad = x.grad.data.clone()
x.grad.data.zero_()
w[n:2 * n] += grad - grad_prev
grad_prev = grad
if bf.rank() == 0:
mse.append(torch.norm(x.data - w_opt, p=2))
bf.barrier()
w = bf.win_update_then_collect(name="w_buff")
x.data = w[:n] / w[-1]
return x, mse
