1bit: rm wd mom

jasperzhong · jasperzhong · commit f888c8d8f9e8 · 2020-06-20T09:10:08.000Z
diff --git a/byteps/mxnet/__init__.py b/byteps/mxnet/__init__.py
@@ -276,12 +276,6 @@ def _register_compressor(self, params, optimizer_params, compression_params):
 
         # change
         if compression_params.get("momentum"):
-            # 1bit compressor use an additional momentum for weight decay
-            if compressor == "onebit" and "wd" in optimizer_params:
-                intra_compressor = Compression.wdmom(
-                    intra_compressor, optimizer_params["momentum"], optimizer_params["wd"])
-                del optimizer_params["wd"]
-
             del optimizer_params['momentum']
 
         return intra_compressor
@@ -308,7 +302,7 @@ def _allreduce_grads(self):
                 byteps_push_pull(compressed, is_average=False,
                                  name="gradient_" + str(i), priority=-i)
                 param._grad[0] = self._intra_compressors[i].decompress(
-                    compressed, ctx,  x=param._data[0])
+                    compressed, ctx)
 
     def _init_params(self):
         tensors = []
diff --git a/byteps/mxnet/compression.py b/byteps/mxnet/compression.py
@@ -65,41 +65,6 @@ def decompress(self, tensor, ctx, *args, **kwargs):
         return tensor_decompressed
 
 
-class WeightDecayMomentum(Compressor):
-    """For 1bit compression."""
-
-    def __init__(self, compressor, mu, wd, *args, **kwargs):
-        self.compressor = compressor
-        self.mom = None
-        self.cache = None
-        self.mu = mu
-        self.wd = wd
-
-    def compress(self, tensor, *args, **kwargs):
-        """Returns the tensor unmodified."""
-        return self.compressor.compress(tensor)
-
-    def decompress(self, tensor, ctx, *args, **kwargs):
-        """Returns the tensor added with additional momentum for wd
-            m_t = \mu * m_{t-1} + wd * x_t
-            x_{t+1} = x_t - \eta_t (tensor + \mu m_t + wd * x_t)
-        """
-        if "x" not in kwargs:
-            return self.compressor.decompress(tensor, ctx)
-
-        x = kwargs["x"]
-        
-        if self.mom is None:
-            self.mom = nd.zeros_like(tensor)
-            self.cache = nd.zeros_like(tensor)
-
-        nd._internal._mul_scalar(x, self.wd, out=self.cache)
-        self.mom += self.cache
-        nd._internal._mul_scalar(self.mom, self.mu, out=self.mom)
-        tensor += self.mom + self.cache
-        return self.compressor.decompress(tensor, ctx)
-
-
 class Compression(object):
     """Optional gradient compression algorithm used during push_pull."""
 
@@ -109,9 +74,6 @@ class Compression(object):
     """Compress all floating point gradients to 16-bit."""
     fp16 = FP16Compressor()
 
-    """Additional Momentum for weight decay. This is only for 1bit. This is a wrapper."""
-    wdmom = WeightDecayMomentum
-
 
 # if __name__ == "__main__":
 #     x = WeightDecayMomentum(Compression.none, 0.9, 1e-4)
