def _run_sharded_linear(self, spec, input_size, linear_size, sharded_dim,
dtype):
# Use same seed.
torch.manual_seed(0)
local_linear = torch.nn.Linear(*linear_size,
dtype=dtype).cuda(self.rank)
sharded_linear = torch.nn.Linear(*linear_size, dtype=dtype)
# Copy the weights and bias from local linear
sharded_linear.weight = clone_module_parameter(local_linear, "weight")
sharded_linear.bias = clone_module_parameter(local_linear, "bias")
# Shard the parameter.
shard_parameter(sharded_linear, "weight", spec)
# Run sharded computation
torch.manual_seed(self.rank) # inputs different on each rank
inp = torch.rand(*input_size, dtype=dtype).cuda(self.rank)
reshard_spec = copy.deepcopy(spec)
reshard_spec.dim = 0
reshard_spec.placements.sort(key=lambda placement: placement.rank())
sharded_linear = _collect_local_shard(
_reshard_output(sharded_linear, reshard_spec))
sharded_output = sharded_linear(inp)
# Run local computation
local_output = local_linear(inp)
# Verify
self.assertEqual(local_output, sharded_output, atol=1e-3, rtol=1e-3)
# Validate for torch.nn.functional.linear version.
local_output = torch.nn.functional.linear(inp, local_linear.weight,
local_linear.bias)
sharded_output = torch.nn.functional.linear(inp, sharded_linear.weight,
sharded_linear.bias)
sharded_output = sharded_output.reshard(reshard_spec).local_tensor()
# When local tensor only has one dimension, we increase one more dimension
# for reshard. We need to squeeze the # of dimensions manually.
if inp.dim() == 1:
sharded_output = sharded_output.squeeze(reshard_spec.dim)
self.assertEqual(local_output, sharded_output, atol=1e-3, rtol=1e-3)
# Compute loss and run backward pass.
local_output.sum().backward()
sharded_output.sum().backward()
local_grad = local_linear.weight.grad
# Verify that both weight and bias in the sharded linear has non-None grad.
sharded_weight = sharded_linear.weight.local_tensor()
self.assertNotEqual(sharded_linear.bias.grad, None)
self.assertNotEqual(sharded_weight.grad, None)
# Shard the local linear's weight grad so that we can compare.
dist.all_reduce(local_grad)
(start_pos,
chunk_size) = generate_local_weight_sharding_params_for_test(
local_linear.weight, sharded_dim, TEST_GPU_NUM, spec, self.rank)
local_grad_narrowed = local_grad.narrow(sharded_dim, start_pos,
chunk_size)
local_bias_grad = local_linear.bias.grad
dist.all_reduce(local_bias_grad)
# Test backward gradient calculation.
self.assertEqual(sharded_linear.bias.grad,
local_bias_grad,
atol=1e-3,
rtol=1e-3)
self.assertEqual(sharded_weight.grad,
local_grad_narrowed,
atol=1e-3,
rtol=1e-3)
# Test optimizer.
previous = local_linear.weight.clone().detach()
optim = torch.optim.SGD(local_linear.parameters(), lr=0.1)
optim.step()
self.assertNotEqual(previous, local_linear.weight)
previous_sharded_weight = sharded_weight.clone()
previous_sharded_bias = sharded_linear.bias.clone()
sharded_optim = ShardedOptimizer(
dict(named_params_with_sharded_tensor(sharded_linear)),
torch.optim.SGD,
lr=0.1,
)
sharded_optim.step()
sharded_weight = sharded_linear.weight.local_tensor()
local_weight_narrowed = local_linear.weight.narrow(
sharded_dim, start_pos, chunk_size)
self.assertEqual(sharded_weight.size(), local_weight_narrowed.size())
self.assertNotEqual(previous_sharded_weight, sharded_weight)
self.assertEqual(sharded_weight,
local_weight_narrowed,
atol=1e-3,
rtol=1e-3)
self.assertNotEqual(previous_sharded_bias, sharded_linear.bias)
self.assertEqual(sharded_linear.bias,
local_linear.bias,
atol=1e-3,
rtol=1e-3)
def _run_sharded_embedding(
self,
spec,
input_size,
num_embeddings,
embedding_dim,
max_norm=None,
norm_type=2.0,
padding_idx=None,
):
# Use same seed.
torch.manual_seed(0)
local_embedding = torch.nn.Embedding(
num_embeddings,
embedding_dim,
max_norm=max_norm,
norm_type=norm_type,
padding_idx=padding_idx,
).cuda(self.rank)
sharded_embedding = torch.nn.Embedding(
num_embeddings,
embedding_dim,
max_norm=max_norm,
norm_type=norm_type,
padding_idx=padding_idx,
)
# Copy the weights from local embedding
sharded_embedding.weight = clone_module_parameter(
local_embedding, "weight")
# Shard the parameter.
shard_parameter(sharded_embedding, "weight", spec)
# Run sharded computation
torch.manual_seed(self.rank) # inputs different on each rank
inp = torch.randint(0, num_embeddings,
tuple(input_size)).cuda(self.rank)
sharded_output = sharded_embedding(inp)
# If max_norm is set, we need to ensure that the renorm has been applied across
# inputs from all ranks.
if max_norm is not None:
gathered_inputs = [
torch.zeros_like(inp) for _ in range(TEST_GPU_NUM)
]
dist.all_gather(gathered_inputs, inp)
unique_inp = torch.unique(torch.cat(gathered_inputs))
local_embedding(unique_inp)
# Run local computation
local_output = local_embedding(inp)
# Compare local weight and shared one to ensure the renorm
# as expected.
if max_norm is not None:
sharded_dim = spec.dim
sharded_weight = sharded_embedding.weight.local_shards()[0].tensor
(start_pos,
chunk_size) = generate_local_weight_sharding_params_for_test(
local_embedding.weight, sharded_dim, TEST_GPU_NUM, spec,
self.rank)
local_weight_narrowed = local_embedding.weight.narrow(
sharded_dim, start_pos, chunk_size)
self.assertEqual(local_weight_narrowed, sharded_weight)
# Verify
self.assertEqual(local_output, sharded_output)
# Validate for torch.nn.functional.embedding version.
local_output = torch.nn.functional.embedding(
inp,
local_embedding.weight,
max_norm=max_norm,
norm_type=norm_type,
padding_idx=padding_idx,
)
sharded_output = torch.nn.functional.embedding(
inp,
sharded_embedding.weight,
max_norm=max_norm,
norm_type=norm_type,
padding_idx=padding_idx,
)
self.assertEqual(local_output, sharded_output)
def _weight_override(module_dst, module_src):
module_dst.fc1.weight = clone_module_parameter(module_src.fc1, "weight")
module_dst.fc1.bias = clone_module_parameter(module_src.fc1, "bias")
module_dst.fc2.weight = clone_module_parameter(module_src.fc2, "weight")
module_dst.fc2.bias = clone_module_parameter(module_src.fc2, "bias")
def _run_sharded_embedding_bag(
self,
spec,
input_size,
num_embeddings,
embedding_dim,
mode,
sharded_dim=None,
include_last_offset=False,
offset_size=None,
max_norm=None,
norm_type=2.0,
padding_idx=None,
):
# Use same seed.
torch.manual_seed(0)
local_embedding_bag = torch.nn.EmbeddingBag(
num_embeddings,
embedding_dim,
mode=mode,
max_norm=max_norm,
norm_type=norm_type,
include_last_offset=include_last_offset,
padding_idx=padding_idx,
).cuda(self.rank)
sharded_embedding_bag = torch.nn.EmbeddingBag(
num_embeddings,
embedding_dim,
mode=mode,
max_norm=max_norm,
norm_type=norm_type,
include_last_offset=include_last_offset,
padding_idx=padding_idx,
)
# Copy the weights from local embedding bag.
sharded_embedding_bag.weight = clone_module_parameter(
local_embedding_bag, "weight")
# Shard the parameter.
shard_parameter(sharded_embedding_bag, "weight", spec)
# Run sharded computation
torch.manual_seed(self.rank) # inputs different on each rank
inp = torch.randint(0, num_embeddings,
tuple(input_size)).cuda(self.rank)
per_sample_weights = None
if mode == "sum":
per_sample_weights = torch.rand(*input_size).cuda(self.rank)
offsets = None
if len(input_size) == 1:
# We need to generate certain length offset for each rank.
# The current implementation and dist API does not support
# the case when the offset has different lengths.
# input_size[0] >> offset_size, so the while loop will not
# for too long.
while offsets is None or (offsets.size(0) != offset_size):
offsets = torch.randint(input_size[0], (offset_size, ))
offsets[0] = 0
if include_last_offset:
offsets[-1] = input_size[0]
offsets = (torch.unique(
offsets, sorted=True).contiguous().cuda(self.rank))
# If max_norm is set, we need to ensure that the renorm has been applied across
# inputs from all ranks.
if max_norm is not None:
gathered_inputs = [
torch.zeros_like(inp) for _ in range(TEST_GPU_NUM)
]
dist.all_gather(gathered_inputs, inp)
unique_inp = torch.unique(torch.cat(gathered_inputs))
offsets_dummy = torch.tensor([len(unique_inp) // 2
]).cuda(self.rank)
local_embedding_bag(unique_inp, offsets=offsets_dummy)
sharded_output = sharded_embedding_bag(
inp,
offsets=offsets,
per_sample_weights=per_sample_weights,
)
# Run local computation
local_output = local_embedding_bag(
inp,
offsets=offsets,
per_sample_weights=per_sample_weights,
)
# Compare local weight and shared one to ensure the renorm
# as expected.
if max_norm is not None:
sharded_weight = sharded_embedding_bag.weight.local_shards(
)[0].tensor
(start_pos,
chunk_size) = generate_local_weight_sharding_params_for_test(
local_embedding_bag.weight, sharded_dim, TEST_GPU_NUM, spec,
self.rank)
local_weight_narrowed = local_embedding_bag.weight.narrow(
sharded_dim, start_pos, chunk_size)
self.assertEqual(local_weight_narrowed, sharded_weight)
# Verify
self.assertEqual(local_output, sharded_output)
# Validate for torch.nn.functional.embedding_bag version.
local_output = torch.nn.functional.embedding_bag(
inp,
local_embedding_bag.weight,
offsets=offsets,
mode=mode,
per_sample_weights=per_sample_weights,
include_last_offset=include_last_offset,
max_norm=max_norm,
norm_type=norm_type,
padding_idx=padding_idx,
)
sharded_output = torch.nn.functional.embedding_bag(
inp,
sharded_embedding_bag.weight,
offsets=offsets,
mode=mode,
per_sample_weights=per_sample_weights,
include_last_offset=include_last_offset,
max_norm=max_norm,
norm_type=norm_type,
padding_idx=padding_idx,
)
self.assertEqual(local_output, sharded_output)