Skip to content

Commit 4d9718d

Browse files
yeounohyeounoh
committed
[SPMD] Copy sharding from XLATensor to XLATensor while copying
1 parent e75516e commit 4d9718d

3 files changed

Lines changed: 31 additions & 3 deletions

File tree

‎test/test_xla_sharding.py‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ class XlaShardingTest(unittest.TestCase):
2323

2424
class SimpleLinear(nn.Module):
2525
def __init__(self):
26-
super(SimpleLinear, self).__init__()
26+
super(XlaShardingTest.SimpleLinear, self).__init__()
2727
self.fc1 = nn.Linear(128, 64)
2828
self.relu = nn.ReLU()
2929
self.fc2 = nn.Linear(64, 1)
@@ -47,6 +47,8 @@ def _get_mesh(self, mesh_shape, device_ids=None):
4747
assert len(device_ids) == self.n_devices
4848
return xs.Mesh(device_ids, mesh_shape)
4949

50+
class BasicShardingTest(XlaShardingTest):
51+
5052
def test_xla_sharded_tensor(self):
5153
partition_spec = (0, 1)
5254
xt1 = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]],
@@ -117,6 +119,15 @@ def test_deep_copy(self):
117119
torch_xla._XLAC._get_xla_sharding_spec(xt),
118120
torch_xla._XLAC._get_xla_sharding_spec(xt2))
119121

122+
def test_clone(self):
123+
xt = torch.randn(2, 4, 8, 16).to(xm.xla_device())
124+
xs.mark_sharding(xt, self._get_mesh((1, 1, 1, self.n_devices)),
125+
(0, 1, 2, 3))
126+
xt2 = xt.clone()
127+
self.assertEqual(
128+
torch_xla._XLAC._get_xla_sharding_spec(xt),
129+
torch_xla._XLAC._get_xla_sharding_spec(xt2))
130+
120131
def test_mark_step_with_sharding(self):
121132
xt = torch.ones(2, 2).to(xm.xla_device())
122133
xs.mark_sharding(xt, self._get_mesh((1, self.n_devices)), (0, 1))
@@ -133,10 +144,11 @@ def test_optimizer_step_with_sharding(self):
133144
optimizer = optim.SGD(model.parameters(), lr=0.1)
134145
data = torch.randn(128, 128).to(xm.xla_device())
135146
target = torch.zeros(128).to(xm.xla_device())
147+
loss_fn = nn.CrossEntropyLoss()
136148
for i in range(5):
137149
optimizer.zero_grad()
138150
output = model(data)
139-
loss = nn.CrossEntropy(output, target)
151+
loss = loss_fn(output, target)
140152
loss.backward()
141153
optimizer.step()
142154
xm.mark_step()

‎torch_xla/csrc/tensor.cpp‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,14 @@ XLATensorPtr XLATensor::Create(
7979
c10::optional<at::ScalarType> logical_element_type) {
8080
XLATensorPtr xtensor = c10::make_intrusive<XLATensor>(
8181
XLATensor(std::move(ir_value), device, logical_element_type));
82+
// Preserve sharding if a new tensor is created from a sharded IR node.
83+
if (ir_value) {
84+
auto* xla_node = dynamic_cast<XlaNode*>(ir_value.node.get());
85+
if (xla_node->GetSharding() != nullptr) {
86+
ShardingSpec sharding = ShardingSpec{*xla_node->GetSharding()};
87+
xtensor->SetShardingSpec(sharding);
88+
}
89+
}
8290
XLAGraphExecutor::Get()->RegisterTensor(xtensor->data());
8391
if (UseEagerDebugMode()) {
8492
std::vector<XLATensorPtr> xtensors({xtensor});
@@ -449,6 +457,7 @@ at::Tensor XLATensor::ToTensor(bool detached) {
449457
void XLATensor::ShallowCopyTo(XLATensorPtr dest) const {
450458
dest->SetScalarType(data()->logical_element_type);
451459
dest->SetIrValue(GetIrValue(), /*inplace=*/false);
460+
dest->SetShardingSpec(*sharding_spec());
452461
}
453462

454463
void XLATensor::SetScalarType(

‎torch_xla/csrc/tensor_methods.cpp‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -884,7 +884,9 @@ XLATensorPtr clamp(const XLATensorPtr& input,
884884
}
885885

886886
XLATensorPtr clone(const XLATensorPtr& input) {
887-
return input->CreateFrom(input->GetIrValue());
887+
XLATensorPtr cloned = input->CreateFrom(input->GetIrValue());
888+
cloned->SetShardingSpec(*input->sharding_spec());
889+
return cloned;
888890
}
889891

890892
XLATensorPtr constant_pad_nd(const XLATensorPtr& input,
@@ -2139,6 +2141,11 @@ void copy_(XLATensorPtr& input, XLATensorPtr& src) {
21392141
}
21402142
input->UpdateFromTensor(std::move(src_tensor), /*sync=*/false);
21412143
}
2144+
2145+
// Preserves sharding when copying.
2146+
if (src->sharding_spec() != nullptr) {
2147+
input->SetShardingSpec(*src->sharding_spec());
2148+
}
21422149
}
21432150

21442151
XLATensorPtr scatter(const XLATensorPtr& input, int64_t dim,

0 commit comments

Comments
 (0)