Add GPU memory usage tracker (#326)

This commit is contained in:
Wenkai Du
2021-03-06 20:32:30 -08:00
committed by GitHub
parent 8e180cf087
commit f60b76c67a
3 changed files with 25 additions and 1 deletions
+4 -1
View File
@@ -58,6 +58,8 @@ NCCL_PARAM(CheckPointers, "CHECK_POINTERS", 0);
ncclNet_t* ncclNet = NULL;
ncclCollNet_t* ncclCollNet = NULL;
struct allocationTracker allocTracker[MAX_ALLOC_TRACK_NGPU] = {};
// Returns ncclInternalError if anything fails, causing that network to be ignored.
ncclResult_t initNet(ncclNet_t* net) {
int ndev;
@@ -1141,7 +1143,7 @@ ncclResult_t ncclCommInitRankSync(ncclComm_t* newcomm, int nranks, ncclUniqueId
NCCLCHECKGOTO(initTransportsRank(*newcomm, &commId), res, cleanup);
NCCLCHECKGOTO(devCommSetup(*newcomm), res, cleanup);
INFO(NCCL_INIT,"comm %p rank %d nranks %d cudaDev %d busId %lx - Init COMPLETE", *newcomm, myrank, nranks, (*newcomm)->cudaDev, (*newcomm)->busId);
INFO(NCCL_INIT,"comm %p rank %d nranks %d cudaDev %d busId %lx used %ld bytes - Init COMPLETE", *newcomm, myrank, nranks, (*newcomm)->cudaDev, (*newcomm)->busId, allocTracker[(*newcomm)->cudaDev].totalAllocSize);
return ncclSuccess;
cleanup:
@@ -1161,6 +1163,7 @@ static ncclResult_t ncclCommInitRankDev(ncclComm_t* newcomm, int nranks, ncclUni
NCCLCHECKGOTO(ncclInit(), res, end);
if (myrank == 0) showVersion();
memset(allocTracker+cudaDev, 0, sizeof(struct allocationTracker));
// Make sure the CUDA runtime is initialized.
CUDACHECKGOTO(hipFree(NULL), res, end);