-
Notifications
You must be signed in to change notification settings - Fork 2
Expand file tree
/
Copy pathtest_doublebackprop.py
More file actions
182 lines (161 loc) · 6.08 KB
/
Copy pathtest_doublebackprop.py
File metadata and controls
182 lines (161 loc) · 6.08 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
import pytest
import torch
import torch.nn as nn
from utils import cleanup_parallel_strategy, fp32_allclose
from distconv import DCTensor, DistConvDDP, ParallelStrategy
@pytest.fixture(scope="module")
def parallel_strategy(device: torch.device):
ps = ParallelStrategy(num_shards=2, device_type=device.type)
yield ps
cleanup_parallel_strategy(ps)
def generate_configs():
configs = []
for ndims in [1, 2, 3]:
for shard_dim in range(ndims):
for kernel_size in [1, 3, 5]:
for num_shards in [2]:
configs.append((ndims, shard_dim, kernel_size, num_shards))
return "ndims,shard_dim,kernel_size,num_shards", configs
@pytest.mark.parametrize(*generate_configs())
def test_double_backprop_gradientloss(
parallel_strategy: ParallelStrategy,
ndims: int,
shard_dim: int,
kernel_size: int,
num_shards: int,
device: torch.device,
):
"""
Test distributed convolution with different number of dimensions and shard dimensions.
Also consider hybrid spatial-data parallelism.
Checks the output and gradients of the distributed convolution against the reference DDP
convolution.
Args:
ndims (int): Number of dimensions for the convolution (1, 2, or 3).
shard_dim (int): Dimension along which the tensor is sharded.
kernel_size (int): Size of the convolution kernel.
num_shards (int): Number of spatial partitions for data
device (torch.device): Torch device to run test with.
"""
# Set the shard dimension for the parallel strategy
parallel_strategy.shard_dim = shard_dim + 2
conv_kwargs = dict(
kernel_size=kernel_size,
padding=kernel_size // 2,
bias=False,
stride=1,
padding_mode="circular",
)
# Initialize the input tensor and convolution layer
shape = [1, 4] + [16] * ndims
x = torch.randn(*shape, device=device, requires_grad=True)
conv_class = getattr(nn, f"Conv{ndims}d")
conv = conv_class(4, 8, **conv_kwargs).to(device).requires_grad_(False)
conv.requires_grad_(True)
# Perform forward and backward pass for reference (non-distributed) convolution
conv.zero_grad()
ref_y = conv(x)
# find gradient wrt input
ref_grads = torch.autograd.grad(
outputs=[ref_y.sum()], inputs=[x], create_graph=True
)[0]
# find all losses
ref_loss_grad = ref_grads.mean()
ref_loss = ref_loss_grad
ref_loss.backward()
ref_conv_grad = conv.weight.grad.clone()
# Perform forward and backward pass for distributed convolution
conv.zero_grad()
ddp_conv = DistConvDDP(conv, parallel_strategy=parallel_strategy)
dcx = DCTensor.distribute(x, parallel_strategy)
dcy = ddp_conv(dcx)
ddpy = dcy.to_replicate()
# find gradient wrt input
dc_grads = torch.autograd.grad(
outputs=[ddpy.sum()], inputs=[dcx], create_graph=True
)[0]
dc_grads_rep = dc_grads.to_replicate()
# find all losses
dc_loss_grad = dc_grads_rep.mean()
dc_loss = dc_loss_grad
dc_loss.backward()
dc_conv_grad = ddp_conv.module.weight.grad
# Validate the results
assert fp32_allclose(ref_loss, dc_loss)
assert fp32_allclose(ref_y, ddpy)
assert fp32_allclose(ref_grads, dc_grads_rep)
assert fp32_allclose(ref_conv_grad, dc_conv_grad)
@pytest.mark.parametrize(*generate_configs())
def test_double_backprop_combinedloss(
parallel_strategy: ParallelStrategy,
ndims: int,
shard_dim: int,
kernel_size: int,
num_shards: int,
device: torch.device,
):
"""
Test distributed convolution with different number of dimensions and shard dimensions.
Also consider hybrid spatial-data parallelism.
Checks the output and gradients of the distributed convolution against the reference DDP
convolution.
Args:
ndims (int): Number of dimensions for the convolution (1, 2, or 3).
shard_dim (int): Dimension along which the tensor is sharded.
kernel_size (int): Size of the convolution kernel.
num_shards (int): Number of spatial partitions for data
device (torch.device): Torch device to run test with.
"""
# Set the shard dimension for the parallel strategy
parallel_strategy.shard_dim = shard_dim + 2
conv_kwargs = dict(
kernel_size=kernel_size,
padding=kernel_size // 2,
bias=False,
stride=1,
padding_mode="circular",
)
# Initialize the input tensor and convolution layer
shape = [1, 4] + [16] * ndims
x = torch.randn(*shape, device=device, requires_grad=True)
conv_class = getattr(nn, f"Conv{ndims}d")
conv = conv_class(4, 8, **conv_kwargs).to(device).requires_grad_(False)
conv.requires_grad_(True)
# Perform forward and backward pass for reference (non-distributed) convolution
conv.zero_grad()
ref_y = conv(x)
# find gradient wrt input
ref_grads = torch.autograd.grad(
outputs=[ref_y.sum()], inputs=[x], create_graph=True
)[0]
# find all losses
ref_loss_y = ref_y.square().norm()
ref_loss_grad = ref_grads.mean()
ref_loss = ref_loss_y + ref_loss_grad
ref_loss.backward()
ref_x_grad = x.grad
ref_conv_grad = conv.weight.grad.clone()
# Perform forward and backward pass for distributed convolution
conv.zero_grad()
ddp_conv = DistConvDDP(conv, parallel_strategy=parallel_strategy)
dcx = DCTensor.distribute(x, parallel_strategy)
dcy = ddp_conv(dcx)
ddpy = dcy.to_replicate()
# find gradient wrt input
dc_grads = torch.autograd.grad(
outputs=[ddpy.sum()], inputs=[dcx], create_graph=True
)[0]
dc_grads_rep = dc_grads.to_replicate()
# find all losses
dc_loss_y = ddpy.square().norm()
dc_loss_grad = dc_grads_rep.mean()
dc_loss = dc_loss_y + dc_loss_grad
dc_loss.backward()
x_grad = dcx.grad.to_replicate()
dc_conv_grad = ddp_conv.module.weight.grad
# Validate the results
assert fp32_allclose(ref_loss, dc_loss)
assert fp32_allclose(ref_y, ddpy)
assert fp32_allclose(ref_grads, dc_grads_rep)
assert fp32_allclose(ref_x_grad, x_grad)
assert fp32_allclose(ref_conv_grad, dc_conv_grad)