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)