HCCL: Collective Communication for Meta Training and Inference Accelerators

Before any work is run on the compute or communication portion of the MTIA 300 chip, it is first submitted through a stream to the CPU-C to ensure it is ready to be dispatched and that the place where it will be executed is ready to receive it. The CPU-C (shown in purple in Figure 1 inside the CCP module) is responsible for maintaining ordering within the stream and via dependencies expressed across streams. Each of these streams can hold a mix of both compute and communication operations. For communication, this means the CPU-C checks the description of each work packet and schedules each subgraph on the requested ME (see Section III-B for more details about collective construction).

The MTIA 300 chip contains two groups of eight message engines which are used to manage the chip’s network communications and collective compute. MEs bridge the gap between the subgraph representation of collectives and the representation used by the NICs with RDMA queue pairs and work requests. Each ME contains three distinct areas to accomplish different functions, the details of which can be seen in Figure 2 and are discussed below.

The network chiplet contains all of the backend off-device I/O capabilities and is integrated directly to avoid PCIe traffic when reading or writing data to caches or HBM. MTIA 300 contains two network chiplets and each chiplet has six custom 800 Gb/s RDMA NICs. Here we describe some details of the NICs.

The HCCL control path manages all of the network resources that are used by the data path later for communication. It uses standard RDMA verbs interfaces to manage these resources (e.g., ibv_create_qp and ibv_reg_mr). Each of these resources is created on-demand as the application requires them. Before constructing a collective (see Section III-B2 for more details), HCCL needs to determine which QP resources will be needed. The control path re-uses existing idle QPs and allocates incremental QPs to fulfill the request, and then provides their descriptors to the data path. Dynamic QP allocation is a well-known optimization (other collective communication libraries employ a similar control-path design) – but in HCCL it serves an additional purpose: the express doorbell design described in Section II-A3 limits the total number of queue pairs that the system can support, making resource reuse important. While this can theoretically require HCCL to harvest existing QPs to allow them to be reallocated for other purposes, in practice we have not seen a need for this as we can support enough QPs without requiring larger communicator sizes. Instead we focus on keeping the number of QPs required for a collective algorithm to a minimum and using optimizations other than QP scaling, which can be common in other collective libraries, but do not provide significant benefit in HCCL.

The HCCL data path covers the components that are responsible for generating and executing the data movement instructions themselves. HCCL differentiates itself from traditional communication libraries such as many MPI [14] implementations as well as other collective communication libraries by offloading not only the data transfer on the NIC itself, but also the execution and dependency processing of the entire communication operation.

The same streaming semantics are used both for communications and compute allowing both to be inserted into the same stream (or they could use cross-stream dependencies if desired). A graphical representation of the construction of communication instructions (work packets, subgraphs, and WQEs) is in Figure 4. Streams operate at the granularity of work packets (the equivalent of a compute kernel) which are executed sequentially within a stream. Each work packet can be broken down into multiple subgraphs which are parallelized across available MEs (see Section II-A2) to maximize throughput. Subgraphs contain a list of work queue elements (WQEs) which describe individual operations to be executed sequentially and include metadata about dependencies between each of the WQEs.

There are a number of WQE types, but each of them can be put into one of three categories:

RDMA Transfers

These include Send, Receive, Write, WriteWithImmediate, etc. and describe data movement through the RDMA NIC. They contain all of the information necessary to construct an RDMA work request (e.g., QP number, lkey (local key), rkey (remote key), addresses, etc.).

Compute operations

These describe offloaded compute operations that occur in the NMC (see Section II-A2) which can include sums and copies. The copy operation can also function as a copy engine for the communication path.

Subgraph coordination

These describe operations to wait for a condition on a memory location (e.g. the value of a short larger than four) or update / set a value in memory. For instance, two subgraphs can synchronize with each other by using two memory location as a signal pair (using monotonic increasing values) to indicate when the other subgraph can continue to next phase.

Figure 5 depicts a sample subgraph that uses these WQEs. This shows how a series of send, receive, and compute operations can be used to construct a simple ring algorithm for AllReduce [24]. This also demonstrates how WQEs can express dependencies on other WQEs. There are five types of dependencies:

Fence

Do not start the next WQE until the current WQE has completed.

Sync

Wait for all previous WQEs to complete before starting the current WQE. This introduces a barrier in the subgraph execution.

WQE Sync

Wait for a specific WQE to complete before starting the current WQE. Typically used for checking for valid input (pointing at Recv WQE) or buffer release (pointing at Send WQE).

Receive Sync

Wait for all outstanding receive WQEs to complete before starting the current WQE (i.e., all data for next pass is available).

Send Sync

Wait for all outstanding send WQEs to complete before starting the current WQE - aka all (temporal) buffers from previous pass are released and available for reuse.

To facilitate multiple dependencies for scheduling a WQE (e.g., a tree based NMC reduction) we have also introduced a NOOP WQE containing only a dependency.

Using send-receive semantics, if there is no matching receive posted when data is incoming to the remote NIC, it will trigger an RNR (receiver not ready) NAK back to the sender, forcing a back-off (wait) with progressively longer retry times. The retry time is programmable, so there is a compromise between retry time and global/application time-out defined by the max skew between processes.

To mitigate this, HCCL has a special mode (SendSync) which relies on responder-driven signaling. The requester waits on the signal (using a WAIT WQE) before posting the SEND WQE.

A further optimization of this is to use the SendSync protocol only on the first transfer in the collective, and let the algorithm guarantee that sufficient receives are posted for any execution scenario, to guarantee that no send will become unmatched by a receive. An example of this would be a bidirectional ring with number of receives exceeding pipeline depth.

One-sided communication APIs have been available in many forms including remote memory access APIs in MPI-2 [12] and again in MPI-3 [13], various SHMEM [3, 23, 22] libraries, and others. Some of the relevant APIs available to us through these interfaces are put (issue an RDMA write directly) and synchronization semantics such as flush or fence to ensure local completion and ordering as well as a way of signaling a remote process that data has landed (e.g., signal and wait). These interfaces are often used in fused kernels that combine compute and collective operations in a single kernel that runs on the PE grid. This minimizes latency by eliminating the computational cost of launching additional kernels.

For MTIA 300, implementing one-sided communication presents some design and performance tradeoffs. As previously described, communication flows through the MEs, which are responsible for executing subgraphs and dispatching work to the NICs and NMCs. For one-sided communication, each communication round is a single NIC WQE, so this model no longer makes sense due to the computational cost of issuing each as a subgraph. On the other hand, because of the shared completion queue described in Section II-A2, it is difficult to efficiently monitor completion of individual WQEs on the PE grid. To manage this, we use a hybrid model where the PE submits WQEs directly through the express doorbell, but the ME receives the completion notification and manages a counter of the completions that the PE can read to know when its messages have been delivered. The host is responsible for setting up the communication control path (i.e., queue pairs and memory regions) and handing those over to the device, which is managed through HCCL APIs that can have an adaptable implementation from one hardware generation to the next as capabilities and requirements shift. An example of this workflow is found in Figure 6 under the column ”One-Sided”.

A common pattern of compute and communication in inference workloads for mixture of experts (MoE) kernels is to have a routing kernel that determines which expert(s) to use, followed by an AllToAllv collective whose metadata (i.e., sizes and offsets) depend on the output of the routing kernel. Traditionally each of these would be submitted and executed eagerly, either in the same stream to express the data dependency or across two streams with an explicit cross stream dependency. However, this adds significant latency because HCCL cannot construct the subgraphs for the AllToAllv collective until the previous kernel is complete. Then the collective is submitted, copied, and scheduled on the device, all of which would be exposed as they are no longer overlapped with the compute kernel that precedes them. To address this, a variant collective was proposed called AllToAllvDynamic [25] which, in addition the existing parameters of AllToAllv, takes pointers to device memory where the final metadata will be written. This allows the collective library to modify the WQEs before they are sent out without requiring the entire host-side software stack to be involved. Figure 6 demonstrates this workflow in the column labeled ”Device-Resident”.

AllToAllvDynamic removes the computational overhead of generating the collective and copying it to the device, however, it still leaves the scheduling latency exposed. For compute kernels, this is typically hidden by fusing together multiple compute kernels in graph mode allowing them to execute as one large kernel which avoids the scheduling latency that would normally occur between each. However, though compute kernels and collective operations can occur in the same stream, they cannot directly be fused because they are not executed by the same hardware.

We explored another concept called device-triggered collectives where the caller can attach a memory address and comparator to a collective to indicate that the collective should not begin until the memory address fulfills the requirement of the comparator (e.g., a memory address becomes non-zero). Additionally, the user can provide a second address and value to be set when the collective is done. With these tools, compute kernels can be compiled in graph mode and still incorporate accelerated communication on the MEs by scheduling the collective on another stream where it is immediately executed, but then blocked until the compute kernel signals that it is ready. In the subgraph, this is implemented by adding additional WQEs where necessary to perform the indicated comparator and wait for it to be true. In the case of multiple subgraphs, these WQEs will need to be added to each of the subgraphs that are immediately scheduled.

Upon completion, the CPU-M needs to indicate back to the PE grid that the next compute work is unblocked. This can either be done in an existing subgraph via a WQE to set the corresponding value or via a dedicated subgraph if there are multiple subgraphs that need to avoid prematurely writing to a shared value. In the case of multiple subgraphs, HCCL uses internal semaphores to ensure ordering is preserved. Figure 6 visually compares this method (labeled ”Device-Triggered”) to the other two methods described here.

The primary benefit of this optimization is to maintain a single compute kernel when executing in graph mode. Graph submission and launching is a non-negligible computational overhead that can be avoided if the compute and communication work can be executed in a single kernel launch. By embedding the collective semaphore into the compute kernel and pre-launching the collective in a separate stream, we can avoid graph breaks that would force multiple CPU-C scheduling computational overheads.

Scheduling the collective on a parallel stream also enables the MEs to pre-post parts of the work until the compute has signaled the ME to start collective execution. Upon scheduling, the ME can also pre-post receive WQEs before the wait WQEs mentioned above since these do not incur in any data exchange until a matching send is posted on the remote node. This enables a two-fold optimization: 1) we save execution time in the ME since fewer WQEs need to be posted in the exposed time of the collective and 2) we mitigate the occurrence of ”Receiver not Ready” NAKs since, for a sufficiently large time window between scheduling and signaling, we can expect all ranks to have pre-posted their receive WQEs.

We design an end-to-end benchmark to compare the computational overhead of issuing WQEs from the PE vs ME, shown in Figure 13. We observe a static computational overhead of around $450~\mathrm{ns}$ when submitting WQEs directly from the PE grid, which is due to the computational overhead of building WQEs on the device, bookkeeping for operation correctness, and obtaining completions routed to the PE through an ME. This demonstrates the efficacy of submitting work from the PE grid to enable fused kernels.

Sections IV-B2 and IV-B3 describe optimizations that allow collectives to be executed in fused kernels with compute operations in graph mode. We measured the performance benefit of these two optimizations together as they are less impactful separately. In Figure 15, we can see an AllToAllvDynamic collective that is scheduled simultaneously with a series of fused kernels. The collective itself appears to execute for $123\mu\mathrm{s}$ on the trace, but in the compute stream, it is actually only executing for $31\mu\mathrm{s}$ as most of the execution time of the collective is actually blocked while waiting for the compute kernel to signal readiness. This avoids all of the computational overheads demonstrated in Figure 10 that would be injected by having to schedule and execute the collective after the compute kernel is finished. Figure 14 shows a performance comparison of device-triggered collectives against their eagerly executed counterparts showing significant latency improvement from avoiding the compute necessary for setup.

We implement AllReduce and AllToAllv kernels on the PE grid to be used as a direct replacement for HCCL offloaded collectives. We take this approach based on the numbers presented in Figure 10, which indicate some computational overhead from scheduling to the ME itself. Figure 16 shows the performance of these kernels when running end-to-end and at the job sizes within the scale-up domain for running inference workloads. We achieve sub-$6\mu\mathrm{s}$ collective execution for both AllToAllv and AllReduce, with AllReduce containing additional synchronization requirements due to the reduction step.

Both collectives are implemented in an “all-to-all” style, with all ranks writing to all ranks over the network. This also enables different cores to issue parallel messages automatically without synchronization, which helps mitigate skew and message enqueue computational overhead for such fine-grained workloads.

At this stage, our solution is mostly bound by the RTT time of the network, with little optimization to be done at the SW level. Further improvements on these numbers would likely rely on updates to the network itself to improve end-to-end message delivery time. While not mentioned in this work, this functionality also enables fused-kernel design which would reduce the control core computational overhead and further cut execution time during end-to-end inference (along with further optimizations described next).