def test_unsqueeze(self):
local_tensor = torch.rand(3, 3, device=self.rank)
replicated_tensor = ReplicatedTensor(local_tensor)
unsqueezed_replicated_tensor = replicated_tensor.unsqueeze(0)
unsqueezed_local_tensor = local_tensor.unsqueeze(0)
self.assertIsInstance(unsqueezed_replicated_tensor, ReplicatedTensor)
self.assertIsInstance(torch.unsqueeze(replicated_tensor, 0), ReplicatedTensor)
self.assertEqual(unsqueezed_local_tensor, unsqueezed_replicated_tensor)
self.assertEqual(torch.unsqueeze(replicated_tensor, 0), unsqueezed_replicated_tensor)
def test_replicated_tensor_inter_op_replicated_tensor(self):
local_tensor = torch.ones(3, 3, device=f"cuda:{self.rank}")
replica_tensor1 = ReplicatedTensor(local_tensor * 4)
replica_tensor2 = ReplicatedTensor(local_tensor * 6)
new_tensor = replica_tensor1 * replica_tensor2
self.assertIsInstance(new_tensor, ReplicatedTensor)
self.assertEqual(new_tensor, torch.ones(3, 3) * 24)
# test replicated tensor inter-op with different pgs
new_pg = dist.new_group(ranks=[1, 2, 3])
replica_tensor_new_group = ReplicatedTensor(local_tensor * 3, process_group=new_pg)
with self.assertRaisesRegex(RuntimeError, 'must be in the same'):
replica_tensor_new_group * replica_tensor1
def test_replicated_tensor_basics(self):
local_tensor = torch.ones(3, 3, device=f"cuda:{self.rank}") * 4
replica_tensor = ReplicatedTensor(local_tensor)
# validate it's a replicated tensor by checking values on all rank
validated = replica_tensor.validate()
self.assertEqual(validated, True)
res = replica_tensor + 2
self.assertIsInstance(res, torch.Tensor)
self.assertNotIsInstance(res, ReplicatedTensor)
self.assertEqual(res, torch.ones(3, 3) * 6)
# modify local tensor on certain rank, and test if validation raise
if self.rank == 2:
local_tensor += 3
with self.assertRaisesRegex(ValueError, 'have different values'):
replica_tensor.validate()
def test_getitem(self):
local_tensor = torch.rand(3, 3, device=self.rank)
replicated_tensor = ReplicatedTensor(local_tensor)
replicated_tensor_view = replicated_tensor[0]
local_tensor_view = local_tensor[0]
self.assertIsInstance(replicated_tensor_view, ReplicatedTensor)
self.assertEqual(local_tensor_view, replicated_tensor_view)
def test_replicated_tensor_inter_op_tensor(self):
local_tensor = torch.ones(3, 3, device=f"cuda:{self.rank}") * 4
replica_tensor = ReplicatedTensor(local_tensor)
local_rand_tensor = torch.randn(3, 3, device=f"cuda:{self.rank}")
new_tensor = replica_tensor + local_rand_tensor
self.assertIsInstance(new_tensor, torch.Tensor)
self.assertNotIsInstance(new_tensor, ReplicatedTensor)
self.assertEqual(new_tensor, local_tensor + local_rand_tensor)
def test_replicated_tensor_inter_op_sharded_tensor(self):
torch.manual_seed(self.rank)
local_tensor1 = torch.rand(12, 3, device=f"cuda:{self.rank}") * 4
local_tensor2 = torch.ones(12, 3, device=f"cuda:{self.rank}") * 4
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
st = _shard_tensor(local_tensor1, spec, src_rank=0)
replica_tensor = ReplicatedTensor(local_tensor2)
ops = [
"torch.add", "torch.sub", "torch.mul", "torch.div", "+", "-", "*",
"/"
]
for op in ops:
binary_op = gen_binary_op_func(op)
res = binary_op(st, replica_tensor)
self.assertIsInstance(res, sharded_tensor.ShardedTensor)
self.assertNotIsInstance(res, ReplicatedTensor)
output = torch.empty((12, 3)) if self.rank == 0 else None
res.gather(dst=0, out=output)
if self.rank == 0:
local_output = binary_op(local_tensor1, local_tensor2)
self.assertEqual(output, local_output)
# reflective
reflect_res = binary_op(replica_tensor, st)
self.assertIsInstance(reflect_res, sharded_tensor.ShardedTensor)
self.assertNotIsInstance(reflect_res, ReplicatedTensor)
reflect_output = torch.empty((12, 3)) if self.rank == 0 else None
reflect_res.gather(dst=0, out=reflect_output)
if self.rank == 0:
reflect_local_output = binary_op(local_tensor2, local_tensor1)
self.assertEqual(reflect_output, reflect_local_output)
def test_replicated_tensor_implicit_broadcasting(self):
# use same seed
torch.manual_seed(self.rank)
# test implicit broadcasting
local_tensor1 = torch.rand(12, 3, device=f"cuda:{self.rank}") * 4
# we use size (3) to trigger the implicit broadcasting logic
# and it will fail if implicit broadcasting not happen.
local_tensor2 = torch.ones(3, device=f"cuda:{self.rank}")
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
st = _shard_tensor(local_tensor1, spec, src_rank=0)
replica_tensor = ReplicatedTensor(local_tensor2)
ops = [
"torch.add", "torch.sub", "torch.mul", "torch.div", "+", "-", "*",
"/"
]
for op in ops:
binary_op = gen_binary_op_func(op)
# replicated tensor should automatically broadcasted
res = binary_op(st, replica_tensor)
self.assertIsInstance(res, sharded_tensor.ShardedTensor)
output = torch.empty((12, 3)) if self.rank == 0 else None
res.gather(dst=0, out=output)
if self.rank == 0:
local_output = binary_op(local_tensor1, local_tensor2)
self.assertEqual(output, local_output)
def test_replicated_tensor_inter_op_sharded_tensor_errors(self):
local_tensor = torch.ones(3, 3, device=f"cuda:{self.rank}") * 4
replica_tensor = ReplicatedTensor(local_tensor)
torch.manual_seed(self.rank)
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
st1 = sharded_tensor.rand(spec, (20, 3, 3))
st2 = sharded_tensor.rand(spec, (30, 3, 3))
with self.assertRaisesRegex(RuntimeError, 'Implicit broadcasting'):
st1 + st2
with self.assertRaisesRegex(RuntimeError, 'not supported for ShardedTensor'):
st1 % replica_tensor
def test_with_ddp(self):
# Test Replicated params for DDP
replica_tensor = ReplicatedTensor(torch.rand(4, 8, device=self.rank))
model = torch.nn.Linear(8, 2).cuda(self.rank)
optim = torch.optim.SGD(model.parameters(), lr=0.1)
ddp = DDP(model)
# Test module.parameters.
params = list(ddp.parameters())
self.assertEqual(2, len(params))
self.assertEqual(ddp.module.weight, params[0])
self.assertEqual(ddp.module.bias, params[1])
params = list(model.parameters())
self.assertEqual(2, len(params))
self.assertEqual(model.weight, params[0])
self.assertEqual(model.bias, params[1])
# Validate output
out = ddp(replica_tensor)
self.assertIsInstance(out, ReplicatedTensor)
# Test backward and optimizer.
# Validate backward.
out.sum().backward()
self.assertIsNotNone(model.weight.grad)
self.assertIsNotNone(model.bias.grad)
self.assertIsNotNone(ddp.module.weight.grad)
self.assertIsNotNone(ddp.module.bias.grad)
original_params = []
for param_group in optim.param_groups:
for original_param in param_group['params']:
self.assertIsNotNone(original_param.grad)
original_params.append(original_param)
self.assertEqual(model.weight.grad, original_params[0].grad)
self.assertEqual(model.bias.grad, original_params[1].grad)
self.assertEqual(model.weight.grad, ddp.module.weight.grad)
self.assertEqual(model.bias.grad, ddp.module.bias.grad)
# Validate optimizer.
optim.step()
self.assertEqual(model.weight, ddp.module.weight)
self.assertEqual(model.weight, original_params[0])
self.assertEqual(model.bias, ddp.module.bias)
self.assertEqual(model.bias, original_params[1])
# Validate zero_grad
optim.zero_grad()
self.assertEqual(model.weight.grad,
torch.zeros_like(model.weight.grad))
self.assertEqual(model.weight.grad, ddp.module.weight.grad)
self.assertEqual(model.weight.grad, original_params[0].grad)
self.assertEqual(model.bias.grad, torch.zeros_like(model.bias.grad))
self.assertEqual(model.bias.grad, ddp.module.bias.grad)
self.assertEqual(model.bias.grad, original_params[1].grad)
# Validate zero_grad set_to_none
optim.zero_grad(set_to_none=True)
self.assertIsNone(model.weight.grad)
self.assertEqual(model.weight.grad, ddp.module.weight.grad)
self.assertEqual(model.weight.grad, original_params[0].grad)
self.assertIsNone(model.bias.grad)
self.assertEqual(model.bias.grad, ddp.module.bias.grad)
self.assertEqual(model.bias.grad, original_params[1].grad)
# Multiple forward passes.
for _ in range(5):
out = ddp(replica_tensor)
self.assertIsInstance(out, ReplicatedTensor)
# Test with context manager.
from torch.nn.parallel._replicated_tensor_ddp_utils import _ddp_replicated_tensor
with _ddp_replicated_tensor(False):
for _ in range(5):
with _ddp_replicated_tensor(True):
ddp = DDP(model)
out = ddp(replica_tensor)
self.assertIsInstance(out, ReplicatedTensor)
# Test save and load.
with _ddp_replicated_tensor(False):
ddp = DDP(model)
expected_state_dict = ddp.state_dict()
buffer = io.BytesIO()
torch.save(ddp, buffer)
buffer.seek(0)
obj = torch.load(buffer)
self.assertEqual(expected_state_dict, obj.state_dict())
with _ddp_replicated_tensor(True):
ddp = DDP(model)
buffer = io.BytesIO()
torch.save(ddp, buffer)
buffer.seek(0)
obj = torch.load(buffer)
self.assertEqual(expected_state_dict, obj.state_dict())
def forward(ctx, inp, process_group=None):
# set_materialize_grads(False) will ensure that None gradients stay as
# None and are not filled with zeros.
ctx.set_materialize_grads(False)
return ReplicatedTensor(inp, process_group)