def test_excepts_mismatched_partitions(barrier_fence_fixture, comm_split_fixture): import numpy as np from distdl.backends.mpi.partition import MPIPartition from distdl.nn.transpose import DistributedTranspose # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) in_shape = (1, 4, 1, 1) out_shape = (1, 2) in_size = np.prod(in_shape) out_size = np.prod(out_shape) # Create the partitions P_x_base = P_world.create_partition_inclusive(np.arange(0, in_size)) P_x = P_x_base.create_cartesian_topology_partition(in_shape) P_y_base = P_world.create_partition_inclusive(np.arange(P_world.size-out_size, P_world.size)) P_y = P_y_base.create_cartesian_topology_partition(out_shape) with pytest.raises(ValueError) as e_info: # noqa: F841 DistributedTranspose(P_x, P_y)
def test_linear_adjoint_input(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, P_w_ranks, P_w_shape, x_global_shape, y_global_shape): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.linear import DistributedLinear from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) P_w_base = P_world.create_partition_inclusive(P_w_ranks) P_w = P_w_base.create_cartesian_topology_partition(P_w_shape) x_global_shape = np.asarray(x_global_shape) y_global_shape = np.asarray(y_global_shape) layer = DistributedLinear(P_x, P_y, P_w, x_global_shape[1], y_global_shape[1], bias=False) x = zero_volume_tensor(x_global_shape[0]) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.Tensor(np.random.randn(*x_local_shape)) x.requires_grad = True y = layer(x) dy = zero_volume_tensor(x_global_shape[0]) if P_y.active: dy = torch.Tensor(np.random.randn(*y.shape)) y.backward(dy) dx = x.grad x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy)
def test_potentially_deadlocked_send_recv_pairs(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_w_ranks, P_w_shape): from distdl.backends.mpi.partition import MPIPartition from distdl.nn.broadcast import Broadcast device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_w_base = P_world.create_partition_inclusive(P_w_ranks) P_w = P_w_base.create_cartesian_topology_partition(P_w_shape) layer = Broadcast(P_x, P_w) # noqa F841 layer = layer.to(device) P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_w_base.deactivate() P_w.deactivate()
def test_simple_conv2d_adjoint_weight(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, x_global_shape): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.conv_feature import DistributedFeatureConv2d from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) x_global_shape = np.asarray(x_global_shape) layer = DistributedFeatureConv2d(P_x, in_channels=x_global_shape[1], out_channels=10, kernel_size=[3, 3], bias=False) x = zero_volume_tensor(x_global_shape[0]) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.randn(*x_local_shape) x.requires_grad = True y = layer(x) dy = zero_volume_tensor(x_global_shape[0]) if P_x.active: dy = torch.randn(*y.shape) y.backward(dy) W = zero_volume_tensor() dW = zero_volume_tensor() if P_x.active: W = layer.weight.detach() dW = layer.weight.grad.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, W, dW, y, dy) P_world.deactivate() P_x_base.deactivate() P_x.deactivate()
def test_transpose_adjoint(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, x_global_shape): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.transpose import DistributedTranspose from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) # The global tensor size is the same for x and y layer = DistributedTranspose(P_x, P_y, preserve_batch=False) # Forward Input x = zero_volume_tensor() if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.Tensor(np.random.randn(*x_local_shape)) x.requires_grad = True # Adjoint Input dy = zero_volume_tensor() if P_y.active: y_local_shape = compute_subshape(P_y.shape, P_y.index, x_global_shape) dy = torch.Tensor(np.random.randn(*y_local_shape)) # y = F @ x y = layer(x) # dx = F* @ dy y.backward(dy) dx = x.grad x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy)
def test_broadcast_adjoint(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, x_global_shape, transpose_src): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.broadcast import Broadcast from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) # TODO #93: Change this to create a subtensor so we test when local tensors # have different shape. Then, the output size will also be different, which # we will have to get from `y` itself. x_local_shape = np.asarray(x_global_shape) layer = Broadcast(P_x, P_y, transpose_src=transpose_src, preserve_batch=False) x = zero_volume_tensor() if P_x.active: x = torch.Tensor(np.random.randn(*x_local_shape)) x.requires_grad = True dy = zero_volume_tensor() if P_y.active: # Adjoint Input dy = torch.Tensor(np.random.randn(*x_local_shape)) # y = F @ x y = layer(x) # dx = F* @ dy y.backward(dy) dx = x.grad x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy)
def test_excepts_mismatched_nondivisible_tensor(barrier_fence_fixture, comm_split_fixture): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.repartition import Repartition from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # A tensor with size 1 in a dimension cannot be partitioned in that # dimension. (See last dimension of output and tensor.) in_shape = (1, 4, 1, 1) out_shape = (1, 1, 1, 2) x_global_shape = np.array([1, 16, 5, 1]) in_size = np.prod(in_shape) out_size = np.prod(out_shape) # Create the partitions P_x_base = P_world.create_partition_inclusive(np.arange(0, in_size)) P_x = P_x_base.create_cartesian_topology_partition(in_shape) P_y_base = P_world.create_partition_inclusive( np.arange(P_world.size - out_size, P_world.size)) P_y = P_y_base.create_cartesian_topology_partition(out_shape) with pytest.raises(ValueError) as e_info: # noqa: F841 layer = Repartition(P_x, P_y) layer = layer.to(device) # Forward Input x = zero_volume_tensor(device=device) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.randn(*x_local_shape, device=device) x.requires_grad = True layer(x) P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_y_base.deactivate() P_y.deactivate()
def test_conv_class_selection(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, P_w_ranks, P_w_shape, kernel_size, InputLayerType, OutputLayerType): from distdl.backends.mpi.partition import MPIPartition # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) if P_y_ranks is not None: P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) else: P_y_base = None P_y = None if P_w_ranks is not None: P_w_base = P_world.create_partition_inclusive(P_w_ranks) P_w = P_w_base.create_cartesian_topology_partition(P_w_shape) else: P_w_base = None P_w = None layer = InputLayerType(P_x, P_y=P_y, P_w=P_w, in_channels=3, out_channels=3, kernel_size=kernel_size) assert type(layer) == OutputLayerType P_world.deactivate() P_x_base.deactivate() P_x.deactivate() if P_y_ranks is not None: P_y_base.deactivate() P_y.deactivate() if P_w_ranks is not None: P_w_base.deactivate() P_w.deactivate()
def test_simple_conv2d_shape(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, x_global_shape, y_local_shape, padding): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.conv_feature import DistributedFeatureConv2d from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) x_global_shape = np.asarray(x_global_shape) layer = DistributedFeatureConv2d(P_x, in_channels=x_global_shape[1], out_channels=10, kernel_size=[3, 3], padding=padding, bias=False) x = zero_volume_tensor(x_global_shape[0]) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.zeros(*x_local_shape) x.requires_grad = True y = layer(x) if P_x.active: assert(np.array_equal(np.array(y.shape), np.asarray(y_local_shape))) P_world.deactivate() P_x_base.deactivate() P_x.deactivate()
def test_excepts_mismatched_output_partition_tensor(barrier_fence_fixture, comm_split_fixture): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.transpose import DistributedTranspose from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Output partition rank must match tensor rank in_shape = (4, 1, 1) out_shape = (1, 1, 1, 2) x_global_shape = np.array([16, 5, 5]) in_size = np.prod(in_shape) out_size = np.prod(out_shape) # Create the partitions P_x_base = P_world.create_partition_inclusive(np.arange(0, in_size)) P_x = P_x_base.create_cartesian_topology_partition(in_shape) P_y_base = P_world.create_partition_inclusive(np.arange(P_world.size-out_size, P_world.size)) P_y = P_y_base.create_cartesian_topology_partition(out_shape) with pytest.raises(ValueError) as e_info: # noqa: F841 layer = DistributedTranspose(P_x, P_y) # Forward Input x = zero_volume_tensor() if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.Tensor(np.random.randn(*x_local_shape)) x.requires_grad = True layer(x)
def test_excepts_no_match(barrier_fence_fixture, comm_split_fixture): import numpy as np from distdl.backends.mpi.partition import MPIPartition from distdl.nn import DistributedConv2d # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) P2 = P_world.create_partition_inclusive(np.arange(2)) P4 = P_world.create_partition_inclusive(np.arange(4)) ks = [3, 3] P_x = P2.create_cartesian_topology_partition([1, 2, 1, 1]) P_y = P2.create_cartesian_topology_partition([1, 2, 1, 1]) P_w = P4.create_cartesian_topology_partition([2, 2, 1, 1]) with pytest.raises(ValueError) as e_info: # noqa: F841 layer = DistributedConv2d( P_x, P_y=P_y, # noqa: F841 in_channels=3, out_channels=3, kernel_size=ks) with pytest.raises(ValueError) as e_info: # noqa: F841 layer = DistributedConv2d( P_x, P_w=P_w, # noqa: F841 in_channels=3, out_channels=3, kernel_size=ks) P_world.deactivate() P2.deactivate() P4.deactivate() P_x.deactivate() P_y.deactivate() P_w.deactivate()
def test_padnd_adjoint(barrier_fence_fixture, comm_split_fixture, x_local_shape, padding): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.padnd import PadNd # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) x_local_shape = np.asarray(x_local_shape) padding = np.asarray(padding) padded_shape = [t + lpad + rpad for t, (lpad, rpad) in zip(x_local_shape, padding)] layer = PadNd(padding, value=0) x = torch.tensor(np.random.randn(*x_local_shape)) x.requires_grad = True dy = torch.tensor(np.random.randn(*padded_shape)) y = layer(x) y.backward(dy) dx = x.grad x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy)
def test_halo_exchange_adjoint(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, x_global_shape, dtype, kernel_size, stride, padding, dilation, MockKernelStyle): import numpy as np import torch import torch.nn.functional as F from distdl.backends.mpi.partition import MPIPartition from distdl.nn.halo_exchange import HaloExchange from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import distdl_padding_to_torch_padding from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) x_global_shape = np.asarray(x_global_shape) kernel_size = np.asarray(kernel_size) stride = np.asarray(stride) padding = np.asarray(padding) dilation = np.asarray(dilation) halo_shape = None recv_buffer_shape = None send_buffer_shape = None if P_x.active: mockup_layer = MockKernelStyle() exchange_info = mockup_layer._compute_exchange_info( x_global_shape, kernel_size, stride, padding, dilation, P_x.active, P_x.shape, P_x.index) halo_shape = exchange_info[0] recv_buffer_shape = exchange_info[1] send_buffer_shape = exchange_info[2] halo_layer = HaloExchange(P_x, halo_shape, recv_buffer_shape, send_buffer_shape) halo_layer = halo_layer.to(device) x = zero_volume_tensor(x_global_shape[0], device=device) dy = zero_volume_tensor(x_global_shape[0], device=device) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) padding = distdl_padding_to_torch_padding(halo_shape) x = torch.randn(*x_local_shape, device=device).to(dtype) # Pad the input with the halo space. We are only testing the behavior of # the halo exchange so the input must be padded before we can do anything. x = F.pad(x, pad=padding, mode="constant", value=0) # dy is also padded, but we wanted it to start with data inside it. dy = torch.randn(*x.shape, device=device).to(dtype) x.requires_grad = True # Halo Exchange (both fwd and adj) is in-place. So, we copy the input # data and save the original for the adjoint test. Because it is in-place, # the clones themselves are modified. This also prevents issues with us # in-place operations on leaf-nodes. x_clone = x.clone() dy_clone = dy.clone() # x_clone is be modified in place by halo_layer, but we assign y to # reference it for clarity. y and x_clone are the same object. y = halo_layer(x_clone) # dy_clone is modified in place by halo_layer-adjoint, but we assign dx to # reference it for clarity. dx and dy_clone are the same object. # dx is not in the grad field as you might expect because the operation is # in-place. y.backward(dy_clone) dx = dy_clone x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy) P_world.deactivate() P_x_base.deactivate() P_x.deactivate()
def test_mixin(): P_world = MPIPartition(MPI.COMM_WORLD) ranks = np.arange(P_world.size) shape = [1, 1, 4] P_size = np.prod(shape) use_ranks = ranks[:P_size] P_x_base = P_world.create_partition_inclusive(use_ranks) P_x = P_x_base.create_cartesian_topology_partition(shape) rank = P_x.rank layer = MockPoolLayer() x_global_shape = np.array([1, 1, 10]) kernel_size = np.array([2]) stride = np.array([2]) padding = np.array([0]) dilation = np.array([1]) halo_shape, recv_buffer_shape, send_buffer_shape, needed_ranges = \ layer._compute_exchange_info(x_global_shape, kernel_size, stride, padding, dilation, P_x.active, P_x.shape, P_x.index) if P_x.active: if rank == 0: expected_halo_shape = np.array([[0, 0], [0, 0], [0, 1]]) expected_recv_buffer_shape = np.array([[0, 0], [0, 0], [0, 1]]) expected_send_buffer_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_needed_ranges = np.array([[0, 1], [0, 1], [0, 4]]) assert (np.array_equal(halo_shape, expected_halo_shape)) assert (np.array_equal(recv_buffer_shape, expected_recv_buffer_shape)) assert (np.array_equal(send_buffer_shape, expected_send_buffer_shape)) assert (np.array_equal(needed_ranges, expected_needed_ranges)) elif rank == 1: expected_halo_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_recv_buffer_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_send_buffer_shape = np.array([[0, 0], [0, 0], [1, 0]]) expected_needed_ranges = np.array([[0, 1], [0, 1], [1, 3]]) assert (np.array_equal(halo_shape, expected_halo_shape)) assert (np.array_equal(recv_buffer_shape, expected_recv_buffer_shape)) assert (np.array_equal(send_buffer_shape, expected_send_buffer_shape)) assert (np.array_equal(needed_ranges, expected_needed_ranges)) elif rank == 2: expected_halo_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_recv_buffer_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_send_buffer_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_needed_ranges = np.array([[0, 1], [0, 1], [0, 2]]) assert (np.array_equal(halo_shape, expected_halo_shape)) assert (np.array_equal(recv_buffer_shape, expected_recv_buffer_shape)) assert (np.array_equal(send_buffer_shape, expected_send_buffer_shape)) assert (np.array_equal(needed_ranges, expected_needed_ranges)) elif rank == 3: expected_halo_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_recv_buffer_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_send_buffer_shape = np.array([[0, 0], [0, 0], [0, 0]]) expected_needed_ranges = np.array([[0, 1], [0, 1], [0, 2]]) assert (np.array_equal(halo_shape, expected_halo_shape)) assert (np.array_equal(recv_buffer_shape, expected_recv_buffer_shape)) assert (np.array_equal(send_buffer_shape, expected_send_buffer_shape)) assert (np.array_equal(needed_ranges, expected_needed_ranges)) # Inactive ranks should get null results else: assert (halo_shape is None) assert (recv_buffer_shape is None) assert (send_buffer_shape is None) assert (needed_ranges is None) P_world.deactivate() P_x_base.deactivate() P_x.deactivate()
def test_batch_norm_no_training(barrier_fence_fixture, P_x_ranks, P_x_shape, input_shape, num_features, eps, momentum, affine, track_running_stats, comm_split_fixture): from distdl.backends.mpi.partition import MPIPartition device = torch.device('cuda' if use_cuda else 'cpu') torch.manual_seed(0) # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions num_dimensions = len(input_shape) P_in_out_base = P_world.create_partition_inclusive([0]) P_in_out = P_in_out_base.create_cartesian_topology_partition( [1] * num_dimensions) P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) # Create the input if P_world.rank == 0: input_eval = torch.rand(input_shape, dtype=torch.float32, device=device) else: input_eval = zero_volume_tensor(device=device) # Create the sequential network if P_world.rank == 0: seq_net = torch.nn.Sequential( torch.nn.BatchNorm1d(num_features=num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)) seq_net = seq_net.to(device) else: seq_net = None # Evaluate sequential network if P_world.rank == 0: seq_net.eval() seq_out = seq_net(input_eval) seq_out = seq_out.detach().cpu() # Create distributed network dist_net = torch.nn.Sequential( distdl.nn.Repartition(P_in_out, P_x), distdl.nn.DistributedBatchNorm( P_x, num_features=num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats), distdl.nn.Repartition(P_x, P_in_out)) dist_net = dist_net.to(device) # Evaluate distributed network dist_net.eval() dist_out = dist_net(input_eval) dist_out = dist_out.detach().cpu() # Compare the distributed and sequential networks if P_world.rank == 0: assert dist_out.shape == seq_out.shape # Set the absolute tolerance to ~sqrt(e_mach), or the default # Pytorch got their defaults from NumPy, but NumPy defaults to 64-bit # floats, not 32-bit floats as torch does. Consequently, the default # torch atol is actually tighter than one can expect from two fp-equal # floating point numbers. The NumPy default of 1e-8 is closer to # sqrt(e_mach) for 64-bit numbers. So we set the 32-bit tolerance to # a little tighter than sqrt(1e-7), 1e-5. if seq_out.dtype == torch.float64: atol = 1e-8 elif seq_out.dtype == torch.float32: atol = 1e-5 else: # torch default atol = 1e-8 assert torch.allclose(dist_out, seq_out, atol=atol) P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_in_out_base.deactivate() P_in_out.deactivate()
def test_batch_norm_with_training(barrier_fence_fixture, P_x_ranks, P_x_shape, input_shape, num_features, eps, momentum, affine, track_running_stats, affine_workers, comm_split_fixture): from distdl.backends.mpi.partition import MPIPartition torch.manual_seed(0) device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions num_dimensions = len(input_shape) P_in_out_base = P_world.create_partition_inclusive([0]) P_in_out = P_in_out_base.create_cartesian_topology_partition( [1] * num_dimensions) P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_affine = P_world.create_partition_inclusive(affine_workers) # Create the input if P_world.rank == 0: input_train = torch.rand(input_shape, dtype=torch.float32, device=device) input_eval = torch.rand(input_shape, dtype=torch.float32, device=device) exp = torch.rand(input_shape, dtype=torch.float32, device=device) else: input_train = zero_volume_tensor(device=device) input_eval = zero_volume_tensor(device=device) exp = zero_volume_tensor(device=device) # Create the sequential network if len(input_shape) == 2: seq_layer = torch.nn.BatchNorm1d elif len(input_shape) == 3: seq_layer = torch.nn.BatchNorm1d elif len(input_shape) == 4: seq_layer = torch.nn.BatchNorm2d elif len(input_shape) == 5: seq_layer = torch.nn.BatchNorm3d if P_world.rank == 0: seq_bn = seq_layer(num_features=num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats) seq_bn = seq_bn.to(device) # Train sequential network if P_world.rank == 0: seq_bn.train() seq_out1 = seq_bn(input_train) seq_loss = ((seq_out1 - exp)**2).sum() seq_loss.backward() seq_grads = [p.grad.detach().cpu() for p in seq_bn.parameters()] # Do a manual weight update (this is what optimizer does): with torch.no_grad(): for p in seq_bn.parameters(): p.copy_(p + 0.1 * p.grad) # Evaluate sequential network if P_world.rank == 0: seq_bn.eval() seq_out2 = seq_bn(input_eval) seq_out2 = seq_out2.detach().cpu() # Create distributed network tr1 = distdl.nn.Repartition(P_in_out, P_x) tr1 = tr1.to(device) dist_bn = distdl.nn.DistributedBatchNorm( P_x, num_features=num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats) dist_bn = dist_bn.to(device) tr2 = distdl.nn.Repartition(P_x, P_in_out) tr2 = tr2.to(device) # Only rank 0 should have trainable parameters: if P_world.rank in affine_workers: assert len(list(dist_bn.parameters())) == 2 else: assert len(list(dist_bn.parameters())) == 0 # Train distributed network dist_bn.train() dist_out1 = tr2(dist_bn(tr1(input_train))) dist_loss = ((dist_out1 - exp)**2).sum() assert dist_loss.requires_grad dist_loss.backward() # Note: We expect the batch norm gradient to have extra dimensions than PyTorch, # but both ultimately have volume equal to num_features. # So, reshape them now, and gather them onto rank 0 for comparison. if P_world.rank == 0: dist_grads = [] if P_world.rank in affine_workers: for p in dist_bn.parameters(): if affine_workers == [0]: parts = [p.grad.detach().cpu()] else: parts = P_affine._comm.gather(p.grad.detach().cpu(), root=0) if P_world.rank == 0: grad = torch.cat(parts, 1) reshaped = grad.reshape((num_features, )) dist_grads.append(reshaped) # Do a manual weight update (this is what optimizer does): with torch.no_grad(): for p in dist_bn.parameters(): p.copy_(p + 0.1 * p.grad) # Evaluate distributed network dist_bn.eval() dist_out2 = tr2(dist_bn(tr1(input_eval))) dist_out2 = dist_out2.detach().cpu() # Compare the distributed and sequential networks if P_world.rank == 0: # Set the absolute tolerance to ~sqrt(e_mach), or the default # Pytorch got their defaults from NumPy, but NumPy defaults to 64-bit # floats, not 32-bit floats as torch does. Consequently, the default # torch atol is actually tighter than one can expect from two fp-equal # floating point numbers. The NumPy default of 1e-8 is closer to # sqrt(e_mach) for 64-bit numbers. So we set the 32-bit tolerance to # sqrt(1e-7), as our usual choice, 1e-5, is too tight. if seq_out1.dtype == torch.float64: atol = 1e-8 elif seq_out1.dtype == torch.float32: import math atol = math.sqrt(1e-7) else: # torch default atol = 1e-8 assert dist_out1.shape == seq_out1.shape assert torch.allclose(dist_out1, seq_out1, rtol=ERROR_THRESHOLD, atol=atol) assert dist_loss.shape == seq_loss.shape assert torch.allclose(dist_loss, seq_loss, rtol=ERROR_THRESHOLD, atol=atol) for dist_grad, seq_grad in zip(dist_grads, seq_grads): assert dist_grad.shape == seq_grad.shape assert torch.allclose(dist_grad, seq_grad, rtol=ERROR_THRESHOLD, atol=atol) assert dist_out2.shape == seq_out2.shape assert torch.allclose(dist_out2, seq_out2, rtol=ERROR_THRESHOLD, atol=atol) P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_in_out_base.deactivate() P_in_out.deactivate() P_affine.deactivate()
def test_distributed_loss(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, x_global_shape, SequentialLoss, DistributedLoss): import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn import DistributedTranspose from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_0_base = P_x_base.create_partition_inclusive([0]) P_0 = P_0_base.create_cartesian_topology_partition([1] * len(P_x_shape)) scatter = DistributedTranspose(P_0, P_x).to(device) gather = DistributedTranspose(P_x, P_0).to(device) for reduction in DistributedLoss._valid_reductions: distributed_criterion = DistributedLoss(P_x, reduction=reduction).to(device) sequential_criterion = SequentialLoss(reduction=reduction).to(device) with torch.no_grad(): x_g = zero_volume_tensor(device=device) y_g = zero_volume_tensor(device=device) if P_0.active: x_g = torch.rand(x_global_shape, device=device) y_g = torch.rand(x_global_shape, device=device) x_l = scatter(x_g) y_l = scatter(y_g) x_l.requires_grad = True distributed_loss = distributed_criterion(x_l, y_l) # For "none", no reduction is applied so we see if it computed the # same loss as the sequential code by gathering the loss value it to # the root rank. if reduction == "none": distributed_loss = gather(distributed_loss) if P_0.active: x_g.requires_grad = True sequential_loss = sequential_criterion(x_g, y_g) assert (torch.allclose(distributed_loss, sequential_loss)) # For any other reduction, we can compare the loss # value *and* backpropagate through the distributed loss to verify # that it produces the same output. if reduction != "none": distributed_loss.backward() distributed_dx_g = gather(x_l.grad) if P_0.active: sequential_loss.backward() sequential_dx_g = x_g.grad assert (torch.allclose(distributed_dx_g, sequential_dx_g)) P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_0_base.deactivate() P_0.deactivate()
def test_repartition_dtype(barrier_fence_fixture, comm_split_fixture, dtype, test_backward, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, x_global_shape): import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.repartition import Repartition from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) # The global tensor size is the same for x and y layer = Repartition(P_x, P_y, preserve_batch=False) layer = layer.to(device) # Forward Input x = zero_volume_tensor(dtype=dtype, device=device) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = 10 * torch.randn(*x_local_shape, device=device).to(dtype) x.requires_grad = test_backward # y = F @ x y = layer(x) if P_y.active: assert y.dtype == dtype if test_backward: # Adjoint Input dy = zero_volume_tensor(dtype=dtype, device=device) if P_y.active: y_local_shape = compute_subshape(P_y.shape, P_y.index, x_global_shape) dy = 10 * torch.randn(*y_local_shape, device=device).to(dtype) # dx = F* @ dy y.backward(dy) dx = x.grad if P_x.active: assert dx.dtype == dtype P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_y_base.deactivate() P_y.deactivate()
def test_buffer_management_transpose_network(barrier_fence_fixture, comm_split_fixture): import numpy as np import torch import distdl from distdl.backends.mpi.buffer import MPIBufferManager from distdl.backends.mpi.partition import MPIPartition from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return buffer_manager = MPIBufferManager() P_world = MPIPartition(base_comm) # Create the partitions P_1_base = P_world.create_partition_inclusive([0]) P_1 = P_1_base.create_cartesian_topology_partition([1, 1]) P_3_base = P_world.create_partition_inclusive([1, 2, 3]) P_3 = P_3_base.create_cartesian_topology_partition([1, 3]) P_4_base = P_world.create_partition_inclusive([0, 1, 2, 3]) P_4 = P_4_base.create_cartesian_topology_partition([1, 4]) tr1 = distdl.nn.DistributedTranspose(P_1, P_4, buffer_manager=buffer_manager) tr2 = distdl.nn.DistributedTranspose(P_4, P_4, buffer_manager=buffer_manager) tr3 = distdl.nn.DistributedTranspose(P_4, P_3, buffer_manager=buffer_manager) tr4 = distdl.nn.DistributedTranspose(P_3, P_1, buffer_manager=buffer_manager) x = zero_volume_tensor(1) if P_1.active: x = torch.randn(1, 10) x.requires_grad = True # [0 1 2 3 4 5 6 7 8 9] to # [0 1 2] [3 4 5] [6 7] [8 9] x2 = tr1(x) n_buffers_by_rank = (3, 1, 1, 1) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # [0 1 2] [3 4 5] [6 7] [8 9] to # [0 1 2] [3 4 5] [6 7] [8 9] x3 = tr2(x2) n_buffers_by_rank = (3, 1, 1, 1) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # [0 1 2] [3 4 5] [6 7] [8 9] to # [] [0 1 2 3] [4 5 6] [7 8 9] x4 = tr3(x3) n_buffers_by_rank = (3, 2, 2, 1) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # [] [0 1 2 3] [4 5 6] [7 8 9] to # [0 1 2 3 4 5 6 7 8 9] y = tr4(x4) n_buffers_by_rank = (3, 2, 2, 1) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] dy = zero_volume_tensor(1) if P_1.active: dy = torch.randn(1, 10) dy.requires_grad = True y.backward(dy) dx = x.grad # Through the backward call the buffer count do not change n_buffers_by_rank = (3, 2, 2, 1) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # And adjointness is still preserved x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy) P_world.deactivate() P_1_base.deactivate() P_1.deactivate() P_3_base.deactivate() P_3.deactivate() P_4_base.deactivate() P_4.deactivate()
def test_matches_sequential(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, input_dimensions, x_global_shape, kernel_size, padding, stride, dilation, layer_type): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.repartition import Repartition from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) P_world._comm.Barrier() # Create the partitions P_0_base = P_world.create_partition_inclusive(np.arange(1)) P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_0 = P_0_base.create_cartesian_topology_partition([1] * len(P_x_shape)) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) scatter_layer_x = Repartition(P_0, P_x).to(device) scatter_layer_y = Repartition(P_0, P_x).to(device) gather_layer_x = Repartition(P_x, P_0).to(device) gather_layer_y = Repartition(P_x, P_0).to(device) # Create the layers if input_dimensions == 1: if layer_type == 'max': from torch.nn import MaxPool1d as SequentialPoolType from distdl.nn import DistributedMaxPool1d as DistributedPoolType else: from torch.nn import AvgPool1d as SequentialPoolType from distdl.nn import DistributedAvgPool1d as DistributedPoolType elif input_dimensions == 2: if layer_type == 'max': from torch.nn import MaxPool2d as SequentialPoolType from distdl.nn import DistributedMaxPool2d as DistributedPoolType else: from torch.nn import AvgPool2d as SequentialPoolType from distdl.nn import DistributedAvgPool2d as DistributedPoolType elif input_dimensions == 3: if layer_type == 'max': from torch.nn import MaxPool3d as SequentialPoolType from distdl.nn import DistributedMaxPool3d as DistributedPoolType else: from torch.nn import AvgPool3d as SequentialPoolType from distdl.nn import DistributedAvgPool3d as DistributedPoolType # PyTorch AvgPool doesn't support dilation, so skip the test if the combination comes up dilation_is_default = dilation == 1 or all(x == 1 for x in dilation) if layer_type == 'avg' and not dilation_is_default: return layer_kwargs = { 'kernel_size': kernel_size, 'padding': padding, 'stride': stride } # Only max pool layers support dilation if layer_type == 'max': layer_kwargs['dilation'] = dilation dist_layer = DistributedPoolType(P_x, **layer_kwargs).to(device) if P_0.active: seq_layer = SequentialPoolType(**layer_kwargs).to(device) # Forward Input x_ref = zero_volume_tensor(device=device) x_ref.requires_grad = True dy_ref = zero_volume_tensor(device=device) # Construct the inputs to the forward and backward functions as well as the # the outputs of the sequential layer if P_0.active: x_ref = torch.randn(*x_global_shape, device=device) x_ref.requires_grad = True y_ref = seq_layer(x_ref) y_global_shape_calc = y_ref.shape dy_ref = torch.randn(*y_global_shape_calc, device=device) y_ref.backward(dy_ref) dx_ref = x_ref.grad # Ensure that the scatter is not part of the computation we are testing with torch.no_grad(): x = scatter_layer_x(x_ref.detach()) dy = scatter_layer_y(dy_ref.detach()) x.requires_grad = True y = dist_layer(x) y.backward(dy) dx = x.grad # Ensure that the gather is not part of the computation we are testing with torch.no_grad(): dx_comp = gather_layer_x(dx.detach()) y_comp = gather_layer_y(y.detach()) if P_0.active: # Set the absolute tolerance to ~sqrt(e_mach), or the default # Pytorch got their defaults from NumPy, but NumPy defaults to 64-bit # floats, not 32-bit floats as torch does. Consequently, the default # torch atol is actually tighter than one can expect from two fp-equal # floating point numbers. The NumPy default of 1e-8 is closer to # sqrt(e_mach) for 64-bit numbers. So we set the 32-bit tolerance to # a little tighter than sqrt(1e-7), 1e-5. if x_ref.dtype == torch.float64: atol = 1e-8 elif x_ref.dtype == torch.float32: atol = 1e-5 else: # torch default atol = 1e-8 assert torch.allclose(y_ref, y_comp, atol=atol) assert torch.allclose(dx_ref, dx_comp, atol=atol) P_world.deactivate() P_0_base.deactivate() P_0.deactivate() P_x_base.deactivate() P_x.deactivate()
def test_upsample_matches_sequential(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, x_global_shape, scale_factor, y_global_shape, mode, align_corners, use_size): import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.transpose import DistributedTranspose from distdl.nn.upsampling import DistributedUpsample from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return # Test align_corners only in the linear case, otherwise ignore it. if mode != "linear" and align_corners: return torch_mode_map = {3: "linear", 4: "bilinear", 5: "trilinear"} torch_mode = mode if mode == "linear": torch_mode = torch_mode_map[len(x_global_shape)] torch_align_corners = align_corners if mode == "nearest": torch_align_corners = None P_world = MPIPartition(base_comm) # Create the partitions P_0_base = P_world.create_partition_inclusive([0]) P_0 = P_0_base.create_cartesian_topology_partition([1] * len(P_x_shape)) P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) scatter_layer_x = DistributedTranspose(P_0, P_x) scatter_layer_y = DistributedTranspose(P_0, P_x) gather_layer_x = DistributedTranspose(P_x, P_0) gather_layer_y = DistributedTranspose(P_x, P_0) if use_size: dist_layer = DistributedUpsample(P_x, size=y_global_shape, mode=mode, align_corners=align_corners) if P_0.active: seq_layer = torch.nn.Upsample(size=y_global_shape[2:], mode=torch_mode, align_corners=torch_align_corners) else: dist_layer = DistributedUpsample(P_x, scale_factor=scale_factor, mode=mode, align_corners=align_corners) if P_0.active: seq_layer = torch.nn.Upsample(scale_factor=scale_factor, mode=torch_mode, align_corners=torch_align_corners) # Forward Input x_ref = zero_volume_tensor() x_ref.requires_grad = True dy_ref = zero_volume_tensor() # Construct the inputs to the forward and backward functions as well as the # the outputs of the sequential layer if P_0.active: x_ref = torch.randn(*x_global_shape) x_ref.requires_grad = True y_ref = seq_layer(x_ref) y_global_shape_calc = y_ref.shape dy_ref = torch.randn(*y_global_shape_calc) y_ref.backward(dy_ref) dx_ref = x_ref.grad # Ensure that the scatter is not part of the computation we are testing with torch.no_grad(): x = scatter_layer_x(x_ref.detach()) dy = scatter_layer_y(dy_ref.detach()) x.requires_grad = True # Because there is no guarantee that any padding is needed, in this test, # the input x may pass directly to the Halo layer without going through # the padding process. As the halo layer is in-place, that would mean a leaf-node # variable is modified in-place, which PyTorch does not allow. # # Thus, we have to clone it to make the input not a leaf-node. x_clone = x.clone() y = dist_layer(x_clone) y.backward(dy) dx = x.grad # Ensure that the gather is not part of the computation we are testing with torch.no_grad(): dx_comp = gather_layer_x(dx.detach()) y_comp = gather_layer_y(y.detach()) if P_0.active: # Set the absolute tolerance to ~sqrt(e_mach), or the default # Pytorch got their defaults from NumPy, but NumPy defaults to 64-bit # floats, not 32-bit floats as torch does. Consequently, the default # torch atol is actually tighter than one can expect from two fp-equal # floating point numbers. The NumPy default of 1e-8 is closer to # sqrt(e_mach) for 64-bit numbers. So we set the 32-bit tolerance to # a little tighter than sqrt(1e-7), 1e-5. if x_ref.dtype == torch.float64: atol = 1e-8 elif x_ref.dtype == torch.float32: atol = 1e-5 else: # torch default atol = 1e-8 # Test the result of each entry independently assert torch.allclose(y_ref, y_comp, atol=atol) assert torch.allclose(dx_ref, dx_comp, atol=atol) P_world.deactivate() P_0_base.deactivate() P_0.deactivate() P_x_base.deactivate() P_x.deactivate()
def test_transpose_adjoint(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, x_global_shape, balanced): import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.transpose import DistributedTranspose from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) # The global tensor size is the same for x and y layer = DistributedTranspose(P_x, P_y, preserve_batch=False) layer = layer.to(device) # Forward Input x = zero_volume_tensor(device=device) if P_x.active: if balanced: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) else: quotient = np.atleast_1d(x_global_shape) // np.atleast_1d( P_x_shape) remainder = np.atleast_1d(x_global_shape) % np.atleast_1d( P_x_shape) loc = np.where(P_x.index == 0) x_local_shape = quotient.copy() x_local_shape[loc] += remainder[loc] x = torch.randn(*x_local_shape, device=device) x.requires_grad = True # Adjoint Input dy = zero_volume_tensor(device=device) if P_y.active: y_local_shape = compute_subshape(P_y.shape, P_y.index, x_global_shape) dy = torch.randn(*y_local_shape, device=device) # y = F @ x y = layer(x) # dx = F* @ dy y.backward(dy) dx = x.grad x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy) P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_y_base.deactivate() P_y.deactivate()
def test_conv_versus_pytorch(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, input_dimensions, x_global_shape, kernel_size, padding, stride, dilation, bias): import numpy as np import torch from torch.nn import Conv1d from torch.nn import Conv2d from torch.nn import Conv3d from distdl.backends.mpi.partition import MPIPartition from distdl.nn.conv_feature import DistributedFeatureConv1d from distdl.nn.conv_feature import DistributedFeatureConv2d from distdl.nn.conv_feature import DistributedFeatureConv3d from distdl.nn.repartition import Repartition from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) P_world._comm.Barrier() # Create the partitions P_0_base = P_world.create_partition_inclusive(np.arange(1)) P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_0 = P_0_base.create_cartesian_topology_partition([1] * len(P_x_shape)) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) scatter_layer_x = Repartition(P_0, P_x) scatter_layer_x = scatter_layer_x.to(device) scatter_layer_y = Repartition(P_0, P_x) scatter_layer_y = scatter_layer_y.to(device) gather_layer_x = Repartition(P_x, P_0) gather_layer_x = gather_layer_x.to(device) gather_layer_y = Repartition(P_x, P_0) gather_layer_y = gather_layer_y.to(device) # Create the layers if input_dimensions == 1: dist_layer_type = DistributedFeatureConv1d seq_layer_type = Conv1d elif input_dimensions == 2: dist_layer_type = DistributedFeatureConv2d seq_layer_type = Conv2d elif input_dimensions == 3: dist_layer_type = DistributedFeatureConv3d seq_layer_type = Conv3d dist_layer = dist_layer_type(P_x, in_channels=x_global_shape[1], out_channels=10, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation, bias=bias) dist_layer = dist_layer.to(device) if P_0.active: seq_layer = seq_layer_type(in_channels=x_global_shape[1], out_channels=10, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation, bias=bias) # set the weights of both layers to be the same seq_layer = seq_layer.to(device) weight = torch.rand_like(seq_layer.weight, device=device) seq_layer.weight.data = weight dist_layer.weight.data = weight if bias: bias_weight = torch.rand_like(seq_layer.bias, device=device) seq_layer.bias.data = bias_weight dist_layer.bias.data = bias_weight # Forward Input x_ref = zero_volume_tensor(device=device) x_ref.requires_grad = True dy_ref = zero_volume_tensor(device=device) # Construct the inputs to the forward and backward functions as well as the # the outputs of the sequential layer if P_0.active: x_ref = torch.randn(*x_global_shape, device=device) x_ref.requires_grad = True y_ref = seq_layer(x_ref) y_global_shape_calc = y_ref.shape dy_ref = torch.randn(*y_global_shape_calc, device=device) y_ref.backward(dy_ref) dx_ref = x_ref.grad # Ensure that the scatter is not part of the computation we are testing with torch.no_grad(): x = scatter_layer_x(x_ref.detach()) dy = scatter_layer_y(dy_ref.detach()) x.requires_grad = True y = dist_layer(x) y.backward(dy) dx = x.grad # Ensure that the gather is not part of the computation we are testing with torch.no_grad(): dx_comp = gather_layer_x(dx.detach()) y_comp = gather_layer_y(y.detach()) if P_0.active: # Set the absolute tolerance to ~sqrt(e_mach), or the default # Pytorch got their defaults from NumPy, but NumPy defaults to 64-bit # floats, not 32-bit floats as torch does. Consequently, the default # torch atol is actually tighter than one can expect from two fp-equal # floating point numbers. The NumPy default of 1e-8 is closer to # sqrt(e_mach) for 64-bit numbers. So we set the 32-bit tolerance to # a little tighter than sqrt(1e-7), 1e-5. if x_ref.dtype == torch.float64: atol = 1e-8 elif x_ref.dtype == torch.float32: atol = 1e-5 else: # torch default atol = 1e-8 # Test the result of each entry independently assert torch.allclose(y_ref, y_comp, atol=atol) assert torch.allclose(dx_ref, dx_comp, atol=atol) P_world.deactivate() P_0_base.deactivate() P_0.deactivate() P_x_base.deactivate() P_x.deactivate()
def test_linear_adjoint_bias(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, P_w_ranks, P_w_shape, x_global_shape, y_global_shape): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.linear import DistributedLinear from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) P_w_base = P_world.create_partition_inclusive(P_w_ranks) P_w = P_w_base.create_cartesian_topology_partition(P_w_shape) x_global_shape = np.asarray(x_global_shape) y_global_shape = np.asarray(y_global_shape) layer = DistributedLinear(P_x, P_y, P_w, x_global_shape[1], y_global_shape[1], bias=True) x = zero_volume_tensor(x_global_shape[0]) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) # For this test, we are only testing to see if the adjoint works # correctly for the bias term. But the adjoint test only works on the # Jacobian of the linear layer. The Jacobian block for b is 0 for x and # W, so killing x makes the forward operator equal to its Jacobian and # we can test to see that adjoint is computed correctly. x = torch.zeros(*x_local_shape) x.requires_grad = True y = layer(x) dy = zero_volume_tensor(x_global_shape[0]) if P_y.active: dy = torch.Tensor(np.random.randn(*y.shape)) y.backward(dy) b = zero_volume_tensor() db = zero_volume_tensor() if P_w.active and P_w.index[-1] == 0: b = layer.sublinear.bias.detach() db = layer.sublinear.bias.grad.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, b, db, y, dy)
def test_halo_exchange_adjoint(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, x_global_shape, kernel_size, stride, padding, dilation, MockKernelStyle): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.halo_exchange import HaloExchange from distdl.nn.padnd import PadNd from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) x_global_shape = np.asarray(x_global_shape) kernel_size = np.asarray(kernel_size) stride = np.asarray(stride) padding = np.asarray(padding) dilation = np.asarray(dilation) halo_shape = None recv_buffer_shape = None send_buffer_shape = None if P_x.active: mockup_layer = MockKernelStyle() exchange_info = mockup_layer._compute_exchange_info(x_global_shape, kernel_size, stride, padding, dilation, P_x.active, P_x.shape, P_x.index) halo_shape = exchange_info[0] recv_buffer_shape = exchange_info[1] send_buffer_shape = exchange_info[2] pad_layer = PadNd(halo_shape, value=0) halo_layer = HaloExchange(P_x, halo_shape, recv_buffer_shape, send_buffer_shape) x = zero_volume_tensor(x_global_shape[0]) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.tensor(np.random.randn(*x_local_shape)) x = pad_layer.forward(x) x.requires_grad = True dy = zero_volume_tensor(x_global_shape[0]) if P_x.active: dy = torch.tensor(np.random.randn(*x.shape)) x_clone = x.clone() dy_clone = dy.clone() # x_clone is be modified in place by halo_layer, but we assign y to # reference it for clarity y = halo_layer(x_clone) # dy_clone is modified in place by halo_layer-adjoint, but we assign dx to # reference it for clarity y.backward(dy_clone) dx = dy_clone x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy)
def test_buffer_management_mixed_network(barrier_fence_fixture, comm_split_fixture): import numpy as np import torch import distdl from distdl.backends.mpi.buffer import MPIBufferManager from distdl.backends.mpi.partition import MPIPartition from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return buffer_manager = MPIBufferManager() P_world = MPIPartition(base_comm) # Create the partitions P_1_base = P_world.create_partition_inclusive([0]) P_1 = P_1_base.create_cartesian_topology_partition([1, 1, 1, 1]) P_22_base = P_world.create_partition_inclusive([0, 1, 2, 3]) P_22 = P_22_base.create_cartesian_topology_partition([1, 1, 2, 2]) tr1 = distdl.nn.DistributedTranspose(P_1, P_22, buffer_manager=buffer_manager) c1 = distdl.nn.DistributedConv2d(P_22, in_channels=1, out_channels=5, kernel_size=[3, 3], padding=[1, 1], bias=False, buffer_manager=buffer_manager) c2 = distdl.nn.DistributedConv2d(P_22, in_channels=5, out_channels=10, kernel_size=[3, 3], padding=[1, 1], bias=False, buffer_manager=buffer_manager) c3 = distdl.nn.DistributedConv2d(P_22, in_channels=10, out_channels=20, kernel_size=[3, 3], padding=[1, 1], bias=False, buffer_manager=buffer_manager) tr2 = distdl.nn.DistributedTranspose(P_22, P_1, buffer_manager=buffer_manager) x = zero_volume_tensor(1) if P_1.active: x = torch.randn(1, 1, 5, 5) x.requires_grad = True # [[00 01 02 03 04] [[00 01 02] [[03 04] # [10 11 12 13 14] [10 11 12] [13 14]] # [20 21 22 23 24] to [20 21 22]] [[23 24] # [30 31 32 33 34] [[30 31 32] [33 34] # [40 41 42 43 44]] [40 41 42]] [43 44]] x2 = tr1(x) n_buffers_by_rank = (3, 1, 1, 1) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # [[00 01 02] [[03 04] [[00 01 02] [[03 04] # [10 11 12] [13 14]] [10 11 12] [13 14]] # [20 21 22]] [[23 24] to [20 21 22]] [[23 24] # [[30 31 32] [33 34] [[30 31 32] [33 34] # [40 41 42]] [43 44]] [40 41 42]] [43 44]] x3 = c1(x2) n_buffers_by_rank = (4, 4, 4, 4) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # [[00 01 02] [[03 04] [[00 01 02] [[03 04] # [10 11 12] [13 14]] [10 11 12] [13 14]] # [20 21 22]] [[23 24] to [20 21 22]] [[23 24] # [[30 31 32] [33 34] [[30 31 32] [33 34] # [40 41 42]] [43 44]] [40 41 42]] [43 44]] x4 = c2(x3) n_buffers_by_rank = (4, 4, 4, 4) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # [[00 01 02] [[03 04] [[00 01 02] [[03 04] # [10 11 12] [13 14]] [10 11 12] [13 14]] # [20 21 22]] [[23 24] to [20 21 22]] [[23 24] # [[30 31 32] [33 34] [[30 31 32] [33 34] # [40 41 42]] [43 44]] [40 41 42]] [43 44]] x5 = c3(x4) n_buffers_by_rank = (4, 4, 4, 4) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # [[00 01 02] [[03 04] [[00 01 02 03 04] # [10 11 12] [13 14]] [10 11 12 13 14] # [20 21 22]] [[23 24] to [20 21 22 23 24] # [[30 31 32] [33 34] [30 31 32 33 34] # [40 41 42]] [43 44]] [40 41 42 43 44]] y = tr2(x5) n_buffers_by_rank = (4, 4, 4, 4) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] dy = zero_volume_tensor(1) if P_1.active: dy = torch.randn(1, 20, 5, 5) dy.requires_grad = True y.backward(dy) dx = x.grad # Through the backward call the buffer count do not change n_buffers_by_rank = (4, 4, 4, 4) assert len(buffer_manager.buffers_map[np.float32]) == n_buffers_by_rank[ P_world.rank] # And adjointness is still preserved x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy) P_world.deactivate() P_1_base.deactivate() P_1.deactivate() P_22_base.deactivate() P_22.deactivate()
def test_sum_reduce_dtype(barrier_fence_fixture, comm_split_fixture, dtype, test_backward, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, x_global_shape, transpose_src): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.sum_reduce import SumReduce from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) # TODO #93: Change this to create a subtensor so we test when local tensors # have different shape. Then, the output size will also be different, which # we will have to get from `y` itself. x_local_shape = np.asarray(x_global_shape) layer = SumReduce(P_x, P_y, transpose_src=transpose_src, preserve_batch=False) layer = layer.to(device) x = zero_volume_tensor(device=device) if P_x.active: x = 10 * torch.randn(*x_local_shape, device=device).to(dtype) x.requires_grad = test_backward # y = F @ x y = layer(x) # If we are not in the output partition, there is no data to test the type # against. if P_y.active: assert y.dtype == dtype if test_backward: dy = zero_volume_tensor(device=device) if P_y.active: # Adjoint Input dy = 10 * torch.randn(*x_local_shape, device=device).to(dtype) # dx = F* @ dy y.backward(dy) dx = x.grad if P_x.active: assert dx.dtype == dtype P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_y_base.deactivate() P_y.deactivate()
def test_repartition_identity(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, x_global_shape, balanced): import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.repartition import Repartition from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor device = torch.device('cuda' if use_cuda else 'cpu') # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y = P_x # The global tensor size is the same for x and y layer = Repartition(P_x, P_y, preserve_batch=False) layer = layer.to(device) # Forward Input x = zero_volume_tensor(device=device) if P_x.active: if balanced: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) else: quotient = np.atleast_1d(x_global_shape) // np.atleast_1d( P_x_shape) remainder = np.atleast_1d(x_global_shape) % np.atleast_1d( P_x_shape) loc = np.where(P_x.index == 0) x_local_shape = quotient.copy() x_local_shape[loc] += remainder[loc] x = torch.randn(*x_local_shape, device=device) x.requires_grad = True # Adjoint Input dy = zero_volume_tensor(device=device) if P_y.active: y_local_shape = compute_subshape(P_y.shape, P_y.index, x_global_shape) dy = torch.randn(*y_local_shape, device=device) # y = F @ x y = layer(x) # In the balanced case, this should be a true identity, so there should # be no communication performed, just self-copies. if balanced: for sl, sz, p in layer.P_x_to_y_overlaps: assert p == "self" or (sl, sz, p) == (None, None, None) for sl, sz, p in layer.P_y_to_x_overlaps: assert p == "self" or (sl, sz, p) == (None, None, None) # dx = F* @ dy y.backward(dy) dx = x.grad x = x.detach() dx = dx.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, x, dx, y, dy) P_world.deactivate() P_x_base.deactivate() P_x.deactivate()
from distdl.nn.halo_exchange import HaloExchange from distdl.nn.mixins.conv_mixin import ConvMixin from distdl.nn.mixins.halo_mixin import HaloMixin from distdl.nn.padnd import PadNd from distdl.utilities.debug import print_sequential from distdl.utilities.misc import DummyContext from distdl.utilities.slicing import compute_subshape class MockConvLayer(HaloMixin, ConvMixin): pass torch.set_printoptions(linewidth=200) P_world = MPIPartition(MPI.COMM_WORLD) ranks = np.arange(P_world.size) shape = [1, 1, 2, 2] P_size = np.prod(shape) use_ranks = ranks[:P_size] P_x_base = P_world.create_partition_inclusive(use_ranks) P_x = P_x_base.create_cartesian_topology_partition(shape) rank = P_x.rank cart_comm = P_x.comm x_global_shape = np.array([1, 1, 10, 12]) if P_x.active: mockup_conv_layer = MockConvLayer()
def test_general_conv1d_adjoint_bias(barrier_fence_fixture, comm_split_fixture, P_x_ranks, P_x_shape, P_y_ranks, P_y_shape, P_w_ranks, P_w_shape, x_global_shape): import numpy as np import torch from distdl.backends.mpi.partition import MPIPartition from distdl.nn.conv_general import DistributedGeneralConv1d from distdl.utilities.slicing import compute_subshape from distdl.utilities.torch import zero_volume_tensor # Isolate the minimum needed ranks base_comm, active = comm_split_fixture if not active: return P_world = MPIPartition(base_comm) # Create the partitions P_x_base = P_world.create_partition_inclusive(P_x_ranks) P_x = P_x_base.create_cartesian_topology_partition(P_x_shape) P_y_base = P_world.create_partition_inclusive(P_y_ranks) P_y = P_y_base.create_cartesian_topology_partition(P_y_shape) P_w_base = P_world.create_partition_inclusive(P_w_ranks) P_w = P_w_base.create_cartesian_topology_partition(P_w_shape) x_global_shape = np.asarray(x_global_shape) layer = DistributedGeneralConv1d(P_x, P_y, P_w, in_channels=x_global_shape[1], out_channels=10, kernel_size=[3], bias=True) x = zero_volume_tensor(x_global_shape[0]) if P_x.active: x_local_shape = compute_subshape(P_x.shape, P_x.index, x_global_shape) x = torch.zeros(*x_local_shape) x.requires_grad = True y = layer(x) dy = zero_volume_tensor(x_global_shape[0]) if P_y.active: dy = torch.randn(*y.shape) y.backward(dy) b = zero_volume_tensor() db = zero_volume_tensor() if P_w.active and layer.stores_bias: b = layer.bias.detach() db = layer.bias.grad.detach() dy = dy.detach() y = y.detach() check_adjoint_test_tight(P_world, b, db, y, dy) P_world.deactivate() P_x_base.deactivate() P_x.deactivate() P_y_base.deactivate() P_y.deactivate() P_w_base.deactivate() P_w.deactivate()