What does torch_xla.runtime.use_spmd() actually do?

It sets the XLA_USE_SPMD environment variable to "1", and everything else follows from that. The C++ side reads that variable in ShouldUseVirtualDevice, so from then on the default device for the process is a single virtual device whose name is the literal string SPMD:0 rather than one of the real chips. If tensors were already created on real devices before you called it, use_spmd() notices (the SPMD config gets locked the first time any device is initialised), warns you, and forces the switch by pulling each of those tensors back to the host and re-uploading it to the virtual device. That is a real one-time cost, which is why the recommendation is to call use_spmd() at the very top of the program. Passing auto=True additionally turns on auto-sharding, which lets the compiler choose the shardings instead of you.

What is a Mesh and a partition spec in PyTorch/XLA?

A Mesh is a flat list of device ids plus a shape to fold them into, and optionally a name per axis. Mesh(np.arange(8), (4, 2), ("data", "model")) says: take the eight devices, lay them out as four rows by two columns, call the row axis "data" and the column axis "model". The list has to cover every device the runtime reports, and the ids have to be unique. A partition spec then has exactly one entry per dimension of the tensor you are annotating, and each entry says which mesh axis that tensor dimension is split along: an axis index, an axis name, a tuple of axes if one tensor dimension is split across several, or None to replicate that dimension everywhere. So (0, None) on an 8x32 tensor cuts the rows four ways along the "data" axis and gives every device a full copy of all 32 columns.

What does mark_sharding do to a tensor?

It attaches an XLA sharding annotation, and if the tensor already has bytes on the device it also re-lays those bytes out. The Python side translates your partition spec into an OpSharding proto, caches it (the translation is expensive as the mesh grows, so it is memoised per spec), and calls into C++. There the path forks. If the tensor is the result of pending operations rather than plain data, the annotation becomes a custom-sharding node in the graph and nothing moves. If the tensor is real data, torch_xla takes the host copy, slices it into one shard per device by the tile assignment, zero-pads any shard that came up short, uploads the shards, and replaces its single-device buffer with a sharded handle. Annotating the same tensor twice with different shardings is refused: you get "Existing annotation must be cleared first".

Why does my whole program compile as sharded once SPMD is on?

Because the executor does not ask whether this particular graph has sharded tensors in it. When it compiles, it sets is_sharded if the graph is on the virtual device or if SPMD is on at all, and that flag flows straight into the runtime client, which turns on SPMD partitioning, sets the number of partitions to the number of devices and the number of replicas to one, and hands the compiler a device assignment covering every chip. A graph with no annotation at all still goes through the partitioner; it simply comes out replicated. The practical consequence is that mixing SPMD and non-SPMD execution inside one process is not a thing you can do.

How do I read a sharded tensor back?

Two ways, and they answer different questions. Printing it or calling .cpu() gives you the whole logical tensor, gathered from every device, exactly as if it had never been split. If you want to see the split itself, mark_sharding hands back an XLAShardedTensor, and its local_shards property gives you one XLAShard per addressable device, each carrying the shard data, the device it came from, the index range it covers in the global tensor, and a replica id. Those shards are snapshots on the host, detached from the tensor, and they include any zero padding the split added, so a shard can be slightly larger than the slice it represents.

SPMD in torch_xla: A Mesh, an Annotation, and One Virtual Device

Your model outgrew one chip. The usual PyTorch answer is to rewrite it into a parallel model; the SPMD answer is to keep writing single-device code and tell the compiler how a few tensors are cut across the chips. This walks the whole path in PyTorch/XLA at a pinned commit: what use_spmd() switches on and why every tensor then lands on a device called SPMD:0, what a Mesh and a partition spec really are, the permutation that turns a spec into a tile assignment, what mark_sharding does to the bytes you already uploaded, how the annotation reaches the HLO, and why one step becomes ExecuteReplicated behind a single device lock.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of SPMD in torch_xla: A Mesh, an Annotation, and One Virtual Device — built to be understood, not skimmed.

Questions

What does torch_xla.runtime.use_spmd() actually do?: It sets the XLA_USE_SPMD environment variable to "1", and everything else follows from that. The C++ side reads that variable in ShouldUseVirtualDevice, so from then on the default device for the process is a single virtual device whose name is the literal string SPMD:0 rather than one of the real chips. If tensors were already created on real devices before you called it, use_spmd() notices (the SPMD config gets locked the first time any device is initialised), warns you, and forces the switch by pulling each of those tensors back to the host and re-uploading it to the virtual device. That is a real one-time cost, which is why the recommendation is to call use_spmd() at the very top of the program. Passing auto=True additionally turns on auto-sharding, which lets the compiler choose the shardings instead of you.
What is a Mesh and a partition spec in PyTorch/XLA?: A Mesh is a flat list of device ids plus a shape to fold them into, and optionally a name per axis. Mesh(np.arange(8), (4, 2), ("data", "model")) says: take the eight devices, lay them out as four rows by two columns, call the row axis "data" and the column axis "model". The list has to cover every device the runtime reports, and the ids have to be unique. A partition spec then has exactly one entry per dimension of the tensor you are annotating, and each entry says which mesh axis that tensor dimension is split along: an axis index, an axis name, a tuple of axes if one tensor dimension is split across several, or None to replicate that dimension everywhere. So (0, None) on an 8x32 tensor cuts the rows four ways along the "data" axis and gives every device a full copy of all 32 columns.
What does mark_sharding do to a tensor?: It attaches an XLA sharding annotation, and if the tensor already has bytes on the device it also re-lays those bytes out. The Python side translates your partition spec into an OpSharding proto, caches it (the translation is expensive as the mesh grows, so it is memoised per spec), and calls into C++. There the path forks. If the tensor is the result of pending operations rather than plain data, the annotation becomes a custom-sharding node in the graph and nothing moves. If the tensor is real data, torch_xla takes the host copy, slices it into one shard per device by the tile assignment, zero-pads any shard that came up short, uploads the shards, and replaces its single-device buffer with a sharded handle. Annotating the same tensor twice with different shardings is refused: you get "Existing annotation must be cleared first".
Why does my whole program compile as sharded once SPMD is on?: Because the executor does not ask whether this particular graph has sharded tensors in it. When it compiles, it sets is_sharded if the graph is on the virtual device or if SPMD is on at all, and that flag flows straight into the runtime client, which turns on SPMD partitioning, sets the number of partitions to the number of devices and the number of replicas to one, and hands the compiler a device assignment covering every chip. A graph with no annotation at all still goes through the partitioner; it simply comes out replicated. The practical consequence is that mixing SPMD and non-SPMD execution inside one process is not a thing you can do.
How do I read a sharded tensor back?: Two ways, and they answer different questions. Printing it or calling .cpu() gives you the whole logical tensor, gathered from every device, exactly as if it had never been split. If you want to see the split itself, mark_sharding hands back an XLAShardedTensor, and its local_shards property gives you one XLAShard per addressable device, each carrying the shard data, the device it came from, the index range it covers in the global tensor, and a replica id. Those shards are snapshots on the host, detached from the tensor, and they include any zero padding the split added, so a shard can be slightly larger than the slice it represents.

SPMD in torch_xla: A Mesh, an Annotation, and One Virtual Device

Questions

Related explainers