def test_named_params_with_sharded_tensor(self): rowwise_spec = ChunkShardingSpec( dim=0, placements=[ "rank:0/cuda:0", "rank:1/cuda:1", "rank:2/cuda:2", "rank:3/cuda:3", ], ) sharded_model = MyShardedModel(spec=rowwise_spec).cuda(self.rank) sharded_model_params = dict(named_params_with_sharded_tensor(sharded_model)) param_keys = list(sharded_model_params.keys()) self.assertEqual(len(param_keys), 2) self.assertTrue("param" in param_keys) self.assertTrue("sharded_param" in param_keys) sharded_linear = MyShardedLinear(rank=self.rank).cuda(self.rank) sharded_linear.shard_parameter() sharded_linear_params = dict(named_params_with_sharded_tensor(sharded_linear)) param_keys = list(sharded_linear_params.keys()) self.assertEqual(len(param_keys), 4) self.assertTrue("linear1.bias" in param_keys) self.assertTrue("linear2.bias" in param_keys) self.assertTrue("linear1.weight" in param_keys) self.assertTrue("linear2.weight" in param_keys) self.assertFalse("bias" in param_keys)
def _get_toy_module_optim(module, lr): """ Creata a optimizer for sharded tensor by using ShardedOptimizer. """ return ShardedOptimizer( dict(named_params_with_sharded_tensor(module)), torch.optim.SGD, # SGD is only demo purpose, one can use other optims. lr=lr, )
def test_sharded_optim(self): rowwise_spec = ChunkShardingSpec( dim=0, placements=[ "rank:0/cuda:0", "rank:1/cuda:1", "rank:2/cuda:2", "rank:3/cuda:3", ], ) local_model = MyShardedModel().cuda(self.rank) sharded_model = MyShardedModel(spec=rowwise_spec).cuda(self.rank) # copy the parameteres from local model sharded_model.sharded_param.local_shards()[0].tensor = \ local_model.sharded_param.detach().clone().requires_grad_() local_optim = optim.SGD(local_model.parameters(), lr=0.1) sharded_model_params = dict(named_params_with_sharded_tensor(sharded_model)) sharded_optim = ShardedOptimizer(sharded_model_params, optim.SGD, lr=0.1) local_optim.zero_grad() sharded_optim.zero_grad() before_update = deepcopy(sharded_optim.named_params) inp = torch.rand([5, 10]).cuda(self.rank).requires_grad_() # run forward local_output = local_model(inp) sharded_output = sharded_model(inp) # backward local_output.sum().backward() sharded_output.sum().backward() # optimizer update local_optim.step() sharded_optim.step() # make sure the parameters (including sharded param) # get updated by the optimizer, and the updated # local params are the same as the sharded params for key, val in before_update.items(): new_val = sharded_optim.named_params[key] if isinstance(val, sharded_tensor.ShardedTensor): self.assertNotEqual( val.local_shards()[0].tensor, new_val.local_shards()[0].tensor ) self.assertEqual( new_val.local_shards()[0].tensor, local_model.sharded_param ) else: self.assertNotEqual(val, new_val) self.assertEqual(new_val, local_model.param)
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 test_sharding_plan_simple_megatron(self): colwise_sharding_spec = generate_chunk_sharding_specs_for_test(0) rowwise_sharding_spec = generate_chunk_sharding_specs_for_test(1) for spec in zip(colwise_sharding_spec, rowwise_sharding_spec): # test each sharding spec pair and see if we can apply sharding reshard_spec = copy.deepcopy(spec[1]) reshard_spec.placements.sort(key=lambda placement: placement.rank()) reshard_spec.dim = 0 sharding_plan = ShardingPlan( plan={ "fc1.weight": spec[0], "fc2.weight": spec[1] }, output_plan={ "": reshard_spec }, return_local_tensor=[""]) # Use same seed. torch.manual_seed(0) local_megatron_lm = SimpleMegatronLM([[17, 12], [12, 29]]).cuda(self.rank) megatron_lm = copy.deepcopy(local_megatron_lm) # shard the module with the provided sharding plan shard_module(megatron_lm, sharding_plan) # check to make sure the module already been sharded self.assertTrue(isinstance(megatron_lm.fc1.weight, ShardedTensor)) self.assertTrue(isinstance(megatron_lm.fc2.weight, ShardedTensor)) self.assertEqual(megatron_lm.fc1.weight.sharding_spec(), spec[0]) self.assertEqual(megatron_lm.fc2.weight.sharding_spec(), spec[1]) # make sure we can run sharded computation input = torch.rand(22, 17).cuda(self.rank) sharded_output = megatron_lm(input) local_output = local_megatron_lm(input) # verify and make sure local and sharded output matches self.assertEqual(local_output, sharded_output) # Compute loss and run backward pass. local_output.sum().backward() sharded_output.sum().backward() ( local_weight_grad_fc1, local_weight_grad_fc2, ) = local_megatron_lm.get_weight_grads() local_bias_grad_fc1, local_bias_grad_fc2 = local_megatron_lm.get_bias_grads() # Verify that weights in both layers and biases in the sharded linear has non-None grad. ( sharded_weight_fc1, sharded_weight_fc2, ) = megatron_lm.get_weights() bias_grad_fc1, bias_grad_fc2 = megatron_lm.get_bias_grads() self.assertNotEqual(sharded_weight_fc1.grad, None) self.assertNotEqual(sharded_weight_fc2.grad, None) self.assertNotEqual(bias_grad_fc1, None) self.assertNotEqual(bias_grad_fc2, None) # Shard the local linear's weight grad so that we can compare. dist.all_reduce(local_weight_grad_fc1) dist.all_reduce(local_weight_grad_fc2) dist.all_reduce(local_bias_grad_fc1) dist.all_reduce(local_bias_grad_fc2) local_weight_fc1, local_weight_fc2 = local_megatron_lm.get_weights() ( start_pos_fc1, chunk_size_fc1, ) = generate_local_weight_sharding_params_for_test( local_weight_fc1, 0, TEST_GPU_NUM, spec[0], self.rank ) local_grad_narrowed_fc1 = local_weight_grad_fc1.narrow( 0, start_pos_fc1, chunk_size_fc1 ) ( start_pos_fc2, chunk_size_fc2, ) = generate_local_weight_sharding_params_for_test( local_weight_fc2, 1, TEST_GPU_NUM, spec[1], self.rank ) local_grad_narrowed_fc2 = local_weight_grad_fc2.narrow( 1, start_pos_fc2, chunk_size_fc2 ) # Test backward gradient calculation. self.assertEqual(sharded_weight_fc1.grad, local_grad_narrowed_fc1) self.assertEqual(sharded_weight_fc2.grad, local_grad_narrowed_fc2) self.assertEqual(bias_grad_fc1, local_bias_grad_fc1) self.assertEqual(bias_grad_fc2, local_bias_grad_fc2) # Test optimizer. bias_fc1, bias_fc2 = megatron_lm.get_biases() local_bias_fc1, local_bias_fc2 = local_megatron_lm.get_biases() self.assertEqual(bias_fc1, local_bias_fc1) self.assertEqual(bias_fc2, local_bias_fc2) self.assertEqual(bias_fc1.grad, local_bias_fc1.grad) self.assertEqual(bias_fc2.grad, local_bias_fc2.grad) previous_sharded_weight_fc1 = sharded_weight_fc1.clone() previous_sharded_weight_fc2 = sharded_weight_fc2.clone() previous_bias_fc1 = bias_fc1.clone() previous_bias_fc2 = bias_fc2.clone() optim = torch.optim.SGD(local_megatron_lm.parameters(), lr=0.1) optim.step() sharded_optim = ShardedOptimizer( dict(named_params_with_sharded_tensor(megatron_lm)), torch.optim.SGD, lr=0.1, ) sharded_optim.step() local_weight_fc1_narrowed = local_weight_fc1.narrow( 0, start_pos_fc1, chunk_size_fc1 ) local_weight_fc2_narrowed = local_weight_fc2.narrow( 1, start_pos_fc2, chunk_size_fc2 ) # Test weight value after optimizer. self.assertEqual(sharded_weight_fc1.size(), local_weight_fc1_narrowed.size()) self.assertEqual(sharded_weight_fc2.size(), local_weight_fc2_narrowed.size()) self.assertNotEqual(previous_sharded_weight_fc1, sharded_weight_fc1) self.assertNotEqual(previous_sharded_weight_fc2, sharded_weight_fc2) self.assertEqual(sharded_weight_fc1, local_weight_fc1_narrowed) self.assertEqual(sharded_weight_fc2, local_weight_fc2_narrowed) # Test bias value after optimizer. local_bias_fc1, local_bias_fc2 = local_megatron_lm.get_biases() self.assertNotEqual(previous_bias_fc1, bias_fc1) self.assertEqual(bias_fc1, local_bias_fc1) self.assertNotEqual(previous_bias_fc2, bias_fc2) self.assertEqual(bias_fc2, local_bias_fc2)
def _run_megatron_linear(self, spec, input_size, linear_size, dtype): 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 _shard_parameter(module, spec): shard_parameter(module.fc1, "weight", spec[0]) shard_parameter(module.fc2, "weight", spec[1]) # Use same seed. torch.manual_seed(0) local_megatron_lm = SimpleMegatronLM(linear_size, rank=self.rank, dtype=dtype) sharded_megatron_lm = SimpleMegatronLM(linear_size, dtype=dtype) _weight_override(sharded_megatron_lm, local_megatron_lm) # Shard the parameter. First col-wise sharding and then row-wise _shard_parameter(sharded_megatron_lm, spec) # Setup resharding of output. reshard_spec = copy.deepcopy(spec[1]) reshard_spec.placements.sort(key=lambda placement: placement.rank()) reshard_spec.dim = 0 sharded_megatron_lm = _collect_local_shard( _reshard_output(sharded_megatron_lm, reshard_spec) ) torch.manual_seed(self.rank) # inputs different on each rank inp = torch.rand(*input_size, requires_grad=True, device=self.rank, dtype=dtype) # Run local computation local_output = local_megatron_lm(inp) # Compute loss and run backward pass. local_output.sum().backward() # Save and reset input grads. local_input_grad = inp.grad self.assertIsNotNone(inp.grad) inp.grad = None # Run sharded computation sharded_output = sharded_megatron_lm(inp) # Verify local and sharded results self.assertEqual(local_output, sharded_output, atol=1e-3, rtol=1e-6) sharded_output.sum().backward() sharded_input_grad = inp.grad self.assertIsNotNone(inp.grad) # Verify sharded and local grads. self.assertEqual(local_input_grad, sharded_input_grad, atol=1e-3, rtol=1e-6) ( local_weight_grad_fc1, local_weight_grad_fc2, ) = local_megatron_lm.get_weight_grads() local_bias_grad_fc1, local_bias_grad_fc2 = local_megatron_lm.get_bias_grads() # Verify that weights in both layers and biases in the sharded linear has non-None grad. ( sharded_weight_fc1, sharded_weight_fc2, ) = sharded_megatron_lm.get_weights() bias_grad_fc1, bias_grad_fc2 = sharded_megatron_lm.get_bias_grads() self.assertNotEqual(sharded_weight_fc1.grad, None) self.assertNotEqual(sharded_weight_fc2.grad, None) self.assertNotEqual(bias_grad_fc1, None) self.assertNotEqual(bias_grad_fc2, None) # Shard the local linear's weight grad so that we can compare. dist.all_reduce(local_weight_grad_fc1) dist.all_reduce(local_weight_grad_fc2) dist.all_reduce(local_bias_grad_fc1) dist.all_reduce(local_bias_grad_fc2) local_weight_fc1, local_weight_fc2 = local_megatron_lm.get_weights() ( start_pos_fc1, chunk_size_fc1, ) = generate_local_weight_sharding_params_for_test( local_weight_fc1, 0, TEST_GPU_NUM, spec[0], self.rank ) local_grad_narrowed_fc1 = local_weight_grad_fc1.narrow( 0, start_pos_fc1, chunk_size_fc1 ) ( start_pos_fc2, chunk_size_fc2, ) = generate_local_weight_sharding_params_for_test( local_weight_fc2, 1, TEST_GPU_NUM, spec[1], self.rank ) local_grad_narrowed_fc2 = local_weight_grad_fc2.narrow( 1, start_pos_fc2, chunk_size_fc2 ) # Test backward gradient calculation. self.assertEdistNorm(sharded_weight_fc1.grad, local_grad_narrowed_fc1) self.assertEdistNorm(sharded_weight_fc2.grad, local_grad_narrowed_fc2) self.assertEdistNorm(bias_grad_fc1, local_bias_grad_fc1) self.assertEdistNorm(bias_grad_fc2, local_bias_grad_fc2) # Test optimizer. bias_fc1, bias_fc2 = sharded_megatron_lm.get_biases() local_bias_fc1, local_bias_fc2 = local_megatron_lm.get_biases() self.assertEdistNorm(bias_fc1, local_bias_fc1) self.assertEdistNorm(bias_fc2, local_bias_fc2) self.assertEdistNorm(bias_fc1.grad, local_bias_fc1.grad) self.assertEdistNorm(bias_fc2.grad, local_bias_fc2.grad) previous_sharded_weight_fc1 = sharded_weight_fc1.clone() previous_sharded_weight_fc2 = sharded_weight_fc2.clone() previous_bias_fc1 = bias_fc1.clone() previous_bias_fc2 = bias_fc2.clone() optim = torch.optim.SGD(local_megatron_lm.parameters(), lr=0.1) optim.step() sharded_optim = ShardedOptimizer( dict(named_params_with_sharded_tensor(sharded_megatron_lm)), torch.optim.SGD, lr=0.1, ) sharded_optim.step() local_weight_fc1_narrowed = local_weight_fc1.narrow( 0, start_pos_fc1, chunk_size_fc1 ) local_weight_fc2_narrowed = local_weight_fc2.narrow( 1, start_pos_fc2, chunk_size_fc2 ) # Test weight value after optimizer. self.assertEqual(sharded_weight_fc1.size(), local_weight_fc1_narrowed.size()) self.assertEqual(sharded_weight_fc2.size(), local_weight_fc2_narrowed.size()) self.assertNotEqual(previous_sharded_weight_fc1, sharded_weight_fc1) self.assertNotEqual(previous_sharded_weight_fc2, sharded_weight_fc2) self.assertEdistNorm(sharded_weight_fc1, local_weight_fc1_narrowed) self.assertEdistNorm(sharded_weight_fc2, local_weight_fc2_narrowed) # Test bias value after optimizer. local_bias_fc1, local_bias_fc2 = local_megatron_lm.get_biases() self.assertNotEqual(previous_bias_fc1, bias_fc1) self.assertEdistNorm(bias_fc1, local_bias_fc1) self.assertNotEqual(previous_bias_fc2, bias_fc2) self.assertEdistNorm(bias_fc2, local_bias_fc2)
def _run_sharded_linear(self, spec, input_size, linear_size, sharded_dim): # Use same seed. torch.manual_seed(0) local_linear = torch.nn.Linear(*linear_size).cuda(self.rank) sharded_linear = torch.nn.Linear(*linear_size) # Copy the weights and bias from local linear sharded_linear.weight = torch.nn.Parameter( local_linear.weight.detach().clone()) sharded_linear.bias = torch.nn.Parameter( local_linear.bias.detach().clone()) # 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).cuda(self.rank) sharded_output = sharded_linear(inp) # Run local computation local_output = local_linear(inp) # Verify self.assertEqual(local_output, sharded_output) # 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) self.assertEqual(local_output, sharded_output) # 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_shards()[0].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) # Test backward gradient calculation. self.assertEqual(sharded_linear.bias.grad, local_linear.bias.grad) self.assertEqual(sharded_weight.grad, local_grad_narrowed) # 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_shards()[0].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) self.assertNotEqual(previous_sharded_bias, sharded_linear.bias) self.assertEqual(sharded_linear.bias, local_linear.bias)