Add FP16 and FP64 Supports to PureNcclComunicator #187
Conversation
… support_fp16_fp64
shu65
changed the title
shu65
changed the title
chainermn/communicators/__init__.py
Outdated
| @@ -41,6 +45,10 @@ def create_communicator( | |||
| import mpi4py.MPI | |||
| mpi_comm = mpi4py.MPI.COMM_WORLD | |||
|
|
|||
| if communicator_name != 'pure_nccl' and allreduce_grad_dtype is not None: | |||
| raise ValueError( | |||
| 'allreduce_grad_dtype is not supported except for \'pure_nccl\'.') | |||
There was a problem hiding this comment.
'allreduce_grad_dtype is only available at 'pure_nccl' communicator.'
There was a problem hiding this comment.
Thank you. I will fix this error message according to your comment.
chainermn/communicators/__init__.py
Outdated
| @@ -31,6 +34,7 @@ def create_communicator( | |||
| ``hierarchical``, ``two_dimensional``, ``pure_nccl``, or | |||
| ``single_node``) | |||
| mpi_comm: MPI4py communicator | |||
| allreduce_grad_dtype: Data type of gradient used in All-Reduce | |||
There was a problem hiding this comment.
I think allowed value here is only np.float32, np.float16 or None. Maybe we can add value check somewhere around here, before actually creating communicator?
There was a problem hiding this comment.
Sorry there are already value check in PureNcclCommunicator, so may be we just can add documentation here, like "Allowed types are numpy.float16, 32, 64 or None....".
There was a problem hiding this comment.
So the behaviour when None is passed should be described here, which is to use float type of models.
There was a problem hiding this comment.
OK. I will add the description.
| @@ -70,8 +70,8 @@ def ptr(self): | |||
| def buffer(self, size): | |||
| return self.ffi.buffer(self.ffi.cast('void *', self.memory.ptr), size) | |||
|
|
|||
| def array(self, shape, offset=0): | |||
| return cp.ndarray(shape, memptr=self.memory + offset, dtype=cp.float32) | |||
| def array(self, shape, offset=0, dtype=np.float32): | |||
There was a problem hiding this comment.
So if we don't want cupy's default behaviour setting dtype as float64 when None is passed, we might better add non-None assertion here.
There was a problem hiding this comment.
I I also think so, I will add the assertion.
This PR adds FP16 and FP64 support to allreduce_grad of PureNcclCommunicator.