chore: import upstream snapshot with attribution

tests/unit/v1/zero/test_zero_functorch_linear.py

@@ -0,0 +1,219 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""Regression: ZeRO-3 linear autograd.Function must work with torch.func transforms.
+
+ZeRO Stage 3 uses ``LinearFunctionForZeroStage3`` (via ``zero3_linear_wrap``) as
+the memory-efficient linear path. After ``deepspeed.initialize``, global
+``torch.nn.functional.linear`` is often the built-in again, so tests call
+``zero3_linear_wrap`` directly-the same ``autograd.Function`` as when the patch
+is active. Legacy ``forward(ctx, ...)`` + ``ctx.save_for_backward`` in forward
+raises on strict functorch builds::
+
+    RuntimeError: In order to use an autograd.Function with functorch
+    transforms ... it must override the setup_context staticmethod.
+"""
+
+import pytest
+import torch
+import torch.nn as nn
+
+import deepspeed
+from deepspeed.accelerator import get_accelerator
+from deepspeed.runtime.zero.linear import zero3_linear_wrap
+
+from unit.common import DistributedTest
+
+
+def _zero3_functorch_config():
+    config = {
+        "train_micro_batch_size_per_gpu": 1,
+        "gradient_accumulation_steps": 1,
+        "steps_per_print": 2147483647,
+        "zero_optimization": {
+            "stage": 3,
+            "stage3_param_persistence_threshold": 0,
+        },
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 1e-3
+            },
+        },
+    }
+    acc = get_accelerator()
+    if acc.is_bf16_supported():
+        config["bf16"] = {"enabled": True}
+    elif acc.is_fp16_supported():
+        config["fp16"] = {"enabled": True, "initial_scale_power": 8}
+    return config
+
+
+class TestZeroFunctorchLinearRegression(DistributedTest):
+    """``torch.func.grad_and_value`` over ``zero3_linear_wrap`` / LinearFunctionForZeroStage3."""
+
+    world_size = 1
+
+    def test_grad_and_value_over_patched_functional_linear(self):
+        if not hasattr(torch, "func"):
+            pytest.skip("torch.func not available")
+
+        model = nn.Linear(8, 8, bias=True)
+        engine, _, _, _ = deepspeed.initialize(
+            model=model,
+            config=_zero3_functorch_config(),
+            model_parameters=model.parameters(),
+        )
+
+        device = engine.device
+        dtype = engine.module.weight.dtype
+        weight = torch.randn(8, 8, device=device, dtype=dtype, requires_grad=True)
+        inp = torch.randn(2, 8, device=device, dtype=dtype, requires_grad=True)
+
+        with torch.enable_grad():
+            probe = zero3_linear_wrap(inp, weight, None)
+        assert "LinearFunctionForZeroStage3" in type(probe.grad_fn).__name__
+
+        def loss_fn(w, x):
+            return zero3_linear_wrap(x, w, None).sum()
+
+        grads, value = torch.func.grad_and_value(loss_fn, argnums=(0, 1))(weight, inp)
+        assert torch.isfinite(value)
+        assert grads[0] is not None and torch.isfinite(grads[0]).all()
+        assert grads[1] is not None and torch.isfinite(grads[1]).all()
+
+
+class TestZeroLinearAutocast(DistributedTest):
+    """Verify autocast state is correctly propagated through forward and backward."""
+
+    world_size = 1
+
+    def _run_forward_backward(self, device, use_autocast, dtype=None):
+        """Run zero3_linear_wrap forward+backward, optionally inside autocast."""
+        weight = torch.randn(4, 4, device=device, dtype=torch.float32, requires_grad=True)
+        inp = torch.randn(2, 4, device=device, dtype=torch.float32, requires_grad=True)
+        bias = torch.randn(4, device=device, dtype=torch.float32, requires_grad=True)
+
+        if use_autocast:
+            with torch.amp.autocast(device_type=device.type, dtype=dtype):
+                out = zero3_linear_wrap(inp, weight, bias)
+        else:
+            out = zero3_linear_wrap(inp, weight, bias)
+
+        loss = out.sum()
+        loss.backward()
+        return out, weight.grad, inp.grad, bias.grad
+
+    def test_backward_without_autocast(self):
+        """Backward without autocast should produce float32 gradients."""
+        model = nn.Linear(4, 4)
+        engine, _, _, _ = deepspeed.initialize(
+            model=model,
+            config=_zero3_functorch_config(),
+            model_parameters=model.parameters(),
+        )
+        device = engine.device
+
+        out, w_grad, i_grad, b_grad = self._run_forward_backward(device, use_autocast=False)
+        assert out.dtype == torch.float32
+        assert w_grad.dtype == torch.float32
+        assert i_grad.dtype == torch.float32
+        assert b_grad.dtype == torch.float32
+
+    def test_backward_with_autocast(self):
+        """Backward with autocast should produce float32 gradients (autocast only affects forward)."""
+        acc = get_accelerator()
+        if acc.is_bf16_supported():
+            amp_dtype = torch.bfloat16
+        elif acc.is_fp16_supported():
+            amp_dtype = torch.float16
+        else:
+            pytest.skip("No half-precision support")
+
+        model = nn.Linear(4, 4)
+        engine, _, _, _ = deepspeed.initialize(
+            model=model,
+            config=_zero3_functorch_config(),
+            model_parameters=model.parameters(),
+        )
+        device = engine.device
+
+        out, w_grad, i_grad, b_grad = self._run_forward_backward(device, use_autocast=True, dtype=amp_dtype)
+        # Forward output should be in reduced precision
+        assert out.dtype == amp_dtype
+        # Gradients accumulate in float32 (master weights)
+        assert w_grad.dtype == torch.float32
+        assert i_grad.dtype == torch.float32
+        assert b_grad.dtype == torch.float32
+
+    def test_no_autocast_leak_into_backward(self):
+        """When forward runs without autocast, an outer autocast during backward must not affect gradient dtype."""
+        model = nn.Linear(4, 4)
+        engine, _, _, _ = deepspeed.initialize(
+            model=model,
+            config=_zero3_functorch_config(),
+            model_parameters=model.parameters(),
+        )
+        device = engine.device
+
+        acc = get_accelerator()
+        if acc.is_bf16_supported():
+            amp_dtype = torch.bfloat16
+        elif acc.is_fp16_supported():
+            amp_dtype = torch.float16
+        else:
+            pytest.skip("No half-precision support")
+
+        weight = torch.randn(4, 4, device=device, dtype=torch.float32, requires_grad=True)
+        inp = torch.randn(2, 4, device=device, dtype=torch.float32, requires_grad=True)
+
+        # Forward WITHOUT autocast
+        out = zero3_linear_wrap(inp, weight, None)
+        assert out.dtype == torch.float32
+
+        # Backward WITH an outer autocast region -- should NOT affect gradient computation
+        # because setup_context captured _fwd_used_autocast=False
+        with torch.amp.autocast(device_type=device.type, dtype=amp_dtype):
+            out.sum().backward()
+
+        assert weight.grad.dtype == torch.float32
+        assert inp.grad.dtype == torch.float32
+
+    def test_setup_context_stores_autocast_attrs(self):
+        """setup_context must store _fwd_used_autocast and _dtype on ctx."""
+        model = nn.Linear(4, 4)
+        engine, _, _, _ = deepspeed.initialize(
+            model=model,
+            config=_zero3_functorch_config(),
+            model_parameters=model.parameters(),
+        )
+        device = engine.device
+
+        weight = torch.randn(4, 4, device=device, dtype=torch.float32, requires_grad=True)
+        inp = torch.randn(2, 4, device=device, dtype=torch.float32, requires_grad=True)
+
+        # Without autocast: setup_context must record that forward did not use autocast
+        out = zero3_linear_wrap(inp, weight, None)
+        grad_fn = out.grad_fn
+        assert hasattr(grad_fn, "_fwd_used_autocast")
+        assert grad_fn._fwd_used_autocast is False
+        assert hasattr(grad_fn, "_dtype")
+        out.sum().backward()
+        assert torch.isfinite(weight.grad).all()
+
+
+class TestLinearFunctionVmap(DistributedTest):
+    """``LinearFunctionForZeroStage3`` must accept ``torch.func.vmap`` directly."""
+
+    world_size = 1
+
+    def test_vmap_over_linear_function(self):
+        from deepspeed.runtime.zero.linear import LinearFunctionForZeroStage3
+        device = get_accelerator().device_name()
+        weight = torch.randn(4, 8, device=device, requires_grad=True)
+        bias = torch.randn(4, device=device, requires_grad=True)
+        xs = torch.randn(3, 8, device=device)
+        y = torch.func.vmap(lambda xi: LinearFunctionForZeroStage3.apply(xi, weight, bias).sum())(xs)
+        ref = torch.func.vmap(lambda xi: (xi @ weight.t() + bias).sum())(xs)
+        assert torch.allclose(y, ref, atol=1e-5)