Repository navigation
Conversation
Signed-off-by: Peter Dykas <wdykas@nvidia.com>
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
|
/te-ci pytorch |
|
/te-ci pytorch |
 [High impact] Replaces symmetric memory cache with a pooling system. The PR needs changes before merging because small reductions can now exhaust GPU memory, and repeated singleton access still clears cached buffers. Findings
SummaryReplaces the size-keyed symmetric-memory cache with reusable buffers grouped by data type, device, and process group.
Diagram%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Symmetric all-reduce] --> B{Pool exists?}
B -- No --> C[Allocate 256M elements]
B -- Yes --> D{Input fits?}
C --> D
D -- No --> E[Allocate larger buffer]
D -- Yes --> F[Take prefix view]
E --> F
F --> G[Copy input into buffer]
G --> H[Run all-reduce]
H --> I[Copy result back]
Reviews (2) · Last reviewed commit: "Merge branch 'main' into symmetric-mem-p..." · Reviewed by Greptile |
| 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 |
There was a problem hiding this comment.
_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.
| 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 |
|
|
||
| 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) |
There was a problem hiding this comment.
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.
| 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__.)
|
@Sunt-ing oh wow that PR got list in the fold at some point. Lets just go with yours mine is like over a year old |
| if pool_key not in self.pools: | ||
| self._create_pool( | ||
| pool_key, self.initial_pool_size, tensor_dtype, tensor_device, tp_group | ||
| ) |
There was a problem hiding this comment.
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.
Description
This is a more memory efficient version for using symmetric memory all reduces. We use a pool of symmetric memory that we grow if we need it to be larger and slice into the pool.
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: