def __init__(self, nccl=False, nccl_fusion=True, hierarchical_nccl=False):
self._nccl = nccl
self._nccl_fusion = nccl_fusion
if self._nccl:
if hierarchical_nccl:
self._group_all_reduce_fn = group_hierarchical_nccl_all_reduce
else:
self._group_all_reduce_fn = group_nccl_all_reduce
else:
self._group_all_reduce_fn = group_all_reduce
_rank, self._num_workers = peer_info()
def test_peer_info():
from kungfu.tensorflow.ops import peer_info
info = peer_info()
with tf.Session() as sess:
rank, np = sess.run(info)
print('rank=%d, np=%d' % (rank, np))
def __init__(self, nccl=False, nccl_fusion=True):
self._nccl = nccl
self._nccl_fusion = nccl_fusion
_rank, self._num_workers = peer_info()
def test_peer_info():
info = peer_info()
with tf.Session() as sess:
rank, np = sess.run(info)
print('rank=%d, np=%d' % (rank, np))