def generate_chunk_sharding_specs_for_test(sharding_dim):
return [
ChunkShardingSpec(
dim=sharding_dim,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
),
# Test different ordering. (Case 1)
ChunkShardingSpec(
dim=sharding_dim,
placements=[
"rank:2/cuda:2",
"rank:3/cuda:3",
"rank:0/cuda:0",
"rank:1/cuda:1",
],
),
# Test different ordering. (Case 2)
ChunkShardingSpec(
dim=sharding_dim,
placements=[
"rank:3/cuda:3",
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
],
),
]
def test_sharded_linear_colwise(self):
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
self._run_sharded_linear(spec, [5, 17], [17, 12], 0)
self._run_sharded_linear(spec, [5, 21], [21, 11], 0)
self._run_sharded_linear(spec, [5, 23], [23, 13], 0)
self._run_sharded_linear(spec, [5, 15], [15, 14], 0)
# Test different ordering.
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:1/cuda:1",
"rank:0/cuda:0",
"rank:3/cuda:3",
"rank:2/cuda:2",
],
)
self._run_sharded_linear(spec, [5, 17], [17, 12], 0)
self._run_sharded_linear(spec, [5, 21], [21, 11], 0)
self._run_sharded_linear(spec, [5, 23], [23, 13], 0)
self._run_sharded_linear(spec, [5, 15], [15, 14], 0)
def test_sharded_linear_rowwise(self):
spec = ChunkShardingSpec(
dim=1,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
# Test even split.
self._run_sharded_linear(spec, [5, 16], [16, 11], 1)
# Test uneven split.
self._run_sharded_linear(spec, [5, 19], [19, 11], 1)
self._run_sharded_linear(spec, [5, 21], [21, 11], 1)
# Test different ordering.
spec = ChunkShardingSpec(
dim=1,
placements=[
"rank:2/cuda:2",
"rank:3/cuda:3",
"rank:0/cuda:0",
"rank:1/cuda:1",
],
)
self._run_sharded_linear(spec, [5, 16], [16, 11], 1)
self._run_sharded_linear(spec, [5, 19], [19, 11], 1)
self._run_sharded_linear(spec, [5, 21], [21, 11], 1)
def test_invalid_sharding(self):
self.init_pg()
spec = ChunkShardingSpec(dim=0, placements=["rank:1/cuda:1"])
pg = dist.new_group(ranks=[2, 3])
if self.rank < 2:
with self.assertRaisesRegex(ValueError, 'not part of process group'):
_sharded_tensor.empty(spec, 10, 20, process_group=pg)
spec = ChunkShardingSpec(dim='H', placements=["rank:1/cuda:1"])
with self.assertRaisesRegex(ValueError, 'needs to be an integer'):
_sharded_tensor.empty(spec, 10, 20)
for dim in [2, 3, 4, -3, -4, -5]:
spec = ChunkShardingSpec(dim=dim, placements=["rank:1/cuda:1"])
with self.assertRaisesRegex(ValueError, 'Invalid sharding dim'):
_sharded_tensor.empty(spec, 10, 20)
spec = ChunkShardingSpec(dim=0, placements=["rank:5/cuda:1"])
with self.assertRaisesRegex(ValueError, 'Invalid rank'):
_sharded_tensor.empty(spec, 10, 20)
spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
sharded_tensor = _sharded_tensor.empty(spec, 10, 20)
tensor = torch.empty(10, 20)
with self.assertRaisesRegex(RuntimeError, "torch function 'add' not supported for ShardedTensor!"):
torch.add(sharded_tensor, tensor)
spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
with self.assertRaisesRegex(ValueError, 'Only torch.strided layout is currently supported'):
_sharded_tensor.empty(spec, 10, 20, layout=torch.sparse)
spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
with self.assertRaisesRegex(ValueError, 'Only torch.contiguous_format memory_format is currently supported'):
_sharded_tensor.empty(spec, 10, 20, memory_format=torch.channels_last)
spec = ChunkShardingSpec(dim=0, placements=["worker0/cuda:1"])
with self.assertRaisesRegex(RuntimeError, 'RPC framework needs to be initialized'):
_sharded_tensor.empty(spec, 10, 20)
spec = ChunkShardingSpec(dim=0, placements=["rank:0/cuda:1"])
with self.assertRaisesRegex(RuntimeError, 'RPC was not initialized'):
st = _sharded_tensor.empty(spec, 10, 20)
st.remote_shards
self.init_rpc()
# ShardedTensor was initialized before RPC.
with self.assertRaisesRegex(RuntimeError, 'RPC was not initialized'):
st.remote_shards
spec = ChunkShardingSpec(dim=0, placements=["workerfoo/cuda:1"])
with self.assertRaisesRegex(ValueError, 'Invalid worker name'):
_sharded_tensor.empty(spec, 10, 20)
def test_sharded_tensor_metadata(self):
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 10, 20)
sharded_tensor_metadata = sharded_tensor.metadata()
self.assertEqual(torch.Size([10, 20]), sharded_tensor_metadata.size)
self.assertEqual(torch.float, sharded_tensor_metadata.dtype)
self.assertEqual(torch.strided, sharded_tensor_metadata.layout)
self.assertEqual(False, sharded_tensor_metadata.requires_grad)
self.assertEqual(torch.contiguous_format,
sharded_tensor_metadata.memory_format)
self.assertEqual(False, sharded_tensor_metadata.pin_memory)
sharded_tensor = _sharded_tensor.empty(spec,
10,
20,
requires_grad=True)
sharded_tensor_metadata = sharded_tensor.metadata()
self.assertEqual(True, sharded_tensor_metadata.requires_grad)
sharded_tensor = _sharded_tensor.empty(spec,
10,
20,
dtype=torch.double)
sharded_tensor_metadata = sharded_tensor.metadata()
self.assertEqual(torch.double, sharded_tensor_metadata.dtype)
# Need CPU for pin_memory
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cpu",
"rank:1/cpu",
"rank:2/cpu",
"rank:3/cpu",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 10, 20, pin_memory=True)
sharded_tensor_metadata = sharded_tensor.metadata()
self.assertEqual(True, sharded_tensor_metadata.pin_memory)
def test_init_sharded_tensor_with_kaiming_uniform(self):
""" Test torch.nn.init.kaiming_uniform_(ShardedTensor, a, mode, nonlinearit) """
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
h, w = 8, 2
expected_h = 2
expected_device = torch.device(f"cuda:{self.rank}")
a, mode, nonlinearity = 0, 'fan_in', 'leaky_relu'
seed = 1234
dtype = torch.double
sharded_tensor = _sharded_tensor.empty(spec, h, w, dtype=dtype)
self.assertEqual(1, len(sharded_tensor.local_shards()))
# Clone local tensor to ensure torch.nn.init starts from the same input
local_tensor_clone = torch.clone(sharded_tensor.local_shards()[0].tensor)
torch.manual_seed(seed)
torch.nn.init.kaiming_normal_(sharded_tensor, a=a, mode=mode, nonlinearity=nonlinearity)
torch.manual_seed(seed)
torch.nn.init.kaiming_normal_(local_tensor_clone, a=a, mode=mode, nonlinearity=nonlinearity)
self.assertEqual(local_tensor_clone, sharded_tensor.local_shards()[0].tensor)
def test_sharded_tensor_state_dict(single_process_pg):
from torch.distributed._sharded_tensor import empty as sharded_tensor_empty
from torch.distributed._sharding_spec import ChunkShardingSpec
class BoringModelWithShardedTensor(BoringModel):
def __init__(self, spec):
super().__init__()
self.sharded_tensor = sharded_tensor_empty(spec, 10, 20)
self.sharded_tensor.local_shards()[0].tensor.fill_(0)
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cpu",
],
)
m_0 = BoringModelWithShardedTensor(spec)
m_0.sharded_tensor.local_shards()[0].tensor.fill_(1)
name_st = ".sharded_tensor" if _TORCH_GREATER_EQUAL_1_11 else "sharded_tensor"
assert name_st in m_0.state_dict(
), 'Expect "sharded_tensor" to appear in the state dict'
m_1 = BoringModelWithShardedTensor(spec)
assert not torch.allclose(
m_1.sharded_tensor.local_shards()[0].tensor,
m_0.sharded_tensor.local_shards()[0].tensor
), "Expect the shards to be different before `m_1` loading `m_0`'s state dict"
m_1.load_state_dict(m_0.state_dict(), strict=False)
assert torch.allclose(
m_1.sharded_tensor.local_shards()[0].tensor,
m_0.sharded_tensor.local_shards()[0].tensor
), "Expect the shards to be same after `m_1` loading `m_0`'s state dict"
def test_multiple_local_shards(self):
self.init_pg()
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 16, 20)
# Validate local shards.
local_shards = sharded_tensor.local_shards()
self.assertEqual(2, len(local_shards))
for local_shard in local_shards:
self.assertEqual(torch.device(f"cuda:{self.rank}"),
local_shard.tensor.device)
self.assertEqual((2, 20), local_shard.tensor.size())
# Validate global metadata.
sharding_metadata = sharded_tensor.sharding_metadata()
self.assertEqual(8, len(sharding_metadata))
for shard_idx, shard_metadata in enumerate(sharding_metadata):
self.assertEqual([shard_idx * 2, 0], shard_metadata.shard_offsets)
self.assertEqual([2, 20], shard_metadata.shard_lengths)
self.assertEqual(f'rank:{shard_idx % 4}/cuda:{shard_idx % 4}',
shard_metadata.placement)
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 test_sharding_columns(self):
self.init_pg()
for dim in [1, -1]:
spec = ChunkShardingSpec(
dim=dim,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 10, 32)
# Validate local shard.
local_shards = sharded_tensor.local_shards()
self.assertEqual(1, len(local_shards))
local_shard = local_shards[0].tensor
self.assertEqual(torch.device(f"cuda:{self.rank}"),
local_shard.device)
self.assertEqual((10, 8), local_shard.size())
# Validate global metadata.
sharding_metadata = sharded_tensor.sharding_metadata()
self.assertEqual(4, len(sharding_metadata))
for rank, shard_metadata in enumerate(sharding_metadata):
self.assertEqual([0, rank * 8], shard_metadata.shard_offsets)
self.assertEqual([10, 8], shard_metadata.shard_lengths)
self.assertEqual(f'rank:{rank}/cuda:{rank}',
shard_metadata.placement)
def test_insufficient_sharding_dims(self):
self.init_pg()
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 2, 20)
# Validate local shard.
local_shards = sharded_tensor.local_shards()
if self.rank <= 1:
self.assertEqual(1, len(local_shards))
local_shard = local_shards[0].tensor
self.assertEqual(torch.device(f"cuda:{self.rank}"),
local_shard.device)
self.assertEqual((1, 20), local_shard.size())
else:
self.assertEqual(0, len(local_shards))
# Validate global metadata.
sharding_metadata = sharded_tensor.sharding_metadata()
self.assertEqual(2, len(sharding_metadata))
for shard_rank, shard_metadata in enumerate(sharding_metadata):
self.assertEqual([shard_rank, 0], shard_metadata.shard_offsets)
self.assertEqual([1, 20], shard_metadata.shard_lengths)
self.assertEqual(f'rank:{shard_rank}/cuda:{shard_rank}',
shard_metadata.placement)
def test_sharded_tensor_sizes(self):
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
# Test with *args
sharded_tensor = _sharded_tensor.empty(spec, 10, 20)
self.assertEqual(torch.Size([10, 20]), sharded_tensor.size())
# Test with single *args
sharded_tensor = _sharded_tensor.empty(spec, 10)
self.assertEqual(torch.Size([10]), sharded_tensor.size())
# Test with list
sharded_tensor = _sharded_tensor.empty(spec, [10, 20])
self.assertEqual(torch.Size([10, 20]), sharded_tensor.size())
# Test with tuple
sharded_tensor = _sharded_tensor.empty(spec, (10, 20))
self.assertEqual(torch.Size([10, 20]), sharded_tensor.size())
with self.assertRaises(TypeError):
sharded_tensor = _sharded_tensor.empty(spec, 'foo')
def test_partial_world_size(self):
self.init_pg()
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 10, 20)
# Validate local shard.
local_shards = sharded_tensor.local_shards()
if self.rank >= 2:
self.assertEqual(1, len(local_shards))
local_shard = local_shards[0].tensor
self.assertEqual(torch.device(f"cuda:{self.rank}"),
local_shard.device)
self.assertEqual((5, 20), local_shard.size())
else:
self.assertEqual(0, len(local_shards))
# Validate global metadata.
sharding_metadata = sharded_tensor.sharding_metadata()
self.assertEqual(2, len(sharding_metadata))
for shard_rank, shard_metadata in enumerate(sharding_metadata):
self.assertEqual([shard_rank * 5, 0], shard_metadata.shard_offsets)
self.assertEqual([5, 20], shard_metadata.shard_lengths)
self.assertEqual(f'rank:{shard_rank + 2}/cuda:{shard_rank + 2}',
shard_metadata.placement)
def test_init_sharded_tensor_with_normal(self):
""" Test torch.nn.init.normal_(ShardedTensor, mean, std) """
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
h, w = 8, 2
expected_h = 2
expected_device = torch.device(f"cuda:{self.rank}")
mean, std = 10, 5
seed = 1234
dtype = torch.double
sharded_tensor = _sharded_tensor.empty(spec, h, w, dtype=dtype)
self.assertEqual(1, len(sharded_tensor.local_shards()))
# Clone local tensor to ensure torch.nn.init starts from the same input
local_tensor_clone = torch.clone(sharded_tensor.local_shards()[0].tensor)
torch.manual_seed(seed)
torch.nn.init.normal_(sharded_tensor, mean=mean, std=std)
torch.manual_seed(seed)
torch.nn.init.normal_(local_tensor_clone, mean=mean, std=std)
self.assertEqual(local_tensor_clone, sharded_tensor.local_shards()[0].tensor)
def test_new_group(self):
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:1/cuda:2",
"rank:2/cuda:3",
],
)
pg = dist.new_group(ranks=[1, 2, 3])
if self.rank >= 1:
sharded_tensor = _sharded_tensor.empty(spec,
10,
20,
process_group=pg)
# Validate local shard.
local_shards = sharded_tensor.local_shards()
if self.rank >= 2:
self.assertEqual(1, len(local_shards))
local_shard = local_shards[0].tensor
self.assertEqual(torch.device(f"cuda:{self.rank}"),
local_shard.device)
self.assertEqual((5, 20), local_shard.size())
else:
self.assertEqual(0, len(local_shards))
# Validate global metadata.
sharded_tensor_metadata = sharded_tensor.metadata()
shards_metadata = sharded_tensor_metadata.shards_metadata
self.assertEqual(2, len(shards_metadata))
for shard_rank, shard_metadata in enumerate(shards_metadata):
self.assertEqual([shard_rank * 5, 0],
shard_metadata.shard_offsets)
self.assertEqual([5, 20], shard_metadata.shard_lengths)
self.assertEqual(
f'rank:{shard_rank + 1}/cuda:{shard_rank + 2}',
shard_metadata.placement)
# Validate remote shards.
remote_shards = sharded_tensor.remote_shards
if self.rank >= 2:
self.assertEqual(1, len(remote_shards))
else:
self.assertEqual(2, len(remote_shards))
for rpc_rank, shards in remote_shards.items():
self.assertEqual(1, len(shards))
for remote_shard in shards:
shard = remote_shard.to_here()
self.assertEqual(rpc_rank, remote_shard.owner().id)
self.assertEqual(f'rank:{rpc_rank - 1}/cuda:{rpc_rank}',
shard.metadata.placement)
self.assertEqual((5, 20), shard.tensor.size())
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 test_complete_world_size(self):
for dim in [0, -2]:
spec = ChunkShardingSpec(
dim=dim,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 10, 20)
# Validate local shard.
local_shards = sharded_tensor.local_shards()
self.assertEqual(1, len(local_shards))
local_shard = local_shards[0].tensor
self.assertEqual(torch.device(f"cuda:{self.rank}"),
local_shard.device)
if self.rank == 3:
self.assertEqual((1, 20), local_shard.size())
else:
self.assertEqual((3, 20), local_shard.size())
# Validate global metadata.
sharded_tensor_metadata = sharded_tensor.metadata()
shards_metadata = sharded_tensor_metadata.shards_metadata
self.assertEqual(4, len(shards_metadata))
for rank, shard_metadata in enumerate(shards_metadata):
self.assertEqual([rank * 3, 0], shard_metadata.shard_offsets)
if rank == 3:
self.assertEqual([1, 20], shard_metadata.shard_lengths)
else:
self.assertEqual([3, 20], shard_metadata.shard_lengths)
self.assertEqual(f'rank:{rank}/cuda:{rank}',
shard_metadata.placement)
# Validate remote shards.
remote_shards = sharded_tensor.remote_shards
self.assertEqual(3, len(remote_shards))
for rpc_rank, shards in remote_shards.items():
self.assertEqual(1, len(shards))
for remote_shard in shards:
self.assertEqual(rpc_rank, remote_shard.owner().id)
shard = remote_shard.to_here()
self.assertEqual(f'rank:{rpc_rank}/cuda:{rpc_rank}',
shard.metadata.placement)
if rpc_rank == 3:
self.assertEqual((1, 20), shard.tensor.size())
else:
self.assertEqual((3, 20), shard.tensor.size())
def shard_parameter(self):
rowwise_sharding_spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
colwise_sharding_spec = ChunkShardingSpec(
dim=1,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor.shard_parameter(self.linear1, "weight", rowwise_sharding_spec)
sharded_tensor.shard_parameter(self.linear2, "weight", colwise_sharding_spec)
def test_invalid_pg_rpc_ranks(self):
self.init_pg()
# Init RPC with different ranks.
rpc_backend_options = rpc.TensorPipeRpcBackendOptions()
rpc_backend_options.init_method = f"file://{self.file_name}"
rank = (self.rank + 1) % self.world_size
rpc.init_rpc(
name=f'worker{rank}',
rank=rank,
world_size=self.world_size,
rpc_backend_options=rpc_backend_options,
)
spec = ChunkShardingSpec(dim=0, placements=["rank:1/cuda:1"])
with self.assertRaisesRegex(ValueError, 'Default ProcessGroup and RPC ranks must be the same'):
_sharded_tensor.empty(spec, 10, 20)
def test_multiple_local_shards(self):
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
sharded_tensor = _sharded_tensor.empty(spec, 16, 20)
# Validate local shards.
local_shards = sharded_tensor.local_shards()
self.assertEqual(2, len(local_shards))
for local_shard in local_shards:
self.assertEqual(torch.device(f"cuda:{self.rank}"),
local_shard.tensor.device)
self.assertEqual((2, 20), local_shard.tensor.size())
# Validate global metadata.
sharded_tensor_metadata = sharded_tensor.metadata()
shards_metadata = sharded_tensor_metadata.shards_metadata
self.assertEqual(8, len(shards_metadata))
for shard_idx, shard_metadata in enumerate(shards_metadata):
self.assertEqual([shard_idx * 2, 0], shard_metadata.shard_offsets)
self.assertEqual([2, 20], shard_metadata.shard_lengths)
self.assertEqual(f'rank:{shard_idx % 4}/cuda:{shard_idx % 4}',
shard_metadata.placement)
# Validate remote shards.
remote_shards = sharded_tensor.remote_shards
self.assertEqual(3, len(remote_shards))
owners = {}
for rpc_rank, shards in remote_shards.items():
self.assertEqual(2, len(shards))
for remote_shard in shards:
shard = remote_shard.to_here()
self.assertEqual((2, 20), shard.tensor.size())
self.assertEqual(rpc_rank, remote_shard.owner().id)
def test_chunked_sharding_spec(self):
# Test valid specs.
ChunkShardingSpec(0, [0, 1])
# Named dimension.
ChunkShardingSpec("N", ["cuda:0", "cuda:1"])
ChunkShardingSpec(0, [torch.device("cuda:0"), torch.device("cuda:1")])
ChunkShardingSpec(-1, ["cuda:0", "cuda:1"])
ChunkShardingSpec(0, ["rank:0/cuda:0", "rank:0/cuda:1"])
ChunkShardingSpec(0, ["rank:0", "rank:1"])
ChunkShardingSpec(0, ["rank:0/cpu", "rank:1/cpu"])
# Test invalid specs
with self.assertRaisesRegex(ValueError, "int or str"):
ChunkShardingSpec(None, ["cuda:0", "cuda:1"])
with self.assertRaisesRegex(ValueError, "int or str"):
ChunkShardingSpec({}, ["cuda:0", "cuda:1"])
with self.assertRaisesRegex(ValueError, "not a valid device"):
ChunkShardingSpec(0, ["random:0", "cuda:1"])
with self.assertRaisesRegex(ValueError, "not a valid device"):
ChunkShardingSpec(0, ["cuda:foo", "cuda:1"])
with self.assertRaisesRegex(ValueError, "not a valid device"):
ChunkShardingSpec(0, ["rank:foo", "cuda:1"])
with self.assertRaisesRegex(ValueError, "not a valid device"):
ChunkShardingSpec(0, ["rank:0/foo", "cuda:1"])
with self.assertRaisesRegex(ValueError, "not a valid device"):
ChunkShardingSpec(0, ["rank:0/random:0", "cuda:1"])
with self.assertRaisesRegex(ValueError, "not a valid device"):
ChunkShardingSpec(0, ["rank:0/cuda:foo", "cuda:1"])
def test_torch_equal(self):
""" Test torch.equal(ShardedTensor, ShardedTensor) """
spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:3/cuda:3",
],
)
alt_spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:1/cuda:1",
"rank:0/cuda:0",
"rank:3/cuda:3",
"rank:2/cuda:2",
],
)
st1, st2 = self.get_random_tensors(spec, spec, 10, 10)
self.assertTrue(torch.equal(st1, st2))
st1, st2 = self.get_random_tensors(spec, spec, 10, 10)
if self.rank == 0:
torch.nn.init.uniform_(st1.local_shards()[0].tensor)
self.assertFalse(torch.equal(st1, st2))
st1 = _sharded_tensor.ones(spec, 10, 10)
st2 = _sharded_tensor.ones(spec, 10, 5)
self.assertFalse(torch.equal(st1, st2))
st1, st2 = self.get_random_tensors(spec, alt_spec, 10, 10)
self.assertFalse(torch.equal(st1, st2))
st1 = _sharded_tensor.ones(spec, 10, 10)
st2 = _sharded_tensor.zeros(spec, 10, 10)
self.assertFalse(torch.equal(st1, st2))
st1 = _sharded_tensor.ones(spec, 10, 10)
st2 = _sharded_tensor.ones(spec, 10, 10, dtype=torch.double)
self.assertFalse(torch.equal(st1, st2))
st1 = _sharded_tensor.ones(spec, 10, 10)
st2 = _sharded_tensor.ones(spec, 10, 10, requires_grad=True)
self.assertFalse(torch.equal(st1, st2))
cpu_spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cpu",
"rank:1/cpu",
"rank:2/cpu",
"rank:3/cpu",
],
)
st1 = _sharded_tensor.ones(cpu_spec, 10, 10)
st2 = _sharded_tensor.ones(cpu_spec, 10, 10, pin_memory=True)
self.assertFalse(torch.equal(st1, st2))
pg = dist.new_group([1, 0, 3, 2])
st1, st2 = self.get_random_tensors(spec, spec, 10, 10, pg2=pg)
self.assertFalse(torch.equal(st1, st2))
pg = dist.new_group([0, 1, 2, 3])
st1, st2 = self.get_random_tensors(spec, spec, 10, 10, pg2=pg)
self.assertFalse(torch.equal(st1, st2))
def test_sharded_linear_errors(self):
for spec in generate_chunk_sharding_specs_for_test(0):
fc1 = torch.nn.Linear(10, 10).cuda(self.rank)
shard_parameter(fc1, "bias", spec)
with self.assertRaisesRegex(
TypeError, 'input and bias need to be torch.Tensor'):
fc1(torch.rand(10, 10).cuda(self.rank))
fc2 = torch.nn.Linear(10, 10).cuda(self.rank)
shard_parameter(fc2, "weight", spec)
with self.assertRaisesRegex(ValueError,
'Input needs to have at least 1 dim'):
fc2(torch.tensor(1).cuda(self.rank))
fc3 = torch.nn.Linear(10, 10).cuda(self.rank)
fc3.weight = torch.nn.Parameter(
torch.rand(10, 10, 10).cuda(self.rank))
shard_parameter(fc3, "weight", spec)
with self.assertRaisesRegex(ValueError,
'Weight needs to have exactly 2 dims'):
fc3(torch.rand(10, 10).cuda(self.rank))
fc4 = torch.nn.Linear(10, 10).cuda(self.rank)
fc4.bias = torch.nn.Parameter(torch.rand(10, 10).cuda(self.rank))
shard_parameter(fc4, "weight", spec)
with self.assertRaisesRegex(ValueError,
'Bias needs to have exactly 1 dim'):
fc4(torch.rand(10, 10).cuda(self.rank))
fc5 = torch.nn.Linear(7, 10).cuda(self.rank)
shard_parameter(fc5, "weight", spec)
with self.assertRaisesRegex(
ValueError,
'Input dim: 13 does not match appropriate weight dim: 7'):
fc5(torch.rand(20, 10, 13).cuda(self.rank))
fc6 = torch.nn.Linear(10, 10).cuda(self.rank)
del fc6.weight
enumerable_spec = EnumerableShardingSpec([
ShardMetadata(
shard_offsets=[0, 0],
shard_sizes=[5, 5],
placement="rank:0/cuda:0",
),
ShardMetadata(
shard_offsets=[0, 5],
shard_sizes=[5, 5],
placement="rank:1/cuda:1",
),
ShardMetadata(
shard_offsets=[5, 0],
shard_sizes=[5, 5],
placement="rank:2/cuda:2",
),
ShardMetadata(
shard_offsets=[5, 5],
shard_sizes=[5, 5],
placement="rank:3/cuda:3",
)
])
fc6.weight = empty(enumerable_spec, 10, 10)
with self.assertRaisesRegex(
ValueError,
'Only ChunkShardingSpec supported for ShardedTensor ops!'):
fc6(torch.rand(10, 10).cuda(self.rank))
fc7 = torch.nn.Linear(10, 80).cuda(self.rank)
multiple_local_shard_spec = ChunkShardingSpec(
dim=0,
placements=[
"rank:0/cuda:0",
"rank:0/cuda:0",
"rank:1/cuda:1",
"rank:1/cuda:1",
"rank:2/cuda:2",
"rank:2/cuda:2",
"rank:3/cuda:3",
"rank:3/cuda:3",
],
)
del fc7.weight
fc7.weight = empty(multiple_local_shard_spec, 80, 10)
with self.assertRaisesRegex(ValueError,
'Only one local shard supported!'):
fc7(torch.rand(10, 10).cuda(self.rank))
def test_chunked_sharding_spec(self):
# Test valid specs.
ChunkShardingSpec(0, [torch.device(0), torch.device(1)])
# Named dimension.
ChunkShardingSpec("N", ["cuda:0", "cuda:1"])
ChunkShardingSpec(0, [torch.device("cuda:0"), torch.device("cuda:1")])
ChunkShardingSpec(-1, ["cuda:0", "cuda:1"])
ChunkShardingSpec(0, ["rank:0/cuda:0", "rank:0/cuda:1"])
ChunkShardingSpec(0, ["rank:0", "rank:1"])
ChunkShardingSpec(0, ["rank:0/cpu", "rank:1/cpu"])
# Test invalid specs
with self.assertRaisesRegex(ValueError, "int or str"):
ChunkShardingSpec(None, ["cuda:0", "cuda:1"])
with self.assertRaisesRegex(ValueError, "int or str"):
ChunkShardingSpec({}, ["cuda:0", "cuda:1"])
with self.assertRaisesRegex(ValueError,
"Could not parse remote_device"):
ChunkShardingSpec(0, ["random:0", "cuda:1"])
with self.assertRaisesRegex(ValueError,
"Could not parse remote_device"):
ChunkShardingSpec(0, ["cuda:foo", "cuda:1"])
with self.assertRaisesRegex(ValueError,
"Could not parse remote_device"):
ChunkShardingSpec(0, ["rank:foo", "cuda:1"])
with self.assertRaisesRegex(RuntimeError, "Expected one of"):
ChunkShardingSpec(0, ["rank:0/foo", "cuda:1"])
with self.assertRaisesRegex(RuntimeError, "Expected one of"):
ChunkShardingSpec(0, ["rank:0/random:0", "cuda:1"])
with self.assertRaisesRegex(RuntimeError, "Invalid device string"):
ChunkShardingSpec(0, ["rank:0/cuda:foo", "cuda:1"])