wehub-resource-sync
173 lines
7.1 KiB
Python
173 lines
7.1 KiB
Python
# Copyright (c) 2023 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.
|
|
|
|
import importlib
|
|
import inspect
|
|
from typing import Any, Union
|
|
|
|
import paddle
|
|
|
|
|
|
def fwd_step_patch(func, output, self, *args, **kwargs):
|
|
"""
|
|
Forward step patch function to handle gradient computation and loss recording during model training.
|
|
If the current model is the last stage and in training mode, it will record the gradient of the output
|
|
to the self._step_losses list. Otherwise, it will not perform any operations on the output.
|
|
|
|
Args:
|
|
func (Callable): The function being called, should be the forward function or any other function that needs to be executed.
|
|
output (Tensor): The output of the model, which should be a tensor.
|
|
self (Any): The model instance, which can be of type nn.Module or other custom model types.
|
|
args (Tuple[Any], optional): Optional arguments passed to func, default is None.
|
|
kwargs (Dict[str, Any], optional): Optional keyword arguments passed to func, default is None.
|
|
|
|
Returns:
|
|
None, no return value, directly modifies the self._step_losses attribute.
|
|
"""
|
|
# Training patch
|
|
if self.training and self.is_pipeline_last_stage():
|
|
if getattr(self, "_step_losses", None):
|
|
self._step_losses.append(output[0].detach())
|
|
else:
|
|
self._step_losses = [output[0].detach()]
|
|
|
|
|
|
def make_wrapper(func, pre_patch=None, post_patch=None):
|
|
"""
|
|
Creates a wrapper function that allows executing additional operations before and after calling the original function.
|
|
|
|
Args:
|
|
func (function): The function to be wrapped.
|
|
pre_patch (Optional[function], optional): The function to be executed before calling the original function, defaults to None.
|
|
The function signature should be `pre_patch(func, None, *args, **kwargs)`.
|
|
post_patch (Optional[function], optional): The function to be executed after calling the original function, defaults to None.
|
|
The function signature should be `post_patch(func, output, *args, **kwargs)`, where `output` is the return value of the original function.
|
|
|
|
Returns:
|
|
function: The wrapped function, which has the same functionality as the original function but executes additional operations before and after the call.
|
|
"""
|
|
|
|
def wrapper(*args, **kwargs):
|
|
if pre_patch is not None:
|
|
pre_patch(func, None, *args, **kwargs)
|
|
output = func(*args, **kwargs)
|
|
if post_patch is not None:
|
|
post_patch(func, output, *args, **kwargs)
|
|
return output
|
|
|
|
return wrapper
|
|
|
|
|
|
funcs = [
|
|
(
|
|
paddle.distributed.fleet.model.PipelineParallel._forward_step,
|
|
fwd_step_patch,
|
|
)
|
|
]
|
|
|
|
for func in funcs:
|
|
fun, patch = func
|
|
module = importlib.import_module(fun.__module__)
|
|
cls_name = fun.__qualname__[: -len(fun.__name__) - 1]
|
|
wrap_fun = make_wrapper(fun, post_patch=patch)
|
|
cls_obj = getattr(module, cls_name)
|
|
setattr(cls_obj, fun.__name__, wrap_fun)
|
|
|
|
|
|
@paddle.no_grad()
|
|
def pad_batches_inputs(inputs, padding_value=0, max_len=None, pad_len=None):
|
|
"""
|
|
Pads the length of tensors shaped [bs, seq_len] to [bs, max(seq_lens)].
|
|
|
|
Args:
|
|
inputs (list of paddle.Tensor or None): List of input tensors or None values.
|
|
padding_value (int or float, optional): The value to pad with, defaults to 0.
|
|
max_len (int, optional): The maximum length to pad to, if not provided it will be calculated.
|
|
pad_len (int or list of int, optional): The length to pad each tensor by, if not provided it will be calculated.
|
|
|
|
Returns:
|
|
list of paddle.Tensor: List of padded input tensors.
|
|
"""
|
|
if pad_len is not None:
|
|
pad_len = [pad_len] * len(inputs) if isinstance(pad_len, int) else pad_len
|
|
elif max_len is None:
|
|
# max_len = max([x.shape[-1] for x in inputs if x is not None])
|
|
max_len = max([x.shape[-1] if isinstance(x, paddle.Tensor) else 0 for x in inputs])
|
|
pad_len = [max_len - x.shape[-1] if isinstance(x, paddle.Tensor) else 0 for x in inputs]
|
|
|
|
for i in range(len(inputs)):
|
|
x = inputs[i]
|
|
# if x is None or x.shape[-1] == max_len:
|
|
if not isinstance(x, paddle.Tensor) or x.shape[-1] == max_len:
|
|
continue
|
|
inputs[i] = paddle.concat(
|
|
[
|
|
x,
|
|
paddle.full([x.shape[0], pad_len[i]], padding_value, dtype=x.dtype),
|
|
],
|
|
axis=-1,
|
|
)
|
|
|
|
return inputs
|
|
|
|
|
|
def get_expected_keys(inputs: dict, keys: list[str]) -> Union[tuple, Any]:
|
|
"""
|
|
Retrieve the expected key-value pairs from the inputs. If the key exists in the inputs,
|
|
return the corresponding value; otherwise, return None. If there is only one key-value pair,
|
|
convert it to a single element.
|
|
|
|
Args:
|
|
inputs (dict): A dictionary containing multiple key-value pairs to search for the expected ones.
|
|
keys (list[str]): A list of keys to be searched for.
|
|
|
|
Returns:
|
|
Union[tuple, Any]: If there is only one key-value pair, return the single element;
|
|
otherwise, return a tuple containing all key-value pairs. If any key does not exist, return None.
|
|
"""
|
|
ret = tuple([inputs.get(k, None) for k in keys if k in inputs])
|
|
if len(ret) == 1:
|
|
ret = ret[0]
|
|
return ret
|
|
|
|
|
|
def fwd_args_to_dict(fun):
|
|
"""
|
|
Converts the function's arguments into a dictionary to support more argument formats in the prediction pipeline step.
|
|
Assumes that no argument is of type inspect.Parameter.VAR_KEYWORD.
|
|
|
|
Args:
|
|
fun (Callable[[Any, Dict[str, Any]], Any]): The function to be converted. Its first argument is an instance of a non-pipeline model class,
|
|
and subsequent arguments can be any format of non-pipeline model forward arguments. The return value is of any type.
|
|
|
|
Returns:
|
|
Callable[[Any, *Any, **Any], Any]: A new function that accepts the same arguments as the original function, but converts all non-self arguments
|
|
into a dictionary and passes it as the second argument to the original function.
|
|
"""
|
|
|
|
def _impl(self, *args, **kwargs):
|
|
try:
|
|
return fun(self, *args, **kwargs)
|
|
except TypeError:
|
|
# otherwise, inputs is any valid format of non_pipe_model forward args,
|
|
# convert to dict, to support more args format in prediction_pipeline_step
|
|
# assume no arg is inspect.Parameter.VAR_KEYWORD
|
|
arg_dict = (
|
|
inspect.signature(self._non_pipe_model_class.forward).bind(*((self,) + args), **kwargs).arguments
|
|
)
|
|
arg_dict.pop("self")
|
|
return fun(self, arg_dict)
|
|
|
|
return _impl
|