wehub-resource-sync
990 lines
37 KiB
Python
990 lines
37 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 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)
|