Make sure proxy threads inherit the CPU affinity.

1 files changed, 6 insertions, 6 deletions

diff --git a/src/init.cc b/src/init.cc
index 867554f..1cc8f7f 100644
--- a/src/init.cc
+++ b/src/init.cc
@@ -798,12 +798,6 @@ static ncclResult_t initTransportsRank(struct ncclComm* comm, ncclUniqueId* comm
   // Compute nChannels per peer for p2p
   NCCLCHECK(ncclTopoComputeP2pChannels(comm));
 
-  // We should have allocated all buffers, collective fifos, ... we can
-  // restore the affinity.
-affinity_restore:
-  sched_setaffinity(0, sizeof(cpu_set_t), &affinitySave);
-  if (ret != ncclSuccess) return ret;
-
   // Compute intra ranks (using AllGather1 data)
   int intraRank0 = -1, intraRank = -1, intraRanks = 0;
   for (int i = 0; i < nranks; i++) {
@@ -828,6 +822,12 @@ affinity_restore:
 
   if (comm->nNodes) NCCLCHECK(ncclProxyCreate(comm));
 
+  // We should have allocated all buffers, collective fifos, ... we can
+  // restore the affinity.
+affinity_restore:
+  sched_setaffinity(0, sizeof(cpu_set_t), &affinitySave);
+  if (ret != ncclSuccess) return ret;
+
   TRACE(NCCL_INIT, "rank %d nranks %d - DONE", rank, nranks);
   return ncclSuccess;
 }