[NFC] Hotfix/format

colossalai/__init__.py

@@ -2,3 +2,4 @@
                          launch_from_slurm, launch_from_torch, get_default_parser)
 
 __version__ = '0.0.1'
+

colossalai/builder/pipeline.py

@@ -251,9 +251,9 @@ def build_pipeline_model(layers: nn.Sequential, num_chunks: int = 1, verbose: bo
     partitions = partition_uniform(len(layers), pipeline_parallel_size, num_chunks)
     module_list = []
     for start, end in partitions[pipeline_rank]:
-        module_list.append(nn.Sequential(*[nn.Identity() for _ in range(start)],
-                                         *layers[start:end],
-                                         *[nn.Identity() for _ in range(len(layers) - end)]))
+        module_list.append(
+            nn.Sequential(*[nn.Identity() for _ in range(start)], *layers[start:end],
+                          *[nn.Identity() for _ in range(len(layers) - end)]))
     if verbose:
         logger = get_dist_logger()
         logger.info(f'Total {len(layers)} layers', ranks=[0])
@@ -264,4 +264,3 @@ def build_pipeline_model(layers: nn.Sequential, num_chunks: int = 1, verbose: bo
                 log_str += '\n'.join([str(layer) for layer in layers[start:end]]) + '\n'
             logger.info(log_str, ranks=[0])
     return nn.ModuleList(module_list) if len(module_list) > 1 else module_list[0]
-

colossalai/kernel/cuda_native/csrc/cpu_adam.cpp

@@ -20,12 +20,14 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE
 */
 #include "cpu_adam.h"
-#include <iostream>
+
 #include <math.h>
-#include <memory>
 #include <omp.h>
 #include <string.h>
 #include <torch/extension.h>
+
+#include <iostream>
+#include <memory>
 #include <type_traits>
 #include <unordered_map>
 
@@ -82,8 +84,7 @@ void Adam_Optimizer::Step_1(float *_params, float *grads, float *_exp_avg,
 
   for (size_t t = 0; t < rounded_size; t += TILE) {
     size_t copy_size = TILE;
-    if ((t + TILE) > rounded_size)
-      copy_size = rounded_size - t;
+    if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
     size_t offset = copy_size + t;
 
 #pragma omp parallel for
@@ -145,8 +146,7 @@ void Adam_Optimizer::Step_1(float *_params, float *grads, float *_exp_avg,
   if (_param_size > rounded_size) {
     for (size_t t = rounded_size; t < _param_size; t += TILE) {
       size_t copy_size = TILE;
-      if ((t + TILE) > _param_size)
-        copy_size = _param_size - t;
+      if ((t + TILE) > _param_size) copy_size = _param_size - t;
       size_t offset = copy_size + t;
 
 #pragma omp parallel for
@@ -235,8 +235,7 @@ void Adam_Optimizer::Step_4(float *_params, float *grads, float *_exp_avg,
 
   for (size_t t = 0; t < rounded_size; t += TILE) {
     size_t copy_size = TILE;
-    if ((t + TILE) > rounded_size)
-      copy_size = rounded_size - t;
+    if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
     size_t offset = copy_size + t;
 
 #pragma omp parallel for
@@ -321,7 +320,6 @@ int create_adam_optimizer(int optimizer_id, float alpha = 1e-3,
   s_optimizers[optimizer_id] = opt;
 
   if (should_log) {
-
     std::string avx_type = "";
 #if defined(__AVX512__)
     avx_type = "AVX512";
@@ -386,8 +384,7 @@ void Adam_Optimizer::Step_8(float *_params, float *grads, float *_exp_avg,
 
   for (size_t t = 0; t < rounded_size; t += TILE) {
     size_t copy_size = TILE;
-    if ((t + TILE) > rounded_size)
-      copy_size = rounded_size - t;
+    if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
     size_t offset = copy_size + t;
 
 #pragma omp parallel for
@@ -463,43 +460,29 @@ void Adam_Optimizer::Step_8(float *_params, float *grads, float *_exp_avg,
            grad_half_precision, loss_scale);
 }
 
-int adam_step(int optimizer_id,
-                 size_t step,
-                 float lr,
-                 float beta1,
-                 float beta2,
-                 float epsilon,
-                 float weight_decay,
-                 bool bias_correction,
-                 torch::Tensor& params,
-                 torch::Tensor& grads,
-                 torch::Tensor& exp_avg,
-                 torch::Tensor& exp_avg_sq,
-                 float loss_scale)
-{
-    auto params_c = params.contiguous();
-    auto grads_c = grads.contiguous();
-    auto exp_avg_c = exp_avg.contiguous();
-    auto exp_avg_sq_c = exp_avg_sq.contiguous();
-
-    float* params_ptr = (float*)params_c.data_ptr();
-    float* grads_ptr = (float*)grads_c.data_ptr();
-    float* exp_avg_ptr = (float*)exp_avg_c.data_ptr();
-    float* exp_avg_sq_ptr = (float*)exp_avg_sq_c.data_ptr();
-    std::shared_ptr<Adam_Optimizer> opt =
-        std::static_pointer_cast<Adam_Optimizer>(s_optimizers[optimizer_id]);
-    opt->IncrementStep(step, beta1, beta2);
-    opt->update_state(lr, epsilon, weight_decay, bias_correction);
-    opt->Step_8(params_ptr,
-                grads_ptr,
-                exp_avg_ptr,
-                exp_avg_sq_ptr,
-                params_c.numel(),
-                (params.options().dtype() == at::kHalf),
-                (grads.options().dtype() == at::kHalf),
-                loss_scale);
-
-    return 0;
+int adam_step(int optimizer_id, size_t step, float lr, float beta1, float beta2,
+              float epsilon, float weight_decay, bool bias_correction,
+              torch::Tensor &params, torch::Tensor &grads,
+              torch::Tensor &exp_avg, torch::Tensor &exp_avg_sq,
+              float loss_scale) {
+  auto params_c = params.contiguous();
+  auto grads_c = grads.contiguous();
+  auto exp_avg_c = exp_avg.contiguous();
+  auto exp_avg_sq_c = exp_avg_sq.contiguous();
+
+  float *params_ptr = (float *)params_c.data_ptr();
+  float *grads_ptr = (float *)grads_c.data_ptr();
+  float *exp_avg_ptr = (float *)exp_avg_c.data_ptr();
+  float *exp_avg_sq_ptr = (float *)exp_avg_sq_c.data_ptr();
+  std::shared_ptr<Adam_Optimizer> opt =
+      std::static_pointer_cast<Adam_Optimizer>(s_optimizers[optimizer_id]);
+  opt->IncrementStep(step, beta1, beta2);
+  opt->update_state(lr, epsilon, weight_decay, bias_correction);
+  opt->Step_8(params_ptr, grads_ptr, exp_avg_ptr, exp_avg_sq_ptr,
+              params_c.numel(), (params.options().dtype() == at::kHalf),
+              (grads.options().dtype() == at::kHalf), loss_scale);
+
+  return 0;
 }
 
 int destroy_adam_optimizer(int optimizer_id) {

colossalai/kernel/cuda_native/csrc/cpu_adam.h

@@ -48,10 +48,10 @@ SOFTWARE
 #define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c)
 #define SIMD_SQRT(x) _mm512_sqrt_ps(x)
 #define SIMD_DIV(x, y) _mm512_div_ps(x, y)
-#define SIMD_LOAD_HALF(x)                                                      \
+#define SIMD_LOAD_HALF(x) \
   _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(x)))
-#define SIMD_STORE_HALF(x, d)                                                  \
-  _mm256_store_ps(                                                             \
+#define SIMD_STORE_HALF(x, d) \
+  _mm256_store_ps(            \
       x, _mm256_castsi256_ps(_mm512_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT)))
 
 #elif defined(__AVX256__) or defined(__AVX2__)
@@ -66,8 +66,8 @@ SOFTWARE
 #define SIMD_SQRT(x) _mm256_sqrt_ps(x)
 #define SIMD_DIV(x, y) _mm256_div_ps(x, y)
 #define SIMD_LOAD_HALF(x) _mm256_cvtph_ps(_mm_loadu_si128((const __m128i *)(x)))
-#define SIMD_STORE_HALF(x, d)                                                  \
-  _mm_store_ps(                                                                \
+#define SIMD_STORE_HALF(x, d) \
+  _mm_store_ps(               \
       x, _mm_castsi128_ps(_mm256_cvtps_ph(d, _MM_FROUND_TO_NEAREST_INT)))
 
 #endif
@@ -83,19 +83,25 @@ union AVX_Data {
 
 #endif
 
-#define STEP(SPAN)                                                             \
-  void Step_##SPAN(float *_params, float *grads, float *_exp_avg,              \
-                   float *_exp_avg_sq, size_t _param_size,                     \
-                   bool param_half_precision = false,                          \
+#define STEP(SPAN)                                                \
+  void Step_##SPAN(float *_params, float *grads, float *_exp_avg, \
+                   float *_exp_avg_sq, size_t _param_size,        \
+                   bool param_half_precision = false,             \
                    bool grad_half_precision = false, float loss_scale = -1);
 
 class Adam_Optimizer {
-public:
+ public:
   Adam_Optimizer(float alpha = 1e-3, float betta1 = 0.9, float betta2 = 0.999,
                  float eps = 1e-8, float weight_decay = 0,
                  bool adamw_mode = true)
-      : _alpha(alpha), _betta1(betta1), _betta2(betta2), _eps(eps),
-        _weight_decay(weight_decay), _betta1_t(1.0), _betta2_t(1.0), _step(0),
+      : _alpha(alpha),
+        _betta1(betta1),
+        _betta2(betta2),
+        _eps(eps),
+        _weight_decay(weight_decay),
+        _betta1_t(1.0),
+        _betta2_t(1.0),
+        _step(0),
         _adamw_mode(adamw_mode) {}
   ~Adam_Optimizer() {}
 
@@ -135,7 +141,7 @@ class Adam_Optimizer {
     }
   }
 
-private:
+ private:
   float _alpha;
   float _betta1;
   float _betta2;

colossalai/kernel/cuda_native/csrc/kernels/cross_entropy.cu

@@ -16,7 +16,7 @@ __global__ void ls_cross_entropy_fw_kernel(
   const int left_idx = block_start + threadIdx.x;
   const int right_idx = (blockIdx.x + 1) * vocab_size;
   float max_input[1] = {REDUCE_FLOAT_INF_NEG};
-  float sum_logits[2] = {0.f, 0.f}; // logit and logit exp
+  float sum_logits[2] = {0.f, 0.f};  // logit and logit exp
   int target_tid = targets[blockIdx.x];
 
   if (target_tid == padding_idx) {

colossalai/kernel/cuda_native/csrc/kernels/dropout_kernels.cu

@@ -1,10 +1,10 @@
+#include <cooperative_groups.h>
+
 #include <chrono>
 #include <ctime>
 
 #include "kernels.h"
 
-#include <cooperative_groups.h>
-
 namespace cg = cooperative_groups;
 
 curandStatePhilox4_32_10_t *curandstate;
@@ -165,8 +165,7 @@ __global__ void ls_dropout_kernel(const int total_count, const float ratio,
   const float scale = 1.f / (1.f - ratio);
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 4 >= total_count)
-    return;
+  if (i * 4 >= total_count) return;
 
   curandStatePhilox4_32_10_t state;
   curand_init(seed, i, 0, &state);
@@ -202,8 +201,7 @@ __global__ void ls_dropout_kernel(const int total_count, const float ratio,
 
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 8 >= total_count)
-    return;
+  if (i * 8 >= total_count) return;
 
   curandStatePhilox4_32_10_t state;
   curand_init(seed, i, 0, &state);
@@ -261,8 +259,7 @@ __global__ void ls_dropout_bwd_kernel(const int total_count, const float ratio,
   const float scale = 1.f / (1.f - ratio);
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 4 >= total_count)
-    return;
+  if (i * 4 >= total_count) return;
 
   uint8_t m[4];
 
@@ -289,8 +286,7 @@ __global__ void ls_dropout_bwd_kernel(const int total_count, const float ratio,
 
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 8 >= total_count)
-    return;
+  if (i * 8 >= total_count) return;
 
   float4 *out4 = reinterpret_cast<float4 *>(out);
   const float4 *vals_float4 = reinterpret_cast<const float4 *>(in);
@@ -380,8 +376,7 @@ __global__ void ls_dropout_res_bias_kernel(
   const float scale = 1.f / (1.f - ratio);
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 4 >= total_count)
-    return;
+  if (i * 4 >= total_count) return;
 
   curandStatePhilox4_32_10_t state;
   curand_init(seed, i, 0, &state);
@@ -424,8 +419,7 @@ __global__ void ls_dropout_res_bias_kernel(
 
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 8 >= total_count)
-    return;
+  if (i * 8 >= total_count) return;
 
   curandStatePhilox4_32_10_t state;
   curand_init(seed, i, 0, &state);
@@ -565,11 +559,9 @@ __global__ void ls_dropout_bias_bwd_kernel(
   }
   __syncthreads();
 
-  for (int i = 1; i < 32; i <<= 1)
-    sum += g.shfl_down(sum, i);
+  for (int i = 1; i < 32; i <<= 1) sum += g.shfl_down(sum, i);
 
-  if (y == 0)
-    tile[0][x] = sum;
+  if (y == 0) tile[0][x] = sum;
   __syncthreads();
 
   if (threadIdx.x < 8) {
@@ -621,11 +613,9 @@ __global__ void ls_dropout_bias_bwd_kernel(
   }
   __syncthreads();
 
-  for (int i = 1; i < WARP_SIZE; i <<= 1)
-    sum += g.shfl_down(sum, i);
+  for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_down(sum, i);
 
-  if (y == 0)
-    tile[0][x] = sum;
+  if (y == 0) tile[0][x] = sum;
   __syncthreads();
 
   if (threadIdx.x < 8) {
@@ -689,8 +679,7 @@ __global__ void ls_dropout_act_bias_kernel(
   const float scale = 1.f / (1.f - ratio);
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 4 >= total_count)
-    return;
+  if (i * 4 >= total_count) return;
 
   curandStatePhilox4_32_10_t state;
   curand_init(seed, i, 0, &state);
@@ -735,8 +724,7 @@ __global__ void ls_dropout_act_bias_kernel(
 
   int i = blockIdx.x * blockDim.x + threadIdx.x;
 
-  if (i * 8 >= total_count)
-    return;
+  if (i * 8 >= total_count) return;
 
   curandStatePhilox4_32_10_t state;
   curand_init(seed, i, 0, &state);
@@ -897,11 +885,9 @@ __global__ void ls_dropout_act_bias_bwd_kernel(
   float sum = tile[threadIdx.y][threadIdx.x];
   __syncthreads();
 
-  for (int i = 1; i < WARP_SIZE; i <<= 1)
-    sum += g.shfl_down(sum, i);
+  for (int i = 1; i < WARP_SIZE; i <<= 1) sum += g.shfl_down(sum, i);
 
-  if (threadIdx.x == 0)
-    tile[0][threadIdx.y] = sum;
+  if (threadIdx.x == 0) tile[0][threadIdx.y] = sum;
   __syncthreads();
 
   if (threadIdx.y == 0) {

colossalai/kernel/cuda_native/csrc/kernels/general_kernels.cu

@@ -1,7 +1,7 @@
-#include "kernels.h"
-
 #include <cooperative_groups.h>
 
+#include "kernels.h"
+
 namespace cg = cooperative_groups;
 
 /**