Skip to content

Commit

Permalink
TL/NCCL: support by version and nb finalize
Browse files Browse the repository at this point in the history
  • Loading branch information
@shimmybalsam
shimmybalsam committed May 15, 2023
1 parent 8e05a95 commit 283fd56
Show file tree
Hide file tree
Showing 3 changed files with 68 additions and 22 deletions.
2 changes: 1 addition & 1 deletion src/components/tl/nccl/tl_nccl.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,8 @@ typedef struct ucc_tl_nccl_team {
ucc_status_t comm_state;
ncclUniqueId *unique_id;
void *oob_req;
ncclComm_t nccl_comm;
int nccl_nb_state;
ncclComm_t nccl_comm;
cudaStream_t stream;
} ucc_tl_nccl_team_t;

Expand Down
80 changes: 64 additions & 16 deletions src/components/tl/nccl/tl_nccl_team.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,14 @@
#include "coll_score/ucc_coll_score.h"
#include "utils/arch/cuda_def.h"

#define NCCL_VERSION_COMM_INIT_NONBLOCKING NCCL_VERSION(2,14,3)

enum {
NCCL_NB_UNUSED,
NCCL_NB_INIT_IN_PROGRESS,
NCCL_NB_FINALIZE_IN_PROGRESS
};

UCC_CLASS_INIT_FUNC(ucc_tl_nccl_team_t, ucc_base_context_t *tl_context,
const ucc_base_team_params_t *params)
{
Expand All @@ -23,7 +31,7 @@ UCC_CLASS_INIT_FUNC(ucc_tl_nccl_team_t, ucc_base_context_t *tl_context,

size = UCC_TL_TEAM_SIZE(self);
self->comm_state = UCC_OK;
self->nccl_nb_state = 0;
self->nccl_nb_state = NCCL_NB_UNUSED;
self->unique_id = ucc_malloc(sizeof(ncclUniqueId) * (size + 1),
"tl_nccl_unique_id");
if (!self->unique_id) {
Expand Down Expand Up @@ -58,38 +66,71 @@ UCC_CLASS_INIT_FUNC(ucc_tl_nccl_team_t, ucc_base_context_t *tl_context,
UCC_CLASS_CLEANUP_FUNC(ucc_tl_nccl_team_t)
{
tl_debug(self->super.super.context->lib, "finalizing tl team: %p", self);
if (self->nccl_comm) {
if (self->comm_state != UCC_OK) {
/* if communication error was detected ncclCommAbort should be used
since ncclCommDestroy could block */
ncclCommAbort(self->nccl_comm);
} else {
ncclCommDestroy(self->nccl_comm);
}
cudaStreamDestroy(self->stream);
}
}

UCC_CLASS_DEFINE_DELETE_FUNC(ucc_tl_nccl_team_t, ucc_base_team_t);
UCC_CLASS_DEFINE(ucc_tl_nccl_team_t, ucc_tl_team_t);

ucc_status_t ucc_tl_nccl_team_destroy(ucc_base_team_t *tl_team)
{
ucc_tl_nccl_team_t *team = ucc_derived_of(tl_team, ucc_tl_nccl_team_t);

#if NCCL_VERSION_CODE >= NCCL_VERSION_COMM_INIT_NONBLOCKING
ncclResult_t nccl_status;

if (team->nccl_nb_state == NCCL_NB_FINALIZE_IN_PROGRESS) {
goto check_finalize;
}
#endif

if (team->nccl_comm) {
if (team->comm_state != UCC_OK) {
/* if communication error was detected ncclCommAbort should be used
since ncclCommDestroy could block */
ncclCommAbort(team->nccl_comm);
} else {
#if NCCL_VERSION_CODE >= NCCL_VERSION_COMM_INIT_NONBLOCKING
ncclCommFinalize(team->nccl_comm);
check_finalize:
ncclCommGetAsyncError(team->nccl_comm, &nccl_status);
if (nccl_status == ncclInProgress) {
team->nccl_nb_state = NCCL_NB_FINALIZE_IN_PROGRESS;
return UCC_INPROGRESS;
}
if (nccl_status != ncclSuccess) {
tl_debug(tl_team->context->lib, "NCCL error %d %s", nccl_status,
ncclGetErrorString(nccl_status));
ncclCommAbort(team->nccl_comm);
return UCC_ERR_NO_MESSAGE;
} else {
ncclCommDestroy(team->nccl_comm);
}
team->nccl_nb_state = NCCL_NB_UNUSED;
#else
ncclCommDestroy(team->nccl_comm);
#endif
}
cudaStreamDestroy(team->stream);
}

UCC_CLASS_DELETE_FUNC_NAME(ucc_tl_nccl_team_t)(tl_team);
return UCC_OK;
}

ucc_status_t ucc_tl_nccl_team_create_test(ucc_base_team_t *tl_team)
{
ucc_tl_nccl_team_t *team = ucc_derived_of(tl_team, ucc_tl_nccl_team_t);
ncclConfig_t nccl_cfg = NCCL_CONFIG_INITIALIZER;
ucc_tl_nccl_team_t *team = ucc_derived_of(tl_team, ucc_tl_nccl_team_t);
ucc_status_t status;
ncclResult_t nccl_status;
ncclUniqueId errorid;

if (team->nccl_nb_state) {
#if NCCL_VERSION_CODE >= NCCL_VERSION_COMM_INIT_NONBLOCKING
ncclConfig_t nccl_cfg = NCCL_CONFIG_INITIALIZER;

if (team->nccl_nb_state == NCCL_NB_INIT_IN_PROGRESS) {
goto ncclInitStage;
}
#endif

status = UCC_TL_TEAM_OOB(team).req_test(team->oob_req);
if (status == UCC_INPROGRESS) {
Expand All @@ -114,7 +155,7 @@ ucc_status_t ucc_tl_nccl_team_create_test(ucc_base_team_t *tl_team)

CUDA_CHECK_GOTO(cudaStreamCreateWithFlags(&team->stream,
cudaStreamNonBlocking), free_unique_id, status);

#if NCCL_VERSION_CODE >= NCCL_VERSION_COMM_INIT_NONBLOCKING
nccl_cfg.blocking = 0;
nccl_status = ncclCommInitRankConfig(&team->nccl_comm,
UCC_TL_TEAM_SIZE(team),
Expand All @@ -127,9 +168,13 @@ ucc_status_t ucc_tl_nccl_team_create_test(ucc_base_team_t *tl_team)
ncclInitStage:
ncclCommGetAsyncError(team->nccl_comm, &nccl_status);
if (nccl_status == ncclInProgress){
team->nccl_nb_state = 1;
team->nccl_nb_state = NCCL_NB_INIT_IN_PROGRESS;
return UCC_INPROGRESS;
}
#else
nccl_status = ncclCommInitRank(&team->nccl_comm, UCC_TL_TEAM_SIZE(team),
team->unique_id[0], UCC_TL_TEAM_RANK(team));
#endif
if (nccl_status != ncclSuccess) {
goto free_stream;
}
Expand All @@ -141,6 +186,9 @@ ucc_status_t ucc_tl_nccl_team_create_test(ucc_base_team_t *tl_team)
tl_debug(tl_team->context->lib, "NCCL error %d %s", nccl_status,
ncclGetErrorString(nccl_status));
status = UCC_ERR_NO_MESSAGE;
#if NCCL_VERSION_CODE >= NCCL_VERSION_COMM_INIT_NONBLOCKING
ncclCommAbort(team->nccl_comm);
#endif
cudaStreamDestroy(team->stream);
free_unique_id:
ucc_free(team->unique_id);
Expand Down
8 changes: 3 additions & 5 deletions test/mpi/test_mpi.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -246,11 +246,9 @@ void UccTestMpi::destroy_team(ucc_test_team_t &team)
ucc_status_t status;

team.free_ee();
while (UCC_INPROGRESS == (status = ucc_team_destroy(team.team))) {
if (UCC_OK != status) {
std::cerr << "ucc_team_destroy failed\n";
break;
}
while (UCC_INPROGRESS == (status = ucc_team_destroy(team.team))) {}
if (UCC_OK != status) {
std::cerr << "ucc_team_destroy failed\n";
}
if (team.comm != MPI_COMM_WORLD) {
MPI_Comm_free(&team.comm);
Expand Down

0 comments on commit 283fd56

Please sign in to comment.