tree 6fdcbc50dfbaa75fec86bf72ff53bf8aabab3054
parent 6a3dce0dcb39c2f43bee1b6468552e22b32afb4d
author Peter Hawkins <phawkins@google.com> 1659131830 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> 1659132114 -0700

[XLA:GPU] Adds a barrier that requires all local participants to call ncclCommInitRank() before allowing any participant to proceed.

Without a barrier, we can experience deadlocks. As best I understand it, the deadlock scenario looks like this:

Thread A:
* calls ncclCommInitRank(), which succeeds,
* issues the collective operation,
* calls an operation that manipulates the device page tables, e.g., copying a device buffer to an unpinned host buffer.
* Since this action manipulates the device page tables, it seems that this action blocks waiting for the device stream.

Thread B:
* calls ncclCommInitRank(), which calls cudaMalloc().
* cudaMalloc() also manipulates device page tables, and cannot proceed without acquiring an internal lock around the device page table state
But thread A already holds this lock, but thread A cannot make progress until thread B issues its collective.

This is a deadlock: neither thread can make progress. We can avoid the problem by requiring a barrier after the calls to ncclCommInitRank(), requiring all GPUs to finish initialization before any of them can issue their collective operation.

Fixes https://github.com/google/jax/issues/11637

PiperOrigin-RevId: 464164328