commit cd5347c06974e4d4d1b65f07eff210e2ebd6280d
author Okwan Kwon <okwan@google.com> Fri Feb 17 08:44:52 2023 -0800
committer GitHub <noreply@github.com> Fri Feb 17 08:44:52 2023 -0800
tree aa8645c5170170a06c8eeda0a76fcba8ed19eb73
parent 16d47a4cc446faff323f1b1e0725070a602d9261

Support mhlo collective ops (#11988)

Add the support for the following mhlo collective ops:
- mhlo.replica_id
- mhlo.all_gather
- mhlo.all_reduce
- mhlo.reduce_scatter

Since NCCL only supports the splitting and concatenation on dim 0 for
all_gather and reduce_scatter, transposes are inserted when the
split/concat dimension is not 0.

To make the implementation simple and incremental, several stages are
planned as follows:

Stage 1 (The current PR):

It assumes a deterministic order of collective operations with 1:1
mapping from replica_id to rank. This means that there is a single
replica group in the mhlo operation and all replicas participate in the
collective operation. Since the order is deterministic and all ranks are
involved in the communication, we can simply use the default channel for
communication. In the MHLO ops, `use_global_device_ids` is set to use
the flattened IDs.

Note that the MHLO collective ops have multiple strategies to interpret
`replica_groups` attribute, such as `flattened_ids`, `cross_replica`,
`cross_partition`, and `cross_replica_and_partition`. (See
https://github.com/openxla/stablehlo/blob/main/docs/spec.md#parallel-execution
for more details of the strategies.)

Stage 2:

Supports multiple channels. This allows us to have multiple replica
groups for collective operations.

Stage 3:

Supports `partition_id` the other strategies: `cross_replica`,
`cross_partition`, and `cross_replica_and_partition`.

Stage 4:

Supports `all_to_all` and `collective_permute`. This would need to
support the NCCL group markers to support multiple collective ops in
parallel since the ops are composite and will need to be lowered into
the existing collective ops.

Stage 5:
PJRT integration and model level testing.