wehub-resource-sync
154 lines
5.0 KiB
Python
154 lines
5.0 KiB
Python
# Copyright (c) 2025 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 os
|
|
|
|
import ml_dtypes
|
|
import numpy as np
|
|
import paddle
|
|
import paddle.nn as nn
|
|
|
|
enable_paddlenlp_monkey_patch = int(os.getenv("FLAGS_enable_paddlenlp_monkey_patch", "1"))
|
|
|
|
np.bfloat16 = ml_dtypes.bfloat16
|
|
np.float8_e5m2 = ml_dtypes.float8_e5m2
|
|
np.float8_e4m3fn = ml_dtypes.float8_e4m3fn
|
|
|
|
|
|
origin_repr = nn.Embedding.extra_repr
|
|
|
|
|
|
def new_repr(self):
|
|
return origin_repr(self) + f", dtype={self.weight.dtype}"
|
|
|
|
|
|
nn.Embedding.extra_repr = new_repr
|
|
|
|
origin_tensort_init = paddle.Tensor.__call__
|
|
origin_to_tensor = paddle.to_tensor
|
|
origin_numpy = paddle.Tensor.numpy
|
|
origin_numel = paddle.Tensor.numel
|
|
origin_set_value = paddle.core.eager.Tensor.set_value
|
|
|
|
|
|
paddle_numpy_mapping = {
|
|
paddle.float8_e5m2: (paddle.int8, np.float8_e5m2),
|
|
paddle.float8_e4m3fn: (paddle.int8, np.float8_e4m3fn),
|
|
# paddle.bfloat16: (paddle.int16, np.bfloat16),
|
|
}
|
|
|
|
numpy_paddle_mapping = {
|
|
np.dtype(np.float8_e5m2): (np.int8, paddle.float8_e5m2),
|
|
np.dtype(np.float8_e4m3fn): (np.int8, paddle.float8_e4m3fn),
|
|
np.dtype(np.bfloat16): (np.uint16, paddle.bfloat16),
|
|
}
|
|
|
|
paddle_numel_mapping = {
|
|
paddle.float8_e5m2: (paddle.int8, None),
|
|
paddle.float8_e4m3fn: (paddle.int8, None),
|
|
}
|
|
|
|
paddle_set_value_mapping = {
|
|
paddle.float8_e5m2: (paddle.int8, None),
|
|
paddle.float8_e4m3fn: (paddle.int8, None),
|
|
# paddle.bfloat16: (paddle.int16, None),
|
|
np.dtype(np.float8_e5m2): (np.int8, paddle.float8_e5m2),
|
|
np.dtype(np.float8_e4m3fn): (np.int8, paddle.float8_e4m3fn),
|
|
}
|
|
|
|
|
|
def enhance_init(*args, **kwargs):
|
|
if len(args) > 0 and isinstance(args[0], np.ndarray) and args[0].dtype in numpy_paddle_mapping:
|
|
inter_dtype, tgt_dtype = numpy_paddle_mapping[args[0].dtype]
|
|
tensor = args[0].view(inter_dtype)
|
|
new_args = (tensor, *args[1:])
|
|
tensor = origin_tensort_init(*new_args, **kwargs)
|
|
return tensor.view(tgt_dtype)
|
|
return origin_tensort_init(*args, **kwargs)
|
|
|
|
|
|
def enhance_to_tensor(*args, **kwargs):
|
|
# Fix with kwargs input
|
|
if len(args) > 0:
|
|
tensor = args[0]
|
|
else:
|
|
tensor = kwargs.get("data", None)
|
|
|
|
if isinstance(tensor, np.ndarray) and tensor.dtype in numpy_paddle_mapping:
|
|
inter_dtype, tgt_dtype = numpy_paddle_mapping[tensor.dtype]
|
|
tensor = tensor.view(inter_dtype)
|
|
if "data" in kwargs:
|
|
new_args = args
|
|
kwargs["data"] = tensor
|
|
else:
|
|
new_args = (tensor, *args[1:])
|
|
tensor = origin_to_tensor(*new_args, **kwargs)
|
|
return tensor.view(tgt_dtype)
|
|
return origin_to_tensor(*args, **kwargs)
|
|
|
|
|
|
def enhance_set_value(self, *args, **kwargs):
|
|
# Fix with kwargs input
|
|
if len(args) > 0:
|
|
tensor = args[0]
|
|
else:
|
|
tensor = kwargs.get("value", None)
|
|
|
|
if isinstance(tensor, np.ndarray) and tensor.dtype in paddle_set_value_mapping:
|
|
inter_dtype, tgt_dtype = paddle_set_value_mapping[tensor.dtype]
|
|
tensor = tensor.view(inter_dtype)
|
|
if "value" in kwargs:
|
|
new_args = args
|
|
kwargs["value"] = tensor
|
|
else:
|
|
new_args = (tensor, *args[1:])
|
|
return origin_set_value(self, *new_args, **kwargs)
|
|
|
|
if isinstance(tensor, paddle.Tensor) and tensor.dtype in paddle_set_value_mapping:
|
|
inter_dtype, _ = paddle_set_value_mapping[tensor.dtype]
|
|
tensor = tensor.view(inter_dtype)
|
|
if "value" in kwargs:
|
|
new_args = args
|
|
kwargs["value"] = tensor
|
|
else:
|
|
new_args = (tensor, *args[1:])
|
|
new_self = self.view(inter_dtype)
|
|
return origin_set_value(new_self, *new_args, **kwargs)
|
|
|
|
return origin_set_value(self, *args, **kwargs)
|
|
|
|
|
|
def _numpy(self, *args, **kwargs):
|
|
if self.dtype in paddle_numpy_mapping:
|
|
inter_pd_dtype, np_dtype = paddle_numpy_mapping[self.dtype]
|
|
tensor = origin_numpy(self.view(inter_pd_dtype), *args, **kwargs)
|
|
return tensor.view(np_dtype)
|
|
return origin_numpy(self, *args, **kwargs)
|
|
|
|
|
|
def _numel(self, *args, **kwargs):
|
|
if self.dtype in paddle_numel_mapping:
|
|
inter_pd_dtype, _ = paddle_numel_mapping[self.dtype]
|
|
ret = origin_numel(self.view(inter_pd_dtype), *args, **kwargs)
|
|
return ret
|
|
return origin_numel(self, *args, **kwargs)
|
|
|
|
|
|
if enable_paddlenlp_monkey_patch:
|
|
paddle.Tensor.numpy = _numpy
|
|
paddle.Tensor.__call__ = enhance_init
|
|
paddle.to_tensor = enhance_to_tensor
|
|
paddle.core.eager.Tensor.set_value = enhance_set_value
|
|
paddle.Tensor.numel = _numel
|