Skip to content

Revert D113962248: Optimizer scratch-buffer stashing in MemoryStashingManager (#4495) #4495

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

Open
doIIarplus wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Revert D113962248: Optimizer scratch-buffer stashing in MemoryStashingManager (#4495) #4495

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
170 changes: 15 additions & 155 deletions torchrec/distributed/memory_stashing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,17 +9,7 @@

import logging
from concurrent.futures import Future, ThreadPoolExecutor
from typing import (
Any,
Callable,
Dict,
List,
Optional,
Protocol,
runtime_checkable,
Tuple,
Union,
)
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
Expand All @@ -42,11 +32,6 @@
_STASH_CHUNK_SIZE_BYTES: int = 32 * 1024**2


@runtime_checkable
class _ScratchBufferOptimizer(Protocol):
def scratch_buffers(self) -> tuple[torch.Tensor, ...]: ...


def _tensor_size_text(tensor: Union[List[torch.Tensor], torch.Tensor]) -> str:
"""Return a human-readable size string (KB / MB / GB) for a tensor.

Expand Down Expand Up @@ -167,7 +152,6 @@ class MemoryStashingManager:
_device_to_host_stream: Optional[torch.cuda.Stream] = None
_embedding_weight_restore_callbacks: List[Callable[..., None]] = []
_optimizer_state_restore_callbacks: List[Callable[..., None]] = []
_optimizer_scratch_buffer_restore_callbacks: List[Callable[..., None]] = []
_emo_cache_restore_callbacks: List[Callable[..., None]] = []
_stash_executor: Optional[ThreadPoolExecutor] = None
_delay_stash: bool = False
Expand Down Expand Up @@ -302,7 +286,6 @@ def reset(cls) -> None:
cls._device_to_host_stream = None
cls._embedding_weight_restore_callbacks.clear()
cls._optimizer_state_restore_callbacks.clear()
cls._optimizer_scratch_buffer_restore_callbacks.clear()
cls._emo_cache_restore_callbacks.clear()
cls._delay_stash = False
cls._stash_chunk_size_bytes = _STASH_CHUNK_SIZE_BYTES
Expand Down Expand Up @@ -376,23 +359,14 @@ def restore_optimizer_state(
cls,
_grad: Optional[torch.Tensor] = None,
sync_event: Optional[torch.cuda.Event] = None,
restore_scratch_buffer: bool = True,
) -> None:
"""Restore copied optimizer state and, optionally, disposable scratch buffer."""
"""Pop and call all optimizer state restore callbacks in reverse order."""
cls._log_ems_once(
"ems_restore_optimizer_state_callbacks",
{
"num_callbacks": str(len(cls._optimizer_state_restore_callbacks)),
"num_scratch_buffer_callbacks": str(
len(cls._optimizer_scratch_buffer_restore_callbacks)
),
},
{"num_callbacks": str(len(cls._optimizer_state_restore_callbacks))},
)
while cls._optimizer_state_restore_callbacks:
cls._optimizer_state_restore_callbacks.pop()(None, sync_event)
if restore_scratch_buffer:
while cls._optimizer_scratch_buffer_restore_callbacks:
cls._optimizer_scratch_buffer_restore_callbacks.pop()(None, sync_event)

@classmethod
def restore_optimizer_state_next(
Expand All @@ -416,18 +390,6 @@ def restore_optimizer_state_next(
)
cls._optimizer_state_restore_callbacks.pop(0)(None, sync_event)

@staticmethod
def _consume_restore_callback(
callback: Callable[..., None],
callbacks: List[Callable[..., None]],
sync_event: Optional[torch.cuda.Event],
) -> None:
try:
callbacks.remove(callback)
except ValueError:
return
callback(None, sync_event)

@classmethod
def _stash_tensors(
cls,
Expand All @@ -437,7 +399,7 @@ def _stash_tensors(
delay: bool = False,
) -> Tuple[
Callable[[Optional[torch.Tensor]], None],
Callable[..., None],
Callable[[Optional[torch.Tensor]], None],
Callable[[Optional[torch.Tensor]], None],
]:
"""
Expand Down Expand Up @@ -1002,77 +964,6 @@ def await_restore(_grad: Optional[torch.Tensor] = None) -> None:
cls._emo_cache_restore_callbacks.append(restore)
return await_restore, restore

@classmethod
def _release_optimizer_scratch_buffers(
cls,
optimizer: torch.optim.Optimizer,
sync_event: Optional[torch.cuda.Event],
) -> Optional[Callable[..., None]]:
if not isinstance(optimizer, _ScratchBufferOptimizer):
return None
if sync_event is not None:
sync_event.synchronize()

scratch_buffer_storages: list[tuple[torch.UntypedStorage, int]] = []
seen_storage_ids: set[int] = set()
for scratch_buffer in optimizer.scratch_buffers():
storage = scratch_buffer.untyped_storage()
storage_id = storage._cdata
storage_size = storage.size()
if storage_size == 0 or storage_id in seen_storage_ids:
continue
seen_storage_ids.add(storage_id)
scratch_buffer_storages.append((storage, storage_size))

released_scratch_buffer_bytes = sum(
storage_size for _, storage_size in scratch_buffer_storages
)
if released_scratch_buffer_bytes == 0:
return None
for storage, _ in scratch_buffer_storages:
storage.resize_(0)

cls._log_ems_once(
"ems_stash_optimizer_scratch_buffer_released",
{
"optimizer": type(optimizer).__name__,
"size_bytes": str(released_scratch_buffer_bytes),
},
)

def restore_scratch_buffer(
_grad: Optional[torch.Tensor] = None,
_sync_event: Optional[torch.cuda.Event] = None,
) -> None:
restored_scratch_buffer_bytes = 0
for storage, storage_size in scratch_buffer_storages:
if storage.size() == 0:
storage.resize_(storage_size)
restored_scratch_buffer_bytes += storage_size
cls._log_ems_once(
"ems_restore_optimizer_scratch_buffer",
{
"optimizer": type(optimizer).__name__,
"size_bytes": str(restored_scratch_buffer_bytes),
},
)

cls._optimizer_scratch_buffer_restore_callbacks.append(restore_scratch_buffer)
return restore_scratch_buffer

@staticmethod
def _collect_optimizer_state_tensors(
optimizer: torch.optim.Optimizer,
) -> List[torch.Tensor]:
"""Collect all stashable CUDA tensors from ``optimizer.state``."""
tensors: List[torch.Tensor] = []
for _param, state_dict in optimizer.state.items():
if not isinstance(state_dict, dict):
continue
for _state_key, state_value in state_dict.items():
tensors.extend(_collect_cuda_tensors_from_value(state_value))
return tensors

@classmethod
def stash_optimizer_state(
cls,
Expand Down Expand Up @@ -1134,7 +1025,12 @@ def stash_optimizer_state(
- Supports nested structures like Shampoo's ShampooKroneckerFactors
"""
# Collect all CUDA tensors from optimizer state
tensors: List[torch.Tensor] = cls._collect_optimizer_state_tensors(optimizer)
tensors: List[torch.Tensor] = []
for _param, state_dict in optimizer.state.items():
if not isinstance(state_dict, dict):
continue
for _state_key, state_value in state_dict.items():
tensors.extend(_collect_cuda_tensors_from_value(state_value))

cls._log_ems_once(
"ems_stash_optimizer_state_collected",
Expand All @@ -1144,33 +1040,11 @@ def stash_optimizer_state(
},
)

scratch_buffer_restore: Callable[..., None] | None = None

if num_slices <= 1:
await_restore, tensor_restore, _execute_stash = cls._stash_tensors(
await_restore, restore, _execute_stash = cls._stash_tensors(
tensors, label="optimizer state", sync_event=sync_event
)
cls._optimizer_state_restore_callbacks.append(tensor_restore)
scratch_buffer_restore = cls._release_optimizer_scratch_buffers(
optimizer, sync_event
)

def restore(
_grad: Optional[torch.Tensor] = None,
restore_sync_event: Optional[torch.cuda.Event] = None,
) -> None:
cls._consume_restore_callback(
tensor_restore,
cls._optimizer_state_restore_callbacks,
restore_sync_event,
)
if scratch_buffer_restore is not None:
cls._consume_restore_callback(
scratch_buffer_restore,
cls._optimizer_scratch_buffer_restore_callbacks,
restore_sync_event,
)

cls._optimizer_state_restore_callbacks.append(restore)
return await_restore, restore

# Gradual (throttled) restore: partition the dense optimizer state into
Expand Down Expand Up @@ -1198,32 +1072,18 @@ def restore(
# maps slice k to the k-th hook to fire.
cls._optimizer_state_restore_callbacks.append(restore_k)

scratch_buffer_restore = cls._release_optimizer_scratch_buffers(
optimizer, sync_event
)

def aggregate_await(_grad: Optional[torch.Tensor] = None) -> None:
"""Gate the current stream on EVERY slice's restore completion."""
for slice_await in slice_awaits:
slice_await(None)

def restore_all(
_grad: Optional[torch.Tensor] = None,
restore_sync_event: Optional[torch.cuda.Event] = None,
sync_event: Optional[torch.cuda.Event] = None,
) -> None:
"""Restore all slices, followed by disposable optimizer scratch buffer."""
"""Restore all slices (used when no per-hook driver is present)."""
for slice_restore in slice_restores:
cls._consume_restore_callback(
slice_restore,
cls._optimizer_state_restore_callbacks,
restore_sync_event,
)
if scratch_buffer_restore is not None:
cls._consume_restore_callback(
scratch_buffer_restore,
cls._optimizer_scratch_buffer_restore_callbacks,
restore_sync_event,
)
slice_restore(None, sync_event)

return aggregate_await, restore_all

Expand Down
94 changes: 0 additions & 94 deletions torchrec/distributed/tests/test_memory_stashing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -512,19 +512,6 @@ def test_is_enabled(self) -> None:
self.assertFalse(MemoryStashingManager.is_enabled())


class ScratchBufferOptimizer(torch.optim.SGD):
def __init__(
self,
params: Any,
scratch_buffer: torch.Tensor,
) -> None:
super().__init__(params, lr=0.01)
self._scratch_buffer = scratch_buffer

def scratch_buffers(self) -> tuple[torch.Tensor, ...]:
return (self._scratch_buffer,)


class TestStashOptimizerState(unittest.TestCase):
"""Tests for MemoryStashingManager.stash_optimizer_state method."""

Expand All @@ -539,87 +526,6 @@ def setUp(self) -> None:
def tearDown(self) -> None:
MemoryStashingManager.reset()

def test_scratch_buffer_only_uses_optimizer_state_stash_restore_api(self) -> None:
model = nn.Linear(10, 10).to(self.device)
scratch_buffer = torch.zeros(1024, dtype=torch.int8, device=self.device)
optimizer = ScratchBufferOptimizer(model.parameters(), scratch_buffer)
scratch_buffer_size = scratch_buffer.untyped_storage().size()

await_restore, _restore = MemoryStashingManager.stash_optimizer_state(optimizer)

self.assertEqual(scratch_buffer.untyped_storage().size(), 0)
self.assertEqual(
len(MemoryStashingManager._optimizer_scratch_buffer_restore_callbacks),
1,
)

MemoryStashingManager.restore_optimizer_state()
await_restore(None)

self.assertEqual(scratch_buffer.untyped_storage().size(), scratch_buffer_size)
self.assertEqual(
len(MemoryStashingManager._optimizer_scratch_buffer_restore_callbacks),
0,
)

def test_returned_restore_consumes_registered_callbacks(self) -> None:
model = nn.Linear(10, 10).to(self.device)
scratch_buffer = torch.zeros(1024, dtype=torch.int8, device=self.device)
optimizer = ScratchBufferOptimizer(model.parameters(), scratch_buffer)
scratch_buffer_size = scratch_buffer.untyped_storage().size()

await_restore, restore = MemoryStashingManager.stash_optimizer_state(optimizer)
restore(None)
await_restore(None)

self.assertEqual(scratch_buffer.untyped_storage().size(), scratch_buffer_size)
self.assertEqual(MemoryStashingManager._optimizer_state_restore_callbacks, [])
self.assertEqual(
MemoryStashingManager._optimizer_scratch_buffer_restore_callbacks,
[],
)

MemoryStashingManager.restore_optimizer_state()
self.assertEqual(scratch_buffer.untyped_storage().size(), scratch_buffer_size)

def test_scratch_buffer_restore_can_be_deferred_until_pre_step_guard(self) -> None:
model = nn.Linear(10, 10).to(self.device)
scratch_buffer = torch.zeros(1024, dtype=torch.int8, device=self.device)
optimizer = ScratchBufferOptimizer(model.parameters(), scratch_buffer)

MemoryStashingManager.stash_optimizer_state(optimizer)
MemoryStashingManager.restore_optimizer_state(restore_scratch_buffer=False)

self.assertEqual(scratch_buffer.untyped_storage().size(), 0)
self.assertEqual(
len(MemoryStashingManager._optimizer_scratch_buffer_restore_callbacks),
1,
)

MemoryStashingManager.restore_optimizer_state()

self.assertGreater(scratch_buffer.untyped_storage().size(), 0)

def test_scratch_buffer_restore_waits_until_all_slices_are_restored(self) -> None:
model = nn.Linear(512, 512).to(self.device)
scratch_buffer = torch.zeros(1024, dtype=torch.int8, device=self.device)
optimizer = ScratchBufferOptimizer(model.parameters(), scratch_buffer)
x = torch.randn(32, 512, device=self.device)
model(x).sum().backward()
optimizer.step()

await_restore, _restore = MemoryStashingManager.stash_optimizer_state(
optimizer, num_slices=2
)
MemoryStashingManager.restore_optimizer_state_next()

self.assertEqual(scratch_buffer.untyped_storage().size(), 0)

MemoryStashingManager.restore_optimizer_state()
await_restore(None)

self.assertGreater(scratch_buffer.untyped_storage().size(), 0)

def test_basic_adam_optimizer_stash_and_restore(self) -> None:
"""Test basic stash and restore with Adam optimizer."""
model = nn.Linear(512, 512).to(self.device)
Expand Down
Loading
Loading