Marking a sharding

Turn SPMD on with xr.use_spmd() and the first thing that changes is a number you can check. The frontend stops reporting eight devices and reports one: _xla_num_devices returns 1 whenever the virtual device is in use, and tensors land on a virtual SPMD:0 instead of on TPU:3. From that point the program is written as though the whole slice were a single large chip, which is the arrangement the mode exists to give you.

The mesh is the object that still remembers there were eight. Building one asserts that the flat list of device ids matches xr.global_runtime_device_count() exactly, with a comment saying XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → requires it, so a mesh is never a subset of the machine and never a superset of it. Reshape those ids however you like; a 512-core slice can be a 16 by 16 by 2 mesh or a flat 512, and the shape you pick is a logical arrangement rather than a claim about the wires.

One latch is worth knowing before the first line runs. UseVirtualDevice sets spmd_config_is_locked true the first time anything queries it, and use_spmd() checks that lock: if tensors were already built on non-virtual devices it warns, calls _xla_force_spmd_device() and waits on device ops, replicating what already existed. The warning names the fix itself, which is to call use_spmd() at the top of the program rather than after the model is on the chip.

    # At the moment, XLA requires that the Mesh uses the global number of
    # devices.
    num_devices = xr.global_runtime_device_count()
    assert num_devices > 0, "This requires XLA supported device(s)."
    assert num_devices == len(
        device_ids
    ), f"Number of device IDs ({len(device_ids)}) must match the global number of devices ({num_devices})"

A partition spec has one entry per tensor dimension, and mark_sharding refuses anything shorter. The assert is deliberate rather than defensive: the comment above it says unspecified dimensions are not filled in with replication, because partial replication builds its group assignment from the spec and the groups would vary with rank if the fill were implicit. So a rank-3 tensor gets a three-entry spec, None included, every time.

What that spec does is reorder the mesh. The tiled entries become a permutation, the untouched mesh axes are appended after them, and the logical mesh is transposed by that permutation. Then any entry that was a tuple collapses its adjacent axes into one, which is how a single tensor dimension gets sharded over two mesh axes at once. The device ordering that falls out is the tile assignment, and it is the whole content of the annotation.

The kind of sharding falls out of the same spec by three cheap tests. Every entry None is replication. Any entry None, with others tiled, is partial replication. Neither of those, and it is tiled. A single device short-circuits all three and gives maximal. The jax path's own lesson at /jax/sharding/mesh-and-spec teaches the mesh and the spec as concepts; the arithmetic below is what the torch side does with them.

  # Flatten the partition spec and ensure that it is fully specified over the
  # mesh for permutation.
  tiled_dims = [x for x in partition_spec if x is not None]
  permutation = np.hstack(tiled_dims).tolist() if tiled_dims else []
  missing_axes = sorted(set(range(len(mesh.shape()))) - set(permutation))
  tile_assignment = mesh.get_logical_mesh().transpose(permutation +
                                                      missing_axes)

  # For any tuples in the partition_spec, the grouped axes will be adjacent
  # after the permutation. Combine these dimensions into a single axis.
  for i, spec in enumerate(tiled_dims):
    if isinstance(spec, tuple):
      shape = tile_assignment.shape
      tile_assignment = tile_assignment.reshape(shape[:i] + (-1,) +
                                                shape[i + len(spec):])

  return tile_assignment

The bridge arc's second lesson lists mark_sharding among the lines that force materialization, with the qualifier that it only does so sometimes. The qualifier is a branch, four lines into XlaMarkSharding, and reading it tells you which of your own annotations are free.

If the tensor's current IR value is anything other than a DeviceData node, the call hands off to XlaAnnotateCustomSharding and returns. That path splices a CustomSharding node into the graph, which lowers to a custom call, and nothing moves. Annotating an activation halfway through a forward pass costs one IR node.

If the IR is a DeviceData node, meaning a parameter or an input, the sharding is physical and the call has to deal with bytes. When the host copy is still around, which is the ordinary case under the virtual device because the initial upload was deferred, it reuses that copy and bumps a VirtualDeviceUsage counter. When it is not, the tensor is pulled back from the device through GetTensors first. Either way the data is re-created against the virtual device as shards.

Two guards sit on that same path. A second annotation with the same spec returns early and does nothing. A second annotation with a different spec is refused outright unless the existing one was replicated or unknown, with the message asking you to clear the old annotation first. Both branches raise their own counter, so a metrics report says which one your program took without you having to reason it out.

  // For Non DeviceData IR values, we directly attach the sharding spec to the
  // xtensor.
  const DeviceData* device_data_node = nullptr;
  if (xtensor->CurrentIrValue()) {
    device_data_node = DeviceData::Cast(xtensor->CurrentIrValue().node.get());
    if (!device_data_node) {
      XlaAnnotateCustomSharding(xtensor, sharding);
      return;
    }
  }

  if (xtensor->CurrentTensorData().has_value()) {
    TORCH_LAZY_COUNTER("VirtualDeviceUsage", 1);
    // When virtual device is enabled for SPMD, we defer the initial
    // data transfer to the device and retain the original data on the
    // host, until the sharded data transfer.
    cpu_tensor = xtensor->CurrentTensorData().value();
  } else {

Up to compile time the sharding is torch-side bookkeeping, mirrored in two places: on the tensor, where a comment calls it the source of truth for every lookup, and on the node, as its output shardings. Neither of those is something XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → can read. One function does the translation, and it runs once per compile, from inside XLAGraphExecutor::Compile.

SetHloSharding walks every emitted output of the lowering context, casts the node back to an XlaNode, reads the sharding for that output index, and writes it straight into the HLO instruction proto. Anything typed unknown is skipped, which is how implicit replication stays implicit rather than becoming an assertion the partitioner has to honour.

Input parameters take a different road entirely, because a parameter is not an emitted output. LoweringContext::GetParameter checks whether the backing data carries a sharding and, when it does, wraps the xla::Parameter call in an xla::XlaScopedShardingAssignment. Same annotation, written at a different moment, which is the underlying reason marking a parameter had to be physical two sections ago.

What the compiler then does with those annotations is not this arc's story. The propagation pass and the collectives it inserts are taught on the xla path at /xla/spmd/the-partitioner, and /xla/collectives/six-shardings-one-matmul captures the collectives for one matmul under six different shardings on this repo's own hardware. Chapter 11's mastery item points you at that comparison for a reason: it is the answer key for the sharding you are about to write.

bool ShardingUtil::SetHloSharding(LoweringContext* lowering_ctx) {
  bool is_sharded = false;
  for (std::pair<torch::lazy::Output, xla::XlaOp> elem :
       lowering_ctx->GetEmittedOutputs()) {
    const torch::lazy::Node* node = elem.first.node;
    const XlaNode* xla_node = dynamic_cast<const XlaNode*>(node);
    xla::HloInstructionProto* instruction =
        XlaBuilderFriend::GetInstruction(elem.second);
    const std::shared_ptr<xla::OpSharding> sharding =
        xla_node->GetSharding(elem.first.index);
    if (sharding != nullptr && sharding->type() != xla::OpSharding::UNKNOWN) {
      *instruction->mutable_sharding() = *sharding;
      is_sharded = true;
    }
  }
  return is_sharded;
}

mark_sharding edits the tensor in place and hands back an XLAShardedTensor view of it. On a parameter or an input that is the whole job. On an activation in the middle of a model it is half of one, because the backward pass will compute a gradient for that activation and nothing has said where the gradient lives.

The docstring on MarkShardingFunction gives the reason in one sentence, quoted here in full: "This is required to guide GSPMD sharding propagation better during the backward pass as during complicated workloads the compiler can introduce extra collectives that can hurt performance." A collective you did not ask for is the failure mode, and it appears in the profile as device time rather than as an error.

The mechanism is an autograd Function whose forward and backward call the same custom op. That op is registered through torch.library, marks a clone rather than the input, and has a fake implementation returning an empty tensor of the same shape, which is what lets it survive AOTAutograd and appear in a dynamo-captured graph. The bridge arc's third lesson explains why forward and backward arrive at the bridge separately in the first place.

Annotating the activation and annotating its gradient are two different instructions to the same partitioner.

  @staticmethod
  def forward(ctx, torch_tensor: torch.Tensor, mesh: Mesh,
              partition_spec: PartitionSpec) -> torch.Tensor:
    o = _aot_mark_sharding(torch_tensor, str(mesh), str(partition_spec))
    ctx.partition_spec = partition_spec
    ctx.mesh = mesh
    return o

  @staticmethod
  def backward(ctx, grad_output: torch.Tensor):  # type: ignore
    partition_spec = ctx.partition_spec
    mesh = ctx.mesh
    o = _aot_mark_sharding(grad_output, str(mesh), str(partition_spec))
    return o, None, None


@torch.library.custom_op("xla::aot_mark_sharding", mutates_args=())
def _aot_mark_sharding(t: torch.Tensor, mesh: str,
                       partition_spec: str) -> torch.Tensor:
  ...
  return xs.mark_sharding(t.clone(), the_mesh,
                          partition_spec_eval).global_tensor