未验证 提交 74412dfe 编写于 作者: S ShenLiang 提交者: GitHub

fix bug of old pp (#44361)

上级 3c074de4
...@@ -2030,6 +2030,10 @@ def alltoall_single(in_tensor, ...@@ -2030,6 +2030,10 @@ def alltoall_single(in_tensor,
return task return task
def _get_group_rank(global_rank, group=None):
return global_rank if group is None else group.get_group_rank(global_rank)
def send(tensor, dst=0, group=None, use_calc_stream=True): def send(tensor, dst=0, group=None, use_calc_stream=True):
""" """
Send a tensor to the receiver. Send a tensor to the receiver.
...@@ -2062,11 +2066,10 @@ def send(tensor, dst=0, group=None, use_calc_stream=True): ...@@ -2062,11 +2066,10 @@ def send(tensor, dst=0, group=None, use_calc_stream=True):
""" """
if group is not None and not group.is_member(): if group is not None and not group.is_member():
return return
dst = _get_group_rank(dst, group)
if in_dygraph_mode(): if in_dygraph_mode():
group = _get_default_group() if group is None else group group = _get_default_group() if group is None else group
group_dst_rank = group.get_group_rank(dst) task = group.process_group.send(tensor, dst)
task = group.process_group.send(tensor, group_dst_rank)
if use_calc_stream: if use_calc_stream:
task.wait() task.wait()
return None return None
...@@ -2126,10 +2129,10 @@ def recv(tensor, src=0, group=None, use_calc_stream=True): ...@@ -2126,10 +2129,10 @@ def recv(tensor, src=0, group=None, use_calc_stream=True):
if group is not None and not group.is_member(): if group is not None and not group.is_member():
return return
src = _get_group_rank(src, group)
if in_dygraph_mode(): if in_dygraph_mode():
group = _get_default_group() if group is None else group group = _get_default_group() if group is None else group
group_src_rank = group.get_group_rank(src) task = group.process_group.recv(tensor, src)
task = group.process_group.recv(tensor, group_src_rank)
if use_calc_stream: if use_calc_stream:
task.wait() task.wait()
return None return None
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册