[contrib/nccl_allocator] Add an interface for the NCCL symmetric registration (#1915)

youngeunkwon0405 · crcrpar · web-flow · commit 2eda0ac9ed59 · 2025-07-12T23:28:40.000+09:00
* add an interface for symmetric reg

Signed-off-by: Youngeun Kwon &lt;youngeunk@nvidia.com&gt;

* update

Signed-off-by: Youngeun Kwon &lt;youngeunk@nvidia.com&gt;

* Update apex/contrib/nccl_allocator/nccl_allocator.py

Co-authored-by: Masaki Kozuki &lt;mkozuki@nvidia.com&gt;

---------

Signed-off-by: Youngeun Kwon &lt;youngeunk@nvidia.com&gt;
Co-authored-by: Masaki Kozuki &lt;mkozuki@nvidia.com&gt;
diff --git a/apex/contrib/nccl_allocator/nccl_allocator.py b/apex/contrib/nccl_allocator/nccl_allocator.py
@@ -7,10 +7,19 @@
 
 __all__ = ["init", "nccl_mem", "create_nccl_mem_pool"]
 
+def get_func_args(func):
+    import inspect
+    sig = inspect.signature(func)
+    return [arg.name for arg in sig.parameters.values()]
 
-def create_nccl_mem_pool():
+def create_nccl_mem_pool(symmetric: bool | None = None) -> torch.cuda.MemPool:
     _allocator = _apex_nccl_allocator.get_nccl_allocator()
-    _pool = torch.cuda.MemPool(_allocator)
+    if symmetric is None:
+        _pool = torch.cuda.MemPool(_allocator)
+    else:
+        assert 'symmetric' in get_func_args(torch.cuda.MemPool), \
+            "symmetric setting with torch.cuda.MemPool requires higher PyTorch version"
+        _pool = torch.cuda.MemPool(_allocator, symmetric=symmetric)
     return _pool
 
 
