wehub-resource-sync
348 lines
12 KiB
Python
348 lines
12 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from __future__ import annotations
|
|
|
|
import random
|
|
from typing import TYPE_CHECKING, Any, TypedDict
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle import framework
|
|
from paddle.autograd import PyLayer
|
|
from paddle.framework import core
|
|
|
|
from ..meta_parallel.parallel_layers.random import get_rng_state_tracker
|
|
from ..meta_parallel.pp_utils import utils
|
|
from .recompute import (
|
|
check_recompute_necessary,
|
|
custom_state_manager,
|
|
detach_variable,
|
|
switch_rng_state_tracker,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
from typing_extensions import NotRequired
|
|
|
|
from paddle.distributed.communication.group import Group
|
|
from paddle.nn import Layer
|
|
|
|
class _Ctx(TypedDict):
|
|
mp_group: Group
|
|
offload: NotRequired[bool]
|
|
partition: NotRequired[bool]
|
|
|
|
|
|
__all__ = []
|
|
|
|
|
|
def _split_activation(tensor, mp_group):
|
|
mp_degree = mp_group.nranks
|
|
mp_rank = mp_group.rank
|
|
if mp_degree < 2:
|
|
return tensor
|
|
|
|
tensor_numel = paddle.numel(tensor)
|
|
assert tensor_numel != 0, "can't recompute zero element"
|
|
assert tensor_numel % mp_degree == 0, (
|
|
f"The capacity of the activation ({tensor_numel}) cannot be divisible by mp_degree({mp_degree})"
|
|
)
|
|
|
|
# use inplace operation to save memory
|
|
data = tensor.flatten_()
|
|
|
|
part_size = tensor_numel // mp_degree
|
|
start = part_size * mp_rank
|
|
end = start + part_size
|
|
return data[start:end]
|
|
|
|
|
|
def _merge_activation(tensor, mp_group):
|
|
mp_degree = mp_group.nranks
|
|
mp_rank = mp_group.rank
|
|
if mp_degree < 2:
|
|
return tensor
|
|
|
|
# adapt to new dygraph
|
|
tensor_shape = list(tensor.shape)
|
|
tensor_shape[0] *= mp_group.nranks
|
|
out = paddle.empty(tensor_shape, tensor.dtype)
|
|
task = mp_group.process_group.all_gather(tensor.cuda(), out)
|
|
task.wait()
|
|
return out
|
|
|
|
|
|
class _HPRecomputeFunction(PyLayer):
|
|
"""
|
|
Compared with paddle.distributed.fleet.utils.recompute, there are the following differences:
|
|
1. In order to support PipeLineParallel, the input of recompute is modified to ensure that the input can be tuple type.
|
|
2. Offload support for activation
|
|
3. Support MP segmentation of activation to further reduce cuda memory
|
|
4. Adapt to the random state of MP
|
|
"""
|
|
|
|
@staticmethod
|
|
def forward(
|
|
ctx,
|
|
run_function,
|
|
all_outputs,
|
|
mp_group,
|
|
offload,
|
|
partition,
|
|
custom_get_state_func,
|
|
custom_set_state_func,
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
# store for recomputing
|
|
ctx.run_function = run_function
|
|
|
|
ctx.kwargs = kwargs
|
|
|
|
# store the rng states
|
|
ctx.fwd_rng_state = paddle.get_rng_state()
|
|
ctx.fwd_rng_state_tracker = get_rng_state_tracker().get_states_tracker()
|
|
ctx.fwd_numpy_state = np.random.get_state()
|
|
ctx.fwd_random_state = random.getstate()
|
|
ctx.fwd_custom_state = custom_get_state_func()
|
|
ctx.custom_get_state_func = custom_get_state_func
|
|
ctx.custom_set_state_func = custom_set_state_func
|
|
|
|
# save config info
|
|
ctx.mp_group = mp_group
|
|
ctx.offload = offload
|
|
ctx.partition = partition
|
|
|
|
# save input for backward
|
|
ctx.inputs = []
|
|
ctx.tensor_indices = []
|
|
ctx.tensor_shapes = []
|
|
tensor_inputs = []
|
|
|
|
cur_device = paddle.get_device()
|
|
assert (
|
|
'gpu:' in paddle.get_device()
|
|
or 'xpu:' in paddle.get_device()
|
|
or cur_device.split(':')[0]
|
|
in paddle.device.get_all_custom_device_type()
|
|
), f"Recompute with RNG is not support current device: {cur_device}."
|
|
|
|
# TODO support AMP
|
|
tracer = framework._dygraph_tracer()
|
|
ctx.is_fw_autocast = (
|
|
False if tracer._amp_level == core.AmpLevel.O0 else True
|
|
)
|
|
if tracer._amp_level == core.AmpLevel.O2:
|
|
ctx.amp_level = 'O2'
|
|
elif tracer._amp_level in (core.AmpLevel.O1, core.AmpLevel.O0):
|
|
ctx.amp_level = 'O1'
|
|
else:
|
|
raise ValueError(f"unsupported amp level: {tracer._amp_level}")
|
|
ctx.amp_dtype = tracer._amp_dtype
|
|
ctx.amp_white_list, ctx.amp_black_list = tracer._get_amp_op_list()
|
|
|
|
with paddle.no_grad():
|
|
outputs = run_function(*args, **kwargs)
|
|
|
|
for i, arg in enumerate(args):
|
|
if paddle.is_tensor(arg):
|
|
state = arg.stop_gradient
|
|
if partition:
|
|
ctx.tensor_shapes.append(arg.shape)
|
|
partition = _split_activation(
|
|
arg.detach(), mp_group
|
|
).clone()
|
|
# TODO(shenliang03) not use calculate stream to D2H to speed
|
|
arg = partition.cpu() if offload else partition
|
|
else:
|
|
arg = arg.cpu() if offload else arg
|
|
arg.stop_gradient = state
|
|
tensor_inputs.append(arg)
|
|
ctx.tensor_indices.append(i)
|
|
ctx.inputs.append(None)
|
|
|
|
# In new dygraph mode, in some cases a subset of outputs is identity to the subset of inputs,
|
|
# which is inplace operating. When the inputs' stop_gradient is True, an
|
|
# error will occurs because the stop_gradient=True and inplace-op are not
|
|
# supported in the same time. The solution is to mark the inputs non_differentiable
|
|
# if its stop_gradient is True.
|
|
# Note:
|
|
# If not marked non_differentiable, all output tensors' attr `stop gradient`
|
|
# will be reset to `False` in c++ backend.
|
|
# See https://github.com/PaddlePaddle/Paddle/blob/9d62efb0e6e5373823039d9eda96cd5905426c0a/paddle/fluid/pybind/eager_py_layer.cc#L388
|
|
if framework.in_dynamic_mode() and state:
|
|
ctx.mark_non_differentiable(arg)
|
|
else:
|
|
ctx.inputs.append(arg)
|
|
|
|
ctx.save_for_backward(*tensor_inputs)
|
|
|
|
if paddle.is_tensor(outputs):
|
|
all_outputs += [outputs]
|
|
return outputs
|
|
else:
|
|
all_outputs += outputs
|
|
return tuple(outputs)
|
|
|
|
@staticmethod
|
|
def backward(ctx, *args):
|
|
with paddle.base.dygraph.guard():
|
|
# Restore inputs
|
|
inputs = list(ctx.inputs)
|
|
tensor_indices = ctx.tensor_indices
|
|
tensor_shapes = ctx.tensor_shapes
|
|
tensors = list(ctx.saved_tensor())
|
|
|
|
device_id = paddle.distributed.ParallelEnv().device_id
|
|
for i, idx in enumerate(tensor_indices):
|
|
if ctx.partition:
|
|
state = tensors[i].stop_gradient
|
|
tensors[i] = (
|
|
_merge_activation(tensors[i], ctx.mp_group)
|
|
.detach()
|
|
.reshape_(tensor_shapes[i])
|
|
)
|
|
tensors[i].stop_gradient = state
|
|
inputs[idx] = (
|
|
tensors[i].cuda(device_id) if ctx.offload else tensors[i]
|
|
)
|
|
|
|
tracer = framework._dygraph_tracer()
|
|
tracer._has_grad = True
|
|
|
|
# need restore auto_cast state as well as w/b list
|
|
with switch_rng_state_tracker(
|
|
ctx.fwd_rng_state,
|
|
ctx.fwd_rng_state_tracker,
|
|
ctx.fwd_numpy_state,
|
|
ctx.fwd_random_state,
|
|
ctx.fwd_custom_state,
|
|
ctx.custom_get_state_func,
|
|
ctx.custom_set_state_func,
|
|
):
|
|
if ctx.is_fw_autocast:
|
|
with paddle.amp.auto_cast(
|
|
enable=ctx.is_fw_autocast,
|
|
custom_white_list=ctx.amp_white_list,
|
|
custom_black_list=ctx.amp_black_list,
|
|
level=ctx.amp_level,
|
|
dtype=ctx.amp_dtype,
|
|
):
|
|
detached_inputs = detach_variable(tuple(inputs))
|
|
outputs = ctx.run_function(
|
|
*detached_inputs, **ctx.kwargs
|
|
)
|
|
else:
|
|
detached_inputs = detach_variable(tuple(inputs))
|
|
outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
|
|
|
|
if isinstance(outputs, core.eager.Tensor):
|
|
outputs = (outputs,)
|
|
assert len(outputs) == len(args)
|
|
|
|
forward_outputs_with_grad = []
|
|
backward_inputs = []
|
|
|
|
for i in range(len(outputs)):
|
|
if (
|
|
isinstance(outputs[i], core.eager.Tensor)
|
|
and not outputs[i].stop_gradient
|
|
):
|
|
forward_outputs_with_grad.append(outputs[i])
|
|
backward_inputs.append(args[i])
|
|
|
|
if len(forward_outputs_with_grad) == 0:
|
|
raise RuntimeError(
|
|
"none of output has stop_gradient=False, this recompute() is not necessary"
|
|
)
|
|
|
|
# actually backward
|
|
paddle.autograd.backward(forward_outputs_with_grad, backward_inputs)
|
|
grads = tuple(
|
|
inp._grad_ivar()
|
|
for inp in detached_inputs
|
|
if isinstance(inp, core.eager.Tensor)
|
|
)
|
|
return grads
|
|
|
|
|
|
def recompute_hybrid(
|
|
ctx: _Ctx, function: Layer | Callable[..., Any], *args: Any, **kwargs: Any
|
|
) -> Any:
|
|
"""
|
|
recompute intermediate activations to save the memory in hybrid parallel scene.
|
|
# NOTE(shenliang03)The current hybrid parallel recompute has limitations.
|
|
# It cannot handle the following situations:
|
|
# 1. The calculation output of recompute, there are tensors that do not require gradients.
|
|
# 2. The forward output tensor has no gradient. This problem can be solved temporarily by detach().
|
|
# 3. Here, we only use float dtype to distinguish whether a gradient is needed in output tensor
|
|
|
|
Parameters:
|
|
ctx(dict): include 'mp_group', 'offload', and 'partition' keys. the key 'mp_group' (Group), represents the activations are splitted
|
|
in which group. the key 'offload' (bool, optional, default=False), represents whether to offload to cpu. the key 'partition' (bool, optional, default=False),
|
|
represents whether to split activations in the mp_group.
|
|
function(paddle.nn.Layer): layer of sequence of layers that describes part of forward pass of the model
|
|
whose intermediate activations will be released to save memory in forward stage and will be recomputed
|
|
in backward stage for gradient calculation.
|
|
*args(Tensor): inputs(tuple) to the function.
|
|
|
|
**kwargs(Dict): inputs(dict) to the function.
|
|
|
|
Returns:
|
|
Output of function on args and kwargs.
|
|
|
|
"""
|
|
mp_group = ctx.get('mp_group', None)
|
|
assert mp_group is not None, (
|
|
"ctx must contains mp_group and mp_group can not be None."
|
|
)
|
|
|
|
offload = ctx.get('offload', False)
|
|
partition = ctx.get('partition', False)
|
|
|
|
if framework._dygraph_tracer()._has_grad:
|
|
check_recompute_necessary(args)
|
|
|
|
if custom_state_manager.custom_get_state_func is None:
|
|
assert custom_state_manager.custom_set_state_func is None
|
|
custom_get_state_func = lambda x=None: None
|
|
custom_set_state_func = lambda x=None: None
|
|
else:
|
|
custom_get_state_func = custom_state_manager.custom_get_state_func
|
|
custom_set_state_func = custom_state_manager.custom_set_state_func
|
|
|
|
all_outputs = []
|
|
_HPRecomputeFunction.apply(
|
|
function,
|
|
all_outputs,
|
|
mp_group,
|
|
offload,
|
|
partition,
|
|
custom_get_state_func,
|
|
custom_set_state_func,
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
|
|
if len(all_outputs) == 1:
|
|
return all_outputs[0]
|
|
else:
|
|
for output in all_outputs:
|
|
if paddle.is_tensor(output) and not utils.is_float_tensor(output):
|
|
output.stop_gradient = True
|
|
|
|
return tuple(all_outputs)
|