From e4bcfd9e27b8b60a852f16722d61a9c468e893f4 Mon Sep 17 00:00:00 2001 From: Masahiro Tanaka Date: Wed, 26 Aug 2026 16:15:53 -0700 Subject: [PATCH] DeepCompile: stabilize ZeRO-3 parameter guards Signed-off-by: Masahiro Tanaka --- deepspeed/compile/patch_fake_tensor.py | 57 +++++++++ deepspeed/runtime/zero/parameter_offload.py | 8 +- tests/unit/compile/test_zero3_grad_dtype.py | 126 ++++++++++++++++++++ 3 files changed, 190 insertions(+), 1 deletion(-) diff --git a/deepspeed/compile/patch_fake_tensor.py b/deepspeed/compile/patch_fake_tensor.py index 2595ceef6d19..72a2195a83ce 100644 --- a/deepspeed/compile/patch_fake_tensor.py +++ b/deepspeed/compile/patch_fake_tensor.py @@ -3,6 +3,8 @@ # DeepSpeed Team +from functools import wraps + import torch from deepspeed.accelerator import get_accelerator @@ -10,7 +12,9 @@ try: from torch._subclasses import FakeTensorMode from torch._subclasses.fake_tensor import unset_fake_temporarily + from torch._dynamo.guards import GuardBuilder, get_verbose_code_parts from torch._dynamo.variables.builder import wrap_to_fake_tensor_and_record + from torch.utils.weak import TensorWeakRef except ImportError: # Unsupported torch version pass @@ -42,6 +46,35 @@ def _get_guard_sizes_strides(t): return t.size(), t.stride() +def _zero3_parameter_guard_metadata(param): + return (type(param), torch._C._dispatch_keys(param), param.dtype, param.device, param.layout, param.requires_grad, + param.ds_id, tuple(param.ds_shape)) + + +def _matches_zero3_parameter_metadata(param, expected): + return isinstance(param, torch.Tensor) and _zero3_parameter_guard_metadata(param) == expected + + +def _get_guard_source_value(builder, guard): + try: + return builder.get(guard) + except TypeError: + # Older PyTorch releases accept the guard's source expression only. + return builder.get(guard.name) + + +def _resolve_zero3_guarded_value(builder, guard, value): + guarded_value = value() if isinstance(value, TensorWeakRef) else value + if not (hasattr(guarded_value, "ds_id") and hasattr(guarded_value, "ds_shape")): + # Some Dynamo paths pass the full-shape dummy/FakeTensor produced + # above instead of the module parameter. Resolve the source so all + # guards for a ZeRO parameter use its stable semantic metadata. + source_value = _get_guard_source_value(builder, guard) + if hasattr(source_value, "ds_id") and hasattr(source_value, "ds_shape"): + guarded_value = source_value + return guarded_value if guarded_value is not None else _get_guard_source_value(builder, guard) + + def patch_fake_tensor(): # dynamo tracer uses wrap_to_fake_tensor_and_record # Wrapping FakeTensorMode.from_tensor is not sufficient as dynamo generates SymbolicContext before calling from_tensor @@ -68,6 +101,30 @@ def wrap_to_fake_tensor_and_record_wrapper(t, *args, **kwargs): torch._dynamo.variables.builder.wrap_to_fake_tensor_and_record = wrap_to_fake_tensor_and_record_wrapper + original_tensor_match = GuardBuilder.TENSOR_MATCH + + @wraps(original_tensor_match) + def tensor_match_wrapper(self, guard, value=None): + guarded_value = _resolve_zero3_guarded_value(self, guard, value) + if hasattr(guarded_value, "ds_id") and hasattr(guarded_value, "ds_shape"): + expected = _zero3_parameter_guard_metadata(guarded_value) + + def semantic_parameter_guard(param): + return _matches_zero3_parameter_metadata(param, expected) + + code = f"DeepCompile ZeRO-3 semantic parameter ds_id={expected[-2]} shape={expected[-1]}" + guard_manager = self.get_guard_manager(guard) + verbose_code = get_verbose_code_parts(code, guard) + try: + guard_manager.add_lambda_guard(semantic_parameter_guard, verbose_code, guard.user_stack) + except TypeError: + # PyTorch 2.10 does not yet accept the user stack argument. + guard_manager.add_lambda_guard(semantic_parameter_guard, verbose_code) + return + return original_tensor_match(self, guard, value) + + GuardBuilder.TENSOR_MATCH = tensor_match_wrapper + # aot_module_simplified uses fake_mode.from_tensor to process inputs original_from_tensor = FakeTensorMode.from_tensor diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py index f9fc7c2c0e64..6d381252a52b 100644 --- a/deepspeed/runtime/zero/parameter_offload.py +++ b/deepspeed/runtime/zero/parameter_offload.py @@ -62,7 +62,13 @@ def __getitem__(self, key): if param is None: return param - is_eager_forward = self._parent_module._parameters._in_forward and not torch.compiler.is_compiling() + # Dynamo traces this getter while lifting module parameters. The + # physical ZeRO status is intentionally volatile across compiled + # forwards and must not become a cache guard. + if torch.compiler.is_compiling(): + return param + + is_eager_forward = self._parent_module._parameters._in_forward fallback = None if hasattr(param, "ds_status") and is_eager_forward: from deepspeed.compile.z3_eager_fallback import get_active_z3_eager_fallback, is_dynamo_guard_evaluation diff --git a/tests/unit/compile/test_zero3_grad_dtype.py b/tests/unit/compile/test_zero3_grad_dtype.py index 38552062a454..30355fe14c5a 100644 --- a/tests/unit/compile/test_zero3_grad_dtype.py +++ b/tests/unit/compile/test_zero3_grad_dtype.py @@ -7,9 +7,13 @@ import torch from deepspeed.compile import backend as backend_mod +import deepspeed.compile.patch_fake_tensor as patch_fake_tensor_mod from deepspeed.compile.init_z3 import _allow_dynamo_dynamic_parameter_shapes_for_z3, _resolve_expected_grad_dtype +from deepspeed.compile.patch_fake_tensor import _resolve_zero3_guarded_value, patch_fake_tensor from deepspeed.compile.patch_compiled_func import (get_backward_inputs, pop_backward_input, register_backward_frame) from deepspeed.runtime.engine import DeepSpeedEngine +from deepspeed.runtime.zero.parameter_offload import ZeROOrderedDict +from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus from deepspeed.utils.torch import required_torch_version @@ -55,6 +59,128 @@ class FakeDynamo: restore() +def test_zero3_tensor_match_delegates_non_zero_tensor(monkeypatch): + from torch._dynamo.guards import GuardBuilder + from torch._subclasses import FakeTensorMode + + calls = [] + sentinel = object() + + def original_tensor_match(builder, guard, value=None): + calls.append((builder, guard, value)) + return sentinel + + monkeypatch.setattr(GuardBuilder, "TENSOR_MATCH", original_tensor_match) + monkeypatch.setattr( + torch._dynamo.variables.builder, + "wrap_to_fake_tensor_and_record", + torch._dynamo.variables.builder.wrap_to_fake_tensor_and_record, + ) + monkeypatch.setattr(FakeTensorMode, "from_tensor", FakeTensorMode.from_tensor) + patch_fake_tensor() + wrapped_tensor_match = GuardBuilder.TENSOR_MATCH + + class Builder: + + def get(self, _guard): + return value + + builder = Builder() + guard = object() + value = torch.ones(1) + + assert wrapped_tensor_match(builder, guard, value) is sentinel + assert calls == [(builder, guard, value)] + + +def test_zero3_semantic_guard_ignores_transient_physical_state_changes(monkeypatch): + + class Accelerator: + + def is_pinned(self, _tensor): + return False + + monkeypatch.setattr(patch_fake_tensor_mod, "get_accelerator", Accelerator) + + class Module(torch.nn.Module): + + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.empty(0)) + self.weight.ds_id = 1 + self.weight.ds_shape = torch.Size((2, 2)) + self.weight.ds_status = ZeroParamStatus.NOT_AVAILABLE + + params = ZeROOrderedDict(parent_module=self) + params.update(self._parameters) + params._in_forward = True + self._parameters = params + + def forward(self, value): + return torch.nn.functional.linear(value, self.weight) + + module = Module() + param = next(module.parameters()) + backend_calls = [] + + def backend(_gm, _inputs): + backend_calls.append(None) + param.data = torch.ones(2, 2, device=param.device) + return lambda _weight, value: (value, ) + + patch_fake_tensor() + compiled = torch.compile(module, backend=backend) + value = torch.ones(1, 2, device=param.device) + + assert torch.equal(compiled(value), value) + param.data = torch.empty(0, device=param.device) + param.ds_status = ZeroParamStatus.AVAILABLE + assert torch.equal(compiled(value), value) + assert len(backend_calls) == 1 + + +def test_zero3_semantic_guard_resolves_real_parameter_from_dummy_value(): + param = torch.nn.Parameter(torch.empty(0)) + param.ds_id = 1 + param.ds_shape = torch.Size((2, 2)) + + class Guard: + name = "param" + + guard = Guard() + + class Builder: + + def get(self, guard_or_name): + if guard_or_name is guard: + raise TypeError("older PyTorch expects the source expression") + assert guard_or_name == guard.name + return param + + dummy = torch.nn.Parameter(torch.empty(2, 2)) + + assert _resolve_zero3_guarded_value(Builder(), guard, dummy) is param + + +def test_zero_ordered_dict_does_not_read_ds_status_while_compiling(monkeypatch): + + class StatusTrap: + + @property + def ds_status(self): + raise AssertionError("volatile ds_status must not be read while compiling") + + module = torch.nn.Module() + params = ZeROOrderedDict(parent_module=module) + param = StatusTrap() + params["weight"] = param + params._in_forward = True + module._parameters = params + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + assert module._parameters["weight"] is param + + @pytest.mark.parametrize("first_owner_to_restore", [0, 1]) def test_zero3_dynamo_config_restores_after_last_overlapping_owner(monkeypatch, first_owner_to_restore):