[nccl_allocator] Adds helper API to create pool (#1877)

syed-ahmed · web-flow · commit 312acb44f9fe · 2025-02-26T17:58:26.000+09:00
diff --git a/apex/contrib/examples/nccl_allocator/allreduce.py b/apex/contrib/examples/nccl_allocator/allreduce.py
@@ -13,8 +13,8 @@
 
 torch.cuda.set_device(local_rank)
 dist.init_process_group(backend="nccl")
-
-with nccl_allocator.nccl_mem():
+pool = nccl_allocator.create_nccl_mem_pool()
+with nccl_allocator.nccl_mem(pool):
     a = torch.ones(1024 * 1024 * 2, device="cuda")
 dist.all_reduce(a)
 
diff --git a/apex/contrib/examples/nccl_allocator/cache.py b/apex/contrib/examples/nccl_allocator/cache.py
@@ -21,7 +21,8 @@ def print_used_mem(string, nvsmi, device_id = 0):
 
 print_used_mem("", nvsmi)
 
-with nccl_allocator.nccl_mem():
+pool = nccl_allocator.create_nccl_mem_pool()
+with nccl_allocator.nccl_mem(pool):
     for i in range(nrep):
       out = torch.randn(1024 * 1024 * 100).cuda() # >= 400 MB
       nccl_mem.append(out)
@@ -42,7 +43,7 @@ def print_used_mem(string, nvsmi, device_id = 0):
 
 del nccl_mem
 nccl_mem2 = []
-with nccl_allocator.nccl_mem():
+with nccl_allocator.nccl_mem(pool):
     for i in range(nrep):
       out = torch.randn(1024 * 1024 * 100).cuda() # >= 400 MB
       nccl_mem2.append(out)
diff --git a/apex/contrib/examples/nccl_allocator/change_cuda_allocator.py b/apex/contrib/examples/nccl_allocator/change_cuda_allocator.py
@@ -3,7 +3,8 @@
 
 nccl_allocator.init()
 nrep = 6
-with nccl_allocator.nccl_mem():
+pool = nccl_allocator.create_nccl_mem_pool()
+with nccl_allocator.nccl_mem(pool):
     for i in range(nrep):
       out = torch.randn(1024).cuda()
 
diff --git a/apex/contrib/examples/nccl_allocator/toy_ddp.py b/apex/contrib/examples/nccl_allocator/toy_ddp.py
@@ -36,7 +36,8 @@ def forward(self, x):
 optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
 
 data_ptrs = []
-with nccl_allocator.nccl_mem():
+pool = nccl_allocator.create_nccl_mem_pool()
+with nccl_allocator.nccl_mem(pool):
     for param in ddp_model.parameters():
         param.grad = torch.empty_like(param)
         data_ptrs.append(param.grad.data_ptr())
diff --git a/apex/contrib/nccl_allocator/nccl_allocator.py b/apex/contrib/nccl_allocator/nccl_allocator.py
@@ -5,11 +5,13 @@
 from contextlib import nullcontext
 
 
-__all__ = ["init", "nccl_mem"]
+__all__ = ["init", "nccl_mem", "create_nccl_mem_pool"]
 
 
-_allocator = _apex_nccl_allocator.get_nccl_allocator()
-_pool = torch.cuda.MemPool(_allocator)
+def create_nccl_mem_pool():
+    _allocator = _apex_nccl_allocator.get_nccl_allocator()
+    _pool = torch.cuda.MemPool(_allocator)
+    return _pool
 
 
 def init() -> None:
@@ -18,10 +20,11 @@ def init() -> None:
 
 
 class nccl_mem:
-    def __init__(self, enabled = True, device = None, group = None):
+    def __init__(self, pool, enabled = True, device = None, group = None):
         self.device = None
         self.group = None
         self.mem_context = None
+        self.pool = pool
 
         if enabled:
             if device is None:
@@ -37,7 +40,7 @@ def __init__(self, enabled = True, device = None, group = None):
             else:
                 self.group = group
 
-            self.mem_context = torch.cuda.use_mem_pool(_pool)
+            self.mem_context = torch.cuda.use_mem_pool(self.pool)
         else:
             self.mem_context = nullcontext()
 
@@ -46,15 +49,15 @@ def __enter__(self):
         if self.group is not None:
             backend = self.group._get_backend(self.device)
             try:
-                backend.deregister_mem_pool(_pool)
+                backend.deregister_mem_pool(self.pool)
             except RuntimeError:
                 pass
 
     def __exit__(self, *args):
         if self.group is not None:
             backend = self.group._get_backend(self.device)
             try:
-                backend.register_mem_pool(_pool)
+                backend.register_mem_pool(self.pool)
             except RuntimeError:
                 pass
         self.mem_context.__exit__(*args)
diff --git a/apex/contrib/optimizers/distributed_fused_adam.py b/apex/contrib/optimizers/distributed_fused_adam.py
@@ -1191,12 +1191,13 @@ def _init_grad_buffer(self) -> None:
                     [buffer_size], dtype=grad_sync_dtype, device=self.device,
                 )
             else:
-                with nccl_allocator.nccl_mem():
+                pool = nccl_allocator.create_nccl_mem_pool()
+                with nccl_allocator.nccl_mem(pool):
                     self._grad_buffers[dtypes] = torch.zeros(
                         [buffer_size], dtype=grad_sync_dtype, device=self.device,
                     )
                 shard_buffer_size = buffer_size // self.distributed_size
-                with nccl_allocator.nccl_mem():
+                with nccl_allocator.nccl_mem(pool):
                     self._shard_grad_buffers[dtypes] = torch.zeros(
                         [shard_buffer_size], dtype=grad_sync_dtype, device=self.device,
                     )

Original file line number	Diff line number	Diff line change
`@@ -1191,12 +1191,13 @@ def _init_grad_buffer(self) -> None:`
`1191`	`1191`	`[buffer_size], dtype=grad_sync_dtype, device=self.device,`
`1192`	`1192`	`)`
`1193`	`1193`	`else:`
`1194`		`- with nccl_allocator.nccl_mem():`
	`1194`	`+ pool = nccl_allocator.create_nccl_mem_pool()`
	`1195`	`+ with nccl_allocator.nccl_mem(pool):`
`1195`	`1196`	`self._grad_buffers[dtypes] = torch.zeros(`
`1196`	`1197`	`[buffer_size], dtype=grad_sync_dtype, device=self.device,`
`1197`	`1198`	`)`
`1198`	`1199`	`shard_buffer_size = buffer_size // self.distributed_size`
`1199`		`- with nccl_allocator.nccl_mem():`
	`1200`	`+ with nccl_allocator.nccl_mem(pool):`
`1200`	`1201`	`self._shard_grad_buffers[dtypes] = torch.zeros(`
`1201`	`1202`	`[shard_buffer_size], dtype=grad_sync_dtype, device=self.device,`
`1202`	`1203`	`)`