support for fp16 in naive comm (#6333)

* support fp16 in naive comm * support for fp16 in bcast
chainer · Mar 4, 2019 · b23f877 · b23f877
1 parent d9bdfc6
commit b23f877
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 2 deletions.
diff --git a/chainermn/communicators/mpi_communicator_base.py b/chainermn/communicators/mpi_communicator_base.py
@@ -611,8 +611,14 @@ def allreduce_obj(self, obj):
     def bcast_data(self, model):
         for _, param in sorted(model.namedparams()):
             if param.data is not None:
-                buf = _memory_utility.array_to_buffer_object(param.data)
+                data = param.data
+                is_float16 = param.data.dtype == numpy.float16
+                if is_float16:
+                    data = data.astype(numpy.float32)
+                buf = _memory_utility.array_to_buffer_object(data)
                 self.mpi_comm.Bcast(buf)
+                if is_float16:
+                    param.data = data.astype(numpy.float16)
 
     # Private methods
     def _init_ranks(self):

diff --git a/chainermn/communicators/naive_communicator.py b/chainermn/communicators/naive_communicator.py
@@ -1,4 +1,5 @@
 import mpi4py.MPI
+import numpy as np
 
 from chainermn.communicators import _memory_utility
 from chainermn.communicators import mpi_communicator_base
@@ -11,6 +12,12 @@ def __init__(self, mpi_comm):
 
     def allreduce_grad(self, model):
         for param in _memory_utility.extract_params_set_grad(model):
-            buf = _memory_utility.array_to_buffer_object(param.grad)
+            grad = param.grad
+            is_float16 = param.grad.dtype == np.float16
+            if is_float16:
+                grad = grad.astype(np.float32)
+            buf = _memory_utility.array_to_buffer_object(grad)
             self.mpi_comm.Allreduce(mpi4py.MPI.IN_PLACE, buf)
+            if is_float16:
+                param.grad = grad.astype(np.float16)
             param.grad /= self.size
diff --git a/tests/chainermn_tests/communicator_tests/test_communicator.py b/tests/chainermn_tests/communicator_tests/test_communicator.py
@@ -64,6 +64,10 @@ def __init__(self, param):
     {
         'communicator_class': NaiveCommunicator,
         'multi_node': True,
+    }, {
+        'communicator_class': NaiveCommunicator,
+        'model_dtype': np.float16,
+        'multi_node': True,
     }, {
         'communicator_class': FlatCommunicator,
         'multi_node': True,