gradient compression updates

.gitignore

@@ -117,3 +117,10 @@ venv.bak/
 # for development
 scripts/
 exps/
+sshlog/
+hosts
+example/mxnet/nohup.out
+example/mxnet/run.sh
+server-hosts
+worker-hosts
+*.sh

byteps/common/common.h

@@ -31,10 +31,9 @@
 #include <vector>
 
 // Add for profiling communication events
-#include <stdio.h>
-#include <stdlib.h>
-
 #include <chrono>
+#include <cstdio>
+#include <cstdlib>
 #include <fstream>
 #include <iostream>
 #include <queue>
@@ -267,7 +266,8 @@ using TensorTable = std::unordered_map<std::string, TensorTableEntry>;
 enum class RequestType {
   kDefaultPushPull,
   kRowSparsePushPull,
-  kCompressedPushPull
+  kCompressedPushPull,
+  kConfigPushPull
 };
 
 int GetCommandType(RequestType requestType, int d);
@@ -278,10 +278,18 @@ ncclDataType_t getNcclDataType(DataType dtype);
 
 int getDataTypeLength(int dtype);
 
-inline size_t Align(size_t size, int dtype) {
-  const size_t min_size =
-      (getDataTypeLength(dtype) * getDataTypeLength(dtype)) * 8;
-  return size + (min_size - size % min_size) % min_size;
+inline size_t Align(size_t size) {
+  constexpr size_t MIN_SIZE = 512;
+  return size + (MIN_SIZE - size % MIN_SIZE) % MIN_SIZE;
+}
+
+// promote low-precision type into float32
+inline void Promote(size_t& size, int& dtype) {
+  size_t ele_size = getDataTypeLength(dtype);
+  if (ele_size < 4) {
+    size *= (4 / ele_size);
+    dtype = BYTEPS_FLOAT32;
+  }
 }
 }  // namespace common
 }  // namespace byteps

byteps/common/compressor/cast.cc

@@ -0,0 +1,32 @@
+// Copyright 2020 Amazon Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+#include "cast.h"
+
+namespace byteps {
+namespace common {
+namespace compressor {
+
+void Cast::Compress(tensor_t grad, tensor_t& output) {
+  _cptr->Compress(CastToFP32(grad), output);
+}
+
+void Cast::Decompress(tensor_t compressed, tensor_t& output) {
+  // directly forward to internal compressor
+  _cptr->Decompress(compressed, output);
+}
+}  // namespace compressor
+}  // namespace common
+}  // namespace byteps

byteps/common/compressor/cast.h

@@ -0,0 +1,62 @@
+// Copyright 2020 Amazon Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+
+#ifndef BYTEPS_COMPRESSOR_CAST_H
+#define BYTEPS_COMPRESSOR_CAST_H
+
+#include "../cpu_reducer.h"
+#include "compressor.h"
+
+namespace byteps {
+namespace common {
+namespace compressor {
+/*!
+ * \brief Wrapper of Compressor to deal with low-precision types
+ *
+ * During summation, data in low-precision suffers from overflow problem.
+ * To solve the issue, we instead use float32 as our data type for inter-
+ * mediate buffers. When intra-node all-reduce is over, locally aggregated
+ * gradients are first transformed into float32 via the wrapper.
+ *
+ * The wrapper has an internel fp32 buffer to store the transformed data.
+ *
+ * \sa Compressor
+ */
+class Cast : public Compressor {
+ public:
+  Cast(size_t size, DataType dtype, std::unique_ptr<Compressor> cptr)
+      : Compressor(size, dtype),
+        _fp32_buf(new byte_t[size]()),
+        _cptr(std::move(cptr)){};
+  ~Cast() override = default;
+
+  void Compress(tensor_t grad, tensor_t& output) final;
+
+  void Decompress(tensor_t compressed, tensor_t& output) final;
+
+ protected:
+  virtual tensor_t CastToFP32(tensor_t grad) = 0;
+
+ protected:
+  std::unique_ptr<byte_t[]> _fp32_buf;
+
+ private:
+  std::unique_ptr<Compressor> _cptr;
+};
+}  // namespace compressor
+}  // namespace common
+}  // namespace byteps
+
+#endif  // BYTEPS_COMPRESSOR_CAST_H

byteps/common/compressor/common.h

@@ -20,31 +20,32 @@
 #if __F16C__
 #include "../half.h"
 using half_t = mshadow::half::half_t;
+#else
+using half_t = void;
 #endif
 
 namespace byteps {
 namespace common {
 namespace compressor {
-typedef char byte_t;
-/*!
- * \brief Tensor type
- */
-typedef struct BPSTensor {
-  byte_t* data;
-  size_t size;
-  int dtype;
-
-  BPSTensor() : data(nullptr), size(0), dtype(0) {}
+using byte_t = char;
+
+struct BPSTensor {
+  byte_t* data{nullptr};
+  size_t size{0};
+  int dtype{0};
+
+  BPSTensor() = default;
   BPSTensor(void* data, size_t size = 0, int dtype = 0)
       : data(reinterpret_cast<byte_t*>(data)), size(size), dtype(dtype) {}
-} tensor_t;
+};
+using tensor_t = BPSTensor;
 
 using kwargs_t = std::unordered_map<std::string, std::string>;
 
 #define COMPRESS_IMPL_SWITCH(dtype, func, dst, src, size)                     \
   switch (dtype) {                                                            \
     case BYTEPS_FLOAT16:                                                      \
-      return func(reinterpret_cast<uint16_t*>(dst),                           \
+      return func(reinterpret_cast<uint32_t*>(dst),                           \
                   reinterpret_cast<const half_t*>(src),                       \
                   size / sizeof(half_t));                                     \
     case BYTEPS_FLOAT32:                                                      \
@@ -58,11 +59,45 @@ using kwargs_t = std::unordered_map<std::string, std::string>;
       BPS_CHECK(0) << "Unsupported data type:" << dtype;                      \
   }
 
+#define COMPRESS_IMPL_SWITCH2(dtype, func, dst, src, size)                    \
+  switch (dtype) {                                                            \
+    case BYTEPS_FLOAT16:                                                      \
+      return func(reinterpret_cast<half_t*>(dst),                             \
+                  reinterpret_cast<const half_t*>(src),                       \
+                  size / sizeof(half_t));                                     \
+    case BYTEPS_FLOAT32:                                                      \
+      return func(reinterpret_cast<float*>(dst),                              \
+                  reinterpret_cast<const float*>(src), size / sizeof(float)); \
+    case BYTEPS_FLOAT64:                                                      \
+      return func(reinterpret_cast<double*>(dst),                             \
+                  reinterpret_cast<const double*>(src),                       \
+                  size / sizeof(double));                                     \
+    default:                                                                  \
+      BPS_CHECK(0) << "Unsupported data type:" << dtype;                      \
+  }
+
+#define COMPRESS_IMPL_SCALAR_SWITCH(dtype, func, dst, src, size)              \
+  switch (dtype) {                                                            \
+    case BYTEPS_FLOAT16:                                                      \
+      return func(reinterpret_cast<index_t*>(dst),                            \
+                  reinterpret_cast<const half_t*>(src),                       \
+                  size / sizeof(half_t));                                     \
+    case BYTEPS_FLOAT32:                                                      \
+      return func(reinterpret_cast<index_t*>(dst),                            \
+                  reinterpret_cast<const float*>(src), size / sizeof(float)); \
+    case BYTEPS_FLOAT64:                                                      \
+      return func(reinterpret_cast<index_t*>(dst),                            \
+                  reinterpret_cast<const double*>(src),                       \
+                  size / sizeof(double));                                     \
+    default:                                                                  \
+      BPS_CHECK(0) << "Unsupported data type:" << dtype;                      \
+  }
+
 #define DECOMPRESS_IMPL_SWITCH(dtype, func, dst, src, compressed_size)      \
   switch (dtype) {                                                          \
     case BYTEPS_FLOAT16:                                                    \
       return func(reinterpret_cast<half_t*>(dst),                           \
-                  reinterpret_cast<const uint16_t*>(src), compressed_size); \
+                  reinterpret_cast<const uint32_t*>(src), compressed_size); \
     case BYTEPS_FLOAT32:                                                    \
       return func(reinterpret_cast<float*>(dst),                            \
                   reinterpret_cast<const uint32_t*>(src), compressed_size); \
@@ -73,13 +108,43 @@ using kwargs_t = std::unordered_map<std::string, std::string>;
       BPS_CHECK(0) << "Unsupported data type:" << dtype;                    \
   }
 
+#define DECOMPRESS_IMPL_SWITCH2(dtype, func, dst, src, compressed_size)  \
+  switch (dtype) {                                                       \
+    case BYTEPS_FLOAT16:                                                 \
+      return func(reinterpret_cast<half_t*>(dst),                        \
+                  reinterpret_cast<const float*>(src), compressed_size); \
+    case BYTEPS_FLOAT32:                                                 \
+      return func(reinterpret_cast<float*>(dst),                         \
+                  reinterpret_cast<const float*>(src), compressed_size); \
+    case BYTEPS_FLOAT64:                                                 \
+      return func(reinterpret_cast<double*>(dst),                        \
+                  reinterpret_cast<const float*>(src), compressed_size); \
+    default:                                                             \
+      BPS_CHECK(0) << "Unsupported data type:" << dtype;                 \
+  }
+
+#define DECOMPRESS_IMPL_SCALAR_SWITCH(dtype, func, dst, src, compressed_size) \
+  switch (dtype) {                                                            \
+    case BYTEPS_FLOAT16:                                                      \
+      return func(reinterpret_cast<half_t*>(dst),                             \
+                  reinterpret_cast<const index_t*>(src), compressed_size);    \
+    case BYTEPS_FLOAT32:                                                      \
+      return func(reinterpret_cast<float*>(dst),                              \
+                  reinterpret_cast<const index_t*>(src), compressed_size);    \
+    case BYTEPS_FLOAT64:                                                      \
+      return func(reinterpret_cast<double*>(dst),                             \
+                  reinterpret_cast<const index_t*>(src), compressed_size);    \
+    default:                                                                  \
+      BPS_CHECK(0) << "Unsupported data type:" << dtype;                      \
+  }
+
 #define FAST_UPDATE_ERROR_IMPL_SWITCH(dtype, func, dst, src1, src2,          \
                                       compressed_size)                       \
   switch (dtype) {                                                           \
     case BYTEPS_FLOAT16:                                                     \
       return func(reinterpret_cast<half_t*>(dst),                            \
                   reinterpret_cast<half_t*>(src1),                           \
-                  reinterpret_cast<const uint16_t*>(src2), compressed_size); \
+                  reinterpret_cast<const uint32_t*>(src2), compressed_size); \
     case BYTEPS_FLOAT32:                                                     \
       return func(reinterpret_cast<float*>(dst),                             \
                   reinterpret_cast<float*>(src1),                            \
@@ -92,6 +157,44 @@ using kwargs_t = std::unordered_map<std::string, std::string>;
       BPS_CHECK(0) << "Unsupported data type:" << dtype;                     \
   }
 
+#define FAST_UPDATE_ERROR_IMPL_SWITCH2(dtype, func, dst, src1, src2,       \
+                                       compressed_size)                    \
+  switch (dtype) {                                                         \
+    case BYTEPS_FLOAT16:                                                   \
+      return func(reinterpret_cast<half_t*>(dst),                          \
+                  reinterpret_cast<half_t*>(src1),                         \
+                  reinterpret_cast<const half_t*>(src2), compressed_size); \
+    case BYTEPS_FLOAT32:                                                   \
+      return func(reinterpret_cast<float*>(dst),                           \
+                  reinterpret_cast<float*>(src1),                          \
+                  reinterpret_cast<const float*>(src2), compressed_size);  \
+    case BYTEPS_FLOAT64:                                                   \
+      return func(reinterpret_cast<double*>(dst),                          \
+                  reinterpret_cast<double*>(src1),                         \
+                  reinterpret_cast<const double*>(src2), compressed_size); \
+    default:                                                               \
+      BPS_CHECK(0) << "Unsupported data type:" << dtype;                   \
+  }
+
+#define FAST_UPDATE_ERROR_IMPL_SCALAR_SWITCH(dtype, func, dst, src1, src2,  \
+                                             compressed_size)               \
+  switch (dtype) {                                                          \
+    case BYTEPS_FLOAT16:                                                    \
+      return func(reinterpret_cast<half_t*>(dst),                           \
+                  reinterpret_cast<half_t*>(src1),                          \
+                  reinterpret_cast<const index_t*>(src2), compressed_size); \
+    case BYTEPS_FLOAT32:                                                    \
+      return func(reinterpret_cast<float*>(dst),                            \
+                  reinterpret_cast<float*>(src1),                           \
+                  reinterpret_cast<const index_t*>(src2), compressed_size); \
+    case BYTEPS_FLOAT64:                                                    \
+      return func(reinterpret_cast<double*>(dst),                           \
+                  reinterpret_cast<double*>(src1),                          \
+                  reinterpret_cast<const index_t*>(src2), compressed_size); \
+    default:                                                                \
+      BPS_CHECK(0) << "Unsupported data type:" << dtype;                    \
+  }
+
 }  // namespace compressor
 }  // namespace common
 }  // namespace byteps