chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,23 @@
|
||||
# Copyright (c) 2022 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 .recompute import ( # noqa: F401
|
||||
custom_state_manager,
|
||||
is_in_recompute,
|
||||
recompute,
|
||||
recompute_sequential,
|
||||
)
|
||||
from .recompute_hybrid import recompute_hybrid # noqa: F401
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,989 @@
|
||||
# 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 contextlib
|
||||
import copy
|
||||
import functools
|
||||
import inspect
|
||||
import random
|
||||
import threading
|
||||
import weakref
|
||||
from typing import TYPE_CHECKING, Any, TypedDict
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle import framework
|
||||
from paddle.autograd import PyLayer
|
||||
from paddle.base.framework import EagerParamBase
|
||||
from paddle.base.wrapped_decorator import copy_signature
|
||||
from paddle.distributed.fleet.meta_parallel.parallel_layers.random import (
|
||||
get_rng_state_tracker,
|
||||
)
|
||||
from paddle.framework import core, in_dynamic_mode
|
||||
from paddle.jit.dy2static.program_translator import StaticFunction
|
||||
|
||||
from ..utils.log_util import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
from typing_extensions import NotRequired
|
||||
|
||||
from paddle.nn import Sequential
|
||||
|
||||
class _Ctx(TypedDict):
|
||||
segments: int = 1
|
||||
preserve_rng_state: NotRequired[bool]
|
||||
|
||||
|
||||
__all__ = []
|
||||
_SIGNATURE_CACHE = weakref.WeakKeyDictionary()
|
||||
|
||||
|
||||
class RecomputeContext:
|
||||
"""
|
||||
A thread-safe context manager and decorator for tracking whether the current
|
||||
execution is inside a recompute phase.
|
||||
|
||||
RecomputeContext uses a thread-local flag to mark when code is running within a
|
||||
recompute region. It can be used as a context manager (``with`` statement) or as
|
||||
a decorator to automatically set and clear the recompute-active state. This allows
|
||||
downstream code to query ``is_in_recompute()`` and adapt its behavior accordingly
|
||||
(e.g., skipping certain logging or side effects during recomputation).
|
||||
|
||||
Parameters:
|
||||
None.
|
||||
|
||||
Returns:
|
||||
RecomputeContext: A recompute context instance that can be used as a context
|
||||
manager or decorator.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> from paddle.distributed.fleet.utils import is_in_recompute
|
||||
|
||||
>>> # Usage as a context manager
|
||||
>>> ctx = RecomputeContext()
|
||||
>>> print(ctx.active)
|
||||
False
|
||||
>>> with ctx:
|
||||
... print(ctx.active)
|
||||
True
|
||||
>>> print(ctx.active)
|
||||
False
|
||||
|
||||
>>> # Usage as a decorator
|
||||
>>> ctx = RecomputeContext()
|
||||
>>> @ctx
|
||||
... def my_forward(x):
|
||||
... return is_in_recompute()
|
||||
>>> print(my_forward(None))
|
||||
True
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._local = threading.local()
|
||||
|
||||
@property
|
||||
def active(self) -> bool:
|
||||
return getattr(self._local, 'active', False)
|
||||
|
||||
def __enter__(self):
|
||||
self._local.active = True
|
||||
return self
|
||||
|
||||
def __exit__(self, *_exc):
|
||||
self._local.active = False
|
||||
return False
|
||||
|
||||
def __call__(self, fn):
|
||||
@functools.wraps(fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
with self:
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
copy_signature(fn, wrapper)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
_recompute_context = RecomputeContext()
|
||||
|
||||
|
||||
def is_in_recompute() -> bool:
|
||||
"""
|
||||
Check whether the current thread is executing inside a recompute context.
|
||||
|
||||
This function inspects the global ``_recompute_context`` to determine if the
|
||||
current thread is within an active recompute phase. It is typically used inside
|
||||
forward computations to detect whether the execution is a normal forward pass
|
||||
or a recompute (re-forward) pass triggered during backpropagation, so that
|
||||
certain operations (e.g., logging, random state management) can be skipped or
|
||||
adjusted accordingly.
|
||||
|
||||
Parameters:
|
||||
None.
|
||||
|
||||
Returns:
|
||||
bool: ``True`` if the current thread is inside a recompute context,
|
||||
``False`` otherwise.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> from paddle.distributed.fleet.utils import is_in_recompute
|
||||
>>> # Outside any recompute context
|
||||
>>> print(is_in_recompute())
|
||||
False
|
||||
|
||||
>>> from paddle.distributed.fleet.utils.__init__ import RecomputeContext
|
||||
>>> ctx = RecomputeContext()
|
||||
>>> with ctx:
|
||||
... print(is_in_recompute())
|
||||
True
|
||||
"""
|
||||
return _recompute_context.active
|
||||
|
||||
|
||||
def _varbase_help(param):
|
||||
state = copy.deepcopy(param.__dict__)
|
||||
new_param = EagerParamBase(
|
||||
shape=param.shape,
|
||||
dtype=param.dtype,
|
||||
trainable=param.trainable,
|
||||
name=param.name,
|
||||
**state,
|
||||
)
|
||||
param._share_buffer_to(new_param)
|
||||
return new_param
|
||||
|
||||
|
||||
def detach_variable(inputs):
|
||||
out = []
|
||||
for inp in inputs:
|
||||
if not isinstance(inp, core.eager.Tensor) and (
|
||||
type(inp) is not tuple or not isinstance(inp[0], core.eager.Tensor)
|
||||
):
|
||||
# the inp is not a tensor or not a tuple of tensors
|
||||
out.append(inp)
|
||||
continue
|
||||
|
||||
if isinstance(inp, EagerParamBase):
|
||||
out.append(_varbase_help(inp))
|
||||
continue
|
||||
|
||||
if type(inp) is tuple:
|
||||
detach_inp = []
|
||||
for i in inp:
|
||||
# detach all tensors in the tuple
|
||||
assert isinstance(i, core.eager.Tensor)
|
||||
|
||||
if isinstance(i, EagerParamBase):
|
||||
detach_inp.append(_varbase_help(i))
|
||||
else:
|
||||
tmp_i = i.detach()
|
||||
tmp_i.stop_gradient = i.stop_gradient
|
||||
detach_inp.append(tmp_i)
|
||||
|
||||
out.append(tuple(detach_inp))
|
||||
continue
|
||||
|
||||
x = inp.detach()
|
||||
x.stop_gradient = inp.stop_gradient
|
||||
out.append(x)
|
||||
return tuple(out)
|
||||
|
||||
|
||||
def check_recompute_necessary(inputs):
|
||||
necessary_for_each_input = []
|
||||
for input_ in inputs:
|
||||
if isinstance(input_, paddle.Tensor):
|
||||
necessary_for_each_input.append(input_.stop_gradient)
|
||||
elif type(input_) is tuple:
|
||||
for i in input_:
|
||||
# traverse all tensors in the tuple
|
||||
if isinstance(i, paddle.Tensor):
|
||||
necessary_for_each_input.append(i.stop_gradient)
|
||||
if all(necessary_for_each_input):
|
||||
logger.warning(
|
||||
"[Recompute]: None of the inputs to current recompute block need grad, "
|
||||
"therefore there is NO need to recompute this block in backward !"
|
||||
)
|
||||
|
||||
|
||||
def _closure_cell_values(run_function):
|
||||
"""Return cell contents of ``run_function``'s ``__closure__`` as a tuple.
|
||||
|
||||
Supports plain functions/lambdas and ``paddle.nn.Layer`` (uses ``forward``).
|
||||
Deep Tensor extraction is done by the C++ side of ``_hold_tensors``.
|
||||
"""
|
||||
fn = (
|
||||
run_function.forward
|
||||
if isinstance(run_function, paddle.nn.Layer)
|
||||
else run_function
|
||||
)
|
||||
closure = getattr(fn, '__closure__', None) or ()
|
||||
values = []
|
||||
for cell in closure:
|
||||
try:
|
||||
values.append(cell.cell_contents)
|
||||
except ValueError: # empty cell
|
||||
pass
|
||||
return tuple(values)
|
||||
|
||||
|
||||
class CustomStatesManager:
|
||||
"""CustomStatesManager"""
|
||||
|
||||
def __init__(self):
|
||||
"""__init__"""
|
||||
self.custom_get_state_func = None
|
||||
self.custom_set_state_func = None
|
||||
|
||||
def set_custom_get_state_func(self, custom_get_state_func):
|
||||
assert_msg = (
|
||||
"The custom_state_manager does not support duplicate settings."
|
||||
)
|
||||
assert self.custom_get_state_func is None, assert_msg
|
||||
self.custom_get_state_func = custom_get_state_func
|
||||
|
||||
def set_custom_set_state_func(self, custom_set_state_func):
|
||||
assert_msg = (
|
||||
"The custom_state_manager does not support duplicate settings."
|
||||
)
|
||||
assert self.custom_set_state_func is None, assert_msg
|
||||
self.custom_set_state_func = custom_set_state_func
|
||||
|
||||
|
||||
custom_state_manager = CustomStatesManager()
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def switch_rng_state_tracker(
|
||||
rng_state,
|
||||
tracker,
|
||||
numpy_state,
|
||||
random_state,
|
||||
custom_state=None,
|
||||
custom_get_state_func=None,
|
||||
custom_set_state_func=None,
|
||||
):
|
||||
orig_rng_state = paddle.get_rng_state()
|
||||
orig_rng_tracker = get_rng_state_tracker().get_states_tracker()
|
||||
paddle.set_rng_state(rng_state)
|
||||
get_rng_state_tracker().set_states_tracker(tracker)
|
||||
|
||||
orig_numpy_state = None
|
||||
orig_random_state = None
|
||||
|
||||
if numpy_state is not None:
|
||||
orig_numpy_state = np.random.get_state()
|
||||
np.random.set_state(numpy_state)
|
||||
if random_state is not None:
|
||||
orig_random_state = random.getstate()
|
||||
random.setstate(random_state)
|
||||
|
||||
if custom_state is not None:
|
||||
assert custom_get_state_func is not None
|
||||
assert custom_set_state_func is not None
|
||||
orig_custom_state = custom_get_state_func()
|
||||
custom_set_state_func(custom_state)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
paddle.set_rng_state(orig_rng_state)
|
||||
get_rng_state_tracker().set_states_tracker(orig_rng_tracker)
|
||||
if orig_numpy_state is not None:
|
||||
np.random.set_state(orig_numpy_state)
|
||||
if orig_random_state is not None:
|
||||
random.setstate(orig_random_state)
|
||||
|
||||
if custom_state is not None:
|
||||
custom_set_state_func(orig_custom_state)
|
||||
|
||||
|
||||
class RecomputeFunction(PyLayer):
|
||||
@staticmethod
|
||||
def forward(
|
||||
ctx,
|
||||
run_function,
|
||||
preserve_rng_state,
|
||||
preserve_external_rng_state,
|
||||
offload_indices,
|
||||
custom_get_state_func,
|
||||
custom_set_state_func,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
# store for recomputing
|
||||
ctx.run_function = run_function
|
||||
ctx.preserve_rng_state = preserve_rng_state
|
||||
ctx.preserve_external_rng_state = preserve_external_rng_state
|
||||
ctx.offload_indices = offload_indices
|
||||
ctx.kwargs = kwargs
|
||||
|
||||
# NOTE the number of outputs of backward() should be equal to the number of tensors in forward()'s input
|
||||
# the order of tensors in backward()'s output should be the same as tensors in forward()'s input
|
||||
# None tensor inputs will be filtered in backward inputs.
|
||||
|
||||
# NOTE recompute with restore RNG only support one scenario where one process for one cuda gpu.
|
||||
# one process with multiple gpu and mix-gpu-cpu scenarios are not support
|
||||
if ctx.preserve_rng_state:
|
||||
ctx.fw_rng_state = paddle.get_rng_state()
|
||||
ctx.fwd_rng_state_tracker = (
|
||||
get_rng_state_tracker().get_states_tracker()
|
||||
)
|
||||
if ctx.preserve_external_rng_state:
|
||||
ctx.fwd_numpy_state = np.random.get_state()
|
||||
ctx.fwd_random_state = random.getstate()
|
||||
else:
|
||||
ctx.fwd_numpy_state = None
|
||||
ctx.fwd_random_state = None
|
||||
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
|
||||
|
||||
# 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}")
|
||||
|
||||
if tracer._amp_dtype == 'float16':
|
||||
ctx.amp_dtype = 'float16'
|
||||
elif tracer._amp_dtype in ('bfloat16', 'float32'):
|
||||
ctx.amp_dtype = 'bfloat16'
|
||||
else:
|
||||
raise ValueError(f"unsupported amp dtype: {tracer._amp_dtype}")
|
||||
|
||||
ctx.amp_white_list, ctx.amp_black_list = tracer._get_amp_op_list()
|
||||
|
||||
with paddle.no_grad(), _recompute_context:
|
||||
outputs = run_function(*args, **kwargs)
|
||||
|
||||
# save input for backward
|
||||
ctx.inputs = []
|
||||
ctx.tensor_indices = []
|
||||
ctx.duplicate_tensor = [False for _ in range(len(args))]
|
||||
tensor_inputs = []
|
||||
for i, arg in enumerate(args):
|
||||
if paddle.is_tensor(arg):
|
||||
if i in ctx.offload_indices:
|
||||
cpu_arg = (
|
||||
arg.pin_memory()
|
||||
if core.is_compiled_with_cuda()
|
||||
else arg.cpu()
|
||||
)
|
||||
cpu_arg._share_buffer_to(arg)
|
||||
tensor_inputs.append(arg)
|
||||
ctx.tensor_indices.append(i)
|
||||
ctx.inputs.append(None)
|
||||
elif type(arg) is tuple:
|
||||
assert i not in ctx.offload_indices, (
|
||||
f"offload_indices should not contain tensor tuple in position{i}"
|
||||
)
|
||||
is_tensors = [paddle.is_tensor(a) for a in arg]
|
||||
if all(is_tensors):
|
||||
# the tuple is a tuple of tensors
|
||||
tensors_stop_gradient = [a.stop_gradient for a in arg]
|
||||
if not all(tensors_stop_gradient) and any(
|
||||
tensors_stop_gradient
|
||||
):
|
||||
# tensors in the tuple have different stop_gradient value, which pylayer doesn't support
|
||||
raise ValueError(
|
||||
"Recompute receive a tuple containing tensor holds different stop gradient."
|
||||
)
|
||||
tensor_inputs.append(arg)
|
||||
ctx.tensor_indices.append(i)
|
||||
# Mark the tuple is a tuple of tensors
|
||||
ctx.duplicate_tensor[i] = True
|
||||
ctx.inputs.append(None)
|
||||
elif any(is_tensors):
|
||||
# the tuple contains tensors and non-tensor values
|
||||
raise ValueError(
|
||||
"Recompute receive a tuple containing tensor and non-tensor at same time."
|
||||
)
|
||||
else:
|
||||
ctx.inputs.append(arg)
|
||||
else:
|
||||
ctx.inputs.append(arg)
|
||||
|
||||
ctx.save_for_backward(*tensor_inputs)
|
||||
|
||||
# Protect tensors captured in run_function's Python __closure__ against
|
||||
# pipeline-parallel _clear_dataptr(); explicit tensor args are already
|
||||
# covered by save_for_backward's tensor_hold_helper.
|
||||
closure_values = _closure_cell_values(run_function)
|
||||
ctx._has_held_tensors = bool(closure_values)
|
||||
if closure_values:
|
||||
ctx._hold_tensors(closure_values)
|
||||
|
||||
return outputs
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, *args):
|
||||
with paddle.base.dygraph.guard():
|
||||
# TODO need to check the recompute calling is valid or not
|
||||
|
||||
# Restore closure-captured tensors potentially emptied by
|
||||
# pipeline-parallel _clear_dataptr() before re-running forward.
|
||||
if getattr(ctx, '_has_held_tensors', False):
|
||||
ctx._restore_held_tensors()
|
||||
|
||||
# Restore inputs
|
||||
inputs = list(ctx.inputs)
|
||||
tensor_indices = ctx.tensor_indices
|
||||
duplicate_tensor = ctx.duplicate_tensor
|
||||
tensors = ctx.saved_tensor()
|
||||
for i, idx in enumerate(tensor_indices):
|
||||
inputs[idx] = (
|
||||
tensors[i].to(
|
||||
paddle.base.framework._current_expected_place()
|
||||
)
|
||||
if i in ctx.offload_indices
|
||||
else tensors[i]
|
||||
)
|
||||
if i in ctx.offload_indices:
|
||||
# NOTE(zhiqiu): tensor.to(device) will set stop_gradient=True, which may break the gragh
|
||||
inputs[idx].stop_gradient = tensors[i].stop_gradient
|
||||
# paddle.enable_grad()
|
||||
tracer = framework._dygraph_tracer()
|
||||
tracer._has_grad = True
|
||||
|
||||
# NOTE support AMP
|
||||
# need restore auto_cast state as well as w/b list
|
||||
if ctx.preserve_rng_state:
|
||||
with (
|
||||
switch_rng_state_tracker(
|
||||
ctx.fw_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,
|
||||
),
|
||||
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,
|
||||
),
|
||||
_recompute_context,
|
||||
):
|
||||
detached_inputs = detach_variable(tuple(inputs))
|
||||
outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
|
||||
else:
|
||||
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,
|
||||
),
|
||||
_recompute_context,
|
||||
):
|
||||
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)
|
||||
|
||||
# run backward() with only tensor that requires grad
|
||||
forward_outputs_with_grad = []
|
||||
# NOTE In Transformer-like network, if user put the attention mask into the recompute segment output,
|
||||
# pylayer will force the stop_gradient of attention mask to be False, which will make the number of
|
||||
# tensor that need grad does not match.
|
||||
# the following backward_inputs_with_grad is used to avoid this case.
|
||||
backward_inputs_with_grad = []
|
||||
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_with_grad.append(args[i])
|
||||
|
||||
if len(forward_outputs_with_grad) == 0:
|
||||
raise RuntimeError(
|
||||
"none of output has requires_grad=True, this recompute() is not necessary"
|
||||
)
|
||||
|
||||
# actually backward
|
||||
with paddle.amp.auto_cast(enable=False):
|
||||
paddle.autograd.backward(
|
||||
forward_outputs_with_grad, backward_inputs_with_grad
|
||||
)
|
||||
|
||||
grads = []
|
||||
for idx, inp in enumerate(detached_inputs):
|
||||
if isinstance(inp, core.eager.Tensor):
|
||||
grads.append(inp._grad_ivar())
|
||||
elif type(inp) is tuple and duplicate_tensor[idx]:
|
||||
# input is a tuple and is a tuple of tensors
|
||||
if all(i.stop_gradient for i in inp):
|
||||
# all tensors in the tuple doesn't need grad, only return a None for the whole tuple
|
||||
grads.append(None)
|
||||
else:
|
||||
# all tensors in the tuple need grad, should return a tuple of grads
|
||||
grads.append(tuple(i._grad_ivar() for i in inp))
|
||||
|
||||
if in_dynamic_mode():
|
||||
grads = tuple(grads)
|
||||
else:
|
||||
grads = list(grads)
|
||||
return grads
|
||||
|
||||
|
||||
def _recompute_without_reentrant(
|
||||
function,
|
||||
custom_get_state_func,
|
||||
custom_set_state_func,
|
||||
preserve_rng_state=True,
|
||||
preserve_external_rng_state=True,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
recompute without reentrant, that means use hook to implement the recompute function rather than re-entrant autograd.
|
||||
"""
|
||||
|
||||
if preserve_rng_state:
|
||||
cur_device = paddle.get_device()
|
||||
if cur_device.startswith('gpu:'):
|
||||
fw_cuda_rng_state = paddle.get_cuda_rng_state()
|
||||
elif 'cpu' in cur_device:
|
||||
fw_cuda_rng_state = paddle.get_rng_state()
|
||||
elif 'xpu:' in cur_device:
|
||||
fw_cuda_rng_state = paddle.get_rng_state()
|
||||
elif (
|
||||
cur_device.split(':')[0]
|
||||
in paddle.device.get_all_custom_device_type()
|
||||
):
|
||||
fw_cuda_rng_state = paddle.get_rng_state(cur_device)
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Recompute with RNG preserve is not support current device: {cur_device}."
|
||||
)
|
||||
fwd_cuda_rng_state_tracker = (
|
||||
get_rng_state_tracker().get_states_tracker()
|
||||
)
|
||||
if preserve_external_rng_state:
|
||||
fwd_numpy_state = np.random.get_state()
|
||||
fwd_random_state = random.getstate()
|
||||
else:
|
||||
fwd_numpy_state = None
|
||||
fwd_random_state = None
|
||||
fwd_custom_state = custom_get_state_func()
|
||||
|
||||
tracer = framework._dygraph_tracer()
|
||||
is_fw_autocast = False if tracer._amp_level == core.AmpLevel.O0 else True
|
||||
if tracer._amp_level == core.AmpLevel.O2:
|
||||
amp_level = 'O2'
|
||||
elif tracer._amp_level in (core.AmpLevel.O1, core.AmpLevel.O0):
|
||||
amp_level = 'O1'
|
||||
|
||||
if tracer._amp_dtype == 'float16':
|
||||
amp_dtype = 'float16'
|
||||
elif tracer._amp_dtype in ('bfloat16', 'float32'):
|
||||
amp_dtype = 'bfloat16'
|
||||
|
||||
amp_white_list, amp_black_list = tracer._get_amp_op_list()
|
||||
|
||||
class Intermediate_Holder:
|
||||
pass
|
||||
|
||||
storage = weakref.WeakKeyDictionary()
|
||||
holder_list = []
|
||||
|
||||
def pack(x):
|
||||
res = Intermediate_Holder()
|
||||
holder_list.append(weakref.ref(res))
|
||||
return res
|
||||
|
||||
def unpack(x):
|
||||
unpack_counter = 0
|
||||
if len(storage) == 0:
|
||||
|
||||
def inner_pack(inner_x):
|
||||
nonlocal unpack_counter
|
||||
unpack_counter += 1
|
||||
|
||||
if holder_list[unpack_counter - 1]() is None:
|
||||
return
|
||||
if inner_x is None:
|
||||
storage[holder_list[unpack_counter - 1]()] = None
|
||||
return
|
||||
if hasattr(inner_x, "main_grad") or inner_x.grad is not None:
|
||||
storage[holder_list[unpack_counter - 1]()] = inner_x
|
||||
else:
|
||||
if inner_x.is_dist():
|
||||
tmp_tensor = core.eager.Tensor(inner_x)
|
||||
else:
|
||||
tmp_tensor = core.eager.Tensor(
|
||||
inner_x.dtype,
|
||||
inner_x.shape,
|
||||
inner_x.name + "cpy",
|
||||
core.VarDesc.VarType.DENSE_TENSOR,
|
||||
inner_x.persistable,
|
||||
)
|
||||
inner_x._unsafe_share_buffer_to(tmp_tensor)
|
||||
storage[holder_list[unpack_counter - 1]()] = tmp_tensor
|
||||
return
|
||||
|
||||
def inner_unpack(inner_x):
|
||||
raise Exception("An unexpected backward called on a tensor!")
|
||||
|
||||
if preserve_rng_state:
|
||||
with (
|
||||
switch_rng_state_tracker(
|
||||
fw_cuda_rng_state,
|
||||
fwd_cuda_rng_state_tracker,
|
||||
fwd_numpy_state,
|
||||
fwd_random_state,
|
||||
fwd_custom_state,
|
||||
custom_get_state_func,
|
||||
custom_set_state_func,
|
||||
),
|
||||
paddle.set_grad_enabled(True),
|
||||
paddle.amp.auto_cast(
|
||||
enable=is_fw_autocast,
|
||||
custom_white_list=amp_white_list,
|
||||
custom_black_list=amp_black_list,
|
||||
level=amp_level,
|
||||
dtype=amp_dtype,
|
||||
),
|
||||
paddle.autograd.saved_tensors_hooks(
|
||||
inner_pack, inner_unpack
|
||||
),
|
||||
):
|
||||
function(*args, **kwargs)
|
||||
else:
|
||||
with (
|
||||
paddle.set_grad_enabled(True),
|
||||
paddle.amp.auto_cast(
|
||||
enable=is_fw_autocast,
|
||||
custom_white_list=amp_white_list,
|
||||
custom_black_list=amp_black_list,
|
||||
level=amp_level,
|
||||
dtype=amp_dtype,
|
||||
),
|
||||
paddle.autograd.saved_tensors_hooks(
|
||||
inner_pack, inner_unpack
|
||||
),
|
||||
):
|
||||
function(*args, **kwargs)
|
||||
|
||||
if x not in storage:
|
||||
raise Exception(
|
||||
"Not supported to retrieve a tensor saved by autograd multiple times that is no need to recompute."
|
||||
)
|
||||
|
||||
return storage.pop(x)
|
||||
|
||||
with paddle.autograd.saved_tensors_hooks(pack, unpack):
|
||||
outputs = function(*args, **kwargs)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def recompute(function, *args, **kwargs):
|
||||
"""
|
||||
recompute intermediate activations to save then memory.
|
||||
|
||||
Parameters:
|
||||
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 to the function.
|
||||
**kwargs(Dict): Kwargs should only contain two kinds of key-value params, the one is part of function's key-value params,
|
||||
and the other contains 'preserve_rng_state', 'preserve_external_rng_state' and 'use_reentrant'.
|
||||
The key-value pair of preserve_rng_state is used to indicate whether to save the forward rng. If it is True,
|
||||
then the last forward rng value will be restored when the forward recalculation of backpropagation is performed,
|
||||
its default value is True.
|
||||
The key-value pair of preserve_external_rng_state is used to indicate whether to save and restore the external
|
||||
random number generator states (numpy.random and python random). If your forward function does not use numpy.random
|
||||
or python random, you can set this to False to improve performance. Its default value is True.
|
||||
The key-value pair of use_reentrant is used to indicate which implementation of recompute you will be used.
|
||||
'use_reentrant=True' means to use the PyLayer implementation of recompute, 'use_reentrant=False' means to
|
||||
use the Hook implementation of recompute, its default value is True.
|
||||
Returns:
|
||||
Output of function on args.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED, env:GPU)
|
||||
>>> import paddle
|
||||
>>> from paddle.distributed.fleet.utils import recompute
|
||||
>>> import random
|
||||
>>> paddle.seed(2023)
|
||||
>>> def get_fc_block(block_idx, input_size, is_last=False):
|
||||
... block_name = "block_" + str(block_idx)
|
||||
... block = paddle.nn.Sequential(
|
||||
... (block_name + "_fc_0", paddle.nn.Linear(input_size, input_size, bias_attr=False)),
|
||||
... (block_name + "_dropout", paddle.nn.Dropout(p=0.5)),
|
||||
... (block_name + "_relu_1", paddle.nn.ReLU()),
|
||||
... (block_name + "_fc_1", paddle.nn.Linear(input_size, input_size, bias_attr=False)),
|
||||
... (block_name + "_relu_2", paddle.nn.ReLU()),
|
||||
... )
|
||||
... if is_last:
|
||||
... block.add_sublayer(
|
||||
... block_name + "_fc_2",
|
||||
... paddle.nn.Linear(input_size, 1, bias_attr=False),
|
||||
... )
|
||||
... else:
|
||||
... block.add_sublayer(
|
||||
... block_name + "_fc_2",
|
||||
... paddle.nn.Linear(input_size, input_size, bias_attr=False),
|
||||
... )
|
||||
... return block
|
||||
|
||||
>>> class Naive_fc_net(paddle.nn.Layer):
|
||||
... def __init__(
|
||||
... self,
|
||||
... input_size=10,
|
||||
... recompute_blocks=[1, 3],
|
||||
... recompute_kwargs={},
|
||||
... ):
|
||||
... super().__init__()
|
||||
... self.recompute_blocks = recompute_blocks
|
||||
... self.recompute_kwargs = recompute_kwargs
|
||||
... self.runfunc0 = get_fc_block(0, input_size, is_last=False)
|
||||
... self.runfunc1 = get_fc_block(1, input_size, is_last=False)
|
||||
... self.runfunc2 = get_fc_block(2, input_size, is_last=False)
|
||||
... self.runfunc3 = get_fc_block(3, input_size, is_last=False)
|
||||
... self.runfunc4 = get_fc_block(4, input_size, is_last=True)
|
||||
... self.total_func = [self.runfunc0, self.runfunc1, self.runfunc2, self.runfunc3, self.runfunc4]
|
||||
...
|
||||
... def forward(self, inputs):
|
||||
... nums = len(self.total_func)
|
||||
... for i in range(nums):
|
||||
... if i in self.recompute_blocks:
|
||||
... inputs = recompute(self.total_func[i], inputs, **{"preserve_rng_state": True})
|
||||
... else:
|
||||
... inputs = self.total_func[i](inputs)
|
||||
... return inputs
|
||||
|
||||
>>> def run_model(cuda_state, recompute_block=[], recompute_kwargs={}):
|
||||
... gen = paddle.seed(10)
|
||||
... gen.manual_seed(10)
|
||||
... random.seed(10)
|
||||
... if cuda_state:
|
||||
... paddle.set_cuda_rng_state(cuda_state)
|
||||
... batch_size, input_size = 1, 10
|
||||
... model = Naive_fc_net(
|
||||
... input_size,
|
||||
... recompute_blocks=recompute_block,
|
||||
... recompute_kwargs=recompute_kwargs,
|
||||
... )
|
||||
... optimizer = paddle.optimizer.SGD(learning_rate=0.01, parameters=model.parameters())
|
||||
... loss_ = []
|
||||
... param_ = []
|
||||
... grad_ = []
|
||||
... for _ in range(5):
|
||||
... x = paddle.rand(shape=[batch_size, input_size], dtype="float32")
|
||||
... y_pred = model(x)
|
||||
... loss = y_pred.mean()
|
||||
... loss_.append(loss.item())
|
||||
... loss.backward()
|
||||
... optimizer.step()
|
||||
... param_.append(model.parameters()[9])
|
||||
... grad_.append(model.parameters()[3]._grad_ivar())
|
||||
... optimizer.clear_grad()
|
||||
... return loss_, param_, grad_
|
||||
|
||||
>>> cuda_state = paddle.get_cuda_rng_state()
|
||||
>>> # without recompute
|
||||
>>> loss_ref, param_ref, grad_ref = run_model(cuda_state, recompute_block=[])
|
||||
|
||||
>>> loss, param, grad = run_model(cuda_state, recompute_block=[1, 2])
|
||||
>>> print("normal_loss: {}, recompute_loss: {}".format(loss_ref, loss))
|
||||
>>> # The result of the recompute_loss should be the same as the normal_loss.
|
||||
normal_loss: [0.0018744759727269411, 0.0, 0.035971127450466156, 0.0, 0.0], recompute_loss: [0.0018744759727269411, 0.0, 0.035971127450466156, 0.0, 0.0]
|
||||
|
||||
"""
|
||||
# Hack to mix *args with **kwargs in a python 2.7-compliant way
|
||||
preserve = kwargs.pop('preserve_rng_state', True)
|
||||
preserve_external_rng_state = kwargs.pop(
|
||||
'preserve_external_rng_state', True
|
||||
)
|
||||
|
||||
# whether to use reentrant method to implement recompute
|
||||
use_reentrant = kwargs.pop('use_reentrant', True)
|
||||
|
||||
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
|
||||
|
||||
if not in_dynamic_mode():
|
||||
from paddle.distributed.auto_parallel.interface import (
|
||||
recompute as static_auto_recompute,
|
||||
)
|
||||
|
||||
return static_auto_recompute(function)(*args, **kwargs)
|
||||
|
||||
if framework._dygraph_tracer()._has_grad:
|
||||
check_args = list(args)
|
||||
check_args.extend(list(kwargs.values()))
|
||||
check_recompute_necessary(check_args)
|
||||
|
||||
if use_reentrant:
|
||||
offload_indices = kwargs.pop('offload_indices', [])
|
||||
if not kwargs: # fast path
|
||||
return RecomputeFunction.apply(
|
||||
function,
|
||||
preserve,
|
||||
preserve_external_rng_state,
|
||||
offload_indices,
|
||||
custom_get_state_func,
|
||||
custom_set_state_func,
|
||||
*args,
|
||||
)
|
||||
|
||||
# rearrange `position-args + keyword-args` into `position-args`
|
||||
target = (
|
||||
function.forward
|
||||
if isinstance(function, paddle.nn.Layer)
|
||||
else function
|
||||
)
|
||||
if isinstance(target, StaticFunction):
|
||||
target = target.dygraph_function
|
||||
|
||||
# Use getattr to get the cached signature. If it doesn't exist, parse and mount it to the target.
|
||||
# This avoids the heavy overhead of inspect.signature during repeated executions.
|
||||
cache_key = getattr(target, "__func__", target)
|
||||
dyfunc_sig = _SIGNATURE_CACHE.get(cache_key)
|
||||
if dyfunc_sig is None:
|
||||
dyfunc_sig = inspect.signature(target)
|
||||
_SIGNATURE_CACHE[cache_key] = dyfunc_sig
|
||||
|
||||
bound_args = dyfunc_sig.bind(*args, **kwargs)
|
||||
bound_args.apply_defaults()
|
||||
input_args = []
|
||||
for arg, param in zip(
|
||||
bound_args.arguments.values(), dyfunc_sig.parameters.values()
|
||||
):
|
||||
if param.kind == param.VAR_POSITIONAL:
|
||||
input_args.extend(arg)
|
||||
elif param.kind in (
|
||||
param.POSITIONAL_ONLY,
|
||||
param.POSITIONAL_OR_KEYWORD,
|
||||
):
|
||||
input_args.append(arg)
|
||||
elif param.kind == param.VAR_KEYWORD:
|
||||
input_args.extend(arg.values())
|
||||
elif param.kind == param.KEYWORD_ONLY:
|
||||
raise ValueError(
|
||||
"Currently, keyword-only arguments are not supported when you want to send kwargs(dict parameter) to function with use_reentrant=True."
|
||||
)
|
||||
else:
|
||||
raise ValueError("Unknown parameter kind.")
|
||||
return RecomputeFunction.apply(
|
||||
function,
|
||||
preserve,
|
||||
preserve_external_rng_state,
|
||||
offload_indices,
|
||||
custom_get_state_func,
|
||||
custom_set_state_func,
|
||||
*input_args,
|
||||
)
|
||||
else:
|
||||
return _recompute_without_reentrant(
|
||||
function,
|
||||
custom_get_state_func,
|
||||
custom_set_state_func,
|
||||
preserve,
|
||||
preserve_external_rng_state,
|
||||
*args,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def recompute_sequential(
|
||||
ctx: _Ctx,
|
||||
functions: Sequential | Sequence[Callable[..., Any]],
|
||||
*args: Any,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""
|
||||
recompute intermediate activations to save the memory for 'Sequential' models. use 'ctx' to transmit some context params, it is similar to 'recompute_hybrid' API.
|
||||
|
||||
Parameters:
|
||||
ctx(dict): include 'segments' and 'preserve_rng_state' keys, the key 'segments' (int, default 1), represents the number of chunks to create in the model,
|
||||
the key 'preserve_rng_state' (bool, optional, default=True) indicate whether to save the forward rng. If it is True, then the last forward rng value will be
|
||||
restored when the forward recalculation of backpropagation is performed.
|
||||
functions(paddle.nn.Sequential): 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.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:DISTRIBUTED)
|
||||
>>> import paddle
|
||||
>>> from paddle.incubate.distributed.fleet import recompute_sequential
|
||||
>>> input = paddle.ones(shape=[8, 10])
|
||||
>>> model = paddle.nn.Sequential(paddle.nn.Linear(10, 10), paddle.nn.Linear(10, 2))
|
||||
>>> output = recompute_sequential({'segments': 1}, model, input)
|
||||
|
||||
"""
|
||||
segments = ctx.get('segments', 1)
|
||||
preserve_rng_state = ctx.get('preserve_rng_state', True)
|
||||
|
||||
def _run_func(begin, end, funcs):
|
||||
def do_run(input):
|
||||
for i in range(begin, end + 1):
|
||||
input = funcs[i](input)
|
||||
return input
|
||||
|
||||
return do_run
|
||||
|
||||
if isinstance(functions, paddle.nn.Sequential):
|
||||
functions = list(functions.children())
|
||||
|
||||
segment_size = len(functions) // segments
|
||||
|
||||
end = -1
|
||||
for begin in range(0, segment_size * (segments - 1), segment_size):
|
||||
end = begin + segment_size - 1
|
||||
args = recompute(
|
||||
_run_func(begin, end, functions),
|
||||
*args,
|
||||
preserve_rng_state=preserve_rng_state,
|
||||
**kwargs,
|
||||
)
|
||||
return _run_func(end + 1, len(functions) - 1, functions)(*args)
|
||||
@@ -0,0 +1,347 @@
|
||||
# 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)
|
||||
Reference in New Issue
Block a user