minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add clear error message when OOM with symmetric memory (#17038)

Browse Source
This commit is contained in:
Nicolas Castet
2026-01-16 20:44:37 +08:00
committed by GitHub
parent daea51385d
commit 9d8bbd4223
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/pynccl_allocator.py
+12 -2
View File
@@ -34,10 +34,20 @@ ncclResult_t ncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, n
ncclResult_t ncclMemAlloc(void** ptr, size_t size); ncclResult_t ncclMemAlloc(void** ptr, size_t size);
ncclResult_t ncclMemFree(void *ptr); ncclResult_t ncclMemFree(void *ptr);
const char* ncclGetErrorString(ncclResult_t result);
#define NCCLCHECK(cmd) do { \
ncclResult_t res = cmd; \
if (res != ncclSuccess) { \
fprintf(stderr, "ERROR: NCCL symmetric memory allocation failed. Most likely out of device memory. '%s'\\n", \
ncclGetErrorString(res)); \
return NULL; \
} \
} while(0)
void* nccl_alloc_plug(size_t size, int device, void* stream) { void* nccl_alloc_plug(size_t size, int device, void* stream) {
void* ptr; void* ptr;
ncclResult_t err = ncclMemAlloc(&ptr, size); NCCLCHECK(ncclMemAlloc(&ptr, size));
const char *str_val = getenv("SGLANG_TMP_NCCL_COMM_VALUE"); const char *str_val = getenv("SGLANG_TMP_NCCL_COMM_VALUE");
char *endptr; char *endptr;
@@ -45,7 +55,7 @@ void* nccl_alloc_plug(size_t size, int device, void* stream) {
ncclComm_t comm = (ncclComm_t)(int_val); ncclComm_t comm = (ncclComm_t)(int_val);
ncclWindow_t win; ncclWindow_t win;
ncclResult_t err2 = ncclCommWindowRegister(comm, ptr, size, &win, NCCL_WIN_COLL_SYMMETRIC); NCCLCHECK(ncclCommWindowRegister(comm, ptr, size, &win, NCCL_WIN_COLL_SYMMETRIC));
return ptr; return ptr;
} }