File tree Expand file tree Collapse file tree 1 file changed +2
-0
lines changed
Expand file tree Collapse file tree 1 file changed +2
-0
lines changed Original file line number Diff line number Diff line change @@ -167,6 +167,7 @@ class Allocator<AllocatorType::CUDA>: public IAllocator {
167167 check_cuda_error (cudaGetDeviceCount (&device_count));
168168 cudaMemPool_t mempool;
169169 check_cuda_error (cudaDeviceGetDefaultMemPool (&mempool, device_id));
170+ #if TM_ENABLE_CUSTOM_ALL_REDUCE
170171 cudaMemAccessDesc desc = {};
171172 int peer_access_available = 0 ;
172173 for (int i = 0 ; i < device_count; i++) {
@@ -184,6 +185,7 @@ class Allocator<AllocatorType::CUDA>: public IAllocator {
184185 desc.flags = cudaMemAccessFlagsProtReadWrite;
185186 check_cuda_error (cudaMemPoolSetAccess (mempool, &desc, 1 ));
186187 }
188+ #endif
187189 // set memory pool threshold to avoid shrinking the pool
188190 uint64_t setVal = UINT64_MAX;
189191 check_cuda_error (cudaMemPoolSetAttribute (mempool, cudaMemPoolAttrReleaseThreshold, &setVal));
You can’t perform that action at this time.
0 commit comments