jax
77c0fd0c - [Mosaic:GPU] Use NCCL API on the host to collective peer addresses.

Commit
11 days ago
[Mosaic:GPU] Use NCCL API on the host to collective peer addresses. This allows to use collective kernels with several processes. Since symmetrical memory returns a peer address of XLA allocation we don't need to subtract offset in param to peer map in order to calculate parameter in multimem address space. However in combined mode when only a subset of parameters are allocated in the collective memory space, we would still need to perform rendezvous and exchange parameters at the host side. In this scenario for parameters used with multimem we still need to subtract the offset. PiperOrigin-RevId: 912522295
Parents
Loading