wehub-resource-sync
161 lines
4.3 KiB
Python
161 lines
4.3 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 paddle
|
|
from paddle import _C_ops
|
|
|
|
__all__ = []
|
|
|
|
|
|
FLOAT_TYPE_DICT = {
|
|
paddle.float16: "float16",
|
|
paddle.float32: "float32",
|
|
paddle.float64: "float64",
|
|
paddle.bfloat16: "bfloat16",
|
|
paddle.bool: "bool",
|
|
}
|
|
|
|
PADDLE_TO_NUMBER = {
|
|
paddle.float16: 0,
|
|
paddle.float32: 1,
|
|
paddle.float64: 2,
|
|
paddle.int32: 3,
|
|
paddle.int64: 4,
|
|
paddle.bfloat16: 5,
|
|
paddle.bool: 6,
|
|
}
|
|
|
|
NUMBER_TO_DTYPE = {
|
|
0: "float16",
|
|
1: "float32",
|
|
2: "float64",
|
|
3: "int32",
|
|
4: "int64",
|
|
5: "bfloat16",
|
|
6: "bool",
|
|
}
|
|
|
|
|
|
def is_float_tensor(tensor):
|
|
"""Is a float tensor"""
|
|
return tensor.dtype in FLOAT_TYPE_DICT.keys()
|
|
|
|
|
|
def get_tensor_dtype(dtype):
|
|
assert dtype in FLOAT_TYPE_DICT.keys()
|
|
return FLOAT_TYPE_DICT[dtype]
|
|
|
|
|
|
def paddle_2_number(dtype):
|
|
assert dtype in PADDLE_TO_NUMBER.keys()
|
|
return PADDLE_TO_NUMBER[dtype]
|
|
|
|
|
|
def number_2_dtype(number):
|
|
assert number in NUMBER_TO_DTYPE.keys()
|
|
return NUMBER_TO_DTYPE[number]
|
|
|
|
|
|
def get_tensor_bytes(tensor):
|
|
"""Get the bytes a tensor occupied."""
|
|
elem_size = None
|
|
if tensor.dtype == paddle.float32:
|
|
elem_size = 4
|
|
elif tensor.dtype == paddle.float64:
|
|
elem_size = 8
|
|
elif tensor.dtype == paddle.int64:
|
|
elem_size = 8
|
|
elif tensor.dtype == paddle.int32:
|
|
elem_size = 4
|
|
elif tensor.dtype == paddle.float16:
|
|
elem_size = 2
|
|
elif tensor.dtype == paddle.int8:
|
|
elem_size = 1
|
|
else:
|
|
raise ValueError(f"unknown data type: {tensor.dtype}")
|
|
return tensor.numel() * elem_size
|
|
|
|
|
|
def _all_gather(tensor, group=None, use_calc_stream=True):
|
|
"""
|
|
The main difference with paddle.distributed.all_gather:
|
|
no need to pass in tensor_list, the returned tensor is spliced
|
|
"""
|
|
if group is not None and not group.is_member():
|
|
return
|
|
ring_id = 0 if group is None else group.id
|
|
nranks = (
|
|
paddle.distributed.collective._get_global_group().nranks
|
|
if group is None
|
|
else group.nranks
|
|
)
|
|
return _C_ops.all_gather(
|
|
tensor,
|
|
ring_id,
|
|
nranks,
|
|
)
|
|
|
|
|
|
def tuple_to_dict_helper(input_tensor):
|
|
# recv tuple -> fwd input dict
|
|
use_dict = False
|
|
if isinstance(input_tensor, tuple):
|
|
use_dict = hasattr(input_tensor[0], "key")
|
|
else: # single tensor
|
|
use_dict = hasattr(input_tensor, "key")
|
|
if use_dict:
|
|
input_tensor = convert_tensor_tuple_to_dict(input_tensor)
|
|
return input_tensor, use_dict
|
|
|
|
|
|
def dict_to_tuple_helper(output_tensor):
|
|
if isinstance(output_tensor, dict):
|
|
output_tensor_tuple = convert_tensor_dict_to_tuple(
|
|
output_tensor_dict=output_tensor
|
|
)
|
|
else: # single tensor or tensor tuple
|
|
output_tensor_tuple = output_tensor
|
|
return output_tensor_tuple
|
|
|
|
|
|
def convert_tensor_dict_to_tuple(output_tensor_dict):
|
|
output_tensor = []
|
|
for key, tensor in output_tensor_dict.items():
|
|
if isinstance(tensor, (list, tuple)):
|
|
for idx, t in enumerate(tensor):
|
|
t.key = key + " " + str(idx)
|
|
output_tensor.append(t)
|
|
else: # single tensor
|
|
tensor.key = key
|
|
output_tensor.append(tensor)
|
|
|
|
return tuple(output_tensor)
|
|
|
|
|
|
def convert_tensor_tuple_to_dict(input_tensor_tuple):
|
|
input_tensor_dict = {}
|
|
for tensor in input_tensor_tuple:
|
|
key = tensor.key
|
|
if " " in key:
|
|
real_key, _ = key.split(" ")
|
|
if real_key in input_tensor_dict.keys():
|
|
input_tensor_dict[real_key].append(tensor)
|
|
else:
|
|
input_tensor_dict[real_key] = [tensor]
|
|
else:
|
|
input_tensor_dict[key] = tensor
|
|
delattr(tensor, "key")
|
|
return input_tensor_dict
|