From bfa3c3dac2035cc681e555b3501b1f8107724436 Mon Sep 17 00:00:00 2001 From: Sean Bell Date: Thu, 29 Oct 2026 23:51:08 +0000 Subject: [PATCH 1/3] ZeRO-1/2: preserve gradients across request boundaries until step Signed-off-by: Stas Bekman --- deepspeed/runtime/zero/stage_1_and_2.py | 21 +-- .../zero/test_zero2_offload_multi_backward.py | 140 ++++++++++++++++-- 2 files changed, 140 insertions(+), 21 deletions(-) diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index c1602f7e6b56..54013a00e311 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -1052,11 +1052,11 @@ def independent_gradient_partition_epilogue(self): device=get_accelerator().current_device_name(), param_group_idx=i, return_tensor_list=True) - # Clear all_grad_tensors after use. With reentrant checkpointing, - # the epilogue may run multiple times per backward pass. Each time, - # we read the cumulative grad_accum (which PyTorch naturally accumulates) - # and the final phase will have all gradients. - self.all_grad_tensors[i] = None + # grad_accum is cumulative across reentrant checkpoint phases, so + # reread it rather than adding it twice. Otherwise keep the running + # sum until step(): more request boundaries may precede the update. + if self.use_grad_accum_attribute: + self.all_grad_tensors[i] = None self._release_ipg_buffers() @@ -2159,11 +2159,12 @@ def copy_grads_in_partition(self, param, group_idx=None): self.clear_grad_attribute(param) return if self.cpu_offload: - # Accumulate when there were prior backwards in this step (restore from - # CPU buffer) or more will follow (save to CPU buffer). Skipping only - # the lone backward of a step preserves the existing fast path for - # ga_steps=1 + single backward. - if self.micro_step_id > 0 or not self.is_gradient_accumulation_boundary(): + # A boundary does not guarantee that step() follows. Refresh the host + # accumulator on the first backward too, or a later request adds the + # previous step's gradients. Low-precision masters already share the + # boundary gradient buffer and need no separate first-backward copy. + if (not self.low_precision_master_weights_and_grads or self.micro_step_id > 0 + or not self.is_gradient_accumulation_boundary()): self.async_accumulate_grad_in_cpu_via_gpu(param) # Record active param so unmanaged finalize skips params unused this window. self._offload_accumulated_param_ids.add(self.get_param_id(param)) diff --git a/tests/unit/v1/zero/test_zero2_offload_multi_backward.py b/tests/unit/v1/zero/test_zero2_offload_multi_backward.py index 0c62b98630d7..5fad98c70371 100644 --- a/tests/unit/v1/zero/test_zero2_offload_multi_backward.py +++ b/tests/unit/v1/zero/test_zero2_offload_multi_backward.py @@ -5,13 +5,20 @@ calls per optimizer step (ga_steps=1, driven via set_gradient_accumulation_boundary). """ +import os +from datetime import timedelta + import pytest import torch +from torch.utils.checkpoint import checkpoint import deepspeed +import deepspeed.comm as dist from unit.common import DistributedTest from unit.simple_model import SimpleModel, random_dataloader from deepspeed.accelerator import get_accelerator +from deepspeed.ops import __compatible_ops__ +from deepspeed.utils import safe_get_full_grad def _base_config(zero_stage, gradient_accumulation_steps=1, cpu_offload=False): @@ -122,18 +129,11 @@ def test_multi_backward_across_multiple_steps(self, zero_stage): test = _run_multi_backward(_base_config(zero_stage, cpu_offload=True), hidden_dim, num_chunks=3, num_steps=3) _assert_params_match(ref, test, label=f"ZeRO-{zero_stage} 3x3") - def test_single_backward_allocates_no_cpu_accumulator(self, zero_stage): + def test_single_backward_across_multiple_steps(self, zero_stage): hidden_dim = 8 - engine = _init_engine(_base_config(zero_stage, cpu_offload=True), hidden_dim) - batch = next( - iter(random_dataloader(model=engine, total_samples=1, hidden_dim=hidden_dim, device=engine.device))) - loss = engine(batch[0], batch[1]) - engine.set_gradient_accumulation_boundary(True) - engine.backward(loss) - engine.step() - populated = len(engine.optimizer.accumulated_grads_in_cpu) - engine.destroy() - assert populated == 0, f"ZeRO-{zero_stage}: ga=1+N=1 populated accumulated_grads_in_cpu ({populated} entries)" + ref = _run_multi_backward(_base_config(zero_stage, cpu_offload=False), hidden_dim, num_chunks=1, num_steps=3) + test = _run_multi_backward(_base_config(zero_stage, cpu_offload=True), hidden_dim, num_chunks=1, num_steps=3) + _assert_params_match(ref, test, label=f"ZeRO-{zero_stage} 3x1") def test_ga_greater_than_one_offload_unchanged(self, zero_stage): hidden_dim = 8 @@ -145,3 +145,121 @@ def test_ga_greater_than_one_offload_unchanged(self, zero_stage): hidden_dim, total_microsteps=ga) _assert_params_match(ref, test, label=f"ZeRO-{zero_stage} ga=4") + + +class RequestBoundaryModel(torch.nn.Module): + + def __init__(self, checkpointed): + super().__init__() + self.projections = torch.nn.ModuleList([torch.nn.Linear(4, 4, bias=False), torch.nn.Linear(4, 2, bias=False)]) + self.checkpointed = checkpointed + + def forward(self, inputs): + loss = inputs.new_zeros(()) + for projection in self.projections: + if self.checkpointed: + outputs = checkpoint(projection, inputs, use_reentrant=True) + else: + outputs = projection(inputs) + loss = loss + outputs.sum() + return loss + + +def _run_request_boundaries(rank, rendezvous, zero_stage, cpu_offload, accumulation_dtype, checkpointed): + os.environ["LOCAL_RANK"] = str(rank) + accelerator = get_accelerator() + accelerator.set_device(rank) + with pytest.MonkeyPatch.context() as patch: + patch.setitem(__compatible_ops__, "deepspeed_shm_comm", False) + dist.init_distributed(accelerator.communication_backend_name(), + auto_mpi_discovery=False, + init_method=rendezvous, + rank=rank, + world_size=2, + timeout=timedelta(seconds=60)) + engine = None + try: + model = RequestBoundaryModel(checkpointed).to(device=accelerator.device_name(), dtype=torch.bfloat16) + with torch.no_grad(): + for param in model.parameters(): + param.fill_(0.5) + references = { + name: torch.nn.Parameter(param.detach().float().cpu().clone()) + for name, param in model.named_parameters() + } + reference_optimizer = torch.optim.SGD(references.values(), lr=1 / 32) + config = { + "train_micro_batch_size_per_gpu": 4, + "gradient_accumulation_steps": 1, + "gradient_clipping": 0, + "bf16": { + "enabled": True, + "bf16_master_weights_and_grads": False, + "bf16_optimizer_states": False, + }, + "data_types": { + "grad_accum_dtype": accumulation_dtype + }, + "zero_allow_untested_optimizer": True, + "zero_force_ds_cpu_optimizer": False, + "zero_optimization": { + "stage": zero_stage, + "reduce_bucket_size": 8 + }, + } + if cpu_offload: + config["zero_optimization"]["offload_optimizer"] = {"device": "cpu", "pin_memory": False} + engine, _, _, _ = deepspeed.initialize(model=model, + optimizer=torch.optim.SGD(model.parameters(), lr=1 / 32), + config=config, + dist_init_required=False) + calls = ["c0", "c1", "c2", "c3"] + windows = [[calls], [["ping"], calls], [["c0"], ["c1"]], [["c0"], ["c1"], ["c2", "c3"]], [["c0"]]] + for step, requests in enumerate(windows): + expected = {name: torch.zeros_like(reference) for name, reference in references.items()} + for request in requests: + for index, name in enumerate(request): + offset = 0 if name == "ping" else calls.index(name) + scale = 0 if name == "ping" else 1 + inputs = (torch.arange(16).reshape(4, 4) + step + rank + offset) / 16 + engine.set_gradient_accumulation_boundary(index == len(request) - 1) + device_inputs = inputs.to(device=engine.device, dtype=torch.bfloat16).requires_grad_(checkpointed) + engine.backward(engine(device_inputs).float() * scale) + # Linear sum-loss gradients are input column sums, averaged over the two ranks. + rank_mean_inputs = inputs + (0.5 - rank) / 16 + for grad in expected.values(): + grad.add_(rank_mean_inputs.sum(dim=0).expand_as(grad) * scale) + for name, param in model.named_parameters(): + torch.testing.assert_close(safe_get_full_grad(param).float().cpu(), + expected[name], + rtol=0, + atol=0, + msg=f"{name}: step={step}, requests={requests}") + references[name].grad = expected[name] + reference_optimizer.step() + engine.step() + for name, param in model.named_parameters(): + torch.testing.assert_close(param.detach().cpu(), + references[name].detach().to(torch.bfloat16), + rtol=0, + atol=0, + msg=f"{name}: update at step={step}, requests={requests}") + finally: + if engine is not None: + engine.destroy() + dist.destroy_process_group() + + +@pytest.mark.skipif(get_accelerator().device_name() not in ("cpu", "cuda") + or (get_accelerator().device_name() == "cuda" and get_accelerator().device_count() < 2), + reason="Requires CPU or two CUDA devices") +@pytest.mark.parametrize("zero_stage", [1, 2]) +@pytest.mark.parametrize("cpu_offload", [False, True], ids=["device-optimizer", "cpu-offload"]) +@pytest.mark.parametrize("accumulation_dtype", ["bf16", "fp32"]) +@pytest.mark.parametrize("checkpointed", [False, True], ids=["plain", "reentrant"]) +def test_multiple_request_boundaries_before_step(tmp_path, zero_stage, cpu_offload, accumulation_dtype, checkpointed): + torch.multiprocessing.spawn(_run_request_boundaries, + args=(f"file://{tmp_path / 'rendezvous'}", zero_stage, cpu_offload, accumulation_dtype, + checkpointed), + nprocs=2, + join=True) From 7e00e1a8e18cbb8e17670883b27873f1bbdaee13 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Fri, 30 Oct 2026 00:14:19 +0000 Subject: [PATCH 2/3] address review Signed-off-by: Stas Bekman --- deepspeed/runtime/zero/stage_1_and_2.py | 29 ++++++++---- .../v1/zero/test_zero2_gradient_safety.py | 46 +++++++++++++++++-- .../zero/test_zero2_offload_multi_backward.py | 22 ++++++--- 3 files changed, 78 insertions(+), 19 deletions(-) diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index 54013a00e311..ef58238e5651 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -1852,7 +1852,9 @@ def copy_gradients_to_cpu(): if param_id not in self.accumulated_grads_in_cpu: self.accumulated_grads_in_cpu[param_id] = buffer_to_accumulate_to_in_cpu() - if self.micro_step_id > 0: + # Other parameters may have run backwards before this parameter's first use. + # Only its participation in this window makes the persistent host buffer current. + if param_id in self._offload_accumulated_param_ids: accumulate_gradients() copy_gradients_to_cpu() self._record_offload_copy() @@ -2159,23 +2161,30 @@ def copy_grads_in_partition(self, param, group_idx=None): self.clear_grad_attribute(param) return if self.cpu_offload: - # A boundary does not guarantee that step() follows. Refresh the host - # accumulator on the first backward too, or a later request adds the - # previous step's gradients. Low-precision masters already share the - # boundary gradient buffer and need no separate first-backward copy. - if (not self.low_precision_master_weights_and_grads or self.micro_step_id > 0 - or not self.is_gradient_accumulation_boundary()): + param_id = self.get_param_id(param) + first_use = param_id not in self._offload_accumulated_param_ids + boundary = self.is_gradient_accumulation_boundary() + if self.low_precision_master_weights_and_grads and first_use and boundary: + # Low-precision masters accumulate in the optimizer's partition grad. + # The boundary copy below initializes that shared buffer, so only + # register its view here; no separate accumulator copy is needed. + [i, _, dest_offset, num_elements] = self.grad_position[param_id] + self.accumulated_grads_in_cpu[param_id] = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow( + 0, dest_offset, num_elements) + else: + # FP32 masters have a separate host accumulator, which must be + # refreshed on first use even at a boundary: step() may come later. self.async_accumulate_grad_in_cpu_via_gpu(param) - # Record active param so unmanaged finalize skips params unused this window. - self._offload_accumulated_param_ids.add(self.get_param_id(param)) - if self.is_gradient_accumulation_boundary(): + if boundary: self.set_norm_for_param_grad_in_gpu(param) self.update_offload_overflow_tracker_for_param_grad(param) self.async_inplace_copy_grad_to_fp32_buffer_from_gpu(param) + # Include the shared-buffer fast path in subsequent accumulation and finalization. + self._offload_accumulated_param_ids.add(param_id) return if self._use_compact_grad_partition_buffer(group_idx): if self.use_grad_accum_attribute: diff --git a/tests/unit/v1/zero/test_zero2_gradient_safety.py b/tests/unit/v1/zero/test_zero2_gradient_safety.py index a6a83101e0da..d58bef8e4fc8 100644 --- a/tests/unit/v1/zero/test_zero2_gradient_safety.py +++ b/tests/unit/v1/zero/test_zero2_gradient_safety.py @@ -138,12 +138,23 @@ def test_oversized_copy_is_independent_and_preserves_routing(dtype, cpu_offload, assert opt.ipg_buckets[torch.float32].params == [(0, 0, 0)] -def test_default_offload_accumulates_across_nonboundary_backwards(): - opt, param, master = make_optimizer() - backward_values(opt, param, [1, 2, 4], [False, False, True]) +@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"]) +@pytest.mark.parametrize("first_boundary", [False, True], ids=["nonboundary-first", "boundary-first"]) +def test_default_offload_accumulates_across_nonboundary_backwards(low_precision, first_boundary): + opt, param, master = make_optimizer(low_precision=low_precision) + backward_values(opt, param, [1, 2, 4], [first_boundary, False, True]) assert torch.equal(master.grad, torch.full_like(master, 7)) +@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"]) +def test_boundary_offload_can_be_finalized_without_another_backward(low_precision): + opt, param, master = make_optimizer(low_precision=low_precision) + backward_values(opt, param, [3], [True]) + opt.finalize_gradient_accumulation_boundary() + opt._wait_for_offload_copies() + assert torch.equal(master.grad, torch.full_like(master, 3)) + + def test_reduction_and_offload_consume_owned_oversized_buffer(): opt, param, master = make_optimizer() opt.is_param_in_current_partition = {0: True} @@ -369,6 +380,35 @@ def test_step_consumes_gradients_after_offload_copies_complete(monkeypatch): assert torch.equal(master, ref) +@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"]) +@pytest.mark.parametrize("first_boundary", [False, True], ids=["nonboundary-first", "boundary-first"]) +def test_offload_first_use_after_another_backward_starts_fresh(monkeypatch, low_precision, first_boundary): + opt, param, master = make_optimizer(low_precision=low_precision) + prepare_step(opt, param, master, monkeypatch) + ref = torch.nn.Parameter(master.detach().clone()) + ref_optimizer = torch.optim.Adam([ref], lr=0.01) + backward_values(opt, param, [8], [True]) + opt.step() + ref.grad = torch.full_like(ref, 8) + ref_optimizer.step() + # Other parameters have already contributed to the new window; this one has not. + opt.micro_step_id = 1 + opt.set_gradient_accumulation_boundary(first_boundary) + param.grad = torch.full_like(param, 2) + opt.copy_grads_in_partition(param) + param.grad = None + opt.micro_step_id = 2 + opt.set_gradient_accumulation_boundary(True) + param.grad = torch.full_like(param, 4) + opt.copy_grads_in_partition(param) + opt._wait_for_offload_copies() + assert torch.equal(master.grad, torch.full_like(master, 6)) + opt.step() + ref.grad = torch.full_like(ref, 6) + ref_optimizer.step() + assert torch.equal(master, ref) + + @pytest.mark.parametrize("operation", ["reset", "legacy_load", "universal_load"]) def test_cpu_buffer_reuse_waits_for_offload_copies(operation): opt, _, _ = make_optimizer() diff --git a/tests/unit/v1/zero/test_zero2_offload_multi_backward.py b/tests/unit/v1/zero/test_zero2_offload_multi_backward.py index 5fad98c70371..4d778d6c10a4 100644 --- a/tests/unit/v1/zero/test_zero2_offload_multi_backward.py +++ b/tests/unit/v1/zero/test_zero2_offload_multi_backward.py @@ -154,9 +154,10 @@ def __init__(self, checkpointed): self.projections = torch.nn.ModuleList([torch.nn.Linear(4, 4, bias=False), torch.nn.Linear(4, 2, bias=False)]) self.checkpointed = checkpointed - def forward(self, inputs): + def forward(self, inputs, projection_indices=(0, 1)): loss = inputs.new_zeros(()) - for projection in self.projections: + for projection_index in projection_indices: + projection = self.projections[projection_index] if self.checkpointed: outputs = checkpoint(projection, inputs, use_reentrant=True) else: @@ -214,8 +215,15 @@ def _run_request_boundaries(rank, rendezvous, zero_stage, cpu_offload, accumulat config=config, dist_init_required=False) calls = ["c0", "c1", "c2", "c3"] - windows = [[calls], [["ping"], calls], [["c0"], ["c1"]], [["c0"], ["c1"], ["c2", "c3"]], [["c0"]]] - for step, requests in enumerate(windows): + windows = [ + (requests, {}) + for requests in [[calls], [["ping"], calls], [["c0"], ["c1"]], [["c0"], ["c1"], ["c2", "c3"]], [["c0"]]] + ] + # Each projection first joins a later backward after contributing to the preceding update. + for early, late in [(0, 1), (1, 0)]: + windows.append(([["c0"], ["c1", "c2"], ["c3"]], {"c0": (early, ), "c1": (late, ), "c2": (late, )})) + reference_names = list(references) + for step, (requests, selections) in enumerate(windows): expected = {name: torch.zeros_like(reference) for name, reference in references.items()} for request in requests: for index, name in enumerate(request): @@ -224,10 +232,12 @@ def _run_request_boundaries(rank, rendezvous, zero_stage, cpu_offload, accumulat inputs = (torch.arange(16).reshape(4, 4) + step + rank + offset) / 16 engine.set_gradient_accumulation_boundary(index == len(request) - 1) device_inputs = inputs.to(device=engine.device, dtype=torch.bfloat16).requires_grad_(checkpointed) - engine.backward(engine(device_inputs).float() * scale) + projection_indices = selections.get(name, (0, 1)) + engine.backward(engine(device_inputs, projection_indices).float() * scale) # Linear sum-loss gradients are input column sums, averaged over the two ranks. rank_mean_inputs = inputs + (0.5 - rank) / 16 - for grad in expected.values(): + for projection_index in projection_indices: + grad = expected[reference_names[projection_index]] grad.add_(rank_mean_inputs.sum(dim=0).expand_as(grad) * scale) for name, param in model.named_parameters(): torch.testing.assert_close(safe_get_full_grad(param).float().cpu(), From e5d68d99955c308e0060ba06f590b28109f54393 Mon Sep 17 00:00:00 2001 From: Stas Bekman Date: Fri, 30 Oct 2026 00:53:46 +0000 Subject: [PATCH 3/3] ZeRO-1/2: allocate CPU gradient accumulation lazily --- deepspeed/runtime/zero/stage_1_and_2.py | 32 ++++--- .../v1/zero/test_zero2_gradient_safety.py | 90 +++++++++++++++++-- .../zero/test_zero2_offload_multi_backward.py | 7 +- 3 files changed, 109 insertions(+), 20 deletions(-) diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index ef58238e5651..e0f19da0ff77 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -650,6 +650,8 @@ def _enforce_cpu_offload(): self.grad_position = {} # Param ids reduced since the last step(); used to finalize only active params in unmanaged mode. self._offload_accumulated_param_ids = set() + # FP32-master gradients first seen at a boundary stay in the partition until another backward. + self._offload_partition_only_param_ids = set() self.temp_grad_buffer_for_cpu_offload = torch.zeros(largest_param_numel, device=self.device, dtype=self.dtype) @@ -1099,6 +1101,8 @@ def _finalize_cpu_offload_gradient_accumulation(self): continue if self.get_param_id(param) not in self._offload_accumulated_param_ids: continue + if self.get_param_id(param) in self._offload_partition_only_param_ids: + continue if self._use_compact_grad_partition_buffer(group_idx): self.copy_grads_in_partition(param, group_idx) continue @@ -1829,7 +1833,13 @@ def buffer_to_accumulate_to_in_cpu(): #accumulate gradients into param.grad_accum or parts of it that belongs to this partition def accumulate_gradients(): grad_accum = self.get_param_gradient_attribute(param) - if not self.low_precision_master_weights_and_grads: + if param_id in self._offload_partition_only_param_ids: + # Recover only this rank's fragment, converting back to the accumulation dtype. + previous = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow(0, dest_offset, num_elements) + local_buffer = dest_buffer.narrow(0, source_offset, num_elements) + local_buffer.copy_(previous, non_blocking=True) + grad_accum.data.view(-1).narrow(0, source_offset, num_elements).add_(local_buffer) + elif not self.low_precision_master_weights_and_grads: dest_buffer.copy_(self.accumulated_grads_in_cpu[param_id].view(-1), non_blocking=True) grad_accum.data.view(-1).add_(dest_buffer) else: @@ -1857,6 +1867,7 @@ def copy_gradients_to_cpu(): if param_id in self._offload_accumulated_param_ids: accumulate_gradients() copy_gradients_to_cpu() + self._offload_partition_only_param_ids.discard(param_id) self._record_offload_copy() def set_norm_for_param_grad(self, param): @@ -2164,16 +2175,16 @@ def copy_grads_in_partition(self, param, group_idx=None): param_id = self.get_param_id(param) first_use = param_id not in self._offload_accumulated_param_ids boundary = self.is_gradient_accumulation_boundary() - if self.low_precision_master_weights_and_grads and first_use and boundary: - # Low-precision masters accumulate in the optimizer's partition grad. - # The boundary copy below initializes that shared buffer, so only - # register its view here; no separate accumulator copy is needed. - [i, _, dest_offset, num_elements] = self.grad_position[param_id] - self.accumulated_grads_in_cpu[param_id] = self.single_partition_of_fp32_groups[i].grad.view(-1).narrow( - 0, dest_offset, num_elements) + if first_use and boundary: + if self.low_precision_master_weights_and_grads: + # Low-precision masters accumulate in the optimizer's partition grad. + [i, _, dest_offset, num_elements] = self.grad_position[param_id] + self.accumulated_grads_in_cpu[param_id] = self.single_partition_of_fp32_groups[i].grad.view( + -1).narrow(0, dest_offset, num_elements) + else: + # The boundary copy is sufficient unless another backward follows before step(). + self._offload_partition_only_param_ids.add(param_id) else: - # FP32 masters have a separate host accumulator, which must be - # refreshed on first use even at a boundary: step() may come later. self.async_accumulate_grad_in_cpu_via_gpu(param) if boundary: @@ -3028,6 +3039,7 @@ def step(self, closure=None): self._wait_for_offload_copies() if self.cpu_offload: self._offload_accumulated_param_ids = set() + self._offload_partition_only_param_ids = set() if not self.compute_grad_norm: self._global_grad_norm = None diff --git a/tests/unit/v1/zero/test_zero2_gradient_safety.py b/tests/unit/v1/zero/test_zero2_gradient_safety.py index d58bef8e4fc8..e430d3691b1d 100644 --- a/tests/unit/v1/zero/test_zero2_gradient_safety.py +++ b/tests/unit/v1/zero/test_zero2_gradient_safety.py @@ -50,6 +50,7 @@ def make_optimizer(dtype=torch.bfloat16, device=None, low_precision=False, cpu_o opt.micro_step_id = 0 opt._pending_offload_events = {} opt._offload_accumulated_param_ids = set() + opt._offload_partition_only_param_ids = set() opt.accumulated_grads_in_cpu = {} opt.norm_for_param_grads = {} opt.local_overflow = False @@ -151,10 +152,60 @@ def test_boundary_offload_can_be_finalized_without_another_backward(low_precisio opt, param, master = make_optimizer(low_precision=low_precision) backward_values(opt, param, [3], [True]) opt.finalize_gradient_accumulation_boundary() + opt.finalize_gradient_accumulation_boundary() opt._wait_for_offload_copies() assert torch.equal(master.grad, torch.full_like(master, 3)) +@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"]) +def test_single_boundary_offload_avoids_auxiliary_storage_and_copy(monkeypatch, low_precision): + opt, param, master = make_optimizer(low_precision=low_precision) + prepare_step(opt, param, master, monkeypatch) + ref = torch.nn.Parameter(master.detach().clone()) + ref_optimizer = torch.optim.Adam([ref], lr=0.01) + original_copy = torch.Tensor.copy_ + copies = [] + + def count_copy(destination, source, *args, **kwargs): + copies.append(destination.data_ptr()) + return original_copy(destination, source, *args, **kwargs) + + for value in [2, 3, 4]: + copies.clear() + with monkeypatch.context() as patch: + patch.setattr(torch.Tensor, "copy_", count_copy) + backward_values(opt, param, [value], [True]) + # Pin the resource regression: one partition copy, no additional host storage. + assert copies == [master.grad.data_ptr()] + for buffer in opt.accumulated_grads_in_cpu.values(): + assert buffer.untyped_storage().data_ptr() == master.grad.untyped_storage().data_ptr() + assert torch.equal(master.grad, torch.full_like(master, value)) + opt.finalize_gradient_accumulation_boundary() + opt.step() + ref.grad = torch.full_like(ref, value) + ref_optimizer.step() + assert torch.equal(master, ref) + + +@pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"]) +@pytest.mark.parametrize("source_offset,dest_offset,num_elements", [(0, 0, 8), (2, 1, 3), (5, 4, 3)]) +def test_boundary_first_offload_preserves_owned_fragment_and_accumulation_dtype(low_precision, source_offset, + dest_offset, num_elements): + opt, param, master = make_optimizer(low_precision=low_precision) + opt.grad_position[0] = [0, source_offset, dest_offset, num_elements] + master.grad.fill_(123) + reference = torch.nn.Parameter(torch.zeros_like(param)) + # BF16 rounds 256 + 1 back to 256; accumulating in the FP32 master would differ. + for value, boundary in [(256, True), (1, False), (2, True)]: + (reference * value).sum().backward() + backward_values(opt, param, [value], [boundary]) + opt.finalize_gradient_accumulation_boundary() + opt._wait_for_offload_copies() + expected = torch.full_like(master, 123) + expected.narrow(0, dest_offset, num_elements).copy_(reference.grad.narrow(0, source_offset, num_elements)) + assert torch.equal(master.grad, expected) + + def test_reduction_and_offload_consume_owned_oversized_buffer(): opt, param, master = make_optimizer() opt.is_param_in_current_partition = {0: True} @@ -317,11 +368,19 @@ def test_record_stream_is_separate_from_readiness(monkeypatch): assert not stream.waited -def test_cpu_accumulator_waits_before_pageable_source_is_read(): +@pytest.mark.parametrize("first_boundary", [False, True], ids=["host-accumulator", "optimizer-partition"]) +def test_cpu_accumulator_waits_before_pageable_source_is_read(first_boundary): opt, param, master = make_optimizer() - backward_values(opt, param, [1], [False]) - # The earlier copy into the CPU accumulator only lands once its event completes. - event = FakeEvent(lambda: opt.accumulated_grads_in_cpu[0].fill_(3)) + backward_values(opt, param, [1], [first_boundary]) + + # The earlier CPU copy only lands once its event completes. + def complete_copy(): + if first_boundary: + master.grad.fill_(3) + if 0 in opt.accumulated_grads_in_cpu: + opt.accumulated_grads_in_cpu[0].fill_(3) + + event = FakeEvent(complete_copy) opt._pending_offload_events["producer"] = event opt.micro_step_id = 1 opt.set_gradient_accumulation_boundary(True) @@ -382,15 +441,27 @@ def test_step_consumes_gradients_after_offload_copies_complete(monkeypatch): @pytest.mark.parametrize("low_precision", [False, True], ids=["fp32-master", "bf16-master"]) @pytest.mark.parametrize("first_boundary", [False, True], ids=["nonboundary-first", "boundary-first"]) -def test_offload_first_use_after_another_backward_starts_fresh(monkeypatch, low_precision, first_boundary): +@pytest.mark.parametrize("overflow", [False, True], ids=["successful-step", "overflow-step"]) +def test_offload_first_use_after_another_backward_starts_fresh(monkeypatch, low_precision, first_boundary, overflow): opt, param, master = make_optimizer(low_precision=low_precision) prepare_step(opt, param, master, monkeypatch) ref = torch.nn.Parameter(master.detach().clone()) ref_optimizer = torch.optim.Adam([ref], lr=0.01) - backward_values(opt, param, [8], [True]) + # Warm the persistent auxiliary buffer before testing first use in a later window. + backward_values(opt, param, [8], [False]) + opt.finalize_gradient_accumulation_boundary() opt.step() ref.grad = torch.full_like(ref, 8) ref_optimizer.step() + assert torch.equal(master, ref) + if overflow: + # A boundary-only overflow must clear partition-only state without reusing the warmed buffer. + backward_values(opt, param, [7], [True]) + opt.finalize_gradient_accumulation_boundary() + opt.local_overflow = True + opt.step() + assert opt.overflow + assert torch.equal(master, ref) # Other parameters have already contributed to the new window; this one has not. opt.micro_step_id = 1 opt.set_gradient_accumulation_boundary(first_boundary) @@ -407,6 +478,13 @@ def test_offload_first_use_after_another_backward_starts_fresh(monkeypatch, low_ ref.grad = torch.full_like(ref, 6) ref_optimizer.step() assert torch.equal(master, ref) + # A later boundary-only window must not reuse the warmed auxiliary sum. + backward_values(opt, param, [3], [True]) + opt.finalize_gradient_accumulation_boundary() + opt.step() + ref.grad = torch.full_like(ref, 3) + ref_optimizer.step() + assert torch.equal(master, ref) @pytest.mark.parametrize("operation", ["reset", "legacy_load", "universal_load"]) diff --git a/tests/unit/v1/zero/test_zero2_offload_multi_backward.py b/tests/unit/v1/zero/test_zero2_offload_multi_backward.py index 4d778d6c10a4..c04e8b9e2303 100644 --- a/tests/unit/v1/zero/test_zero2_offload_multi_backward.py +++ b/tests/unit/v1/zero/test_zero2_offload_multi_backward.py @@ -215,10 +215,9 @@ def _run_request_boundaries(rank, rendezvous, zero_stage, cpu_offload, accumulat config=config, dist_init_required=False) calls = ["c0", "c1", "c2", "c3"] - windows = [ - (requests, {}) - for requests in [[calls], [["ping"], calls], [["c0"], ["c1"]], [["c0"], ["c1"], ["c2", "c3"]], [["c0"]]] - ] + windows = [(requests, {}) + for requests in [[calls], [["ping"], calls], [["c0"], ["c1"]], [["c0"], ["c1"], ["c2", "c3"]], + [["c0"], ["c1", "c2", "c3"]], [["c0"]]]] # Each projection first joins a later backward after contributing to the preceding update. for early, late in [(0, 1), (1, 0)]: windows.append(([["c0"], ["c1", "c2"], ["c3"]], {"c0": (early, ), "c1": (late, ), "c2": (late, )}))