mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-01-16 23:52:00 +00:00
0fedef4f3c
* integrated parallel layers for ease of building models * integrated 2.5d layers * cleaned codes and unit tests * added log metric by step hook; updated imagenet benchmark; fixed some bugs * reworked initialization; cleaned codes Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
17 lines
867 B
Python
17 lines
867 B
Python
from .collective import all_gather, reduce_scatter, all_reduce, broadcast, reduce
|
|
from .p2p import (send_forward, send_forward_recv_forward,
|
|
send_backward_recv_forward, send_backward,
|
|
send_backward_recv_backward, send_forward_recv_backward,
|
|
send_forward_backward_recv_forward_backward, recv_forward,
|
|
recv_backward)
|
|
from .ring import ring_forward
|
|
from .utils import send_tensor_meta, recv_tensor_meta
|
|
|
|
__all__ = [
|
|
'all_gather', 'reduce_scatter', 'all_reduce', 'broadcast', 'reduce',
|
|
'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',
|
|
'ring_forward', 'send_tensor_meta', 'recv_tensor_meta'
|
|
] |