wehub/deepspeedai--deepspeed
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deepspeed/runtime/zero/linear.py
+203
View File
@@ -0,0 +1,203 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
#Linear Module to use with ZeRO Stage 3 to allow for parameter memory release
#after the module execution during forward
#Instead of saving variables using save_for_backward, we save variable ids
#Allowing us to retrieve the variable without creating pointer to it
#Which allows for underlying tensor to be garbage collected
#When partitioned as needed by the Zero Stage 3 optimizer
#TODO instead of patching Linear module, we could patch the ctx.save_for_backward
#ctx.saved_tensors so that this approach works for all nn modules that are built upon
#torch.nn.function. However the issue is that many modules uses C++ implementations
#which does not have pytorch implementation. Eg torch.addmm which acts as a functional
#when implemented outside of torch.autograd.Function
import math
import functools
import torch
from torch import Tensor
from torch.nn.parameter import Parameter
from torch.nn import init
from torch.nn.modules.module import Module
from deepspeed.runtime.utils import noop_decorator
from deepspeed import comm as dist
from deepspeed.accelerator import get_accelerator
def print_rank_0(message, debug=False, force=False):
if dist.get_rank() == 0 and (debug or force):
print(message)
def _get_legacy_autocast_decorators(device_type):
legacy_amp = getattr(getattr(torch, device_type, None), 'amp', None)
custom_fwd = getattr(legacy_amp, 'custom_fwd', None)
custom_bwd = getattr(legacy_amp, 'custom_bwd', None)
if custom_fwd is not None and custom_bwd is not None:
return custom_fwd, custom_bwd
return noop_decorator, noop_decorator
def _get_autocast_decorators():
amp = getattr(torch, 'amp', None)
custom_fwd = getattr(amp, 'custom_fwd', None)
custom_bwd = getattr(amp, 'custom_bwd', None)
if custom_fwd is not None and custom_bwd is not None:
device_type = get_accelerator().device_name()
return functools.partial(custom_fwd, device_type=device_type), functools.partial(custom_bwd,
device_type=device_type)
return _get_legacy_autocast_decorators(get_accelerator().device_name())
autocast_custom_fwd, autocast_custom_bwd = _get_autocast_decorators()
def _is_autocast_enabled(device_type):
try:
return torch.is_autocast_enabled(device_type)
except TypeError:
legacy_getter = getattr(torch, f'is_autocast_{device_type}_enabled', None)
if legacy_getter is not None:
return legacy_getter()
return torch.is_autocast_enabled()
def _get_autocast_dtype(device_type):
try:
return torch.get_autocast_dtype(device_type)
except TypeError:
legacy_getter = getattr(torch, f'get_autocast_{device_type}_dtype', None)
if legacy_getter is not None:
return legacy_getter()
return None
class LinearFunctionForZeroStage3(torch.autograd.Function):
generate_vmap_rule = True
@staticmethod
# bias is an optional argument
def forward(input, weight, bias=None):
if input.dim() == 2 and bias is not None:
# fused op is marginally faster
ret = torch.addmm(bias, input, weight.t())
else:
output = input.matmul(weight.t())
if bias is not None:
output += bias
ret = output
return ret
@staticmethod
def setup_context(ctx, inputs, output):
device_type = get_accelerator().device_name()
ctx._dtype = _get_autocast_dtype(device_type)
ctx._fwd_used_autocast = _is_autocast_enabled(device_type)
input, weight, bias = inputs[0], inputs[1], inputs[2] if len(inputs) > 2 else None
ctx.save_for_backward(input, weight, bias)
# This function has only a single output, so it gets only one gradient
@staticmethod
def backward(ctx, grad_output):
# Match @custom_bwd semantics: always run backward under the same
# autocast state as forward — including explicitly disabling autocast
# when forward did not use it, to guard against outer autocast regions.
device_type = get_accelerator().device_name()
with torch.autocast(device_type=device_type, enabled=ctx._fwd_used_autocast, dtype=ctx._dtype):
input, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
dim = grad_output.dim()
if ctx.needs_input_grad[0]:
grad_input = grad_output.matmul(weight)
if ctx.needs_input_grad[1]:
if dim > 2:
grad_weight = grad_output.reshape(-1, grad_output.shape[-1]).t().matmul(
input.reshape(-1, input.shape[-1]))
else:
grad_weight = grad_output.t().matmul(input)
if bias is not None and ctx.needs_input_grad[2]:
if dim > 2:
grad_bias = grad_output.sum([i for i in range(dim - 1)])
else:
grad_bias = grad_output.sum(0)
return grad_input, grad_weight, grad_bias
def zero3_linear_wrap(input, weight, bias=None):
return LinearFunctionForZeroStage3.apply(input, weight, bias)
class LinearModuleForZeroStage3(Module):
r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.
The weights are pre-transposed and stored as A^T instead of transposing during each
forward. Memory savings proportional to the parameter size.
Args:
in_features: size of each input sample
out_features: size of each output sample
bias: If set to ``False``, the layer will not learn an additive bias.
Default: ``True``
Shape:
- Input: :math:`(N, *, H_{in})` where :math:`*` means any number of
additional dimensions and :math:`H_{in} = \text{in\_features}`
- Output: :math:`(N, *, H_{out})` where all but the last dimension
are the same shape as the input and :math:`H_{out} = \text{out\_features}`.
Attributes:
weight: the learnable weights of the module of shape
:math:`(\text{out\_features}, \text{in\_features})`. The values are
initialized from :math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})`, where
:math:`k = \frac{1}{\text{in\_features}}`
bias: the learnable bias of the module of shape :math:`(\text{out\_features})`.
If :attr:`bias` is ``True``, the values are initialized from
:math:`\mathcal{U}(-\sqrt{k}, \sqrt{k})` where
:math:`k = \frac{1}{\text{in\_features}}`
Examples::
>>> m = nn.Linear(20, 30)
>>> input = torch.randn(128, 20)
>>> output = m(input)
>>> print(output.size())
torch.Size([128, 30])
"""
__constants__ = ['in_features', 'out_features']
in_features: int
out_features: int
weight: Tensor
def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
super(LinearModuleForZeroStage3, self).__init__()
print("Building ZeRO module")
self.in_features = in_features
self.out_features = out_features
self.weight = Parameter(torch.Tensor(out_features, in_features))
if bias:
self.bias = Parameter(torch.Tensor(out_features))
else:
self.register_parameter('bias', None)
self.reset_parameters()
def reset_parameters(self) -> None:
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
def forward(self, input: Tensor) -> Tensor:
return LinearFunctionForZeroStage3.apply(input, self.weight, self.bias)
def extra_repr(self) -> str:
return 'in_features={}, out_features={}, bias={}'.format(self.in_features, self.out_features, self.bias
is not None)