feat: implement wake-sleep API for trainer

Author: DaniilSergeev17

d9d/core/offload/__init__.py

@@ -0,0 +1,15 @@
+"""Offloading of GPU-resident training state to host memory and back."""
+
+from .api import DEFAULT_SLEEP_TAGS, Offloadable, OffloadContext, OnloadContext, SleepTag
+from .tensor import OffloadedTensor, offload_tensor, onload_tensor
+
+__all__ = [
+    "DEFAULT_SLEEP_TAGS",
+    "OffloadContext",
+    "Offloadable",
+    "OffloadedTensor",
+    "OnloadContext",
+    "SleepTag",
+    "offload_tensor",
+    "onload_tensor",
+]

d9d/core/offload/api.py

@@ -0,0 +1,85 @@
+import dataclasses
+from enum import StrEnum
+from typing import Protocol, runtime_checkable
+
+from d9d.core.dist_context import DistributedContext
+
+
+class SleepTag(StrEnum):
+    """
+    Subsystem selector for Trainer.sleep and Trainer.wake.
+
+    Attributes:
+        TENSOR_STATES: All GPU tensor state - model parameters and buffers, optimizer state,
+            gradient buckets and the residual loss accumulator. Always offloaded as a unit.
+        COMMS: NCCL process groups. Opt-in; its implementation is deferred to a second phase,
+            so requesting it currently raises NotImplementedError.
+    """
+
+    TENSOR_STATES = "tensor_states"
+    COMMS = "comms"
+
+
+DEFAULT_SLEEP_TAGS = frozenset({SleepTag.TENSOR_STATES})
+"""The default tag set for Trainer.sleep and Trainer.wake: tensor state only, no comms."""
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class OffloadContext:
+    """
+    Context passed to Offloadable.offload.
+
+    Attributes:
+        dist_context: The distributed context the subsystem was built under.
+        pin_memory: Whether to allocate the host buffer in pinned memory.
+    """
+
+    dist_context: DistributedContext
+    pin_memory: bool
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class OnloadContext:
+    """
+    Context passed to Offloadable.onload.
+
+    Attributes:
+        dist_context: The distributed context the subsystem was built under.
+    """
+
+    dist_context: DistributedContext
+
+
+@runtime_checkable
+class Offloadable(Protocol):
+    """
+    Protocol for subsystems that own GPU-resident state and can release it to host memory.
+
+    An "offload" followed by an "onload" must be observationally a no-op: parameter identities,
+    optimizer state keys, DTensor wrapper instances, placements and dtypes are all preserved
+    across the round trip. Only the underlying device storages are reallocated.
+    """
+
+    def offload(self, ctx: OffloadContext) -> None:
+        """
+        Releases the GPU memory owned by this subsystem, moving its state to host memory.
+
+        Args:
+            ctx: Context for this operation.
+        """
+
+    def onload(self, ctx: OnloadContext) -> None:
+        """
+        Restores GPU residency of the state previously released by "offload".
+
+        Args:
+            ctx: Context for this operation.
+        """
+
+    def is_offloaded(self) -> bool:
+        """
+        Reports whether this subsystem currently has its state on host memory.
+
+        Returns:
+            True if the subsystem is offloaded, False otherwise.
+        """

d9d/core/offload/tensor.py

@@ -0,0 +1,68 @@
+import dataclasses
+
+import torch
+from torch.distributed.tensor import DTensor
+
+
+@dataclasses.dataclass(slots=True, frozen=True)
+class OffloadedTensor:
+    """
+    Handle to a tensor whose local storage has been swapped to host memory.
+
+    Produced by "offload_tensor" and consumed by "onload_tensor". Subsystems hold these as
+    opaque handles between an offload and the matching onload.
+
+    Attributes:
+        host: The host-memory mirror that currently backs the offloaded tensor's local storage.
+    """
+
+    host: torch.Tensor
+
+
+def _local_storage_holder(tensor: torch.Tensor) -> torch.Tensor:
+    """Returns the storage-bearing tensor: the DTensor's local shard, or the tensor itself."""
+    return tensor._local_tensor if isinstance(tensor, DTensor) else tensor  # noqa: SLF001
+
+
+def offload_tensor(tensor: torch.Tensor, *, pin_memory: bool) -> OffloadedTensor:
+    """
+    Swaps "tensor"'s local storage for a host-memory mirror, in place.
+
+    For a DTensor, the wrapper instance is preserved across the swap - only the local shard's
+    underlying storage is rebound, via "_local_tensor.data". "device_mesh", "placements",
+    global "shape" and global "stride" live on the wrapper and never leave it, so any
+    external reference to the DTensor keeps pointing at the same object. For a plain tensor,
+    "tensor.data" itself is rebound.
+
+    Args:
+        tensor: The device tensor to offload. May be a plain tensor or a DTensor.
+        pin_memory: Whether to allocate the host buffer in pinned memory.
+
+    Returns:
+        A handle holding the host buffer; pass back to "onload_tensor" with the same tensor.
+    """
+
+    local = _local_storage_holder(tensor)
+    host = torch.empty_like(local, device="cpu", pin_memory=pin_memory)
+    host.copy_(local, non_blocking=True)
+    local.data = host
+    return OffloadedTensor(host=host)
+
+
+def onload_tensor(tensor: torch.Tensor, offloaded: OffloadedTensor, *, device: torch.device) -> None:
+    """
+    Restores "tensor"'s local storage to "device", in place, from a host buffer.
+
+    The tensor object (and, for a DTensor, its wrapper) is the same instance as before the
+    offload; only the underlying device storage is freshly allocated.
+
+    Args:
+        tensor: The same tensor previously passed to "offload_tensor".
+        offloaded: The handle returned by "offload_tensor".
+        device: The accelerator device to restore the local storage onto.
+    """
+
+    local = _local_storage_holder(tensor)
+    fresh = torch.empty_like(local, device=device)
+    fresh.copy_(offloaded.host, non_blocking=True)
+    local.data = fresh

d9d/loop/component/__init__.py

@@ -13,6 +13,7 @@
 from .stepper import Stepper
 from .task_operator import ForwardResult, InferenceTaskOperator, TrainTaskOperator
 from .timeout_manager import TimeoutManager
+from .train_sleeper import TrainSleeper
 
 __all__ = [
     "BatchMaths",
@@ -34,5 +35,6 @@
     "Stepper",
     "TimeoutManager",
     "TrackedModules",
+    "TrainSleeper",
     "TrainTaskOperator",
 ]

d9d/loop/component/gradient_manager.py

@@ -4,6 +4,7 @@
 from torch.distributed.tensor import DTensor
 
 from d9d.core.dist_context import DistributedContext
+from d9d.core.offload import Offloadable, OffloadContext, OnloadContext
 from d9d.internals.grad_sync import GradientSynchronizer
 from d9d.loop.config import GradientManagerConfig
 from d9d.metric.impl.aggregation import WeightedMeanMetric
@@ -12,7 +13,7 @@
 from .model_stage_factory import TrackedModules
 
 
-class GradientManager:
+class GradientManager(Offloadable):
     """
     Manages the lifecycle of gradients during the training loop.
 
@@ -53,6 +54,10 @@ def __init__(
         )
         self._grads_to_scale: list[torch.Tensor] | None = None
 
+        self._installed = False
+        self._offloaded = False
+        self._in_flight_count = 0
+
     def _setup_grad_dtype(self):
         if self._config.grad_dtype is None:
             return
@@ -85,6 +90,15 @@ def _scale_grads(self):
         if len(self._grads_to_scale) > 0:
             torch._foreach_mul_(self._grads_to_scale, scale_factor)
 
+    def _bind(self):
+        self._setup_grad_dtype()
+        self._grad_sync.bind()
+        self._bind_grads_to_scale()
+
+    def _unbind(self):
+        self._unbind_grads_to_scale()
+        self._grad_sync.unbind()
+
     @contextmanager
     def install(self):
         """
@@ -95,12 +109,11 @@ def install(self):
         as the boundary for the accumulation phase.
         """
 
-        self._setup_grad_dtype()
-        self._grad_sync.bind()
-        self._bind_grads_to_scale()
+        self._bind()
+        self._installed = True
         yield
-        self._unbind_grads_to_scale()
-        self._grad_sync.unbind()
+        self._installed = False
+        self._unbind()
 
     def add_loss_with_weight(self, loss: torch.Tensor, loss_weight: torch.Tensor):
         """
@@ -112,6 +125,7 @@ def add_loss_with_weight(self, loss: torch.Tensor, loss_weight: torch.Tensor):
         """
 
         self._loss.update(loss, loss_weight)
+        self._in_flight_count += 1
 
     def sync_and_scale(self):
         """
@@ -151,3 +165,62 @@ def zero_grad(self):
 
         self._grad_sync.zero_grad()
         self._loss.reset()
+        self._in_flight_count = 0
+
+    @property
+    def has_in_flight_gradients(self) -> bool:
+        """
+        Checks whether a gradient accumulation is currently in flight.
+
+        The counter is raised by "add_loss_with_weight" and reset by "zero_grad". While it is
+        non-zero, partial accumulation state lives in the synchronizer buckets, so offloading
+        the gradient state would lose it.
+        """
+
+        return self._in_flight_count > 0
+
+    def offload(self, ctx: OffloadContext) -> None:
+        """
+        Releases the GPU memory of the gradient state to host memory.
+
+        The synchronizer bucket buffers are released and the residual loss accumulator is reset.
+
+        Args:
+            ctx: Context for this operation.
+
+        Raises:
+            RuntimeError: If the gradient state is already offloaded, or if the manager has
+                not been installed.
+        """
+
+        if self._offloaded:
+            raise RuntimeError("GradientManager is already offloaded.")
+        if not self._installed:
+            raise RuntimeError("GradientManager must be installed before it can be offloaded.")
+
+        self._unbind()
+        self._loss.reset()
+        self._offloaded = True
+
+    def onload(self, ctx: OnloadContext) -> None:
+        """
+        Restores GPU residency of the gradient state released by "offload".
+
+        Args:
+            ctx: Context for this operation.
+
+        Raises:
+            RuntimeError: If the gradient state is not offloaded.
+        """
+
+        if not self._offloaded:
+            raise RuntimeError("GradientManager is not offloaded.")
+
+        if self._installed:
+            self._bind()
+
+        self._offloaded = False
+
+    def is_offloaded(self) -> bool:
+        """Reports whether the gradient state is currently on host memory."""
+        return self._offloaded