fix(zero3): async grad offload + pinned offload buffers by default

accelerator/abstract_accelerator.py

@@ -260,9 +260,16 @@ def IntTensor(self):
     def LongTensor(self):
         ...
 
-    @abc.abstractmethod
     def pin_memory(self, tensor, align_bytes=1):
-        ...
+        from deepspeed.utils.pin_memory_tracker import track_pinned_memory
+        track_pinned_memory(tensor.nbytes)
+        return self._pin_memory(tensor, align_bytes)
+
+    def _pin_memory(self, tensor, align_bytes=1):
+        """Device-specific pinning hook. Accelerators that need custom pinning
+        behavior should override this method rather than ``pin_memory`` so that
+        the pinned-memory accounting in ``pin_memory`` is preserved."""
+        return tensor.pin_memory()
 
     @abc.abstractmethod
     def is_pinned(self, tensor):

accelerator/cpu_accelerator.py

@@ -279,6 +279,9 @@ def LongTensor(self):
         return torch.LongTensor
 
     def pin_memory(self, tensor, align_bytes=1):
+        # Overrides pin_memory directly (not _pin_memory) to bypass the ABC's
+        # pinned-memory accounting: this is a no-op, nothing is page-locked, so
+        # counting would mislead OOM diagnostics. Do not rename to _pin_memory.
         return tensor
 
     def is_pinned(self, tensor):

accelerator/cuda_accelerator.py

@@ -326,9 +326,6 @@ def IntTensor(self):
     def LongTensor(self):
         return functools.partial(torch.tensor, dtype=torch.long, device='cuda')
 
-    def pin_memory(self, tensor, align_bytes=1):
-        return tensor.pin_memory()
-
     def is_pinned(self, tensor):
         return tensor.is_pinned()
 

accelerator/hpu_accelerator.py

@@ -231,7 +231,7 @@ def IntTensor(self):
     def LongTensor(self):
         return functools.partial(torch.tensor, dtype=torch.long, device='hpu')
 
-    def pin_memory(self, tensor, align_bytes=1):
+    def _pin_memory(self, tensor, align_bytes=1):
         return tensor.pin_memory(self.device())
 
     def is_pinned(self, tensor):

accelerator/mlu_accelerator.py

@@ -223,9 +223,6 @@ def IntTensor(self):
     def LongTensor(self):
         return functools.partial(torch.tensor, dtype=torch.long, device='mlu')
 
-    def pin_memory(self, tensor, align_bytes=1):
-        return tensor.pin_memory()
-
     def is_pinned(self, tensor):
         return tensor.is_pinned()
 

accelerator/mps_accelerator.py

@@ -214,9 +214,6 @@ def IntTensor(self):
     def LongTensor(self):
         return
 
-    def pin_memory(self, tensor, align_bytes=1):
-        return tensor.pin_memory()
-
     def is_pinned(self, tensor):
         return tensor.is_pinned()
 

accelerator/npu_accelerator.py

@@ -222,9 +222,6 @@ def IntTensor(self):
     def LongTensor(self):
         return torch.npu.LongTensor
 
-    def pin_memory(self, tensor, align_bytes=1):
-        return tensor.pin_memory()
-
     def is_pinned(self, tensor):
         return tensor.is_pinned()
 

accelerator/sdaa_accelerator.py

@@ -251,9 +251,6 @@ def IntTensor(self):
     def LongTensor(self):
         return functools.partial(torch.tensor, dtype=torch.long, device='sdaa')
 
-    def pin_memory(self, tensor, align_bytes=1):
-        return tensor.pin_memory()
-
     def is_pinned(self, tensor):
         return tensor.is_pinned()
 

accelerator/supa_accelerator.py

@@ -214,9 +214,6 @@ def IntTensor(self):
     def LongTensor(self):
         return torch.supa.LongTensor
 
-    def pin_memory(self, tensor, align_bytes=1):
-        return tensor.pin_memory()
-
     def is_pinned(self, tensor):
         return tensor.is_pinned()
 

accelerator/xpu_accelerator.py

@@ -229,7 +229,7 @@ def IntTensor(self):
     def LongTensor(self):
         return functools.partial(torch.tensor, dtype=torch.long, device=self._name)
 
-    def pin_memory(self, tensor, align_bytes=1):
+    def _pin_memory(self, tensor, align_bytes=1):
         if align_bytes == 1:
             return tensor.pin_memory(device=self.current_device_name())
         elif align_bytes == 0:

deepspeed/runtime/zero/offload_config.py

@@ -42,10 +42,13 @@ class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel):
     NVMe is enabled.
     """
 
-    pin_memory: bool = False
+    pin_memory: bool = True
     """
-    Offload to page-locked CPU memory. This could boost throughput at the cost
-    of extra memory overhead.
+    Offload to page-locked (pinned) CPU memory. Required for asynchronous,
+    full-bandwidth GPU<->CPU transfers and for overlap of grad/param offload
+    with compute. Defaults to True. Disable only on hosts with tight memlock
+    limits (ulimit -l) or very limited resident RAM, since pinned memory
+    cannot be paged out.
     """
 
 
@@ -69,10 +72,13 @@ class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel):
     gradient, momentum, and variance).
     """
 
-    pin_memory: bool = False
+    pin_memory: bool = True
     """
-    Offload to page-locked CPU memory. This could boost throughput at the cost
-    of extra memory overhead.
+    Offload to page-locked (pinned) CPU memory. Required for asynchronous,
+    full-bandwidth GPU<->CPU transfers and for overlap of grad/param offload
+    with compute. Defaults to True. Disable only on hosts with tight memlock
+    limits (ulimit -l) or very limited resident RAM, since pinned memory
+    cannot be paged out.
     """
 
     pipeline_read: bool = False

deepspeed/runtime/zero/parameter_offload.py

@@ -15,6 +15,7 @@
 from deepspeed.runtime.zero.partitioned_param_coordinator import PartitionedParameterCoordinator, InflightParamRegistry, iter_params
 from deepspeed.accelerator import get_accelerator
 from deepspeed import utils
+from deepspeed.utils.pin_memory_tracker import pinned_memory_summary
 
 FWD_MODULE_STACK = list()
 
@@ -225,6 +226,7 @@ def __init__(
             force=False)
 
         see_memory_usage("DeepSpeedZeRoOffload initialize [end]", force=False)
+        pinned_memory_summary("ZeRO-3 parameter offload init")
 
     @instrument_w_nvtx
     def partition_all_parameters(self):

deepspeed/runtime/zero/stage3.py

@@ -18,6 +18,7 @@
 from deepspeed.runtime.base_optimizer import ZeROOptimizer
 from deepspeed.utils import logger
 from deepspeed.utils.torch import register_grad_hook, required_torch_version
+from deepspeed.utils.pin_memory_tracker import pinned_memory_summary
 from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler
 from deepspeed.runtime.torch_autocast import get_autocast_dtype, get_all_comm_dtypes, is_autocast_initialized, sort_dtypes
 from deepspeed.runtime.comm.coalesced_collectives import reduce_scatter_coalesced, all_to_all_quant_reduce, all_to_all_loco_quant_reduce
@@ -719,6 +720,8 @@ def _setup_for_real_optimizer(self):
                 0, offset, param.partition_numel())
             offset += param.partition_numel()
 
+        pinned_memory_summary("ZeRO-3 optimizer init")
+
     def _link_all_hp_params(self):
         for p in self.module.parameters():
             p._z3_optimizer = self
@@ -1778,14 +1781,6 @@ def set_norm_for_param_grad_in_gpu(self, param):
         #Using a more memory efficient version
         self.norm_for_param_grads[param_id] = self._constant_buffered_norm2(param.grad)
 
-    def async_inplace_copy_grad_to_fp32_buffer_from_gpu(self, param, fp32_grad_tensor):
-        with get_accelerator().stream(self.copy_grad_stream):
-            param_id = self.get_param_id(param)
-            src_tensor = param.grad.view(-1).to(dtype=self.master_weights_and_grads_dtype)
-            #print(f"src_tensor {src_tensor.size()} and fp32 grad {fp32_grad_tensor.size()}")
-            fp32_grad_tensor.copy_(src_tensor, non_blocking=True)
-            param.grad = None
-
     def complete_grad_norm_calculation_for_cpu_offload(self, params):
         self._assert_same_partition_group(params)
         process_group = self._get_param_partition_group(params[0])
@@ -1864,7 +1859,8 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L
                     else:
                         fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow(
                             0, dest_offset, grad_buffer.numel())
-                        fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype))
+                        fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype),
+                                               non_blocking=True)
 
             # free the gradient
             if not get_accelerator().is_synchronized_device():

deepspeed/utils/pin_memory_tracker.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+
+from deepspeed.utils.logging import logger
+
+_GB = 1024**3
+# Emit an INFO checkpoint each time cumulative pinned memory reaches the next
+# power-of-two multiple of this base. Doubling keeps the number of checkpoints
+# logarithmic in the total, so large offload paths (e.g. per-parameter shards
+# pinned at init) produce a few readable milestones instead of flooding the log.
+_CHECKPOINT_BASE_GB = 32
+
+
+def _fmt_bytes(num_bytes: int) -> str:
+    kb = 1024
+    mb = kb * 1024
+    gb = mb * 1024
+    if num_bytes >= gb:
+        return f"{num_bytes / gb:.3f} GB"
+    if num_bytes >= mb:
+        return f"{num_bytes / mb:.2f} MB"
+    if num_bytes >= kb:
+        return f"{num_bytes / kb:.1f} KB"
+    return f"{num_bytes} B"
+
+
+class _PinnedMemoryTracker:
+    """Process-wide total of host memory pinned through the accelerator's
+    ``pin_memory``. Pinned memory is page-locked: it cannot be swapped out and
+    counts against the host memlock limit (``ulimit -l``). The running total is
+    a useful hint when diagnosing host out-of-memory errors, which often surface
+    far from the call site that consumed the resident-RAM budget.
+    """
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self) -> None:
+        self._bytes = 0
+        self._calls = 0
+        self._next_checkpoint = _CHECKPOINT_BASE_GB * _GB
+
+    def track(self, num_bytes: int) -> None:
+        self._bytes += num_bytes
+        self._calls += 1
+        logger.debug(f"pin_memory: +{_fmt_bytes(num_bytes)} "
+                     f"(call #{self._calls}, running total: {_fmt_bytes(self._bytes)})")
+        while self._bytes >= self._next_checkpoint:
+            msg = (f"[pinned-memory checkpoint] crossed {_fmt_bytes(self._next_checkpoint)}: "
+                   f"{_fmt_bytes(self._bytes)} pinned across {self._calls} allocations")
+            logger.info(msg)
+            self._next_checkpoint *= 2
+
+    def log_summary(self, tag: str = "") -> None:
+        prefix = f"[pinned-memory {tag}] " if tag else "[pinned-memory] "
+        logger.info(f"{prefix}{_fmt_bytes(self._bytes)} pinned across {self._calls} allocations")
+
+
+_tracker = _PinnedMemoryTracker()
+
+
+def track_pinned_memory(num_bytes: int) -> None:
+    _tracker.track(num_bytes)
+
+
+def pinned_memory_summary(tag: str = "") -> None:
+    _tracker.log_summary(tag)

docs/_pages/config-json.md

@@ -657,7 +657,7 @@ Note that if the value of "device" is not specified or not supported, an asserti
 
 | Description                                                                                          | Default |
 | ---------------------------------------------------------------------------------------------------- | ------- |
-| Offload to page-locked CPU memory. This could boost throughput at the cost of extra memory overhead. | `false` |
+| Offload to page-locked (pinned) CPU memory. Pinning enables asynchronous, full-bandwidth CPU<->GPU DMA so parameter fetches during forward/backward overlap with compute. Pinned memory is non-swappable and counts against the host memlock limit (`ulimit -l`); on hosts with tight memlock limits this may fail at init or cause out-of-memory errors elsewhere — set to `false` in that case. | `true` |
 
 ***buffer_count***: [integer]
 
@@ -707,7 +707,10 @@ Note that if the value of "device" is not specified or not supported, an asserti
 
 | Description                                                                                          | Default |
 | ---------------------------------------------------------------------------------------------------- | ------- |
-| Offload to page-locked CPU memory. This could boost throughput at the cost of extra memory overhead. | `false` |
+| Offload to page-locked (pinned) CPU memory. Pinning is required for the asynchronous GPU->CPU gradient offload to run as a full-bandwidth DMA that overlaps with backward compute (needs `overlap_comm: true`). Pinned memory is non-swappable and counts against the host memlock limit (`ulimit -l`); on hosts with tight memlock limits this may fail at init or cause out-of-memory errors elsewhere — set to `false` in that case. | `true` |
+
+**Note:** `pin_memory` now defaults to `true` for both `offload_param` and `offload_optimizer` (previously `false`). If you see out-of-memory errors after upgrading — especially on hosts with a low memlock limit (`ulimit -l`) — explicitly set `"pin_memory": false`.
+{: .notice--warning}
 
 ***ratio***: [float]
 

docs/code-docs/source/memory.rst

@@ -265,22 +265,20 @@ Note about gradients: While gradients are stored in fp16 (2 bytes), during the w
 
 **Pinned Memory**
 
-Pinned general RAM is included in normal general RAM allocations (i.e. this is not extra memory allocations but simply shows how much of the general RAM is pinned)
+Pinned general RAM is included in normal general RAM allocations (i.e. this is not extra memory allocations but simply shows how much of the general RAM is pinned). Pinning is controlled by the ``pin_memory`` field of ``offload_optimizer`` / ``offload_param`` (both default to ``true``); set to ``false`` on hosts with tight memlock limits (``ulimit -l``).
 
-* ZeRO-2: can't be controlled
+* ZeRO-1/2: controlled by ``offload_optimizer.pin_memory``
 
 * ZeRO-3
 
-To enable add: ``"cpu_offload_use_pin_memory" : true``
+With pinning enabled there are 2 sub-cases:
 
-Now there are 2 sub-cases:
-
-1. ``"cpu_offload_params": true``:
+1. ``offload_param`` enabled (``device: cpu``):
 
    - 6 * params (2b for fp16 params + 4b for fp32 gradients)
    - if ``gradient_accumulation_steps > 1`` an additional 2b for fp16 gradients are pinned
 
-2. ``"cpu_offload_params": false``:
+2. ``offload_param`` not enabled:
 
    - 4b for fp32 gradients
 

tests/unit/utils/test_pin_memory_tracker.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# DeepSpeed Team
+
+import logging
+
+import torch
+
+from deepspeed.utils.pin_memory_tracker import (
+    _fmt_bytes,
+    _tracker,
+    pinned_memory_summary,
+    track_pinned_memory,
+)
+
+
+def test_track_accumulates_and_resets():
+    _tracker.reset()
+    track_pinned_memory(100)
+    track_pinned_memory(2**30)
+    assert _tracker._bytes == 100 + 2**30
+    assert _tracker._calls == 2
+    _tracker.reset()
+    assert _tracker._bytes == 0 and _tracker._calls == 0
+
+
+def test_summary_does_not_raise():
+    _tracker.reset()
+    track_pinned_memory(2**30)
+    pinned_memory_summary("unit-test")
+    _tracker.reset()
+
+
+def test_fmt_bytes():
+    assert _fmt_bytes(512) == "512 B"
+    assert _fmt_bytes(2048) == "2.0 KB"
+    assert _fmt_bytes(2**20) == "1.00 MB"
+    assert _fmt_bytes(2**30).endswith("GB")
+
+
+def test_torch_tensor_nbytes_is_consistent():
+    t = torch.zeros(1024, dtype=torch.float32)
+    track_pinned_memory(t.nbytes)
+    assert _tracker._bytes == 4096
+    _tracker.reset()
+
+
+def test_checkpoint_thresholds_double_from_32gb():
+    _tracker.reset()
+    gb = 1024**3
+    assert _tracker._next_checkpoint == 32 * gb
+    track_pinned_memory(30 * gb)  # below 32 GB -> no crossing
+    assert _tracker._next_checkpoint == 32 * gb
+    track_pinned_memory(10 * gb)  # 40 GB -> crosses 32
+    assert _tracker._next_checkpoint == 64 * gb
+    track_pinned_memory(100 * gb)  # 140 GB -> crosses 64 and 128 in one call
+    assert _tracker._next_checkpoint == 256 * gb
+    _tracker.reset()
+
+
+def test_checkpoint_emits_info(caplog):
+    # The DeepSpeed logger does not propagate, so flip propagation so caplog
+    # (root-based) can observe the checkpoint INFO records.
+    _tracker.reset()
+    ds_logger = logging.getLogger("DeepSpeed")
+    old_prop = ds_logger.propagate
+    ds_logger.propagate = True
+    try:
+        caplog.clear()
+        with caplog.at_level(logging.INFO, logger="DeepSpeed"):
+            track_pinned_memory(33 * (1024**3))  # crosses the 32 GB checkpoint
+            track_pinned_memory(5 * (1024**3))  # 38 GB, no new checkpoint
+        checkpoints = [r.message for r in caplog.records if "checkpoint" in r.message]
+        assert len(checkpoints) == 1
+        assert "32" in checkpoints[0]
+    finally:
+        ds_logger.propagate = old_prop
+        _tracker.reset()