From 65e749c523edb86501e5a3e598f1d70ef66f628a Mon Sep 17 00:00:00 2001 From: Peter Dykas Date: Wed, 6 Aug 2025 13:11:05 -0700 Subject: [PATCH 1/4] initial commit Signed-off-by: Peter Dykas --- transformer_engine/pytorch/distributed.py | 184 +++++++++++++--------- 1 file changed, 107 insertions(+), 77 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 868fc3a27ae..d7e555dbebc 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1453,57 +1453,94 @@ def gather_along_first_dim( return out, handle -# Global cache to store symmetric memory tensors -symmetric_mem_cache = {} +# Global buffer pool for symmetric memory +class SymmetricMemoryBufferPool: + _instance: Optional['SymmetricMemoryBufferPool'] = None + _lock = threading.Lock() + + def __new__(cls): + if cls._instance is None: + with cls._lock: + # Double-checked locking pattern + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._initialized = False + return cls._instance + def __init__(self): + # Ensure __init__ is only called once + if self._initialized: + return -def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_group, tag=None): - """ - Gets or creates a symmetric memory tensor with specified properties. - - Reuses cached tensors when available to avoid redundant creation and rendezvous operations. - - Note: This function always returns a 1D tensor. - - Parameters - ---------- - tensor_numel : int - Number of elements in the tensor. - tensor_dtype : torch.dtype - Data type of the tensor. - tensor_device : torch.device - Device on which to allocate the tensor. - tp_group : dist_group_type - Process group for rendezvous operation. - tag : Any, optional - Optional identifier to further distinguish tensors. - - Returns - ------- - torch.Tensor - A symmetric memory tensor with the specified properties. - """ - # Create a cache key based on tensor properties and group - cache_key = (tensor_numel, tensor_dtype, tensor_device, tp_group.group_name, tag) - - # Check if we already have a symmetric memory tensor for this configuration - if cache_key not in symmetric_mem_cache: - # Create a new symmetric memory tensor if not in cache - msg = symm_mem.empty( - tensor_numel, - dtype=tensor_dtype, - device=tensor_device, - ) - # Perform the rendezvous once for this tensor - symm_mem.rendezvous(msg, group=tp_group) - # Store in cache - symmetric_mem_cache[cache_key] = msg - else: - # Reuse the existing symmetric memory tensor - msg = symmetric_mem_cache[cache_key] - - return msg - + self.pools: Dict[Tuple, torch.Tensor] = {} # (dtype, device, group_name) -> buffer + self.pool_sizes: Dict[Tuple, int] = {} # Track allocated sizes + self.initial_pool_size = 1024 * 1024 * 256 # 256M elements initially + self.growth_factor = 2.0 + + @classmethod + def get_instance(cls) -> 'SymmetricMemoryBufferPool': + """Alternative way to get the singleton instance.""" + return cls() + + def get_buffer_slice(self, tensor_numel: int, tensor_dtype: torch.dtype, + tensor_device: torch.device, tp_group, tag=None) -> torch.Tensor: + """ + Get a slice from the buffer pool that can hold tensor_numel elements. + Grows the pool if necessary. + """ + pool_key = (tensor_dtype, tensor_device, tp_group.group_name) + + # Initialize pool if it doesn't exist + if pool_key not in self.pools: + self._create_pool(pool_key, self.initial_pool_size, tensor_dtype, + tensor_device, tp_group) + + # Check if we need to grow the pool + current_size = self.pool_sizes[pool_key] + if tensor_numel > current_size: + # Grow to at least accommodate the requested size + new_size = max(int(current_size * self.growth_factor), tensor_numel) + self._grow_pool(pool_key, new_size, tensor_dtype, tensor_device, tp_group) + + # Return a view/slice of the buffer + buffer = self.pools[pool_key] + return buffer[:tensor_numel] + + def _create_pool(self, pool_key: Tuple, size: int, dtype: torch.dtype, + device: torch.device, tp_group): + """Create a new buffer pool.""" + buffer = symm_mem.empty(size, dtype=dtype, device=device) + symm_mem.rendezvous(buffer, group=tp_group) + self.pools[pool_key] = buffer + self.pool_sizes[pool_key] = size + + def _grow_pool(self, pool_key: Tuple, new_size: int, dtype: torch.dtype, + device: torch.device, tp_group): + """Grow an existing buffer pool.""" + old_buffer = self.pools[pool_key] + + # Create new larger buffer + new_buffer = symm_mem.empty(new_size, dtype=dtype, device=device) + symm_mem.rendezvous(new_buffer, group=tp_group) + + # Replace the old buffer + self.pools[pool_key] = new_buffer + self.pool_sizes[pool_key] = new_size + + # The old buffer will be garbage collected + del old_buffer + + def get_memory_usage(self) -> Dict[str, float]: + """Get memory usage statistics for all pools.""" + stats = {} + for key, buffer in self.pools.items(): + dtype, device, group = key + size_mb = (buffer.numel() * buffer.element_size()) / (1024 * 1024) + stats[f"{dtype}_{device}_{group}"] = size_mb + return stats + +# Global buffer pool instance +buffer_pool = SymmetricMemoryBufferPool() def symmetric_all_reduce( inp: torch.Tensor, @@ -1513,10 +1550,10 @@ def symmetric_all_reduce( ): """ Performs an all-reduce operation across multiple processes using symmetric memory. - If the input tensor is already in the symmetric memory cache we can avoid copy - overheads by just directly using the input tensor for all reduce. Externally - created symmetric memory tensors not in the cache currently will not be able to - avoid the extra copies. + Currently we do a copy into the symmetric memory buffer in the future we would like + to be able to recognize if a tensor is already in symmetric memory so that we can + call all-reduce with no copy. Pytorch does not have functionality currently to + recognize if a tensor is already symmetric memory. Parameters ---------- @@ -1572,30 +1609,23 @@ def symmetric_all_reduce( tensor_dtype = inp.dtype tensor_device = inp.device - input_id = id(inp) - is_cached = any(id(cached_tensor) == input_id for cached_tensor in symmetric_mem_cache.values()) - # Check if the input tensor is already in the symmetric memory cache. If it is we can avoid copy overheads. - if is_cached: - all_reduce_impl( - inp, - "sum", - group_name, - ) - else: - # Get symmetric memory tensor. Build or retrieve from cache. - msg = get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_group) - - msg.copy_(inp.reshape(-1)) - - all_reduce_impl( - msg, - "sum", - group_name, - ) - - # Copy the result back to the input tensor - inp.copy_(msg.reshape(tensor_shape)) - + # Get a slice from the buffer pool + buffer_slice = buffer_pool.get_buffer_slice(tensor_numel, tensor_dtype, + tensor_device, tp_group) + + # Copy input to buffer slice + buffer_slice.copy_(inp.reshape(-1)) + + # Perform all-reduce on the buffer slice + all_reduce_impl( + buffer_slice, + "sum", + group_name, + ) + + # Copy result back to input tensor + inp.copy_(buffer_slice.reshape(tensor_shape)) + return inp, None From f1d42a0f2b5ef420a5959fa8bc0845f00fba068b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 6 Aug 2025 20:13:21 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/distributed.py | 66 +++++++++++++---------- 1 file changed, 38 insertions(+), 28 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d7e555dbebc..8f8d40ce5e0 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1455,9 +1455,9 @@ def gather_along_first_dim( # Global buffer pool for symmetric memory class SymmetricMemoryBufferPool: - _instance: Optional['SymmetricMemoryBufferPool'] = None + _instance: Optional["SymmetricMemoryBufferPool"] = None _lock = threading.Lock() - + def __new__(cls): if cls._instance is None: with cls._lock: @@ -1478,58 +1478,67 @@ def __init__(self): self.growth_factor = 2.0 @classmethod - def get_instance(cls) -> 'SymmetricMemoryBufferPool': + def get_instance(cls) -> "SymmetricMemoryBufferPool": """Alternative way to get the singleton instance.""" return cls() - - def get_buffer_slice(self, tensor_numel: int, tensor_dtype: torch.dtype, - tensor_device: torch.device, tp_group, tag=None) -> torch.Tensor: + + def get_buffer_slice( + self, + tensor_numel: int, + tensor_dtype: torch.dtype, + tensor_device: torch.device, + tp_group, + tag=None, + ) -> torch.Tensor: """ Get a slice from the buffer pool that can hold tensor_numel elements. Grows the pool if necessary. """ pool_key = (tensor_dtype, tensor_device, tp_group.group_name) - + # Initialize pool if it doesn't exist if pool_key not in self.pools: - self._create_pool(pool_key, self.initial_pool_size, tensor_dtype, - tensor_device, tp_group) - + self._create_pool( + pool_key, self.initial_pool_size, tensor_dtype, tensor_device, tp_group + ) + # Check if we need to grow the pool current_size = self.pool_sizes[pool_key] if tensor_numel > current_size: # Grow to at least accommodate the requested size new_size = max(int(current_size * self.growth_factor), tensor_numel) self._grow_pool(pool_key, new_size, tensor_dtype, tensor_device, tp_group) - + # Return a view/slice of the buffer buffer = self.pools[pool_key] return buffer[:tensor_numel] - - def _create_pool(self, pool_key: Tuple, size: int, dtype: torch.dtype, - device: torch.device, tp_group): + + def _create_pool( + self, pool_key: Tuple, size: int, dtype: torch.dtype, device: torch.device, tp_group + ): """Create a new buffer pool.""" buffer = symm_mem.empty(size, dtype=dtype, device=device) symm_mem.rendezvous(buffer, group=tp_group) self.pools[pool_key] = buffer self.pool_sizes[pool_key] = size - - def _grow_pool(self, pool_key: Tuple, new_size: int, dtype: torch.dtype, - device: torch.device, tp_group): + + def _grow_pool( + self, pool_key: Tuple, new_size: int, dtype: torch.dtype, device: torch.device, tp_group + ): """Grow an existing buffer pool.""" old_buffer = self.pools[pool_key] - + # Create new larger buffer new_buffer = symm_mem.empty(new_size, dtype=dtype, device=device) symm_mem.rendezvous(new_buffer, group=tp_group) - + # Replace the old buffer self.pools[pool_key] = new_buffer self.pool_sizes[pool_key] = new_size - + # The old buffer will be garbage collected del old_buffer - + def get_memory_usage(self) -> Dict[str, float]: """Get memory usage statistics for all pools.""" stats = {} @@ -1539,9 +1548,11 @@ def get_memory_usage(self) -> Dict[str, float]: stats[f"{dtype}_{device}_{group}"] = size_mb return stats + # Global buffer pool instance buffer_pool = SymmetricMemoryBufferPool() + def symmetric_all_reduce( inp: torch.Tensor, tp_group: Optional[dist_group_type] = None, @@ -1550,7 +1561,7 @@ def symmetric_all_reduce( ): """ Performs an all-reduce operation across multiple processes using symmetric memory. - Currently we do a copy into the symmetric memory buffer in the future we would like + Currently we do a copy into the symmetric memory buffer in the future we would like to be able to recognize if a tensor is already in symmetric memory so that we can call all-reduce with no copy. Pytorch does not have functionality currently to recognize if a tensor is already symmetric memory. @@ -1610,22 +1621,21 @@ def symmetric_all_reduce( tensor_device = inp.device # Get a slice from the buffer pool - buffer_slice = buffer_pool.get_buffer_slice(tensor_numel, tensor_dtype, - tensor_device, tp_group) - + buffer_slice = buffer_pool.get_buffer_slice(tensor_numel, tensor_dtype, tensor_device, tp_group) + # Copy input to buffer slice buffer_slice.copy_(inp.reshape(-1)) - + # Perform all-reduce on the buffer slice all_reduce_impl( buffer_slice, "sum", group_name, ) - + # Copy result back to input tensor inp.copy_(buffer_slice.reshape(tensor_shape)) - + return inp, None From c152c510bf9fd6818b9b921e45b72a35ae6294b6 Mon Sep 17 00:00:00 2001 From: Peter Dykas Date: Thu, 7 Aug 2025 06:39:46 -0700 Subject: [PATCH 3/4] refactor Signed-off-by: Peter Dykas --- transformer_engine/pytorch/distributed.py | 42 ++++++++++++----------- 1 file changed, 22 insertions(+), 20 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d7e555dbebc..24d76081f2b 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -10,6 +10,7 @@ from functools import lru_cache from dataclasses import dataclass import math +import threading from typing import Any, Callable, Dict, List, Optional, Tuple, Union import warnings @@ -1455,9 +1456,10 @@ def gather_along_first_dim( # Global buffer pool for symmetric memory class SymmetricMemoryBufferPool: + """A memory pool of Pytorch Symmetric memory to be used by symmetric all-reduce""" _instance: Optional['SymmetricMemoryBufferPool'] = None _lock = threading.Lock() - + def __new__(cls): if cls._instance is None: with cls._lock: @@ -1481,55 +1483,55 @@ def __init__(self): def get_instance(cls) -> 'SymmetricMemoryBufferPool': """Alternative way to get the singleton instance.""" return cls() - - def get_buffer_slice(self, tensor_numel: int, tensor_dtype: torch.dtype, - tensor_device: torch.device, tp_group, tag=None) -> torch.Tensor: + + def get_buffer_slice(self, tensor_numel: int, tensor_dtype: torch.dtype, + tensor_device: torch.device, tp_group) -> torch.Tensor: """ Get a slice from the buffer pool that can hold tensor_numel elements. Grows the pool if necessary. """ pool_key = (tensor_dtype, tensor_device, tp_group.group_name) - + # Initialize pool if it doesn't exist if pool_key not in self.pools: - self._create_pool(pool_key, self.initial_pool_size, tensor_dtype, + self._create_pool(pool_key, self.initial_pool_size, tensor_dtype, tensor_device, tp_group) - + # Check if we need to grow the pool current_size = self.pool_sizes[pool_key] if tensor_numel > current_size: # Grow to at least accommodate the requested size new_size = max(int(current_size * self.growth_factor), tensor_numel) self._grow_pool(pool_key, new_size, tensor_dtype, tensor_device, tp_group) - + # Return a view/slice of the buffer buffer = self.pools[pool_key] return buffer[:tensor_numel] - - def _create_pool(self, pool_key: Tuple, size: int, dtype: torch.dtype, + + def _create_pool(self, pool_key: Tuple, size: int, dtype: torch.dtype, device: torch.device, tp_group): """Create a new buffer pool.""" buffer = symm_mem.empty(size, dtype=dtype, device=device) symm_mem.rendezvous(buffer, group=tp_group) self.pools[pool_key] = buffer self.pool_sizes[pool_key] = size - + def _grow_pool(self, pool_key: Tuple, new_size: int, dtype: torch.dtype, device: torch.device, tp_group): """Grow an existing buffer pool.""" old_buffer = self.pools[pool_key] - + # Create new larger buffer new_buffer = symm_mem.empty(new_size, dtype=dtype, device=device) symm_mem.rendezvous(new_buffer, group=tp_group) - + # Replace the old buffer self.pools[pool_key] = new_buffer self.pool_sizes[pool_key] = new_size - + # The old buffer will be garbage collected del old_buffer - + def get_memory_usage(self) -> Dict[str, float]: """Get memory usage statistics for all pools.""" stats = {} @@ -1610,22 +1612,22 @@ def symmetric_all_reduce( tensor_device = inp.device # Get a slice from the buffer pool - buffer_slice = buffer_pool.get_buffer_slice(tensor_numel, tensor_dtype, + buffer_slice = buffer_pool.get_buffer_slice(tensor_numel, tensor_dtype, tensor_device, tp_group) - + # Copy input to buffer slice buffer_slice.copy_(inp.reshape(-1)) - + # Perform all-reduce on the buffer slice all_reduce_impl( buffer_slice, "sum", group_name, ) - + # Copy result back to input tensor inp.copy_(buffer_slice.reshape(tensor_shape)) - + return inp, None From fd71b00bd71faf1b2f154c683237ac16bc9e37a9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 7 Aug 2025 13:49:50 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/distributed.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 827d1ceade0..c22f0e0859c 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1457,30 +1457,31 @@ def gather_along_first_dim( # Global buffer pool for symmetric memory class SymmetricMemoryBufferPool: """This class manages a pool of reusable symmetric memory buffers for symmetric - all-reduce operations. It implements automatic buffer growth and reuse to + all-reduce operations. It implements automatic buffer growth and reuse to minimize memory allocation overhead during distributed communication operations. - + The pool maintains separate buffers for different combinations of data type, device, and process group, automatically growing buffers when larger sizes are requested. - + Attributes ---------- pools : Dict[Tuple, torch.Tensor] Dictionary mapping (dtype, device, group_name) tuples to allocated buffers. Each buffer is a contiguous symmetric memory tensor that can be sliced for individual operations. - + pool_sizes : Dict[Tuple, int] Tracks the allocated size (in elements) for each buffer pool. - + initial_pool_size : int Initial number of elements to allocate for new buffer pools. Default is 256M elements (1024 * 1024 * 256). - + growth_factor : float Multiplicative factor for growing buffers when current size is insufficient. Default is 2.0 (doubles the buffer size on each growth).""" + _instance: Optional["SymmetricMemoryBufferPool"] = None _lock = threading.Lock()