DeepCompile: stabilize ZeRO-3 parameter guards

deepspeed/compile/patch_fake_tensor.py

@@ -3,14 +3,18 @@
 
 # DeepSpeed Team
 
+from functools import wraps
+
 import torch
 
 from deepspeed.accelerator import get_accelerator
 
 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
 

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

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):