Custom Collectives — infra.bea.ai

Standard NCCL collectives are general-purpose — they work well for most patterns. But frontier training systems (DeepSeek, Megatron-LM) sometimes need custom communication patterns that NCCL can’t express efficiently. This chapter covers the tools for going beyond standard collectives.

1. When Standard NCCL Isn’t Enough

Scenario	Problem with standard NCCL	Custom solution
MoE expert routing	All-to-All doesn’t match irregular token distribution	Custom dispatch with per-expert buffers
Pipeline + TP hybrid	Need overlapped send/recv + AllReduce simultaneously	MSCCL programmable schedules
Cross-node TP	AllReduce latency too high for per-layer sync	NVIDIA UB (fuse communication into compute)
Hierarchical AllReduce	Flat ring doesn’t exploit intra-node NVLink advantage	Two-level: intra-node NVLink ring + inter-node IB tree

2. MSCCL: Programmable Collectives

MSCCL (Microsoft Collective Communication Library) lets you define custom collective algorithms as XML programs:

<!-- Custom 2-level AllReduce: intra-node ring + inter-node tree -->
<algo name="hierarchical_allreduce">
  <step type="ReduceScatter" gpus="0-7" algo="ring"/>    <!-- NVLink -->
  <step type="AllReduce" gpus="0,8,16,24" algo="tree"/>  <!-- IB, one per node -->
  <step type="AllGather" gpus="0-7" algo="ring"/>         <!-- NVLink -->
</algo>

Why MSCCL Exists

NCCL selects from a fixed set of algorithms (ring, tree). For specific topologies (e.g., Azure’s NDv4 with 8 A100s per node and 8 IB NICs), a custom hierarchical algorithm can be 20-40% faster than NCCL’s default by exploiting the exact topology.

Using MSCCL with NCCL

# MSCCL compiles to NCCL-compatible algorithms
export MSCCL_ALGO_DIR=/path/to/custom_algos
export NCCL_ALGO=MSCCL  # tell NCCL to use MSCCL algorithms

MSCCL is primarily used by cloud providers (Azure) and large training teams. Most users should stick with standard NCCL.

3. NVIDIA Unified Bus (UB)

Unified Bus fuses communication directly into compute kernels — instead of launching a compute kernel and then a communication kernel, a single kernel does both.

The Problem UB Solves

In tensor parallelism, every transformer layer requires:

Compute (GEMM) → AllReduce → Compute (GEMM) → AllReduce → ...

With standard NCCL, this is:

[GEMM kernel] [NCCL AllReduce kernel] [GEMM kernel] [NCCL AllReduce kernel]
              ↑ GPU idle while waiting for AllReduce

With Unified Bus:

[GEMM kernel that also does AllReduce as it produces output]
  ↑ communication overlapped with the tail of computation

How It Works

UB enables GEMM kernels to write output tiles directly to remote GPUs’ memory as they’re computed — using NVLink DMA within the compute kernel. No separate communication launch.

Availability

Integrated into Megatron-LM for tensor-parallel GEMM
Requires NVLink (intra-node only)
Used by NVIDIA’s internal training infrastructure and DeepSeek-V3

UB is cutting-edge — most teams don’t need it. It matters for >100B models where tensor parallelism’s AllReduce overhead is the training bottleneck.

4. Hierarchical Communication

Hierarchical AllReduce

For large clusters, flat algorithms are suboptimal. Hierarchical approaches exploit the bandwidth asymmetry:

Step 1: Intra-node AllReduce (NVLink, 900 GB/s)
  GPU 0-7: ReduceScatter → AllGather (fast, NVLink)

Step 2: Inter-node AllReduce (InfiniBand, 50 GB/s per NIC)
  Node 0 GPU 0 ↔ Node 1 GPU 0 ↔ Node 2 GPU 0 ↔ ...
  (one representative per node, slower IB link)

Step 3: Intra-node broadcast (NVLink, fast)
  GPU 0 broadcasts to GPUs 1-7

Benefit: The expensive inter-node communication only transfers $1/8$ th of the data (one GPU’s share per node). The cheap intra-node NVLink handles the rest.

NCCL does this automatically for multi-node setups — but MSCCL lets you customize the exact strategy.

5. InfiniBand SHARP: Hardware AllReduce

SHARP (Scalable Hierarchical Aggregation and Reduction Protocol) offloads AllReduce to the InfiniBand switches themselves:

Without SHARP: GPU → NIC → Switch → NIC → GPU (data traverses full path)
With SHARP:    GPU → NIC → Switch(reduces here) → NIC → GPU
                              ↑ switch does the sum!

The switch performs the reduction in-network — reducing data volume by 2× and halving latency.

Availability: Requires Mellanox Quantum InfiniBand switches with SHARP support. Not available on Ethernet/RoCE.

Self-Check

Name a scenario where standard NCCL is insufficient
Explain what MSCCL does and when to use it vs standard NCCL
Explain Unified Bus and why fusing communication into compute helps TP
Explain hierarchical AllReduce and why it exploits NVLink/IB asymmetry
Explain SHARP and how it reduces AllReduce to in-network computation