What do replica_groups mean on an XLA all-reduce?

They are a list of lists of integers naming who participates, but the list alone does not say which devices talk. Two optional fields on the instruction decide how the integers are read, and there are exactly four readings. With no channel id the mode is cross-replica and the integers are replica ids, with each group formed inside one partition. With a channel id and no use_global_device_ids field the mode is cross-partition and the integers are partition ids, with each group formed inside one replica. With a channel id and use_global_device_ids set to false the mode is cross-replica-and-partition, where the integers are still replica ids but the group sweeps every partition of each listed replica. And with a channel id and use_global_device_ids set to true the mode is flattened-id, where the integers are positions in one flat enumeration of the whole replica-by-partition grid. Present-and-false differs from absent because the field was added after the mode it disambiguates, so an older module with no field has to keep meaning cross-partition. An empty list means all participants, but which participants depends on the mode, and flattened-id may not be empty at all.

Why does one replica group turn into several device groups?

Because in cross-replica mode the written group is read once per partition. The function that does the translation loops over the written groups on the outside and over every partition on the inside, building a fresh participant list each time, so two replicas across four partitions with an empty group list produce four groups of two devices rather than one group of eight. That is the correct behaviour for data parallelism inside a model-parallel program: each partition averages independently with the corresponding partition on the other replicas, because partition 0 of replica 0 holds a different slice of the model than partition 1 does and averaging across them would be meaningless. Cross-partition mode does the mirror image, looping over every replica for each written group. Cross-replica-and-partition is the one that collapses instead of multiplying, since each written group becomes a single device group containing every partition of every listed replica. Flattened-id produces exactly one device group per written group, which is why the partitioner prefers it.

Is a flattened id the same as a global device id in XLA?

No, and on any mesh with more than one replica and more than one partition they disagree. A flattened id walks the grid replica-major: the formula, which appears identically in three places in the same file, is replica_id times partition_count plus partition_id. A global device id comes from the device assignment, and the default placement walks the grid the other way, writing computation times replica_count plus replica into cell (replica, computation), where computation is the partition. So on two replicas by two partitions, replica 1 partition 0 has flattened id 2 and device number 1. The consequence you actually see is that a flattened group written as {0,1} can resolve to devices {0,2}, which is why a replica group in a dump sometimes looks shuffled. When either count is one the two numberings coincide, which is why the divergence usually appears for the first time when somebody combines data and model parallelism.

What is a clique in XLA, and how does it relate to a communicator?

A communicator is one participant handle, one instance per device in this process, and the collective operations are its virtual methods: all-reduce, all-gather, reduce-scatter, broadcast, collective-permute, all-to-all, and point-to-point send and receive. A clique is all the communicators that make up one collective, stored in a map sorted by rank, with a comment saying the sorting exists to guarantee deterministic traversal order. Asking a clique for a rank returns an optional, and the empty case is not an error: it means that rank exists but is not local to this process, which is normal on multi-host jobs where a thirty-two device clique has eight communicators in each of four processes. Above both sits the factory, which creates communicators from a clique key plus an optional rendezvous id plus a list of device-and-rank pairs, and which can also split an existing set of communicators into sub-groups by colour without a fresh bootstrap. The whole surface is defined without naming a vendor, and the GPU, CPU and TPU implementations all satisfy it.

Why does an XLA distributed job hang instead of erroring?

Because collectives match by arrival order rather than by name. Nothing on the wire identifies which collective a call belongs to, so there is nothing to compare and no mismatch to detect. Each rank enters the next collective on its communicator and is paired with whatever the other ranks entered next, so two ranks issuing the same two operations in opposite orders each end up waiting for a partner that is already waiting for them. The source is blunt about this: the clique header opens by saying it is notoriously easy to get a deadlock, which is why communicators are grouped into cliques with a well defined order of operations. In practice the common causes are a rank that never arrives because a host crashed or was never started, a guest list that differs between processes because the mesh or the device assignment differed, and control flow that diverges so one device issues a collective the others never reach. The compiler does have a pass that forces a total order on collectives, but on GPU it is enabled to work around executable divergence caused by online autotuning of convolutions, not as general deadlock insurance.

What does async collective conversion do in XLA, and when does it help?

A pass rewrites each collective into a start instruction that issues the operation and a done instruction that waits for it, so the scheduler can put independent arithmetic in between and hide the network latency behind work that was going to run anyway. Whether a given instruction is converted is one expression: the config predicate for that opcode, and then either a flag that skips the size check or the output being at least a minimum size in bytes. Only all-reduce, all-gather and reduce-scatter consult a size at all; collective-permute, all-to-all, ragged-all-to-all and collective-broadcast convert whenever their predicate says yes. The threshold exists because a start and a done are two instructions where there was one, and on a small buffer the overhead exceeds the overlap. The GPU pipeline converts everything unconditionally with the thresholds at zero, then annotates only the ones that should really run asynchronously, filtered by a debug option. After scheduling, one more pass converts any pair with nothing overlapping between them back to the plain synchronous instruction, so the speculation is refunded where it did not pay.

XLA Collectives: One Guest List, Then One Order

A collective is the one instruction in a compiled program that cannot be executed by looking at its own operands, because it needs values from machines that do not share memory with this one. This reads the whole subsystem at a pinned commit of openxla/xla, from both sides of the line. On the compiler side: the list of integers an instruction carries, the two optional fields that decide whether those integers are replica ids, partition ids or positions in a flat enumeration of the device grid, and the switch with four arms that multiplies one written group out against a mesh into the sets of global device numbers the runtime will actually open communicators for. Then the flattened id and the device number, which walk the same grid in opposite directions and disagree on any mesh with more than one of each. On the runtime side: ranks, communicators, cliques and the factory that makes them, all defined without naming a vendor; the future that means launched rather than finished, and means two different things on GPU and CPU; the rendezvous that agrees on membership before any bytes move; and the three separate places where the answer to how do we avoid a deadlock is the same answer. The signature exhibit runs the real group arithmetic and the real verifier checks over a mesh you dial.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of XLA Collectives: One Guest List, Then One Order — built to be understood, not skimmed.

Questions

What do replica_groups mean on an XLA all-reduce?: They are a list of lists of integers naming who participates, but the list alone does not say which devices talk. Two optional fields on the instruction decide how the integers are read, and there are exactly four readings. With no channel id the mode is cross-replica and the integers are replica ids, with each group formed inside one partition. With a channel id and no use_global_device_ids field the mode is cross-partition and the integers are partition ids, with each group formed inside one replica. With a channel id and use_global_device_ids set to false the mode is cross-replica-and-partition, where the integers are still replica ids but the group sweeps every partition of each listed replica. And with a channel id and use_global_device_ids set to true the mode is flattened-id, where the integers are positions in one flat enumeration of the whole replica-by-partition grid. Present-and-false differs from absent because the field was added after the mode it disambiguates, so an older module with no field has to keep meaning cross-partition. An empty list means all participants, but which participants depends on the mode, and flattened-id may not be empty at all.
Why does one replica group turn into several device groups?: Because in cross-replica mode the written group is read once per partition. The function that does the translation loops over the written groups on the outside and over every partition on the inside, building a fresh participant list each time, so two replicas across four partitions with an empty group list produce four groups of two devices rather than one group of eight. That is the correct behaviour for data parallelism inside a model-parallel program: each partition averages independently with the corresponding partition on the other replicas, because partition 0 of replica 0 holds a different slice of the model than partition 1 does and averaging across them would be meaningless. Cross-partition mode does the mirror image, looping over every replica for each written group. Cross-replica-and-partition is the one that collapses instead of multiplying, since each written group becomes a single device group containing every partition of every listed replica. Flattened-id produces exactly one device group per written group, which is why the partitioner prefers it.
Is a flattened id the same as a global device id in XLA?: No, and on any mesh with more than one replica and more than one partition they disagree. A flattened id walks the grid replica-major: the formula, which appears identically in three places in the same file, is replica_id times partition_count plus partition_id. A global device id comes from the device assignment, and the default placement walks the grid the other way, writing computation times replica_count plus replica into cell (replica, computation), where computation is the partition. So on two replicas by two partitions, replica 1 partition 0 has flattened id 2 and device number 1. The consequence you actually see is that a flattened group written as {0,1} can resolve to devices {0,2}, which is why a replica group in a dump sometimes looks shuffled. When either count is one the two numberings coincide, which is why the divergence usually appears for the first time when somebody combines data and model parallelism.
What is a clique in XLA, and how does it relate to a communicator?: A communicator is one participant handle, one instance per device in this process, and the collective operations are its virtual methods: all-reduce, all-gather, reduce-scatter, broadcast, collective-permute, all-to-all, and point-to-point send and receive. A clique is all the communicators that make up one collective, stored in a map sorted by rank, with a comment saying the sorting exists to guarantee deterministic traversal order. Asking a clique for a rank returns an optional, and the empty case is not an error: it means that rank exists but is not local to this process, which is normal on multi-host jobs where a thirty-two device clique has eight communicators in each of four processes. Above both sits the factory, which creates communicators from a clique key plus an optional rendezvous id plus a list of device-and-rank pairs, and which can also split an existing set of communicators into sub-groups by colour without a fresh bootstrap. The whole surface is defined without naming a vendor, and the GPU, CPU and TPU implementations all satisfy it.
Why does an XLA distributed job hang instead of erroring?: Because collectives match by arrival order rather than by name. Nothing on the wire identifies which collective a call belongs to, so there is nothing to compare and no mismatch to detect. Each rank enters the next collective on its communicator and is paired with whatever the other ranks entered next, so two ranks issuing the same two operations in opposite orders each end up waiting for a partner that is already waiting for them. The source is blunt about this: the clique header opens by saying it is notoriously easy to get a deadlock, which is why communicators are grouped into cliques with a well defined order of operations. In practice the common causes are a rank that never arrives because a host crashed or was never started, a guest list that differs between processes because the mesh or the device assignment differed, and control flow that diverges so one device issues a collective the others never reach. The compiler does have a pass that forces a total order on collectives, but on GPU it is enabled to work around executable divergence caused by online autotuning of convolutions, not as general deadlock insurance.
What does async collective conversion do in XLA, and when does it help?: A pass rewrites each collective into a start instruction that issues the operation and a done instruction that waits for it, so the scheduler can put independent arithmetic in between and hide the network latency behind work that was going to run anyway. Whether a given instruction is converted is one expression: the config predicate for that opcode, and then either a flag that skips the size check or the output being at least a minimum size in bytes. Only all-reduce, all-gather and reduce-scatter consult a size at all; collective-permute, all-to-all, ragged-all-to-all and collective-broadcast convert whenever their predicate says yes. The threshold exists because a start and a done are two instructions where there was one, and on a small buffer the overhead exceeds the overlap. The GPU pipeline converts everything unconditionally with the thresholds at zero, then annotates only the ones that should really run asynchronously, filtered by a debug option. After scheduling, one more pass converts any pair with nothing overlapping between them back to the plain synchronous instruction, so the speculation is refunded where it did not pay.

XLA Collectives: One Guest List, Then One Order

Questions

Related explainers