wehub-resource-sync
85 lines
3.4 KiB
Python
85 lines
3.4 KiB
Python
# Copyright (c) 2024 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.
|
|
|
|
# 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.distributed.auto_parallel.ring_attention import shard_seq_load_balance
|
|
from paddle.distributed.fleet import fleet
|
|
|
|
|
|
def split_inputs_sequence_dim_load_balance(inputs, rank=None, degree=None):
|
|
if degree is None and rank is None:
|
|
_hcg = fleet.get_hybrid_communicate_group()
|
|
degree = _hcg.get_sep_parallel_world_size()
|
|
rank = _hcg.get_sep_parallel_rank()
|
|
assert isinstance(degree, int) and isinstance(
|
|
rank, int
|
|
), f"degree:{type(degree)} and rank:{type(rank)} must be int"
|
|
if degree <= 1:
|
|
return inputs
|
|
|
|
def do_split_sequence_dim_load_balance(data, rank, degree):
|
|
if data is None:
|
|
return None
|
|
assert isinstance(data, paddle.Tensor), f"data should be paddle.Tensor, but is type:{type(data)}"
|
|
assert len(data.shape) == 2, f"data dims should be 2, but shaped: {data.shape}"
|
|
sliced_datas = paddle.split(data, num_or_sections=degree * 2, axis=-1)
|
|
sliced_data0, sliced_data1 = sliced_datas[rank], sliced_datas[degree * 2 - 1 - rank]
|
|
return paddle.concat([sliced_data0, sliced_data1], axis=-1)
|
|
|
|
if isinstance(inputs, paddle.Tensor):
|
|
return do_split_sequence_dim_load_balance(inputs, rank, degree)
|
|
elif isinstance(inputs, dict):
|
|
res = {}
|
|
for k, tensor in inputs.items():
|
|
res[k] = do_split_sequence_dim_load_balance(tensor, rank, degree)
|
|
elif isinstance(inputs, list):
|
|
res = []
|
|
for tensor in inputs:
|
|
res.append(do_split_sequence_dim_load_balance(tensor, rank, degree))
|
|
else:
|
|
raise ValueError(f"the inputs should be a list or a dict, but is type: {type(inputs)}")
|
|
return res
|
|
|
|
|
|
def auto_split_sequence_dim_load_balance(inputs):
|
|
"""
|
|
for auto_parallel mode
|
|
"""
|
|
if isinstance(inputs, paddle.Tensor):
|
|
return shard_seq_load_balance(inputs, 1)
|
|
elif isinstance(inputs, dict):
|
|
res = {}
|
|
for k, tensor in inputs.items():
|
|
res[k] = shard_seq_load_balance(tensor, 1)
|
|
elif isinstance(inputs, list):
|
|
res = []
|
|
for tensor in inputs:
|
|
res.append(shard_seq_load_balance(tensor, 1))
|
|
else:
|
|
raise ValueError(f"the inputs should be a list or a dict, but is type: {type(inputs)}")
|
|
return res
|