RFC - Using a Handshake to enable Fat Collectives

Handshake

Introduction

Skinny collectives are the torch_tpu solution for deadlocks in XLA that would get caused by generating differing HLOs across ranks. This comes at a major performance cost as XLA is unable to overlap communication and computation. In order to achieve optimal performance we need a reliable way of enabling fat collectives.

Proposed solution

A TCP handshake on graph fingerprints between ranks at configurable stages of the execution.

The coordinator rank

runs a TCP server
receives graph fingerprints from participating ranks
sends a split decision

All other ranks

send a fingerprint
receive the split decision
- if all graphs match compile full graph
- otherwise split and compile with skinny collectives

The expectation here is that most graphs should be the same across ranks so this should safely ensure enablement of fat collectives in most cases.

Handshake stages

We will provide users with three options for handshake stages

OFF (no Handshake)
COMPILE_STAGE (before compilation)
DISPATCH_STAGE (at dispatch time, this is intended to be default)

COMPILE_STAGE

We run the Handshake only once before compiling.

[UPDATE] We intend to use this as default as it is a good tradeoff for correctness and performance.

DISPATCH_STAGE

~~This mode will be default because we anticipate that a recompilation might occur on some ranks while not on others.~~

[UPDATE] We plan to keep this as an option for debugging but as of right now it looks too defensive as it would only be needed for graph specific recompilations which seems like a stretch.

Dummy example:

iteration = 0
divergent_rank = ...

def divergence_after_iteration(x, y):
    z = torch.add(x, y)
    if rank == divergent_rank and iteration == 1:
      z = torch.mul(z, 5.0)
    dist.all_reduce(z)
    return z

compiled_fn = torch.compile(divergence_after_iteration, backend=backend)
res0 = compiled_fn(x, y)
iteration += 1
res1 = compiled_fn(x, y)      # recompilation happens only on divergent rank

Note that there are worfklow in VLLM and SGLang that apparently have purposely diverging graphs.

If the handshake doesn't help with enabling fat collectives, we can use it to gather data points in divergences between graphs.

@basioli-k

Thanks for this proposal @basioli-k!

I have a question about the post-dispatch handshake. I need to know what 'dispatch' actually refers to to pose the question though. Are we assuming that compile is done with 'fullgraph=true', or can it have graph-breaks? if there are graph-breaks on rank1 but no graph-breaks on rank0, does that mean rank1 has more dispatches than rank0?

Basically, I'm concerned about cases where we end up with a hang or timeout at the 'handshake' apparatus instead of at the HLO runtime, due to a mismatch in the number of either compilations or dispatches. Do you have a way to guard against that?

@wconstab

@wconstab

I need to know what 'dispatch' actually refers to to pose the question though.

Dispatch would mean before setting of a pjrt executable - basically before invoking this __call__

torch_tpu/torch_tpu/_internal/compile/torch_tpu_compiled_executable.py

Line 288 in 2282137

def __call__(

Are we assuming that compile is done with 'fullgraph=true', or can it have graph-breaks? if there are graph-breaks on rank1 but no graph-breaks on rank0, does that mean rank1 has more dispatches than rank0?

We don't need fullgraph=True. The way we designed the handshake protocol it can handle splits.

Here's a bit of a deeper dive in the protocol which hopefully clarifies this

Implementation details

Important observations

dispatch/compile order for modules with collectives matches across ranks -> this is reasonable as collectives have to be dispatched in the same order even on CUDA
[UPDATE] We keep track of collective counts per process group per process and snapshot the state when a new Handshake is dispatched
we know which frame is the last frame of the torch.compile region -> we use graph_break=False from GraphCompileReason https://github.com/pytorch/pytorch/blob/a910eb0f598f557d12de8def6768ca0b8889ec6b/torch/_dynamo/output_graph.py#L340

[UPDATE] For this example

even_ranks = [r for r in range(world_size) if r % 2 == 0]
even_group = dist.new_group(ranks=even_ranks)

def foo(x, y):
    dist.all_reduce(x, op=dist.ReduceOp.SUM)
    dist.all_reduce(x, op=dist.ReduceOp.SUM)

    if rank in even_ranks:
        dist.all_reduce(y, op=dist.ReduceOp.SUM, group=even_group)

the collective counts for this module would be:

For even ranks:
global subgroup: {"collective_count_before": 0, "num_collectives_in_graph": 2}
even subgroup: {"collective_count_before": 0, "num_collectives_in_graph": 1}

For odd ranks 
global subgroup: {"collective_count_before": 0, "num_collectives_in_graph": 2}

[UPDATE] Why we made this switch?
We realized that the last_frame logic falls apart if there are process subgroups (i.e. if there is a non global subgroup) as there is not enough information on which Handshakes are comparable.

What each participant sends

all graphs with collectives
[UPDATE] information on collectives contained inside the module
~~the last frame of the torch.compile region that CONTAINED a collective (even if that frame doesn't contain it)~~

Consensus algorithm

Receive FIRST Handshake from ANY participant
Receive Handshakes from all participants in this Handshake
Decide on the consensus and broadcast to participants
[UPDATE] If any participant runs ahead within some process group we flush other participants until we are synchronized
~~If ANY participant sent a last frame while some didn't we flush the Handshakes for other participants until we reach the last frame (see What each participant sends point 2)~~

Example

Rank 0 generated Frame1, Frame2, Frame3. For rank 0 the last compilation frame is Frame3.
Rank 1 generated Frame1’, Frame2’. For rank 1 the last compilation frame is Frame2’.

Each rank will initiate a Handshake for every Frame.

Frame1 and Frame1’will participate in the first handshake and get a consensus.
Frame2 and Frame2’will participate in the second handshake and get a consensus.
Because Frame2’ is the last frame in a frame for rank 1 the following happens:
1. We poll the queues of participating ranks (in this case rank 0) for the next frame
2. Frame3 arrives via a Handshake initiated by rank 0
3. We immediately send a negative result for Frame3 as it doesn’t match any other frame.
  1. If this frame wasn’t the last frame we go back to step a
  2. Otherwise we finish and go to our full handshake.

[UPDATE] Another example of ranks progressing through the global PG:

Implement CollectiveHandshakeRequest class and serialization.

Implement _CollectiveHandshakeConsensus message queues.

Implement _select_handshake_port utility

Implement _CollectiveHandshakeConsensus message queues.

Implement _select_handshake_port utility

Implement matches and collective_count APIs in RankCollectiveCounts.

Implement should_handshake utility function.

Expose executable fingerprint and fingerprint calculation to python API.

Implement should_handshake utility function.

Expose executable fingerprint and fingerprint calculation to python API.

Implement should_handshake utility function.

Expose executable fingerprint and fingerprint calculation to python API.

Initial implementation has landed.

Current status

handshake is turned off by default
there have been some benchmarks with vllm and torchtitan suggesting that the performance is good
Handshake works only on modules containing the global process group

The implementation differ from what was described above.

Instead of comparing graphs containing only collectives in the order they are arriving, we are comparing them based on their sequence number within a process group.

Will follow up with a more detailed high level overview of the algorithm with examples but the code can be found in

torch_tpu/torch_tpu/_internal/distributed/handshake.py

Line 1181 in 83fc580

async def _consensus_loop(self) -> None:

with unit tests in https://github.com/google-pytorch/torch_tpu/blob/main/tests/distributed/handshake_test.py and e2e tests in https://github.com/google-pytorch/torch_tpu/blob/main/tests/distributed/handshake_e2e_test.py

I updated the original comment #2886 (comment) with the current implementation state (I crossed out the old to maintain history and annotated new text with [UPDATE]).

As said previously the initial implementation has landed.

Current State

Consensus algorithm works for global process group

Next steps

Make consensus algorithm work for arbitrary process group
Add a debug mode which uses print_readable to automatically warn the user what they should change to get a sucessfull handshake
Maybe automate finding the largest common subgraph (more compute expensive)

Benchmarks

Torchtitan

Llama 70B
world size 256

Configuration	Compile Time (Step 1)	TPS (Stabilized)	TFLOPS	MFU	End-to-End Time	Peak Memory	Splits > 1	Successful Handshakes
Baseline fat collectives with spmd_safe (`fat_spmd_safe`)	36.77s	1,081	520.28	45.10%	30.32s	14.91 GiB (15.53%)	0	0
Baseline skinny collectives (`skinny_baseline`)	46.85s	946	455.46	39.49%	34.64s	4.91 GiB (5.12%)	2	0
Baseline fat collectives (`fat_baseline`)	37.87s	1,080	519.99	45.08%	30.34s	14.91 GiB (15.53%)	0	0
Fat collectives enabled by handshake - compile time (`handshake_compilation_stage`)	37.44s	1,081	520.39	45.11%	30.32s	14.91 GiB (15.53%)	0	6
Fat collectives enabled by handshake - dispatch time (`handshake_dispatch_stage`)	37.59s	1,075	517.42	44.86%	30.49s	14.91 GiB (15.53%)	0	3,330

VLLM

I measured on Qwen 30B and the Handshake failed (as expected).

All of the failure reasons were trivial:

[Difference at line 4]
'ge: "b8[8192]" = torch.ops.aten.ge.Scalar(arg0_1, 0)'       (in worker_rank_0.log)
'ge: "b8[8192]" = torch.ops.aten.ge.Scalar(arg0_1, 18992)'   (in worker_rank_1.log)
...

[Difference at line 5]
'lt: "b8[8192]" = torch.ops.aten.lt.Scalar(arg0_1, 18992)'   (in worker_rank_0.log)
'lt: "b8[8192]" = torch.ops.aten.lt.Scalar(arg0_1, 37984)'   (in worker_rank_1.log)
...

[Difference at line 18]
'mul: "i64[8192]" = torch.ops.aten.mul.Tensor(bitwise_and, 0)' (in worker_rank_0.log)
'mul: "i64[8192]" = torch.ops.aten.mul.Tensor(bitwise_and, 18992)' (in worker_rank_1.log)
...

and using this method enabled me to find the issue was in vllm/model_executor/layers/vocab_parallel_embedding.py#L171.

I manually fixed the discrepancy and confirmed I am able to get the same TPS as when skinny collectives are manually turned off.

@wconstab

@wconstab @cbasile-g

Here's additional benchmarking for results for models using VLLM.

TLDR

Even when handshaking at dispatch time we perform very well. I do expect it is a bit nosy but we can fairly clearly see that the handshake consistently performs at the level of fat collectives

Llama 70B Summary Of Results

Scaling up to 1024 chips (world size 2048) scales almost perfectly.

Problems

At 2048 chips due to the zeroth dimension of wqkv[10240, 8192] not being evenly divisible by world size (4096) the fat collective run hung while the handshake runs fell back to skinny collectives.

Attempt at workaround

Sharding across dimension 1 (8192 is divisible by 4096!).

Here we saw a perf degradation when using the dispatch handshake this is likely because the model is small and the handshake cannot be hidden behind device execution

Result Summary

Navigation Menu

Repository navigation

RFC - Using a Handshake to enable Fat Collectives #2886

Handshake

Introduction

Proposed solution

Handshake stages

COMPILE_STAGE

DISPATCH_STAGE

Implementation details

Important observations

What each participant sends

Consensus algorithm

Example

36 remaining items

Current State

Next steps

Benchmarks

Torchtitan

VLLM

Llama 70B Summary Of Results

Problems

Attempt at workaround

Result Summary

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

RFC - Using a Handshake to enable Fat Collectives #2886

Description

Handshake

Introduction

Proposed solution

Handshake stages

COMPILE_STAGE

DISPATCH_STAGE

Activity

basioli-k commented on Aug 11, 2026

wconstab commented on Aug 11, 2026

basioli-k commented on Aug 12, 2026

Implementation details

Important observations

What each participant sends

Consensus algorithm

Example

36 remaining items

basioli-k commented on Sep 9, 2026

basioli-k commented on Sep 11, 2026

Current State

Next steps

Benchmarks

Torchtitan

VLLM

basioli-k commented on Sep 11, 2026

basioli-k commented on Oct 6, 2026

basioli-k commented on Oct 26, 2026

Llama 70B Summary Of Results

Problems

Attempt at workaround

Result Summary

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions