From d6f43c5dcef2ff9145af5880d7c1f14a41f9b38f Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Mon, 3 Aug 2026 17:07:33 +0800 Subject: [PATCH 1/4] fix(zero3): async grad offload + pinned offload buffers by default The ZeRO-3 grad GPU->CPU offload copy in partition_grads() was a blocking copy_() without non_blocking, and its destination buffer defaulted to pageable host memory. Together these forced the grad offload onto a synchronous, low-bandwidth (staged pageable) path with no overlap against backward compute, even though the copy already runs on the dedicated reduce_and_partition_stream. - offload_config.py: default offload_optimizer/offload_param pin_memory to True. Pinned (page-locked) host memory is required for asynchronous, full-bandwidth DMA; the prior False default silently selected the slow staged pageable copy. Disable only on hosts with tight memlock limits. - stage3.py: issue the grad offload copy with non_blocking=True. The copy stays on reduce_and_partition_stream, and correctness is preserved by the existing reduce_and_partition_stream.synchronize() barriers ahead of step(). Validated on 4x RTX 4080-SUPER (autotp=2, offload_optimizer, cpu_adam): BWD -18% to -28% (1.5B/3B), no memory regression, FWD/STEP unchanged. Signed-off-by: Guokai Ma --- deepspeed/runtime/zero/offload_config.py | 18 ++++++++++++------ deepspeed/runtime/zero/stage3.py | 3 ++- 2 files changed, 14 insertions(+), 7 deletions(-) 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/stage3.py b/deepspeed/runtime/zero/stage3.py index 62cdf6cb9dad..956f5a05ce9f 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1858,7 +1858,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(): From d537ae376415d6c1b8ca32fe9691aba87b2e8d4b Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Mon, 3 Aug 2026 18:13:25 +0800 Subject: [PATCH 2/4] refactor(zero3): remove orphaned async grad copy stub `async_inplace_copy_grad_to_fp32_buffer_from_gpu` in stage3.py has been dead code since #1453 (2022-01, "Various ZeRO Stage3 Optimizations + Improvements"). That rewrite moved grad handling from the per-param path to the batched partition_grads() path and deleted both the call site and the `self.copy_grad_stream` initialization, but left this function definition behind. It now (a) has zero callers in stage3 and (b) references `self.copy_grad_stream`, an attribute that no longer exists in stage3, so it would raise AttributeError if ever invoked. Keeping it is a maintenance hazard: it reads like working async-offload infrastructure and obscures the fact that stage3's grad offload was actually synchronous (fixed in the previous commit). The live, called version in stage_1_and_2.py (still using non_blocking, called at ~line 1590) is unaffected and intentionally left in place. Signed-off-by: Guokai Ma --- deepspeed/runtime/zero/stage3.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 956f5a05ce9f..b70ab26a0158 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -1772,14 +1772,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]) From 12bdbc5176b4df354d8995081e9a25b6e9c9bdd9 Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Tue, 4 Aug 2026 23:23:46 +0800 Subject: [PATCH 3/4] docs(zero): document pin_memory default, mechanism, and memlock caveat Expand the pin_memory descriptions for offload_param and offload_optimizer in config-json.md to cover the mechanism (async full-bandwidth DMA + compute overlap), the precondition (overlap_comm for the grad offload path), and the failure mode (pinned memory is non-swappable, counts against ulimit -l; may cause init failures or indirect OOM on hosts with tight memlock limits). Update the default to true and add a release notice. Fix memory.rst Pinned Memory section: replace the deprecated cpu_offload_use_pin_memory flag with the new offload_optimizer / offload_param pin_memory fields, correct the stale 'ZeRO-2 can't be controlled' claim, and note the new default. Signed-off-by: Guokai Ma --- docs/_pages/config-json.md | 7 +++++-- docs/code-docs/source/memory.rst | 12 +++++------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 426eac8f51a1..4f3ca56485b2 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 From fbb02770ed5c40828fb125ac2d88b6f7dd3dffa8 Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Wed, 5 Aug 2026 09:19:15 +0800 Subject: [PATCH 4/4] feat(accel): track pinned host memory for OOM diagnostics Centralize pin_memory accounting in the accelerator ABC so every get_accelerator().pin_memory() call records the bytes pinned. Pinned (page-locked) memory is non-swappable and counts against the host memlock limit (ulimit -l); the running total helps diagnose host OOM that surfaces far from the pinning call site. abstract_accelerator.py: pin_memory is now a concrete template that records tensor.nbytes then delegates to a new _pin_memory hook (concrete default tensor.pin_memory()). The hook is NOT abstract, so custom accelerators that only override pin_memory keep working (they just bypass accounting). cuda/npu/sdaa/mlu/mps/supa drop their now-redundant override; xpu/hpu rename pin_memory->_pin_memory; cpu is unchanged (its no-op shadows and correctly skips accounting). Per-call logging is at DEBUG; pinned_memory_summary() logs the total at INFO, emitted at the end of ZeRO-3 init (stage3.py, parameter_offload.py). Signed-off-by: Guokai Ma --- accelerator/abstract_accelerator.py | 11 ++- accelerator/cpu_accelerator.py | 3 + accelerator/cuda_accelerator.py | 3 - accelerator/hpu_accelerator.py | 2 +- accelerator/mlu_accelerator.py | 3 - accelerator/mps_accelerator.py | 3 - accelerator/npu_accelerator.py | 3 - accelerator/sdaa_accelerator.py | 3 - accelerator/supa_accelerator.py | 3 - accelerator/xpu_accelerator.py | 2 +- deepspeed/runtime/zero/parameter_offload.py | 2 + deepspeed/runtime/zero/stage3.py | 3 + deepspeed/utils/pin_memory_tracker.py | 67 ++++++++++++++++++ tests/unit/utils/test_pin_memory_tracker.py | 77 +++++++++++++++++++++ 14 files changed, 163 insertions(+), 22 deletions(-) create mode 100644 deepspeed/utils/pin_memory_tracker.py create mode 100644 tests/unit/utils/test_pin_memory_tracker.py 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/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 b70ab26a0158..e9be49681d3f 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 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/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()