[NFC] Hotfix/format (#984)

* [NFC] Polish colossalai/kernel/cuda_native/csrc/multi_tensor_lamb.cu code style. (#937)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/cuda_util.h code style (#939)

* [NFC] polish colossalai/kernel/cuda_native/csrc/cpu_adam.cpp code style (#936)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/block_reduce.h code style (#938)

* [NFC] polish moe_cuda_kernel.cu code style (#940)
Co-authored-by: Xiao Ye <xiaoye2@illinois.edu>

* [NFC] polish pre-commit run --files colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax_cuda.cu code style (#943)

* [NFC] polish colossalai/kernel/cuda_native/csrc/moe_cuda.cpp code style (#942)

* [NFC] polish colossalai/kernel/cuda_native/csrc/cpu_adam.h code style (#945)

* [NFC] polish colossalai/kernel/jit/bias_gelu.py code style (#946)
Co-authored-by: jnbai <897086360@qq.com>

* [NFC] polish colossalai/kernel/cuda_native/csrc/scaled_masked_softmax_cuda.cu code style (#949)
Co-authored-by: Jiatong <jiatong.han@u.nus.edu>

* [NFC] polish colossalai/builder/pipeline.py code style (#951)

* [NFC] polish colossalai/kernel/cuda_native/csrc/multihead_attention_1d.cpp code style (#952)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/cross_entropy.cu code style (#953)
Co-authored-by: 何晓昕 <cautious@hexiaoxins-MacBook-Pro.local>

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/softmax_kernels.cu code style (#954)

* [NFC] polish colossalai/kernel/cuda_native/scaled_softmax.py  code style (#955)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/context.h code style (#956)
Co-authored-by: RichardoLuo <14049555596@qq.com>

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/cross_entropy_layer.h code style (#957)

* [NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_l2norm_kernel.cu code style (#958)

* [NFC] polish colossalai/kernel/cuda_native/csrc/multihead_attention_1d.h code style (#962)

* [NFC] polish colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.cpp code style (#959)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/general_kernels.cu code style (#963)
Co-authored-by: “Arsmart123 <202476410arsmart@gmail.com>

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/softmax.h code style (#964)

* [NFC] polish __init__.py code style (#965)

* [NFC] polish colossalai/nn/layer/parallel_3d/layers.py code style (#966)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/feed_forward.h (#968)

code style

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/include/dropout.h code style (#970)

* [NFC] polish colossalai/nn/layer/parallel_2p5d/layers.py code style (#972)

* [NFC] polish colossalai/kernel/cuda_native/csrc/layer_norm_cuda.cpp code style (#973)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/normalize_kernels.cu code style (#974)

* [NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_scale_kernel.cu code style (#977)

* [NFC] polish colossalai/nn/layer/parallel_2d/layers.py code style (#976)

* [NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_sgd_kernel.cu code style (#978)

* [NFC] polish colossalai/kernel/cuda_native/csrc/kernels/dropout_kernels.cu code style (#979)

* [NFC] polish colossalai/kernel/cuda_native/layer_norm.py code style (#980)

* [NFC] polish colossalai/nn/layer/utils/common.py code style (#983)
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
Co-authored-by: yuxuan-lou <83441848+yuxuan-lou@users.noreply.github.com>
Co-authored-by: Geng Zhang <34452939+zxgx@users.noreply.github.com>
Co-authored-by: Maruyama_Aya <38985202+MaruyamaAya@users.noreply.github.com>
Co-authored-by: XYE <92607131+Itok2000u@users.noreply.github.com>
Co-authored-by: Xiao Ye <xiaoye2@illinois.edu>
Co-authored-by: HaoyuQin <79465534+coder-chin@users.noreply.github.com>
Co-authored-by: wky <64853922+wangkuangyi@users.noreply.github.com>
Co-authored-by: bajiaoyu517 <59548007+bajiaoyu517@users.noreply.github.com>
Co-authored-by: luoling-LC <105470086+luoling-LC@users.noreply.github.com>
Co-authored-by: jnbai <897086360@qq.com>
Co-authored-by: JT.Han <59948448+JThh@users.noreply.github.com>
Co-authored-by: Jiatong <jiatong.han@u.nus.edu>
Co-authored-by: xyupeng <99191637+xyupeng@users.noreply.github.com>
Co-authored-by: Sze-qq <68757353+Sze-qq@users.noreply.github.com>
Co-authored-by: Cautiousss <48676630+Cautiousss@users.noreply.github.com>
Co-authored-by: 何晓昕 <cautious@hexiaoxins-MacBook-Pro.local>
Co-authored-by: Luxios22 <67457897+Luxios22@users.noreply.github.com>
Co-authored-by: Wangbo Zhao(黑色枷锁) <56866854+wangbo-zhao@users.noreply.github.com>
Co-authored-by: RichardoLuo <50363844+RichardoLuo@users.noreply.github.com>
Co-authored-by: RichardoLuo <14049555596@qq.com>
Co-authored-by: doubleHU <98150031+huxin711@users.noreply.github.com>
Co-authored-by: runluo <68489000+run-qiao@users.noreply.github.com>
Co-authored-by: MaxT <854721132@qq.com>
Co-authored-by: superhao1995 <804673818@qq.com>
Co-authored-by: ziyu huang <huang0ziyu@gmail.com>
Co-authored-by: “Arsmart123 <202476410arsmart@gmail.com>
Co-authored-by: Yuer867 <62204893+Yuer867@users.noreply.github.com>
Co-authored-by: lucasliunju <lucasliunju@gmail.com>
Co-authored-by: LuGY <74758262+Gy-Lu@users.noreply.github.com>
Co-authored-by: ExtremeViscent <zhangyiqi55732@sina.com>
Co-authored-by: Xu Kai <xukai16@foxmail.com>
Co-authored-by: Zirui Zhu <zhuzr21@gmail.com>
Co-authored-by: Ofey Chan <ofey206@gmail.com>
Co-authored-by: DouJS <dujiangsu@163.com>
Co-authored-by: Jie Zhu <chore.08-protist@icloud.com>
Co-authored-by: shenggan <csg19971016@gmail.com>
Co-authored-by: Kai Wang (Victor Kai) <37533040+kaiwang960112@users.noreply.github.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: Ziheng Qin <37519855+henryqin1997@users.noreply.github.com>

colossalai/kernel/cuda_native/csrc/multi_tensor_lamb.cu

@@ -15,7 +15,8 @@
 #define BLOCK_SIZE 512
 #define ILP 4
 
-template <typename T> __device__ __forceinline__ bool is_aligned(T *p) {
+template <typename T>
+__device__ __forceinline__ bool is_aligned(T *p) {
   return ((uint64_t)p) % (ILP * sizeof(T)) == 0;
 }
 
@@ -32,16 +33,17 @@ typedef enum {
   MOMENT_MODE_1 = 1   // Decoupled weight decay mode
 } adamMode_t;
 
-std::tuple<at::Tensor, at::Tensor>
-multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
+    int chunk_size, at::Tensor noop_flag,
     std::vector<std::vector<at::Tensor>> tensor_lists,
     at::optional<bool> per_tensor_python);
 
 using MATH_T = float;
 
-template <typename T> struct LAMBStage1Functor {
-  __device__ __forceinline__ void
-  operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<4> &tl,
+template <typename T>
+struct LAMBStage1Functor {
+  __device__ __forceinline__ void operator()(
+      int chunk_size, volatile int *noop_gmem, TensorListMetadata<4> &tl,
       const float beta1, const float beta2, const float beta3,
       const float beta1_correction, const float beta2_correction,
       const float epsilon, adamMode_t mode, const float decay,
@@ -89,8 +91,7 @@ template <typename T> struct LAMBStage1Functor {
            i_start += blockDim.x) {
         // load
         load_store(l_g, g, 0, i_start);
-        if (decay != 0)
-          load_store(l_p, p, 0, i_start);
+        if (decay != 0) load_store(l_p, p, 0, i_start);
         load_store(l_m, m, 0, i_start);
         load_store(l_v, v, 0, i_start);
         // unpack
@@ -204,12 +205,12 @@ template <typename T> struct LAMBStage1Functor {
 
 // Step 2 reads in 'update' value and per-tensor param_norm and update_norm.
 // It computes new parameter value.
-template <typename T> struct LAMBStage2Functor {
-  __device__ __forceinline__ void
-  operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<2> &tl,
-             const float *per_tensor_param_norm,
-             const float *per_tensor_update_norm, const float learning_rate,
-             const float decay, bool use_nvlamb) {
+template <typename T>
+struct LAMBStage2Functor {
+  __device__ __forceinline__ void operator()(
+      int chunk_size, volatile int *noop_gmem, TensorListMetadata<2> &tl,
+      const float *per_tensor_param_norm, const float *per_tensor_update_norm,
+      const float learning_rate, const float decay, bool use_nvlamb) {
     // I'd like this kernel to propagate infs/nans.
     // if(*noop_gmem == 1)
     //   return;
@@ -310,8 +311,7 @@ void multi_tensor_lamb_cuda(int chunk_size, at::Tensor noop_flag,
 
   // Handle grad averaging mode
   float beta3 = 1.0f;
-  if (grad_averaging == 1)
-    beta3 = 1 - beta1;
+  if (grad_averaging == 1) beta3 = 1 - beta1;
 
   std::vector<std::vector<at::Tensor>> grad_list(tensor_lists.begin(),
                                                  tensor_lists.begin() + 1);

colossalai/kernel/cuda_native/csrc/multi_tensor_scale_kernel.cu

@@ -15,7 +15,8 @@
 #define BLOCK_SIZE 512
 #define ILP 4
 
-template <typename T> __device__ __forceinline__ bool is_aligned(T *p) {
+template <typename T>
+__device__ __forceinline__ bool is_aligned(T *p) {
   return ((uint64_t)p) % (ILP * sizeof(T)) == 0;
 }
 
@@ -27,7 +28,8 @@ __device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset,
   ((LT *)dst)[dst_offset] = ((LT *)src)[src_offset];
 }
 
-template <typename in_t, typename out_t> struct ScaleFunctor {
+template <typename in_t, typename out_t>
+struct ScaleFunctor {
   __device__ __forceinline__ void operator()(int chunk_size,
                                              volatile int *noop_gmem,
                                              TensorListMetadata<2> &tl,
@@ -76,8 +78,7 @@ template <typename in_t, typename out_t> struct ScaleFunctor {
         for (int ii = 0; ii < ILP; ii++) {
           r_in[ii] = 0;
           int i = i_start + threadIdx.x + ii * blockDim.x;
-          if (i < n && i < chunk_size)
-            r_in[ii] = in[i];
+          if (i < n && i < chunk_size) r_in[ii] = in[i];
         }
         // note for clarification to future michael:
         // From a pure memory dependency perspective, there's likely no point
@@ -93,8 +94,7 @@ template <typename in_t, typename out_t> struct ScaleFunctor {
 #pragma unroll
         for (int ii = 0; ii < ILP; ii++) {
           int i = i_start + threadIdx.x + ii * blockDim.x;
-          if (i < n && i < chunk_size)
-            out[i] = r_out[ii];
+          if (i < n && i < chunk_size) out[i] = r_out[ii];
         }
       }
     }

colossalai/kernel/cuda_native/csrc/multi_tensor_sgd_kernel.cu

-// modified from https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_sgd_kernel.cu
+// modified from
+// https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_sgd_kernel.cu
 #include <ATen/ATen.h>
 #include <ATen/AccumulateType.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/Exceptions.h>
-#include "multi_tensor_apply.cuh"
-#include "compat.h"
-
 #include <assert.h>
 #include <cuda_runtime.h>
 
+#include "compat.h"
+#include "multi_tensor_apply.cuh"
+
 #define BLOCK_SIZE 512
 #define ILP 4
 
@@ -28,24 +29,13 @@
  * wd_after_momentum : apply weight decay _after_ momentum instead of before
  **/
 template <int N, typename T_grad, typename T_weight>
-struct SGDFunctor
-{
+struct SGDFunctor {
   __device__ __forceinline__ void operator()(
-        int chunk_size,
-        volatile int *noop_gmem,
-        TensorListMetadata<N> &tl,
-        float wd,
-        float momentum,
-        float dampening,
-        float lr,
-        bool nesterov,
-        bool first_run,
-        bool wd_after_momentum,
-        float scale)
-    {
+      int chunk_size, volatile int *noop_gmem, TensorListMetadata<N> &tl,
+      float wd, float momentum, float dampening, float lr, bool nesterov,
+      bool first_run, bool wd_after_momentum, float scale) {
     // Early exit if we don't need to do anything
-        if (*noop_gmem)
-            return;
+    if (*noop_gmem) return;
 
     int tensor_loc = tl.block_to_tensor[blockIdx.x];
     int chunk_idx = tl.block_to_chunk[blockIdx.x];
@@ -61,8 +51,7 @@ struct SGDFunctor
     mom_in += chunk_idx * chunk_size;
 
     at::Half *model_weights_out = nullptr;
-        if (N == 4)
-        {
+    if (N == 4) {
       model_weights_out = (at::Half *)tl.addresses[3][tensor_loc];
       model_weights_out += chunk_idx * chunk_size;
     }
@@ -73,19 +62,15 @@ struct SGDFunctor
     float incoming_grads[ILP];
     float incoming_weights[ILP];
     float incoming_moms[ILP];
-        for (int i_start = 0;
-             i_start < n && i_start < chunk_size;
-             i_start += blockDim.x * ILP)
-        {
+    for (int i_start = 0; i_start < n && i_start < chunk_size;
+         i_start += blockDim.x * ILP) {
 #pragma unroll
-            for (int ii = 0; ii < ILP; ii++)
-            {
+      for (int ii = 0; ii < ILP; ii++) {
         incoming_grads[ii] = 0;
         incoming_weights[ii] = 0;
         incoming_moms[ii] = 0;
         int i = i_start + threadIdx.x + ii * blockDim.x;
-                if (i < n && i < chunk_size)
-                {
+        if (i < n && i < chunk_size) {
           incoming_grads[ii] = static_cast<float>(grad_in[i]) * scale;
           incoming_weights[ii] = static_cast<float>(weight_in[i]);
           incoming_moms[ii] = static_cast<float>(mom_in[i]);
@@ -98,19 +83,17 @@ struct SGDFunctor
 // Put another way, the STGs are dependent on the LDGs, but not on each other.
 // There is still compute ILP benefit from unrolling the loop though.
 #pragma unroll
-            for (int ii = 0; ii < ILP; ii++)
-            {
+      for (int ii = 0; ii < ILP; ii++) {
         int i = i_start + threadIdx.x + ii * blockDim.x;
-                if (i < n && i < chunk_size)
-                {
+        if (i < n && i < chunk_size) {
           // apply weight decay before momentum if necessary
           if (wd != 0.f && !wd_after_momentum)
             incoming_grads[ii] += wd * incoming_weights[ii];
 
-                    if (momentum != 0.f)
-                    {
+          if (momentum != 0.f) {
             if (!first_run)
-                            incoming_moms[ii] = incoming_moms[ii] * momentum + (1.f - dampening) * incoming_grads[ii];
+              incoming_moms[ii] = incoming_moms[ii] * momentum +
+                                  (1.f - dampening) * incoming_grads[ii];
             else  // initialize momentums to current incoming grads
               incoming_moms[ii] = incoming_grads[ii];
 
@@ -132,27 +115,18 @@ struct SGDFunctor
             model_weights_out[i] = static_cast<at::Half>(weight_in[i]);
 
           // also write out the new momentum
-                    if (momentum != 0.f)
-                        mom_in[i] = incoming_moms[ii];
+          if (momentum != 0.f) mom_in[i] = incoming_moms[ii];
         }
       }
     }
   }
 };
 
-void multi_tensor_sgd_cuda(
-    int chunk_size,
-    at::Tensor noop_flag,
+void multi_tensor_sgd_cuda(int chunk_size, at::Tensor noop_flag,
                            std::vector<std::vector<at::Tensor>> tensor_lists,
-    float wd,
-    float momentum,
-    float dampening,
-    float lr,
-    bool nesterov,
-    bool first_run,
-    bool wd_after_momentum,
-    float scale)
-{
+                           float wd, float momentum, float dampening, float lr,
+                           bool nesterov, bool first_run,
+                           bool wd_after_momentum, float scale) {
   auto num_tensors = tensor_lists.size();
   auto grad_type = tensor_lists[0][0].scalar_type();
   auto weight_type = tensor_lists[1][0].scalar_type();
@@ -162,7 +136,8 @@ void multi_tensor_sgd_cuda(
       TORCH_CHECK(tensor_lists[3][i].scalar_type() == at::ScalarType::Half,
                   "Additional output tensors should always be fp16.");
 
-    TORCH_CHECK(noop_flag.device() == tensor_lists[0][0].device(), "expected noop flag to be on the same device as tensors");
+  TORCH_CHECK(noop_flag.device() == tensor_lists[0][0].device(),
+              "expected noop flag to be on the same device as tensors");
 
   // We have 3 possibilities to handle here, in terms of
   // grad_type, param_type, momentum_type, requires_fp16_copy
@@ -176,22 +151,10 @@ void multi_tensor_sgd_cuda(
 
   // Case 1. fp16, fp16, fp16, No
   if (grad_type == at::ScalarType::Half &&
-        weight_type == at::ScalarType::Half &&
-        num_tensors == 3)
-    {
-        multi_tensor_apply<3>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<3, at::Half, at::Half>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+      weight_type == at::ScalarType::Half && num_tensors == 3) {
+    multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<3, at::Half, at::Half>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
   }
   // Case 2. fp16, fp32, fp32, No
@@ -214,68 +177,33 @@ void multi_tensor_sgd_cuda(
   // }
   // Case 2. fp32, fp32, fp32, No
   else if (grad_type == at::ScalarType::Float &&
-             weight_type == at::ScalarType::Float &&
-             num_tensors == 3)
-    {
-        multi_tensor_apply<3>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<3, float, float>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+           weight_type == at::ScalarType::Float && num_tensors == 3) {
+    multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<3, float, float>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
   }
   // Case 3. fp16, fp32, fp32, Yes
   else if (grad_type == at::ScalarType::Half &&
-             weight_type == at::ScalarType::Float &&
-             num_tensors == 4)
-    {
-        multi_tensor_apply<4>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<4, at::Half, float>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+           weight_type == at::ScalarType::Float && num_tensors == 4) {
+    multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<4, at::Half, float>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
   }
   // Case 4. fp32, fp32, fp32, Yes
   else if (grad_type == at::ScalarType::Float &&
-             weight_type == at::ScalarType::Float &&
-             num_tensors == 4)
-    {
-        multi_tensor_apply<4>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<4, float, float>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+           weight_type == at::ScalarType::Float && num_tensors == 4) {
+    multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<4, float, float>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
-    }
-    else
-    {
-        AT_ERROR("multi_tensor_sgd only supports some combinations of gradient & weight types. Given: ",
-                 "gradient: ", grad_type, ", weight: ", weight_type, ", num_lists: ", num_tensors);
+  } else {
+    AT_ERROR(
+        "multi_tensor_sgd only supports some combinations of gradient & weight "
+        "types. Given: ",
+        "gradient: ", grad_type, ", weight: ", weight_type,
+        ", num_lists: ", num_tensors);
   }
 
   AT_CUDA_CHECK(cudaGetLastError());

colossalai/kernel/cuda_native/csrc/multihead_attention_1d.h

@@ -19,21 +19,25 @@
 template <typename T>
 class MultiHeadAttention {
  public:
-  MultiHeadAttention(int layer_id, int max_batch_tokens, int _max_seq_len, int hidden_size,
-                     int num_heads, float attn_dropout_ratio, float hidden_output_dropout_ratio,
+  MultiHeadAttention(int layer_id, int max_batch_tokens, int _max_seq_len,
+                     int hidden_size, int num_heads, float attn_dropout_ratio,
+                     float hidden_output_dropout_ratio,
                      bool pre_or_postLayerNorm);
 
   virtual ~MultiHeadAttention();
 
   void Forward(const T *input_ptr, const T *input_mask_ptr, T *out_ptr);
 
-  void Backward(const T *grad_output_ptr, const T *input_ptr, const T *output_ptr,
-                const T *input_mask_ptr, T *grad_input_ptr);
+  void Backward(const T *grad_output_ptr, const T *input_ptr,
+                const T *output_ptr, const T *input_mask_ptr,
+                T *grad_input_ptr);
 
-  void attn_layer_fw(const T *input_ptr, const T *input_mask_ptr, T *output_ptr, T *buffer);
+  void attn_layer_fw(const T *input_ptr, const T *input_mask_ptr, T *output_ptr,
+                     T *buffer);
 
-  void attn_layer_bw(const T *input_ptr, const T *input_mask_ptr, const T *output_ptr,
-                     const T *grad_output_ptr, T *grad_input_attn_layer_bwptr, T *buffer);
+  void attn_layer_bw(const T *input_ptr, const T *input_mask_ptr,
+                     const T *output_ptr, const T *grad_output_ptr,
+                     T *grad_input_attn_layer_bwptr, T *buffer);
 
   void set_cur_batch_shape(int batch_size, int seq_len) {
     _batch_size = batch_size;
@@ -83,12 +87,15 @@ class MultiHeadAttention {
     }
 
     _qkv_ptr = cuda_malloc<T>(_max_batch_tokens * _hidden_size * 3);
-    _soft_out_ptr = cuda_malloc<T>(_max_batch_tokens * _heads / pg_size * _max_seq_len);
-    _ctx_bufB_ptr = cuda_malloc<T>(_max_batch_tokens * _heads / pg_size * _max_seq_len);
+    _soft_out_ptr =
+        cuda_malloc<T>(_max_batch_tokens * _heads / pg_size * _max_seq_len);
+    _ctx_bufB_ptr =
+        cuda_malloc<T>(_max_batch_tokens * _heads / pg_size * _max_seq_len);
     _attn_o_inp_ptr = cuda_malloc<T>(_max_batch_tokens * _hidden_size);
 
     // buffer size needed by attn bw
-    size_t smem_size = 4 * _max_batch_tokens * _hidden_size / pg_size +
+    size_t smem_size =
+        4 * _max_batch_tokens * _hidden_size / pg_size +
         std::max(3 * _max_batch_tokens * _hidden_size / pg_size,
                  _max_batch_tokens * _heads / pg_size * _max_seq_len);
 

colossalai/kernel/cuda_native/csrc/scaled_masked_softmax_cuda.cu

@@ -2,12 +2,13 @@
  *     with minor changes. */
 
 #include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <cuda.h>
-#include <cuda_runtime.h>
 #include <cuda_fp16.h>
 #include <cuda_profiler_api.h>
-#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
 #include <torch/extension.h>
+
 #include "scaled_masked_softmax.h"
 #include "type_shim.h"
 
@@ -15,17 +16,15 @@ namespace multihead_attn {
 namespace fused_softmax {
 namespace scaled_masked_softmax {
 
-int get_batch_per_block_cuda(int query_seq_len, int key_seq_len, int batches, int attn_heads){
+int get_batch_per_block_cuda(int query_seq_len, int key_seq_len, int batches,
+                             int attn_heads) {
   return get_batch_per_block(query_seq_len, key_seq_len, batches, attn_heads);
 }
 
-
-torch::Tensor fwd_cuda(
-    torch::Tensor const& input,
-    torch::Tensor const& mask,
-    float scale_factor)
-{
-  // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]
+torch::Tensor fwd_cuda(torch::Tensor const& input, torch::Tensor const& mask,
+                       float scale_factor) {
+  // input is a 4d tensor with dimensions [batches, attn_heads, seq_len,
+  // seq_len]
   const int batches = input.size(0);
   const int pad_batches = mask.size(0);
   const int attn_heads = input.size(1);
@@ -40,8 +39,8 @@ torch::Tensor fwd_cuda(
 
   // Output
   auto act_options = input.options().requires_grad(false);
-  torch::Tensor softmax_results = 
-      torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options);
+  torch::Tensor softmax_results = torch::empty(
+      {batches, attn_heads, query_seq_len, key_seq_len}, act_options);
 
   // Softmax Intermediate Result Ptr
   void* input_ptr = static_cast<void*>(input.data_ptr());
@@ -49,31 +48,23 @@ torch::Tensor fwd_cuda(
   void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
 
   DISPATCH_HALF_AND_BFLOAT(
-      input.scalar_type(),
-      "dispatch_scaled_masked_softmax_forward",
+      input.scalar_type(), "dispatch_scaled_masked_softmax_forward",
       dispatch_scaled_masked_softmax_forward<scalar_t, scalar_t, float>(
           reinterpret_cast<scalar_t*>(softmax_results_ptr),
           reinterpret_cast<const scalar_t*>(input_ptr),
-	  reinterpret_cast<const uint8_t*>(mask_ptr),
-	  scale_factor,
-	  query_seq_len,
-	  key_seq_len,
-	  batches,
-	  attn_heads,
-	  pad_batches);
-      );
+          reinterpret_cast<const uint8_t*>(mask_ptr), scale_factor,
+          query_seq_len, key_seq_len, batches, attn_heads, pad_batches););
   return softmax_results;
 }
 
-torch::Tensor bwd_cuda(
-    torch::Tensor const& output_grads_, 
+torch::Tensor bwd_cuda(torch::Tensor const& output_grads_,
                        torch::Tensor const& softmax_results_,
                        float scale_factor) {
-	
   auto output_grads = output_grads_.contiguous();
   auto softmax_results = softmax_results_.contiguous();
 
-  //output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]
+  // output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len,
+  // seq_len]
   const int batches = output_grads.size(0);
   const int attn_heads = output_grads.size(1);
   const int query_seq_len = output_grads.size(2);
@@ -81,24 +72,18 @@ torch::Tensor bwd_cuda(
 
   void* output_grads_ptr = static_cast<void*>(output_grads.data_ptr());
 
-  //Softmax Grad
+  // Softmax Grad
   DISPATCH_HALF_AND_BFLOAT(
-      output_grads_.scalar_type(),
-      "dispatch_scaled_masked_softmax_backward",
+      output_grads_.scalar_type(), "dispatch_scaled_masked_softmax_backward",
       dispatch_scaled_masked_softmax_backward<scalar_t, scalar_t, float>(
           reinterpret_cast<scalar_t*>(output_grads_ptr),
           reinterpret_cast<scalar_t*>(output_grads_ptr),
           reinterpret_cast<scalar_t const*>(softmax_results.data_ptr()),
-	  scale_factor,
-	  query_seq_len,
-	  key_seq_len,
-	  batches,
-	  attn_heads);
-			   );
+          scale_factor, query_seq_len, key_seq_len, batches, attn_heads););
 
-  //backward pass is completely in-place
+  // backward pass is completely in-place
   return output_grads;
 }
-}
-}
-}
+}  // namespace scaled_masked_softmax
+}  // namespace fused_softmax
+}  // namespace multihead_attn

colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax.cpp

@@ -3,18 +3,16 @@
 
 #include <cuda_fp16.h>
 #include <torch/extension.h>
+
 #include <vector>
 
 namespace multihead_attn {
 namespace fused_softmax {
 namespace scaled_upper_triang_masked_softmax {
 
-torch::Tensor fwd_cuda(
-    torch::Tensor const& input, 
-    float scale_factor);
+torch::Tensor fwd_cuda(torch::Tensor const& input, float scale_factor);
 
-torch::Tensor bwd_cuda(
-    torch::Tensor const& output_grads, 
+torch::Tensor bwd_cuda(torch::Tensor const& output_grads,
                        torch::Tensor const& softmax_results,
                        float scale_factor);
 
@@ -27,11 +25,8 @@ torch::Tensor fwd(torch::Tensor const& input, float scale_factor) {
   return fwd_cuda(input, scale_factor);
 }
 
-torch::Tensor bwd(
-    torch::Tensor const& output_grads, 
-    torch::Tensor const& softmax_results,
-    float scale_factor) {
-
+torch::Tensor bwd(torch::Tensor const& output_grads,
+                  torch::Tensor const& softmax_results, float scale_factor) {
   AT_ASSERTM(output_grads.dim() == 3, "expected 3D tensor");
   AT_ASSERTM(softmax_results.dim() == 3, "expected 3D tensor");
 

colossalai/kernel/cuda_native/csrc/scaled_upper_triang_masked_softmax_cuda.cu

@@ -2,12 +2,13 @@
  *     with minor changes. */
 
 #include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <cuda.h>
-#include <cuda_runtime.h>
 #include <cuda_fp16.h>
 #include <cuda_profiler_api.h>
-#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
 #include <torch/extension.h>
+
 #include "scaled_upper_triang_masked_softmax.h"
 #include "type_shim.h"
 
@@ -15,10 +16,7 @@ namespace multihead_attn {
 namespace fused_softmax {
 namespace scaled_upper_triang_masked_softmax {
 
-torch::Tensor fwd_cuda(
-    torch::Tensor const& input, 
-    float scale_factor)
-{
+torch::Tensor fwd_cuda(torch::Tensor const& input, float scale_factor) {
   // input is a 3d tensor with dimensions [attn_batches, seq_len, seq_len]
   const int attn_batches = input.size(0);
   const int seq_len = input.size(1);
@@ -36,50 +34,42 @@ torch::Tensor fwd_cuda(
   DISPATCH_HALF_AND_BFLOAT(
       input.scalar_type(),
       "dispatch_scaled_upper_triang_masked_softmax_forward",
-      dispatch_scaled_upper_triang_masked_softmax_forward<scalar_t, scalar_t, float>(
+      dispatch_scaled_upper_triang_masked_softmax_forward<scalar_t, scalar_t,
+                                                          float>(
           reinterpret_cast<scalar_t*>(softmax_results_ptr),
-	  reinterpret_cast<const scalar_t*>(input_ptr),
-	  scale_factor,
-	  seq_len,
-	  seq_len,
-	  attn_batches);
-      );
+          reinterpret_cast<const scalar_t*>(input_ptr), scale_factor, seq_len,
+          seq_len, attn_batches););
   return softmax_results;
 }
 
-
-torch::Tensor bwd_cuda(
-    torch::Tensor const& output_grads_, 
+torch::Tensor bwd_cuda(torch::Tensor const& output_grads_,
                        torch::Tensor const& softmax_results_,
                        float scale_factor) {
-	
   auto output_grads = output_grads_.contiguous();
   auto softmax_results = softmax_results_.contiguous();
 
-  //output grads is a 3d tensor with dimensions [attn_batches, seq_len, seq_len]
+  // output grads is a 3d tensor with dimensions [attn_batches, seq_len,
+  // seq_len]
   const int attn_batches = output_grads.size(0);
   const int seq_len = output_grads.size(1);
   TORCH_INTERNAL_ASSERT(output_grads.size(1) == output_grads.size(2));
 
   void* output_grads_ptr = static_cast<void*>(output_grads.data_ptr());
 
-  //Softmax Grad
+  // Softmax Grad
   DISPATCH_HALF_AND_BFLOAT(
       output_grads_.scalar_type(),
       "dispatch_scaled_upper_triang_masked_softmax_backward",
-      dispatch_scaled_upper_triang_masked_softmax_backward<scalar_t, scalar_t, float>(
+      dispatch_scaled_upper_triang_masked_softmax_backward<scalar_t, scalar_t,
+                                                           float>(
           reinterpret_cast<scalar_t*>(output_grads_ptr),
           reinterpret_cast<scalar_t*>(output_grads_ptr),
           reinterpret_cast<scalar_t const*>(softmax_results.data_ptr()),
-	  scale_factor,
-	  seq_len,
-	  seq_len,
-	  attn_batches);
-      );
+          scale_factor, seq_len, seq_len, attn_batches););
 
-  //backward pass is completely in-place
+  // backward pass is completely in-place
   return output_grads;
 }
-}
-}
-}
+}  // namespace scaled_upper_triang_masked_softmax
+}  // namespace fused_softmax
+}  // namespace multihead_attn

colossalai/kernel/cuda_native/layer_norm.py

@@ -24,8 +24,8 @@ class FusedLayerNormAffineFunction(torch.autograd.Function):
         input_ = input.contiguous()
         weight_ = weight.contiguous()
         bias_ = bias.contiguous()
-        output, mean, invvar = colossal_layer_norm_cuda.forward_affine(
-            input_, ctx.normalized_shape, weight_, bias_, ctx.eps)
+        output, mean, invvar = colossal_layer_norm_cuda.forward_affine(input_, ctx.normalized_shape, weight_, bias_,
+                                                                       ctx.eps)
         ctx.save_for_backward(input_, weight_, bias_, mean, invvar)
 
         return output
@@ -72,8 +72,7 @@ class MixedFusedLayerNorm(torch.nn.Module):
 
     def forward(self, input):
 
-        return FusedLayerNormAffineFunction.apply(input, self.weight, self.bias,
-                                                  self.normalized_shape, self.eps)
+        return FusedLayerNormAffineFunction.apply(input, self.weight, self.bias, self.normalized_shape, self.eps)
 
     def __repr__(self):
         return f'MixedFusedLayerNorm(normalized_shape={self.normalized_shape}, eps={self.eps})'

colossalai/kernel/cuda_native/scaled_softmax.py

@@ -28,9 +28,7 @@ class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
             raise RuntimeError('ScaledUpperTriangMaskedSoftmax requires cuda extensions')
 
         scale_t = torch.tensor([scale])
-        softmax_results = colossal_scaled_upper_triang_masked_softmax.forward(
-            inputs, scale_t[0]
-        )
+        softmax_results = colossal_scaled_upper_triang_masked_softmax.forward(inputs, scale_t[0])
 
         ctx.save_for_backward(softmax_results, scale_t)
         return softmax_results
@@ -43,9 +41,7 @@ class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
             raise RuntimeError('ScaledUpperTriangMaskedSoftmax requires cuda extensions')
 
         softmax_results, scale_t = ctx.saved_tensors
-        input_grads = colossal_scaled_upper_triang_masked_softmax.backward(
-            output_grads, softmax_results, scale_t[0]
-        )
+        input_grads = colossal_scaled_upper_triang_masked_softmax.backward(output_grads, softmax_results, scale_t[0])
 
         return input_grads, None
 
@@ -81,9 +77,7 @@ class ScaledMaskedSoftmax(torch.autograd.Function):
 
         softmax_results, scale_t = ctx.saved_tensors
 
-        input_grads = colossal_scaled_masked_softmax.backward(
-            output_grads, softmax_results, scale_t[0]
-        )
+        input_grads = colossal_scaled_masked_softmax.backward(output_grads, softmax_results, scale_t[0])
         return input_grads, None, None
 
 
@@ -114,9 +108,8 @@ class FusedScaleMaskSoftmax(nn.Module):
         super(FusedScaleMaskSoftmax, self).__init__()
         self.input_in_fp16 = input_in_fp16
         self.input_in_bf16 = input_in_bf16
-        assert not (
-            self.input_in_fp16 and self.input_in_bf16
-        ), "both fp16 and bf16 flags cannot be active at the same time."
+        assert not (self.input_in_fp16
+                    and self.input_in_bf16), "both fp16 and bf16 flags cannot be active at the same time."
         self.input_in_float16 = self.input_in_fp16 or self.input_in_bf16
         self.attn_mask_type = attn_mask_type
         self.scaled_masked_softmax_fusion = scaled_masked_softmax_fusion
@@ -124,9 +117,7 @@ class FusedScaleMaskSoftmax(nn.Module):
         self.softmax_in_fp32 = softmax_in_fp32
         self.scale = scale
 
-        assert (
-            self.scale is None or softmax_in_fp32
-        ), "softmax should be in fp32 when scaled"
+        assert (self.scale is None or softmax_in_fp32), "softmax should be in fp32 when scaled"
 
     def forward(self, input, mask):
         # [b, np, sq, sk]
@@ -140,8 +131,7 @@ class FusedScaleMaskSoftmax(nn.Module):
     def is_kernel_available(self, mask, b, np, sq, sk):
         attn_batches = b * np
 
-        if (
-            self.scaled_masked_softmax_fusion  # user want to fuse
+        if (self.scaled_masked_softmax_fusion    # user want to fuse
                 and self.input_in_float16    # input must be fp16
                 and mask is not None    # mask tensor must not be None
                 and 16 < sk <= 2048    # sk must be 16 ~ 2048

colossalai/kernel/jit/bias_gelu.py

 import torch
 
-
 ###### BIAS GELU FUSION/ NO AUTOGRAD ################
 # 1/sqrt(2*pi)-> 0.3989423
 # 1/sqrt(2)   -> 0.70710678
@@ -9,11 +8,13 @@ import torch
 # actual gelu is:
 # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
 
+
 @torch.jit.script
 def bias_gelu(bias, y):
     x = bias + y
     return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
 
+
 # gradient of tanh approximation of gelu
 # gradient of actual gelu is:
 # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
@@ -23,9 +24,11 @@ def bias_gelu_back(g, bias, y):
     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
     ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
-    return ff*g
+    return ff * g
+
 
 class GeLUFunction(torch.autograd.Function):
+
     @staticmethod
     # bias is an optional argument
     def forward(ctx, input, bias):
@@ -38,4 +41,5 @@ class GeLUFunction(torch.autograd.Function):
         tmp = bias_gelu_back(grad_output, bias, input)
         return tmp, tmp
 
+
 bias_gelu_impl = GeLUFunction.apply

colossalai/nn/layer/parallel_2d/layers.py

@@ -182,7 +182,7 @@ class Linear2D(ParallelLayer):
     def forward(self, x: Tensor) -> Tensor:
         # input: [m/q, n/q, k/q]
         # output: [m/q, n/q, h/q]
-        out_shape = x.shape[:-1] + (self.hidden_size_per_partition, )
+        out_shape = x.shape[:-1] + (self.hidden_size_per_partition,)
 
         output = Matmul_AB_2D.apply(x, self.weight, self.summa_dim, out_shape, self.row_rank, self.col_rank,
                                     ParallelMode.PARALLEL_2D_ROW, ParallelMode.PARALLEL_2D_COL, self.data_parallel_rank,
@@ -1012,7 +1012,7 @@ class Classifier2D(ParallelLayer):
             destination.update(local_state)
 
     def forward(self, input_: Tensor) -> Tensor:
-        out_shape = input_.shape[:-1] + (self.num_classes, )
+        out_shape = input_.shape[:-1] + (self.num_classes,)
 
         return classifier_2d(input_, self.weight, self.bias, self.summa_dim, out_shape, self.row_rank, self.col_rank,
                              ParallelMode.PARALLEL_2D_ROW, ParallelMode.PARALLEL_2D_COL, self.data_parallel_rank,
@@ -1186,7 +1186,7 @@ class VocabParallelClassifier2D(ParallelLayer):
     def forward(self, x: Tensor) -> Tensor:
         # input: [m/q, n/q, k/q]
         # output: [m/q, n/q, h/q]
-        out_shape = x.shape[:-1] + (self.output_size_per_partition, )
+        out_shape = x.shape[:-1] + (self.output_size_per_partition,)
 
         output = Matmul_ABT_2D.apply(x, self.weight, self.summa_dim, out_shape, self.row_rank, self.col_rank,
                                      ParallelMode.PARALLEL_2D_ROW, ParallelMode.PARALLEL_2D_COL,

colossalai/nn/layer/parallel_2p5d/layers.py

@@ -189,7 +189,7 @@ class Linear2p5D(ParallelLayer):
     def forward(self, x: Tensor) -> Tensor:
         # input: [m/dq, n/q, k/q]
         # output: [m/dq, n/q, h/q]
-        out_shape = x.shape[:-1] + (self.hidden_size_per_partition, )
+        out_shape = x.shape[:-1] + (self.hidden_size_per_partition,)
 
         output = Matmul_AB_2p5D.apply(
             x,
@@ -1038,7 +1038,7 @@ class Classifier2p5D(ParallelLayer):
             destination.update(local_state)
 
     def forward(self, input_: Tensor) -> Tensor:
-        out_shape = input_.shape[:-1] + (self.num_classes, )
+        out_shape = input_.shape[:-1] + (self.num_classes,)
 
         return classifier_2p5d(input_, self.weight, self.bias, self.tesseract_dim, out_shape, self.row_rank,
                                self.col_rank, ParallelMode.PARALLEL_2P5D_ROW, ParallelMode.PARALLEL_2P5D_COL,
@@ -1172,7 +1172,7 @@ class VocabParallelClassifier2p5D(ParallelLayer):
     def forward(self, x: Tensor) -> Tensor:
         # input: [m/dq, n/q, k/q]
         # output: [m/dq, n/q, h/q]
-        out_shape = x.shape[:-1] + (self.hidden_size_per_partition, )
+        out_shape = x.shape[:-1] + (self.hidden_size_per_partition,)
 
         output = Matmul_ABT_2p5D.apply(
             x,

colossalai/nn/layer/parallel_3d/layers.py

@@ -53,8 +53,8 @@ class LayerNorm3D(ParallelLayer):
         self.weight = Parameter(
             torch.ones(self.normalized_shape_per_partition, device=get_current_device(), dtype=dtype))
         if bias:
-            self.bias = Parameter(torch.zeros(self.normalized_shape_per_partition,
-                                              device=get_current_device(), dtype=dtype))
+            self.bias = Parameter(
+                torch.zeros(self.normalized_shape_per_partition, device=get_current_device(), dtype=dtype))
         else:
             self.bias = None
         self.variance_epsilon = eps

colossalai/nn/layer/utils/common.py

@@ -13,7 +13,8 @@ from torch import Tensor, nn
 
 
 class CheckpointModule(nn.Module):
-    def __init__(self, checkpoint: bool = True, offload : bool = False):
+
+    def __init__(self, checkpoint: bool = True, offload: bool = False):
         super().__init__()
         self.checkpoint = checkpoint
         self._use_checkpoint = checkpoint
@@ -78,6 +79,7 @@ def get_tensor_parallel_mode():
 
 
 def _ntuple(n):
+
     def parse(x):
         if isinstance(x, collections.abc.Iterable):
             return x