wehub-resource-sync
204 lines
8.0 KiB
Python
204 lines
8.0 KiB
Python
# 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)
|