InternLM/internlm/core/communication/__init__.py

27 lines
646 B
Python

from .p2p import (
recv_backward,
recv_forward,
send_backward,
send_backward_recv_backward,
send_backward_recv_forward,
send_forward,
send_forward_backward_recv_forward_backward,
send_forward_recv_backward,
send_forward_recv_forward,
)
from .utils import recv_obj_meta, send_obj_meta
__all__ = [
"send_forward",
"send_forward_recv_forward",
"send_forward_backward_recv_forward_backward",
"send_backward",
"send_backward_recv_backward",
"send_backward_recv_forward",
"send_forward_recv_backward",
"recv_backward",
"recv_forward",
"send_obj_meta",
"recv_obj_meta",
]