Skip to content

Symmetric Memory Pool #2037

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wdykas wants to merge 9 commits into from
Closed

Symmetric Memory Pool #2037

Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
65e749c
initial commit
wdykas Aug 6, 2025
f1d42a0
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 6, 2025
c152c51
refactor
wdykas Aug 7, 2025
93222f4
refactor
wdykas Aug 7, 2025
fd71b00
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Aug 7, 2025
054d78a
Merge branch 'main' into symmetric-mem-pool
ksivaman Aug 11, 2025
bdee29f
Merge branch 'main' into symmetric-mem-pool
ksivaman Aug 16, 2025
f160c0e
Merge branch 'main' into symmetric-mem-pool
ksivaman Mar 11, 2026
614b176
Merge branch 'main' into symmetric-mem-pool
ksivaman Oct 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
200 changes: 133 additions & 67 deletions transformer_engine/pytorch/distributed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,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

Expand Down Expand Up @@ -1845,56 +1846,129 @@ def gather_along_first_dim(
return output_tensor, handle


# Global cache to store symmetric memory tensors
symmetric_mem_cache = {}
# 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
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.

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.
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()

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

Reuses cached tensors when available to avoid redundant creation and rendezvous operations.
def __init__(self):
# Ensure __init__ is only called once
if self._initialized:
return

Note: This function always returns a 1D tensor.
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
Comment on lines +1889 to +1897

@greptile-apps greptile-apps Bot Mar 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_initialized flag never set to True

The __init__ guard checks self._initialized but never sets it to True after initialization. Since __new__ sets _initialized = False on the instance and __init__ never flips it to True, every subsequent call to SymmetricMemoryBufferPool() or get_instance() will fall through the guard and reinitialize the singleton — resetting self.pools and self.pool_sizes to empty dicts and discarding all cached symmetric memory buffers.

This directly defeats the purpose of the singleton pool. Any code path that calls SymmetricMemoryBufferPool.get_instance() (which calls cls()) will silently blow away the pool.

Suggested change
def __init__(self):
# Ensure __init__ is only called once
if self._initialized:
return
Note: This function always returns a 1D tensor.
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
def __init__(self):
# Ensure __init__ is only called once
if self._initialized:
return
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
self._initialized = True


@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,
) -> 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)

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.
# 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
)
Comment on lines +1918 to +1921

@greptile-apps greptile-apps Bot Oct 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Small reductions reserve large buffers

Every new pool allocates 256M elements before checking the requested size. Even a small all-reduce now reserves 512 MiB for BF16 or 1 GiB for FP32 per device and process group. If a training run has less spare GPU memory, its first symmetric all-reduce can fail even though the old cache allocated only the requested size.

Size the first buffer from tensor_numel. This also avoids allocating an undersized buffer immediately before growing it.


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]
# 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)
Comment on lines 1916 to +1928

@greptile-apps greptile-apps Bot Mar 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pool operations not protected by _lock

_lock is used in __new__ to guard singleton creation, but get_buffer_slice, _create_pool, and _grow_pool are all unprotected. A race between two threads calling get_buffer_slice with the same pool_key for the first time could cause _create_pool to be called twice, triggering a double rendezvous and corrupting the pool state.

While PyTorch training is typically driven from a single thread, adding a per-instance lock (separate from the class-level _lock) around the check-then-act sections would make the implementation safe by construction and easier to reason about.

Suggested change
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.
# 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
)
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]
# 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)
with self._instance_lock:
# 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]

(You would also need to add self._instance_lock = threading.Lock() in __init__.)


return msg
# Return a view/slice of the buffer
buffer = self.pools[pool_key]
return buffer[:tensor_numel]
Comment thread
greptile-apps[bot] marked this conversation as resolved.

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()


_SYMM_MEM_POOL = None
Expand Down Expand Up @@ -2013,10 +2087,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
----------
Expand Down Expand Up @@ -2079,29 +2153,21 @@ 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)
# Get a slice from the buffer pool
buffer_slice = buffer_pool.get_buffer_slice(tensor_numel, tensor_dtype, tensor_device, tp_group)

msg.copy_(inp.reshape(-1))
# Copy input to buffer slice
buffer_slice.copy_(inp.reshape(-1))

all_reduce_impl(
msg,
"sum",
group_name,
)
# Perform all-reduce on the buffer slice
all_reduce_impl(
buffer_slice,
"sum",
group_name,
)

# Copy the result back to the input tensor
inp.copy_(msg.reshape(tensor_shape))
# Copy result back to input tensor
inp.copy_(buffer_slice.reshape(tensor_shape))

return inp, None

Expand Down
Loading