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