chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,15 @@
|
||||
# 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.
|
||||
|
||||
__all__ = []
|
||||
@@ -0,0 +1,133 @@
|
||||
# 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.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import paddle
|
||||
from paddle.distributed.communication.batch_isend_irecv import (
|
||||
P2POp,
|
||||
)
|
||||
|
||||
from .p2p_communication import SendRecvMeta
|
||||
from .utils import (
|
||||
number_2_dtype,
|
||||
paddle_2_number,
|
||||
)
|
||||
|
||||
|
||||
class BatchCommHelper:
|
||||
# NOTE(zhangyuqin1998): Tensors to be sent or received must have a
|
||||
# consistent shape and data type throughout the entire pipeline.
|
||||
def __init__(self, use_cache=True):
|
||||
self._send_recv_meta = SendRecvMeta()
|
||||
self._use_cache = use_cache
|
||||
|
||||
def clear_meta_cache(self):
|
||||
self._send_recv_meta.init_or_erase_meta()
|
||||
|
||||
def _send_meta(self, tensors, group, broadcast=False):
|
||||
self._send_recv_meta.set_send_message(tensors)
|
||||
self._send_recv_meta.send_meta(tensors, group, broadcast=broadcast)
|
||||
self._send_recv_meta.recv_shape_message = (
|
||||
self._send_recv_meta.send_shape_message
|
||||
)
|
||||
self._send_recv_meta.recv_dtype_message = (
|
||||
self._send_recv_meta.send_dtype_message
|
||||
)
|
||||
|
||||
def _recv_meta(self, group, broadcast=False):
|
||||
self._send_recv_meta.recv_meta(group, broadcast=broadcast)
|
||||
|
||||
def _build_from_meta(self):
|
||||
shape_message = self._send_recv_meta.recv_shape_message
|
||||
dtype_message = self._send_recv_meta.recv_dtype_message
|
||||
stop_gradient = self._send_recv_meta.recv_stop_gradient
|
||||
assert (shape_message is not None) and (dtype_message is not None), (
|
||||
"Failed to build from meta."
|
||||
)
|
||||
|
||||
res = []
|
||||
if isinstance(shape_message, tuple):
|
||||
for idx, shape in enumerate(shape_message):
|
||||
tmp = paddle.empty(
|
||||
shape=shape, dtype=number_2_dtype(dtype_message[idx])
|
||||
)
|
||||
tmp.stop_gradient = (
|
||||
stop_gradient[idx] if stop_gradient is not None else False
|
||||
)
|
||||
res.append(tmp)
|
||||
else:
|
||||
tmp = paddle.empty(
|
||||
shape=shape_message, dtype=number_2_dtype(dtype_message)
|
||||
)
|
||||
tmp.stop_gradient = stop_gradient
|
||||
res.append(tmp)
|
||||
return res
|
||||
|
||||
def _check_valid(self, tensors):
|
||||
shape_message = self._send_recv_meta.recv_shape_message
|
||||
dtype_message = self._send_recv_meta.recv_dtype_message
|
||||
|
||||
assert (shape_message is not None) and (dtype_message is not None), (
|
||||
"Failed to build from meta."
|
||||
)
|
||||
|
||||
if isinstance(shape_message, tuple):
|
||||
assert isinstance(tensors, (list, tuple))
|
||||
assert len(tensors) == len(shape_message)
|
||||
for idx, (shape, dtype, tensor) in enumerate(
|
||||
zip(shape_message, dtype_message, tensors)
|
||||
):
|
||||
assert tensor.shape == shape, "Invalid shape."
|
||||
assert number_2_dtype(
|
||||
paddle_2_number(tensor.dtype)
|
||||
) == number_2_dtype(dtype), "Invalid dtype."
|
||||
else:
|
||||
if isinstance(tensors, (list, tuple)):
|
||||
assert len(tensors) == 1
|
||||
tensors = tensors[0]
|
||||
|
||||
assert tensors.shape == shape_message, "Invalid shape."
|
||||
assert number_2_dtype(
|
||||
paddle_2_number(tensors.dtype)
|
||||
) == number_2_dtype(dtype_message), "Invalid dtype."
|
||||
|
||||
def recv_meta_from_head(self, group, need_recv_meta):
|
||||
if not need_recv_meta:
|
||||
return
|
||||
self._recv_meta(group, broadcast=True)
|
||||
|
||||
def append_irecv(self, ops, src, group, alloc_on_comm_stream=False):
|
||||
if alloc_on_comm_stream:
|
||||
send_recv_stream = paddle.device.Stream(
|
||||
stream_base=group.process_group.get_stream(
|
||||
paddle.framework._current_expected_place_()
|
||||
)
|
||||
)
|
||||
with paddle.device.stream_guard(send_recv_stream):
|
||||
tensors = self._build_from_meta()
|
||||
else:
|
||||
tensors = self._build_from_meta()
|
||||
for tensor in tensors:
|
||||
if tensor is not None:
|
||||
ops.append(P2POp(paddle.distributed.irecv, tensor, src, group))
|
||||
return tensors
|
||||
|
||||
def append_isend(self, ops, tensors, dst, group, need_broadcast_meta=False):
|
||||
if need_broadcast_meta:
|
||||
self._send_meta(tensors, group, broadcast=True)
|
||||
self._check_valid(tensors)
|
||||
for tensor in tensors:
|
||||
if tensor is not None:
|
||||
ops.append(P2POp(paddle.distributed.isend, tensor, dst, group))
|
||||
+201
@@ -0,0 +1,201 @@
|
||||
# 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 paddle
|
||||
|
||||
from .utils import dict_to_tuple_helper
|
||||
|
||||
|
||||
class ScheduleChunk:
|
||||
# NOTE(zhangyuqin): ScheduleChunk is the atomic unit of pipeline scheduling.
|
||||
# A ScheduleChunk can contain several ScheduleNodes
|
||||
def __init__(self, nodes):
|
||||
self.nodes = nodes
|
||||
self._check_nodes_valid()
|
||||
|
||||
def forward(self, inputs):
|
||||
for n in self.nodes:
|
||||
inputs = n.forward(inputs)
|
||||
return inputs
|
||||
|
||||
def backward(self, output_grad):
|
||||
for n in reversed(self.nodes):
|
||||
output_grad = n.backward(output_grad)
|
||||
return output_grad
|
||||
|
||||
def _check_nodes_valid(self):
|
||||
for n in self.nodes:
|
||||
assert isinstance(n, (ScheduleNode, ScheduleChunk))
|
||||
|
||||
|
||||
def detach_and_requires_grad(inputs):
|
||||
if isinstance(inputs, (tuple, list)):
|
||||
is_tuple = isinstance(inputs, tuple)
|
||||
ret = []
|
||||
for input in inputs:
|
||||
if isinstance(input, (tuple, list)):
|
||||
ret.append(detach_and_requires_grad(input))
|
||||
elif isinstance(input, paddle.Tensor):
|
||||
tmp = input.detach() if input is not None else None
|
||||
if tmp is not None:
|
||||
tmp.stop_gradient = input.stop_gradient
|
||||
ret.append(tmp)
|
||||
else:
|
||||
ret.append(input)
|
||||
if is_tuple:
|
||||
ret = tuple(ret)
|
||||
return ret
|
||||
elif isinstance(inputs, dict):
|
||||
ret = {}
|
||||
for key in inputs.keys():
|
||||
input = inputs[key]
|
||||
tmp = input.detach() if input is not None else None
|
||||
if tmp is not None:
|
||||
tmp.stop_gradient = input.stop_gradient
|
||||
ret[key] = tmp
|
||||
return ret
|
||||
else:
|
||||
tmp = inputs.detach()
|
||||
tmp.stop_gradient = inputs.stop_gradient
|
||||
return tmp
|
||||
|
||||
|
||||
def clone_and_clear_dataptr(outputs, clear_dataptr=False):
|
||||
if isinstance(outputs, (tuple, list)):
|
||||
is_tuple = isinstance(outputs, tuple)
|
||||
ret = [
|
||||
FakeClone.apply(o)
|
||||
for o in outputs
|
||||
if o is not None and isinstance(o, paddle.Tensor)
|
||||
]
|
||||
|
||||
if clear_dataptr:
|
||||
for o in ret:
|
||||
o._clear_dataptr()
|
||||
if is_tuple:
|
||||
ret = tuple(ret)
|
||||
return ret
|
||||
elif isinstance(outputs, dict):
|
||||
ret = {}
|
||||
for key in outputs.keys():
|
||||
o = outputs[key]
|
||||
if o is not None and isinstance(o, paddle.Tensor):
|
||||
ret[key] = FakeClone.apply(o)
|
||||
if clear_dataptr:
|
||||
for key in ret:
|
||||
ret[key]._clear_dataptr()
|
||||
return ret
|
||||
else:
|
||||
ret = FakeClone.apply(outputs)
|
||||
if clear_dataptr:
|
||||
ret._clear_dataptr()
|
||||
return ret
|
||||
|
||||
|
||||
class FakeClone(paddle.autograd.PyLayer):
|
||||
# NOTE(zhangyuqin): Some input tensors may not be used in the forward function, but their gradients
|
||||
# need to be retained. Therefore, we need a clone here. To avoid the DtoD copy, we need a FakeClone
|
||||
@staticmethod
|
||||
def forward(ctx, input):
|
||||
return paddle.empty_like(input)
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
return grad_output
|
||||
|
||||
|
||||
class ScheduleNode:
|
||||
# NOTE(zhangyuqin): ScheduleNode is a subgraph of the pipeline, capable of independently calling
|
||||
# forward and backward. Users should not use paddle.autograd.backward on the results of ScheduleNode.forward.
|
||||
# Instead, they should use ScheduleNode.backward. Otherwise, resource leakage may occur.
|
||||
def __init__(self, fwd_func, name=""):
|
||||
self.name = name
|
||||
self.fwd_func = fwd_func
|
||||
self.inputs = None
|
||||
self.outputs = None
|
||||
|
||||
self.labels = None
|
||||
self.scale_loss_factor = None
|
||||
|
||||
def forward(self, inputs=(), **kwargs):
|
||||
detached_inputs = detach_and_requires_grad(inputs)
|
||||
self.inputs = detached_inputs
|
||||
if self.labels is not None:
|
||||
outputs = self.fwd_func(self.inputs, self.labels, **kwargs)
|
||||
else:
|
||||
outputs = self.fwd_func(self.inputs, **kwargs)
|
||||
if self.scale_loss_factor is not None:
|
||||
outputs /= self.scale_loss_factor
|
||||
|
||||
# Do not release the loss tensor.
|
||||
clear_dataptr = self.labels is None
|
||||
self.outputs = clone_and_clear_dataptr(outputs, clear_dataptr)
|
||||
return outputs
|
||||
|
||||
def backward(self, output_grad=None, scaler=None):
|
||||
if output_grad is None:
|
||||
if isinstance(self.outputs, (tuple, list)):
|
||||
assert len(self.outputs) == 1
|
||||
outputs = self.outputs[0]
|
||||
else:
|
||||
outputs = self.outputs
|
||||
assert isinstance(outputs, paddle.Tensor)
|
||||
if scaler is not None:
|
||||
paddle.autograd.backward(scaler.scale(outputs))
|
||||
else:
|
||||
paddle.autograd.backward(outputs)
|
||||
else:
|
||||
# Record the original type (tuple or list) to preserve it after filtering
|
||||
is_output_grad_tuple = isinstance(output_grad, tuple)
|
||||
if not isinstance(output_grad, (tuple, list)):
|
||||
is_output_grad_tuple = True # Single value becomes tuple
|
||||
output_grad = (output_grad,)
|
||||
|
||||
outputs = dict_to_tuple_helper(self.outputs)
|
||||
if not isinstance(outputs, (tuple, list)):
|
||||
outputs = (outputs,)
|
||||
outputs = [t for t in outputs if not t.stop_gradient]
|
||||
|
||||
# Filter None values from output_grad
|
||||
output_grad = [grad for grad in output_grad if grad is not None]
|
||||
# Preserve original type (tuple or list)
|
||||
output_grad = (
|
||||
tuple(output_grad)
|
||||
if is_output_grad_tuple
|
||||
else list(output_grad)
|
||||
)
|
||||
|
||||
assert len(outputs) == len(output_grad), (
|
||||
f"{len(outputs)} of {type(outputs[0])} vs {len(output_grad)} of {type(output_grad[0])}"
|
||||
)
|
||||
|
||||
paddle.autograd.backward(outputs, output_grad)
|
||||
|
||||
inputs = dict_to_tuple_helper(self.inputs)
|
||||
if not isinstance(inputs, (tuple, list)):
|
||||
inputs = (inputs,)
|
||||
grad = tuple([e.grad if e is not None else None for e in inputs])
|
||||
# grad = tuple([e.grad if e is not None and not e.stop_gradient else None for e in inputs])
|
||||
self._reset_states()
|
||||
|
||||
# if len(grad) == 1:
|
||||
# grad = grad[0]
|
||||
return grad
|
||||
|
||||
def _reset_states(self):
|
||||
self.inputs = None
|
||||
self.outputs = None
|
||||
self.labels = None
|
||||
self.scale_loss_factor = None
|
||||
+870
@@ -0,0 +1,870 @@
|
||||
# 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.
|
||||
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
|
||||
import paddle
|
||||
from paddle import framework
|
||||
|
||||
from ...utils import timer_helper as timer
|
||||
from ...utils.log_util import logger
|
||||
from .utils import number_2_dtype, paddle_2_number
|
||||
|
||||
_hcg = None
|
||||
_enable_partial_send_recv = True
|
||||
_timers = None
|
||||
|
||||
_xpu_comm_group_started = False
|
||||
|
||||
_sync_send = os.environ.get("PADDLE_P2P_SYNC_SEND", "0")
|
||||
_sync_send = _sync_send.lower() in ['1', 'true']
|
||||
|
||||
|
||||
def _xpu_comm_group_start():
|
||||
if not paddle.is_compiled_with_xpu():
|
||||
return
|
||||
global _xpu_comm_group_started
|
||||
assert not _xpu_comm_group_started
|
||||
framework.core.ProcessGroupBKCL.group_start()
|
||||
_xpu_comm_group_started = True
|
||||
|
||||
|
||||
def _xpu_comm_group_end():
|
||||
if not paddle.is_compiled_with_xpu():
|
||||
return
|
||||
global _xpu_comm_group_started
|
||||
if _xpu_comm_group_started:
|
||||
framework.core.ProcessGroupBKCL.group_end()
|
||||
_xpu_comm_group_started = False
|
||||
|
||||
|
||||
def initialize_p2p_groups(
|
||||
hcg, enable_partial_send_recv=True, enable_timer=False
|
||||
):
|
||||
global _hcg, _enable_partial_send_recv, _timers
|
||||
_hcg = hcg
|
||||
_enable_partial_send_recv = enable_partial_send_recv
|
||||
if enable_timer:
|
||||
_timers = timer.get_timers()
|
||||
(
|
||||
send_next_group,
|
||||
send_prev_group,
|
||||
recv_next_group,
|
||||
recv_prev_group,
|
||||
) = _hcg.get_p2p_groups()
|
||||
|
||||
debug_str = (
|
||||
f"P2pInfo: send_next_group: {send_next_group!r}, send_prev_group: {send_prev_group!r}, "
|
||||
f"recv_next_group: {recv_next_group!r}, recv_prev_group: {recv_prev_group!r}"
|
||||
)
|
||||
logger.info(debug_str)
|
||||
|
||||
|
||||
class SendRecvMeta:
|
||||
"""Mainly used to help p2p communication context information"""
|
||||
|
||||
def __init__(self):
|
||||
self.send_shape_message = None
|
||||
self.send_dtype_message = None
|
||||
|
||||
self.recv_shape_message = None
|
||||
self.recv_dtype_message = None
|
||||
self.recv_stop_gradient = None
|
||||
|
||||
self.has_send_meta = False
|
||||
self.has_recv_meta = False
|
||||
|
||||
def _recv_shape_dtype(self, group):
|
||||
# recv len(shape)
|
||||
dims = paddle.to_tensor([0])
|
||||
src_rank = _hcg._get_p2p_prev_rank()
|
||||
|
||||
paddle.distributed.recv(dims, src=src_rank, group=group)
|
||||
dims = dims.item()
|
||||
|
||||
# recv shape
|
||||
shape = paddle.to_tensor([0] * dims)
|
||||
paddle.distributed.recv(shape, src=src_rank, group=group)
|
||||
|
||||
# recv dtype
|
||||
dtype = paddle.to_tensor([0])
|
||||
paddle.distributed.recv(dtype, src=src_rank, group=group)
|
||||
|
||||
# recv stop_gradient
|
||||
stop_grad = paddle.to_tensor([0])
|
||||
paddle.distributed.recv(stop_grad, src=src_rank, group=group)
|
||||
return shape.tolist(), dtype.item(), stop_grad.item()
|
||||
|
||||
def recv_meta(self, group):
|
||||
tensor_type = paddle.to_tensor([0])
|
||||
src_rank = _hcg._get_p2p_prev_rank()
|
||||
|
||||
paddle.distributed.recv(tensor_type, src=src_rank, group=group)
|
||||
tensor_type = tensor_type.item()
|
||||
|
||||
if tensor_type == 0:
|
||||
shape, dtype, stop_grad = self._recv_shape_dtype(group)
|
||||
self.recv_shape_message = shape
|
||||
self.recv_dtype_message = dtype
|
||||
self.recv_stop_gradient = bool(stop_grad)
|
||||
|
||||
elif tensor_type == 1:
|
||||
num = paddle.to_tensor([0])
|
||||
paddle.distributed.recv(num, src=src_rank, group=group)
|
||||
num = num.item()
|
||||
shapes = []
|
||||
dtypes = []
|
||||
stop_grads = []
|
||||
for i in range(num):
|
||||
shape, dtype, stop_grad = self._recv_shape_dtype(group)
|
||||
shapes.append(shape)
|
||||
dtypes.append(dtype)
|
||||
stop_grads.append(bool(stop_grad))
|
||||
|
||||
self.recv_shape_message = tuple(shapes)
|
||||
self.recv_dtype_message = tuple(dtypes)
|
||||
self.recv_stop_gradient = tuple(stop_grads)
|
||||
|
||||
def _send_dims_shape_dtype(self, tensor, group):
|
||||
# send len(shape)
|
||||
dims = paddle.to_tensor([len(tensor.shape)])
|
||||
dst_rank = _hcg._get_p2p_next_rank()
|
||||
|
||||
paddle.distributed.send(dims, dst=dst_rank, group=group)
|
||||
|
||||
# send shape
|
||||
shape = paddle.to_tensor(tensor.shape)
|
||||
paddle.distributed.send(shape, dst=dst_rank, group=group)
|
||||
|
||||
# send dtype
|
||||
dtype = paddle.to_tensor([paddle_2_number(tensor.dtype)])
|
||||
paddle.distributed.send(dtype, dst=dst_rank, group=group)
|
||||
|
||||
# send trainable
|
||||
stop_grad = paddle.to_tensor([int(tensor.stop_gradient)])
|
||||
paddle.distributed.send(stop_grad, dst=dst_rank, group=group)
|
||||
|
||||
def send_meta(self, tensor, group):
|
||||
dst_rank = _hcg._get_p2p_next_rank()
|
||||
|
||||
if isinstance(tensor, paddle.Tensor):
|
||||
tensor_type = paddle.to_tensor([0])
|
||||
# send tensor type
|
||||
paddle.distributed.send(tensor_type, dst=dst_rank, group=group)
|
||||
|
||||
self._send_dims_shape_dtype(tensor, group)
|
||||
elif isinstance(tensor, tuple):
|
||||
tensor_type = paddle.to_tensor([1])
|
||||
# send tensor type
|
||||
paddle.distributed.send(tensor_type, dst=dst_rank, group=group)
|
||||
|
||||
nums = paddle.to_tensor([len(tensor)])
|
||||
paddle.distributed.send(nums, dst=dst_rank, group=group)
|
||||
|
||||
for d in tensor:
|
||||
assert isinstance(d, paddle.Tensor)
|
||||
self._send_dims_shape_dtype(d, group=group)
|
||||
|
||||
def set_send_message(self, tensor):
|
||||
if isinstance(tensor, paddle.Tensor):
|
||||
self.send_shape_message = tensor.shape
|
||||
self.send_dtype_message = paddle_2_number(tensor.dtype)
|
||||
elif isinstance(tensor, tuple):
|
||||
self.send_shape_message = tuple(
|
||||
[d.shape for d in tensor if not d.stop_gradient]
|
||||
)
|
||||
self.send_dtype_message = tuple(
|
||||
[
|
||||
paddle_2_number(d.dtype)
|
||||
for d in tensor
|
||||
if not d.stop_gradient
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _is_valid_send_recv_partial(tensor, mp_degree):
|
||||
if not _enable_partial_send_recv:
|
||||
return False
|
||||
tensor_numel = np.prod(tensor.shape)
|
||||
assert tensor_numel != 0, "can't send/recv zero element"
|
||||
return mp_degree > 1 and tensor_numel % mp_degree == 0
|
||||
|
||||
|
||||
def _partial_send_op(
|
||||
tensor, group, use_calc_stream, ring_id, dst, nranks, rank_id
|
||||
):
|
||||
dst_rank_in_group = dst if group is None else group.get_group_rank(dst)
|
||||
if framework.in_dynamic_mode():
|
||||
group = (
|
||||
paddle.distributed.collective._get_default_group()
|
||||
if group is None
|
||||
else group
|
||||
)
|
||||
comm_op = (
|
||||
group.process_group.send_partial_on_calc_stream
|
||||
if use_calc_stream
|
||||
else group.process_group.send_partial
|
||||
)
|
||||
return comm_op(tensor, dst_rank_in_group, nranks, rank_id)
|
||||
|
||||
|
||||
def send_partial(
|
||||
tensor, dst=0, nranks=1, rank_id=0, group=None, use_calc_stream=True
|
||||
):
|
||||
# dst: local rank in group
|
||||
if group is not None and not group.is_member():
|
||||
return
|
||||
ring_id = 0 if group is None else group.id
|
||||
|
||||
dst_rank = (
|
||||
_hcg._get_p2p_next_rank() if dst == 1 else _hcg._get_p2p_prev_rank()
|
||||
)
|
||||
|
||||
if _is_valid_send_recv_partial(tensor, nranks):
|
||||
return _partial_send_op(
|
||||
tensor, group, use_calc_stream, ring_id, dst_rank, nranks, rank_id
|
||||
)
|
||||
else:
|
||||
send_op = paddle.distributed.isend
|
||||
return send_op(tensor.detach(), dst=dst_rank, group=group)
|
||||
|
||||
|
||||
def _partial_recv_op(
|
||||
tensor, group, use_calc_stream, ring_id, src, nranks, rank_id
|
||||
):
|
||||
src_rank_in_group = src if group is None else group.get_group_rank(src)
|
||||
group = (
|
||||
paddle.distributed.collective._get_default_group()
|
||||
if group is None
|
||||
else group
|
||||
)
|
||||
comm_op = (
|
||||
group.process_group.recv_partial_on_calc_stream
|
||||
if use_calc_stream
|
||||
else group.process_group.recv_partial
|
||||
)
|
||||
return comm_op(tensor, src_rank_in_group, nranks, rank_id)
|
||||
|
||||
|
||||
def recv_partial(
|
||||
tensor, src=0, nranks=1, rank_id=0, group=None, use_calc_stream=True
|
||||
):
|
||||
# src: local rank in group
|
||||
if group is not None and not group.is_member():
|
||||
return
|
||||
ring_id = 0 if group is None else group.id
|
||||
|
||||
src_rank = (
|
||||
_hcg._get_p2p_prev_rank() if src == 0 else _hcg._get_p2p_next_rank()
|
||||
)
|
||||
|
||||
if _is_valid_send_recv_partial(tensor, nranks):
|
||||
return _partial_recv_op(
|
||||
tensor, group, use_calc_stream, ring_id, src_rank, nranks, rank_id
|
||||
)
|
||||
else:
|
||||
if use_calc_stream:
|
||||
recv_op = paddle.distributed.recv
|
||||
elif framework.in_dynamic_mode():
|
||||
recv_op = paddle.distributed.irecv
|
||||
return recv_op(tensor.detach(), src=src_rank, group=group)
|
||||
|
||||
|
||||
def _partial_allgather_op(
|
||||
tensor, group, use_calc_stream, ring_id, nranks, rank_id
|
||||
):
|
||||
group = (
|
||||
paddle.distributed.collective._get_default_group()
|
||||
if group is None
|
||||
else group
|
||||
)
|
||||
comm_op = (
|
||||
group.process_group.all_gather_partial_on_calc_stream
|
||||
if use_calc_stream
|
||||
else group.process_group.all_gather_partial
|
||||
)
|
||||
return comm_op(tensor, tensor, nranks, rank_id)
|
||||
|
||||
|
||||
def allgather_partial(
|
||||
tensor, nranks=1, rank_id=0, group=None, use_calc_stream=True
|
||||
):
|
||||
if not _is_valid_send_recv_partial(tensor, nranks):
|
||||
return tensor
|
||||
if group is not None and not group.is_member():
|
||||
return
|
||||
ring_id = 0 if group is None else group.id
|
||||
|
||||
return _partial_allgather_op(
|
||||
tensor, group, use_calc_stream, ring_id, nranks, rank_id
|
||||
)
|
||||
|
||||
|
||||
def _p2p_helper(
|
||||
tensor_send_next,
|
||||
tensor_send_prev,
|
||||
recv_prev,
|
||||
recv_next,
|
||||
sync_recv=True,
|
||||
send_recv_meta=None,
|
||||
):
|
||||
global _hcg
|
||||
|
||||
tensor_recv_prev = None
|
||||
tensor_recv_next = None
|
||||
|
||||
# send / recv message
|
||||
assert send_recv_meta is not None, "send_recv_meta should not be None"
|
||||
recv_shape_msg = send_recv_meta.recv_shape_message
|
||||
recv_dtype_msg = send_recv_meta.recv_dtype_message
|
||||
recv_stop_gradient = send_recv_meta.recv_stop_gradient
|
||||
|
||||
send_shape_msg = send_recv_meta.send_shape_message
|
||||
send_dtype_msg = send_recv_meta.send_dtype_message
|
||||
|
||||
# model parallel message
|
||||
mp_group = _hcg.get_model_parallel_group()
|
||||
mp_degree = _hcg.get_model_parallel_world_size()
|
||||
mp_rank = _hcg.get_model_parallel_rank()
|
||||
|
||||
if recv_prev:
|
||||
if isinstance(recv_shape_msg, tuple):
|
||||
tensor_recv_prev = []
|
||||
for idx, shape in enumerate(recv_shape_msg):
|
||||
tmp = paddle.empty(
|
||||
shape=shape, dtype=number_2_dtype(recv_dtype_msg[idx])
|
||||
)
|
||||
tmp.stop_gradient = recv_stop_gradient[idx]
|
||||
tensor_recv_prev.append(tmp)
|
||||
tensor_recv_prev = tuple(tensor_recv_prev)
|
||||
else:
|
||||
tensor_recv_prev = paddle.empty(
|
||||
shape=recv_shape_msg, dtype=number_2_dtype(recv_dtype_msg)
|
||||
)
|
||||
tensor_recv_prev.stop_gradient = recv_stop_gradient
|
||||
|
||||
if recv_next:
|
||||
if isinstance(send_shape_msg, tuple):
|
||||
tensor_recv_next = []
|
||||
for idx, shape in enumerate(send_shape_msg):
|
||||
tensor_recv_next.append(
|
||||
paddle.empty(
|
||||
shape=shape, dtype=number_2_dtype(send_dtype_msg[idx])
|
||||
)
|
||||
)
|
||||
tensor_recv_next = tuple(tensor_recv_next)
|
||||
else:
|
||||
tensor_recv_next = paddle.empty(
|
||||
shape=send_shape_msg, dtype=number_2_dtype(send_dtype_msg)
|
||||
)
|
||||
|
||||
# TODO(Yuang Liu): use batch_isend_irecv replace all these comm ops
|
||||
tasks = []
|
||||
# start to p2p communicate
|
||||
|
||||
if _sync_send:
|
||||
# Some devices(NPU for example) do not support asynchronized send op, So the order is
|
||||
# recv_prev -> send_next -> recv_next -> send_prev
|
||||
# When using this order, the environment variable
|
||||
# 'PADDLE_P2P_SYNC_SEND' should be set True
|
||||
if tensor_recv_prev is not None:
|
||||
if isinstance(tensor_recv_prev, tuple):
|
||||
for d in tensor_recv_prev:
|
||||
task = recv_partial(
|
||||
d,
|
||||
src=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_prev_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
if sync_recv:
|
||||
allgather_partial(
|
||||
d,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
else:
|
||||
task = recv_partial(
|
||||
tensor_recv_prev,
|
||||
src=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_prev_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
|
||||
if sync_recv:
|
||||
allgather_partial(
|
||||
tensor_recv_prev,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
|
||||
if tensor_send_next is not None:
|
||||
if isinstance(tensor_send_next, tuple):
|
||||
for d in tensor_send_next:
|
||||
paddle.distributed.wait(d, use_calc_stream=True)
|
||||
send_partial(
|
||||
d,
|
||||
dst=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_next_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
else:
|
||||
paddle.distributed.wait(tensor_send_next, use_calc_stream=True)
|
||||
send_partial(
|
||||
tensor_send_next,
|
||||
dst=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_next_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
|
||||
if tensor_recv_next is not None:
|
||||
if isinstance(tensor_recv_next, tuple):
|
||||
for d in tensor_recv_next:
|
||||
task = recv_partial(
|
||||
d,
|
||||
src=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_next_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
|
||||
if sync_recv:
|
||||
allgather_partial(
|
||||
d,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
|
||||
else:
|
||||
task = recv_partial(
|
||||
tensor_recv_next,
|
||||
src=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_next_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
if sync_recv:
|
||||
allgather_partial(
|
||||
tensor_recv_next,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
|
||||
if tensor_send_prev is not None:
|
||||
if isinstance(tensor_send_prev, tuple):
|
||||
for d in tensor_send_prev:
|
||||
paddle.distributed.wait(d, use_calc_stream=True)
|
||||
send_partial(
|
||||
d,
|
||||
dst=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_prev_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
else:
|
||||
paddle.distributed.wait(tensor_send_prev, use_calc_stream=True)
|
||||
send_partial(
|
||||
tensor_send_prev,
|
||||
dst=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_prev_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
else:
|
||||
_xpu_comm_group_start()
|
||||
if tensor_send_prev is not None:
|
||||
if isinstance(tensor_send_prev, tuple):
|
||||
for d in tensor_send_prev:
|
||||
paddle.distributed.wait(d, use_calc_stream=True)
|
||||
send_partial(
|
||||
d,
|
||||
dst=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_prev_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
else:
|
||||
paddle.distributed.wait(tensor_send_prev, use_calc_stream=True)
|
||||
send_partial(
|
||||
tensor_send_prev,
|
||||
dst=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_prev_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
|
||||
if tensor_recv_prev is not None:
|
||||
if isinstance(tensor_recv_prev, tuple):
|
||||
for d in tensor_recv_prev:
|
||||
task = recv_partial(
|
||||
d,
|
||||
src=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_prev_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
if sync_recv:
|
||||
_xpu_comm_group_end()
|
||||
allgather_partial(
|
||||
d,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
else:
|
||||
task = recv_partial(
|
||||
tensor_recv_prev,
|
||||
src=0,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_prev_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
|
||||
if sync_recv:
|
||||
_xpu_comm_group_end()
|
||||
allgather_partial(
|
||||
tensor_recv_prev,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
|
||||
if tensor_send_next is not None:
|
||||
if isinstance(tensor_send_next, tuple):
|
||||
for d in tensor_send_next:
|
||||
paddle.distributed.wait(d, use_calc_stream=True)
|
||||
send_partial(
|
||||
d,
|
||||
dst=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_next_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
else:
|
||||
paddle.distributed.wait(tensor_send_next, use_calc_stream=True)
|
||||
send_partial(
|
||||
tensor_send_next,
|
||||
dst=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.send_next_group,
|
||||
use_calc_stream=False,
|
||||
)
|
||||
|
||||
if tensor_recv_next is not None:
|
||||
if isinstance(tensor_recv_next, tuple):
|
||||
for d in tensor_recv_next:
|
||||
task = recv_partial(
|
||||
d,
|
||||
src=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_next_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
|
||||
if sync_recv:
|
||||
_xpu_comm_group_end()
|
||||
allgather_partial(
|
||||
d,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
|
||||
else:
|
||||
task = recv_partial(
|
||||
tensor_recv_next,
|
||||
src=1,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=_hcg.recv_next_group,
|
||||
use_calc_stream=sync_recv,
|
||||
)
|
||||
if sync_recv:
|
||||
_xpu_comm_group_end()
|
||||
allgather_partial(
|
||||
tensor_recv_next,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
else:
|
||||
tasks.append(task)
|
||||
_xpu_comm_group_end()
|
||||
if not sync_recv:
|
||||
if framework.in_dynamic_mode():
|
||||
# wait irecv tasks in eager dygraph mode with new comm library
|
||||
for task in tasks:
|
||||
assert task is not None
|
||||
task.wait()
|
||||
|
||||
tensors_for_all_gather = []
|
||||
if tensor_recv_prev is not None:
|
||||
if isinstance(tensor_recv_prev, tuple):
|
||||
for d in tensor_recv_prev:
|
||||
tensors_for_all_gather.append(d)
|
||||
else:
|
||||
tensors_for_all_gather.append(tensor_recv_prev)
|
||||
if tensor_recv_next is not None:
|
||||
if isinstance(tensor_recv_next, tuple):
|
||||
for d in tensor_recv_next:
|
||||
tensors_for_all_gather.append(d)
|
||||
else:
|
||||
tensors_for_all_gather.append(tensor_recv_next)
|
||||
|
||||
for tensor in tensors_for_all_gather:
|
||||
allgather_partial(
|
||||
tensor,
|
||||
nranks=mp_degree,
|
||||
rank_id=mp_rank,
|
||||
group=mp_group,
|
||||
use_calc_stream=True,
|
||||
)
|
||||
|
||||
return tensor_recv_prev, tensor_recv_next
|
||||
|
||||
|
||||
class P2pHelper:
|
||||
def __init__(self, use_cache=True):
|
||||
self._send_recv_meta = SendRecvMeta()
|
||||
self._use_cache = use_cache
|
||||
|
||||
def _send_meta(self, output_tensor, skip_check_meta=False):
|
||||
if not self._send_recv_meta.has_send_meta:
|
||||
self._send_recv_meta.set_send_message(output_tensor)
|
||||
self._send_recv_meta.send_meta(
|
||||
output_tensor, _hcg.get_pipe_parallel_group()
|
||||
)
|
||||
self._send_recv_meta.has_send_meta = self._use_cache
|
||||
|
||||
def _recv_meta(self):
|
||||
if not self._send_recv_meta.has_recv_meta:
|
||||
self._send_recv_meta.recv_meta(_hcg.get_pipe_parallel_group())
|
||||
self._send_recv_meta.has_recv_meta = self._use_cache
|
||||
|
||||
def recv_forward(self, pp_first_stage, sync_recv=True):
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("recv_forward").start()
|
||||
if pp_first_stage:
|
||||
input_tensor = None
|
||||
else:
|
||||
self._recv_meta()
|
||||
|
||||
input_tensor, _ = _p2p_helper(
|
||||
tensor_send_next=None,
|
||||
tensor_send_prev=None,
|
||||
recv_prev=True,
|
||||
recv_next=False,
|
||||
sync_recv=sync_recv,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("recv_forward").stop()
|
||||
return input_tensor
|
||||
|
||||
def recv_backward(self, pp_last_stage, sync_recv=True):
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("recv_backward").start()
|
||||
if pp_last_stage:
|
||||
output_tensor_grad = None
|
||||
else:
|
||||
_, output_tensor_grad = _p2p_helper(
|
||||
tensor_send_next=None,
|
||||
tensor_send_prev=None,
|
||||
recv_prev=False,
|
||||
recv_next=True,
|
||||
sync_recv=sync_recv,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("recv_backward").stop()
|
||||
return output_tensor_grad
|
||||
|
||||
def send_forward(self, output_tensor, pp_last_stage, skip_check_meta=False):
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("send_forward").start()
|
||||
if not pp_last_stage:
|
||||
self._send_meta(output_tensor, skip_check_meta=skip_check_meta)
|
||||
|
||||
_p2p_helper(
|
||||
tensor_send_next=output_tensor,
|
||||
tensor_send_prev=None,
|
||||
recv_prev=False,
|
||||
recv_next=False,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("send_forward").stop()
|
||||
|
||||
def send_backward(self, input_tensor_grad, pp_first_stage):
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("send_backward").start()
|
||||
if not pp_first_stage:
|
||||
_p2p_helper(
|
||||
tensor_send_next=None,
|
||||
tensor_send_prev=input_tensor_grad,
|
||||
recv_prev=False,
|
||||
recv_next=False,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("send_backward").stop()
|
||||
|
||||
def send_forward_recv_backward(self, output_tensor, pp_last_stage):
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("send_forward_recv_backward").start()
|
||||
if pp_last_stage:
|
||||
output_tensor_grad = None
|
||||
else:
|
||||
_, output_tensor_grad = _p2p_helper(
|
||||
tensor_send_next=output_tensor,
|
||||
tensor_send_prev=None,
|
||||
recv_prev=False,
|
||||
recv_next=True,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("send_forward_recv_backward").stop()
|
||||
return output_tensor_grad
|
||||
|
||||
def send_backward_recv_forward(self, input_tensor_grad, pp_first_stage):
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("send_backward_recv_forward").start()
|
||||
if pp_first_stage:
|
||||
input_tensor = None
|
||||
else:
|
||||
input_tensor, _ = _p2p_helper(
|
||||
tensor_send_next=None,
|
||||
tensor_send_prev=input_tensor_grad,
|
||||
recv_prev=True,
|
||||
recv_next=False,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("send_backward_recv_forward").stop()
|
||||
return input_tensor
|
||||
|
||||
def send_forward_backward_recv_forward_backward(
|
||||
self, output_tensor, input_tensor_grad, recv_prev, recv_next
|
||||
):
|
||||
# always have to send dtype info to downstream
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("send_forward_backward_recv_forward_backward").start()
|
||||
|
||||
self._send_meta(output_tensor)
|
||||
if recv_prev:
|
||||
self._recv_meta()
|
||||
input_tensor, output_tensor_grad = _p2p_helper(
|
||||
tensor_send_next=output_tensor,
|
||||
tensor_send_prev=input_tensor_grad,
|
||||
recv_prev=recv_prev,
|
||||
recv_next=recv_next,
|
||||
sync_recv=False,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("send_forward_backward_recv_forward_backward").stop()
|
||||
return input_tensor, output_tensor_grad
|
||||
|
||||
def send_forward_recv_forward(self, output_tensor, recv_prev):
|
||||
# always have to send dtype info to downstream
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("send_forward_recv_forward").start()
|
||||
|
||||
if output_tensor is not None:
|
||||
self._send_meta(output_tensor)
|
||||
|
||||
if recv_prev:
|
||||
self._recv_meta()
|
||||
|
||||
input_tensor, _ = _p2p_helper(
|
||||
tensor_send_next=output_tensor,
|
||||
tensor_send_prev=None,
|
||||
recv_prev=recv_prev,
|
||||
recv_next=False,
|
||||
sync_recv=False,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("send_forward_recv_forward").stop()
|
||||
return input_tensor
|
||||
|
||||
def send_backward_recv_backward(self, input_tensor_grad, recv_next):
|
||||
global _timers
|
||||
if _timers is not None:
|
||||
_timers("send_backward_recv_backward").start()
|
||||
_, output_tensor_grad = _p2p_helper(
|
||||
tensor_send_next=None,
|
||||
tensor_send_prev=input_tensor_grad,
|
||||
recv_prev=False,
|
||||
recv_next=recv_next,
|
||||
sync_recv=False,
|
||||
send_recv_meta=self._send_recv_meta,
|
||||
)
|
||||
if _timers is not None:
|
||||
_timers("send_backward_recv_backward").stop()
|
||||
return output_tensor_grad
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,41 @@
|
||||
# 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 os
|
||||
|
||||
|
||||
def main():
|
||||
all_record = []
|
||||
all_files = os.listdir('./')
|
||||
all_files = sorted(
|
||||
filter(
|
||||
lambda file: file.startswith("profile_record_tmp_file_for_rank_"),
|
||||
all_files,
|
||||
)
|
||||
)
|
||||
|
||||
for files in all_files:
|
||||
with open(files, 'r') as f:
|
||||
for line in f:
|
||||
all_record.append(line.strip())
|
||||
|
||||
with open('pipeline_profile.json', 'w') as f:
|
||||
f.write('[ ')
|
||||
f.writelines(all_record[i] + ',\n' for i in range(len(all_record) - 1))
|
||||
f.write(all_record[-1])
|
||||
f.write(' ]\n')
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,160 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user