fix

xinhaoc · Jan 26, 2024 · a31f8e9 · a31f8e9
1 parent 01c9d4c
commit a31f8e9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/runtime/optimizer_kernel.cpp b/src/runtime/optimizer_kernel.cpp
@@ -258,7 +258,7 @@ __host__ void AdamOptimizer::nccl_unified_update_task_gpu(
 
   hipStream_t stream;
   checkCUDA(get_legion_stream(&stream));
-  assert(op->reservedWorkSpaceSize < meta->handle.workSpaceSize);
+  // assert(op->reservedWorkSpaceSize < meta->handle.workSpaceSize);
 
   void *workSpace_ptr = meta->handle.workSpace;
 
@@ -275,7 +275,7 @@ __host__ void AdamOptimizer::nccl_unified_update_task_gpu(
   // do allreduce once
   checkNCCL(ncclAllReduce(meta->handle.workSpace,
                           (float *)meta->handle.workSpace,
-                          op->reservedWorkSpaceSize / sizeof(float),
+                          meta->handle.workSpaceSize,
                           ncclFloat,
                           ncclSum,
                           meta->handle.ncclComm,

diff --git a/src/runtime/optimizer_kernel.cu b/src/runtime/optimizer_kernel.cu
@@ -228,7 +228,7 @@ __host__ void AdamOptimizer::nccl_unified_update_task_gpu(
 
   cudaStream_t stream;
   checkCUDA(get_legion_stream(&stream));
-  assert(op->reservedWorkSpaceSize < meta->handle.workSpaceSize);
+  // assert(op->reservedWorkSpaceSize < meta->handle.workSpaceSize);
 
   void *workSpace_ptr = meta->handle.workSpace;
 
@@ -245,7 +245,7 @@ __host__ void AdamOptimizer::nccl_unified_update_task_gpu(
   // do allreduce once
   checkNCCL(ncclAllReduce(meta->handle.workSpace,
                           (float *)meta->handle.workSpace,
-                          op->reservedWorkSpaceSize / sizeof(float),
+                          meta->handle.workSpaceSize,
                           ncclFloat,
                           ncclSum,
                           meta->handle.ncclComm,