Fix dev_runtime memory leaks

[PatriosTheGreat]

Description

Fix CPU memory leak in dev_runtime finalize method.

Without this change we observe heap checker failures on the following code launched on H100 machine:

constexpr int kNumLocalDevices = 2;
ncclComm_t communicators[kNumLocalDevices];
ncclCommInitAll(communicators, kNumLocalDevices,
                /*devlist=*/nullptr);

for (int i = 0; i < kNumLocalDevices; ++i) {
  ncclCommDestroy(communicators[i]);
}

Heap checker shows the following stack trace where unreleased memory is allocated:

ncclDevrInitOnce()
ncclTeamLsa
initTransportsRank()
ncclCommInitRankFunc()
ncclAsyncJobMain()

Considering that lsaRankList is allocated before bigSize is initialized, I assume it should also be deallocated even in bigSize is zero.

Related Issues

No issues opened IIUC

Changes & Impact

No impact

Performance Impact

No impact

[PatriosTheGreat]

I found another leak with the following sample:

constexpr int kNumLocalDevices = 2;
constexpr int kRootDevice = 0;
constexpr int kDataSize = 8;

std::vector<double> host_data(kDataSize);
for (int i = 0; i < kDataSize; i++) host_data[i] = 1;

devices_data_.resize(kNumLocalDevices);

for (int device = 0; device < kNumLocalDevices; device++) {
  cudaSetDevice(device);
  cudaMalloc(&devices_data_[device], kDataSize * sizeof(double));
  cudaMemcpy(devices_data_[device], host_data.data(),  kDataSize * sizeof(double), cudaMemcpyHostToDevice);
 }

ncclCommInitAll(communicators_.data(), kNumLocalDevices,
                /*devlist=*/nullptr);
streams_.resize(kNumLocalDevices);
for (int device = 0; device < kNumLocalDevices; device++) {
  cudaSetDevice(device);
  cudaStreamCreate(&streams_[device]);
}

ncclGroupStart();
for (int device = 0; device < kNumLocalDevices; device++) {
  cudaSetDevice(device);
  ncclReduce(devices_data_[device], devices_data_[device], kDataSize,
             ncclDouble, ncclSum, kRootDevice, communicators_[device],
             streams_[device]);
}
ncclGroupEnd();

for (int device = 0; device < kNumLocalDevices; device++) {
  cudaSetDevice(device);
  cudaStreamSynchronize(streams_[device]);
}

for (int device = 0; device < kNumLocalDevices; device++) {
  cudaSetDevice(device);
  cudaStreamDestroy(streams_[device]);
}
for (int device = 0; device < kNumLocalDevices; device++) {
  ncclCommDestroy(communicators_[device]);
}
for (int device = 0; device < kNumLocalDevices; device++) {
  cudaFree(devices_data_[device]);
}

Apparently ncclDevrInitOnce is getting called twice in this example for each device:

1: ncclCommInitRankFunc->initTransportsRank->ncclTeamLsa->ncclDevrInitOnce

2: groupLaunch->ncclPrepareTasksAndCollPreconnect->ncclPrepareTasks->ncclMakeSymmetricTaskList->ncclDevrInitOnce

But the second object finalize function is never called which leads to memory leak for allocated lsaRankList memory

I guess this happens because in init.cc:1482 bigSize is zeroed

Not sure yet how to fix this

[xiaofanl-nvidia]

++ @kiskra-nvidia @kgioioso to take a look. Feel free to mirror the PR if it looks good.