hotfix: merge two buffer (mentioned in bytedance#285)

jasperzhong · jasperzhong · commit ba60a767f051 · 2020-08-05T17:42:57.000Z
diff --git a/byteps/mxnet/__init__.py b/byteps/mxnet/__init__.py
@@ -217,15 +217,16 @@ def __init__(self, params, optimizer, optimizer_params=None, root_rank=0, compre
         self.root_rank = root_rank
         self._intra_compressors = {}
         for i, param in enumerate(self._params):
-            byteps_declare_tensor("parameter_" + str(i))
             self._intra_compressors[param.name] = type(self._intra_compressor)(
                 **self._intra_compressor.__dict__)
             if param.grad_req != 'null':
                 byteps_params = dict(
                     filter(lambda attr: attr[0].startswith(
                         "byteps_",), param.__dict__.items())
                 )
-                byteps_declare_tensor("gradient_" + str(i), **byteps_params)
+                byteps_declare_tensor("tensor_" + str(i), **byteps_params)
+            else:
+                byteps_declare_tensor("tensor_" + str(i))
 
     def _register_compressor(self, params, optimizer_params, compression_params):
         """Register compressor for BytePS
@@ -313,7 +314,7 @@ def _allreduce_grads(self):
                 compressed, ctx = self._intra_compressors[param.name].compress(
                     param._grad[0])
                 byteps_push_pull(compressed, is_average=False,
-                                 name="gradient_" + str(i), priority=-i)
+                                 name="tensor_" + str(i), priority=-i)
                 param._grad[0][:] = self._intra_compressors[param.name].decompress(
                     compressed, ctx,  x=param._data[0])
 
@@ -328,8 +329,8 @@ def _init_params(self):
 
                 compressed, ctx = self._intra_compressors[param.name].compress(
                     param_arrays[0])
-                byteps_push_pull(param_arrays[0], version=0, priority=0,
-                                 name="parameter_" + str(idx), is_average=False)
+                byteps_push_pull(compressed, version=0, priority=0,
+                                 name="tensor_" + str(idx), is_average=False)
                 param_arrays[0][:] = self._intra_compressors[param.name].decompress(
                     compressed, ctx,  x=param._data[0])
 
