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))