[RFC] Intra-node shared memory (SHM) optimizations for communication operators on CPUs

[jianan-gu]

Overview and Motivation

Tensor parallel (TP) inference is a key optimization and well adopted in LLM inference on CPUs. Most use cases for LLM benchmark/deployment focus on one node host (intra-node), where there are many numbers of CPU sockets/NUMAs/sub-NUMAs available.

Referring to the practices of well-known libraries including DeepSpeed , Sglang, vLLM and TGI, for intra-node TP cases, shared memory (SHM) optimizations bring much benefit for communication operators including allreduce, allgather and allgather_into_tensor.

In this RFC, we propose a low-latency SHM-based optimization for Gloo, intending to cover all key communication operators which are mostly used in LLM inference. And as the CPU default communication backend in PyTorch, this SHM optimization can bring significant improvement when working with PyTorch Tensor parallel solution.

Design:

1. Intra-node check and fallback

# Check if an intra-node case
if  local_size >= 0 and local_size == word_size  
 ->   SHM_lmpl

# Fallback
else ->  Gloo Ring_lmpl

2. OP register and dispatch

3. dtype scope:

In PyTorch integration, Gloo uses datatype defined in PyTorch, including c10::Half, c10::Float and c10::Bfloat16. Thus, our SHM impl is also targeted at c10::Half, c10::Float, and c10::Bfloat16.

4. Shared memory design

In general, there are 3 steps for SHM optimization (take SHM allreduce as example):
a. copy input to shared memory buffer.
b. do sum operation on reduce buffer (or gather for allgather).
c. copy result back to each rank

Compared to ring allreduce, it uses one united shared memory buffer for calculation in intra-node case. 
SHM allreduce won't use other collective primitives such as send, receive like in ring allreduce.
In that case it's faster and more efficient.

Implementation (Work in progress):

• allreduce. Intra-node shared memory (SHM) optimizations for CPU primitives #458 （under review)
• allgather.
• allgather_into_tensor.

[jianan-gu]

Hi @d4l3k ,
Here we are reaching out to you.
We are actively working on this RFC feature proposal. Could you help review? Appreciate any feedback.
Thanks!

[d4l3k]

@jianan-gu this sounds very reasonable to me but I have a few questions:

1. do you intend to have a separate backend for shm or do you want the tcp backend to dynamically dispatch between the two automatically? I'm wondering if it makes sense to make PyTorch instantiate two backends and then have some logic that dispatches between them (a third dispatching backend wrapper?)
2. for ring operations that are both intra and internode is there a hybrid option or its always one or the other?

[jianan-gu]

Hi @d4l3k

Thanks for your reply~

For 1. No, we are not going to add a separate backend to replace current or dispatch between the two.
We introduce shared memory (SHM) communication algorithm, same level as ring and bcube, where SHM will intend to play as a fast path for current ring algorithm. And this fast path will automatically take effect for intra-node cases (by runtime automatical detection whether it is intra-node, or optionally adding a global SHM control flag if it is more appropriate). Thus, for PyTorch, it still instantiates current backend without any change.

For 2. SHM shall replace ring for intra node cases (as the fast path), and it is still ring in internode cases.

Besides, we have implemented the SHM allreduce #458 (mark as draft under our internal review). Then comes for the rest allgather and allgather_into_tensor.