diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py index feb7f897d4f0..c0c0646c6205 100644 --- a/accelerator/abstract_accelerator.py +++ b/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): diff --git a/accelerator/cpu_accelerator.py b/accelerator/cpu_accelerator.py index 4ff0f4dd7527..c089ac07fb69 100644 --- a/accelerator/cpu_accelerator.py +++ b/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): diff --git a/accelerator/cuda_accelerator.py b/accelerator/cuda_accelerator.py index 757560d23e85..c64f2ec11f75 100644 --- a/accelerator/cuda_accelerator.py +++ b/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() diff --git a/accelerator/hpu_accelerator.py b/accelerator/hpu_accelerator.py index e809ae3e1dab..0ce24f72593c 100644 --- a/accelerator/hpu_accelerator.py +++ b/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): diff --git a/accelerator/mlu_accelerator.py b/accelerator/mlu_accelerator.py index 0876116b5208..fa0989d79309 100644 --- a/accelerator/mlu_accelerator.py +++ b/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() diff --git a/accelerator/mps_accelerator.py b/accelerator/mps_accelerator.py index f6600beb779c..c780e9e39253 100644 --- a/accelerator/mps_accelerator.py +++ b/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() diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py index 91515e8112ee..b9c96e24093f 100644 --- a/accelerator/npu_accelerator.py +++ b/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() diff --git a/accelerator/sdaa_accelerator.py b/accelerator/sdaa_accelerator.py index 9940d4a97d17..40d8005c8299 100755 --- a/accelerator/sdaa_accelerator.py +++ b/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() diff --git a/accelerator/supa_accelerator.py b/accelerator/supa_accelerator.py index 5c86b226e28d..783efb7bc15a 100644 --- a/accelerator/supa_accelerator.py +++ b/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() diff --git a/accelerator/xpu_accelerator.py b/accelerator/xpu_accelerator.py index 9f6b21af54ea..a3332df75f06 100644 --- a/accelerator/xpu_accelerator.py +++ b/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: diff --git a/deepspeed/runtime/zero/offload_config.py b/deepspeed/runtime/zero/offload_config.py index ac88d32266f8..a71398e4929f 100644 --- a/deepspeed/runtime/zero/offload_config.py +++ b/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 diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py index 6ef839fa816b..f9fc7c2c0e64 100644 --- a/deepspeed/runtime/zero/parameter_offload.py +++ b/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): diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index e300a1d369c4..93cead30f090 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/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(): diff --git a/deepspeed/utils/pin_memory_tracker.py b/deepspeed/utils/pin_memory_tracker.py new file mode 100644 index 000000000000..f08f42cb5905 --- /dev/null +++ b/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) diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 8b4a066314b2..f09183f36978 100755 --- a/docs/_pages/config-json.md +++ b/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] diff --git a/docs/code-docs/source/memory.rst b/docs/code-docs/source/memory.rst index 28e96955d41f..9f05495001c6 100644 --- a/docs/code-docs/source/memory.rst +++ b/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 diff --git a/tests/unit/utils/test_pin_memory_tracker.py b/tests/unit/utils/test_pin_memory_tracker.py new file mode 100644 index 000000000000..2946801fd603 --- /dev/null +++ b/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()