From 321063284ff9462b77f8f56a81383574d4db46e7 Mon Sep 17 00:00:00 2001
From: ycr0776 <ycr0776@zhihu.com>
Date: Wed, 2 Aug 2023 09:40:49 +0800
Subject: [PATCH 01/27] modified:   bmtrain/optim/adam.py

---
 bmtrain/optim/adam.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/bmtrain/optim/adam.py b/bmtrain/optim/adam.py
index b63a4f51..c39cec9d 100644
--- a/bmtrain/optim/adam.py
+++ b/bmtrain/optim/adam.py
@@ -88,7 +88,7 @@ def step(self, closure=None, scale=1):
                         grad = p.grad
                         
                     if p.dtype == torch.half:
-                        F.adam(
+                        C.f_adam(
                             state["_param_fp32"],    # fp32
                             p,                      # fp16
                             grad,                 # fp16
@@ -101,6 +101,19 @@ def step(self, closure=None, scale=1):
                             group['weight_decay'],
                             state['step']
                         )
+                    elif p.dtype == torch.bfloat16:
+                        C.f_adam_bf16(
+                            state["_param_fp32"],    # fp32
+                            p,                      # bf16
+                            grad,                 # bf16
+                            state['exp_avg'],       # fp32: m
+                            group['betas'][0], group['betas'][1],
+                            group['eps'],
+                            0.0 if state["step"] <= self._hold_steps else group['lr'],
+                            scale,
+                            group['weight_decay'],
+                            state['step']
+                        )
                     else:
                         other_kwargs = {}
                         if 'maximize' in inspect.signature(torch.optim._functional.adam).parameters:

From 4e4d32a45dfe4f723f4353b0e114e731a6f61bd7 Mon Sep 17 00:00:00 2001
From: ycr0776 <ycr0776@zhihu.com>
Date: Wed, 2 Aug 2023 12:32:59 +0800
Subject: [PATCH 02/27] modified:   add bf16 in adam.py

---
 bmtrain/optim/adam.py | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/bmtrain/optim/adam.py b/bmtrain/optim/adam.py
index c39cec9d..6630ffcf 100644
--- a/bmtrain/optim/adam.py
+++ b/bmtrain/optim/adam.py
@@ -40,8 +40,12 @@ def _on_justify_scale(self, old_scale, new_scale):
                 if p in self.state:
                     state = self.state[p]
                     if len(state) > 0:
-                        state['exp_avg'] *= delta
-                        state['exp_avg_sq'] *= delta
+                        #if p belongs to bf16, do not justify scale  
+                        if p.dtype == torch.bfloat16:
+                            continue
+                        else:                 
+                            state['exp_avg'] *= delta
+                            state['exp_avg_sq'] *= delta
 
     @torch.no_grad()
     def step(self, closure=None, scale=1):
@@ -63,19 +67,22 @@ def step(self, closure=None, scale=1):
                 if p.grad is not None and p.requires_grad:
                     if p.grad.is_sparse:
                         raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
-                    if p.dtype not in [torch.float16, torch.float32]:
-                        raise RuntimeError('Adam only supports fp32 or fp16 gradients')
+                    if p.dtype not in [torch.float32, torch.half, torch.bfloat16]:
+                        raise RuntimeError('Adam only supports fp32, fp16 and bf16 gradients')
 
                     state = self.state[p]
                     # Lazy state initialization
                     if len(state) == 0:
                         state['step'] = 0
                         # Exponential moving average of gradient values
-                        state['exp_avg'] = torch.zeros(p.size(), dtype=p.dtype, device=p.device) # on device
+                        if p.dtype == torch.bfloat16:
+                            state['exp_avg'] = torch.zeros(p.size(), dtype=torch.float32, device=p.device) # on device
+                        else:
+                            state['exp_avg'] = torch.zeros(p.size(), dtype=p.dtype, device=p.device) # on device
                         # Exponential moving average of squared gradient values
-                        state['exp_avg_sq'] = torch.zeros(p.size(), dtype=torch.float32, device=p.device)   # on device
-
-                        if p.dtype == torch.half:
+                        state['exp_avg_sq'] = torch.zeros(p.size(), dtype=torch.float32, device=p.device)# on device
+                  
+                        if p.dtype == torch.half or p.dtype == torch.bfloat16:
                             state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, device=p.device)   # on device
                             state['_param_fp32'].copy_(p)
 
@@ -107,6 +114,7 @@ def step(self, closure=None, scale=1):
                             p,                      # bf16
                             grad,                 # bf16
                             state['exp_avg'],       # fp32: m
+                            state["exp_avg_sq"],    # fp32: v
                             group['betas'][0], group['betas'][1],
                             group['eps'],
                             0.0 if state["step"] <= self._hold_steps else group['lr'],
@@ -172,11 +180,11 @@ def load_state_dict(self, state_dict: dict) -> None:
             if k in id_map:
                 param = id_map[k]
 
-                if param.dtype == torch.half and "_param_fp32" not in v:
+                if param.dtype != torch.float32 and "_param_fp32" not in v:
                     v["_param_fp32"] = torch.empty(param.size(), dtype=torch.float32, device=param.device)
                     v["_param_fp32"].copy_(param)
 
-                for name, dtype in [("exp_avg", param.dtype), ("exp_avg_sq", torch.float32), ("_param_fp32", torch.float32)]:
+                for name, dtype in [("exp_avg", torch.float32 if param.dtype == torch.bfloat16 else param.dtype), ("exp_avg_sq", torch.float32), ("_param_fp32", torch.float32)]:
                     if name in v:
                         v[name] = v[name].to(param.device).to(dtype)
 

From ef038ae35a83d3acc9582996870c69c9b708591b Mon Sep 17 00:00:00 2001
From: ycr0776 <ycr0776@zhihu.com>
Date: Wed, 2 Aug 2023 15:55:33 +0800
Subject: [PATCH 03/27] 	modified:   bmtrain/optim/_function.py 	modified:  
 bmtrain/optim/adam.py 	modified:   bmtrain/optim/adam_offload.py 
 modified:   csrc/cuda/adam_cuda.cu 	modified:   csrc/include/adam_cpu.hpp

---
 bmtrain/optim/_function.py    | 37 +++++++++++++++++++++++++++++-
 bmtrain/optim/adam.py         |  6 ++---
 bmtrain/optim/adam_offload.py | 42 ++++++++++++++++++++++++++++++-----
 csrc/cuda/adam_cuda.cu        |  6 ++++-
 csrc/include/adam_cpu.hpp     | 19 ++++++++++++++++
 5 files changed, 99 insertions(+), 11 deletions(-)

diff --git a/bmtrain/optim/_function.py b/bmtrain/optim/_function.py
index ee4b04a7..9777b589 100644
--- a/bmtrain/optim/_function.py
+++ b/bmtrain/optim/_function.py
@@ -1,4 +1,3 @@
-
 from .. import C
 import torch
 CHECK_INPUT = lambda x: x.is_contiguous() and x.is_cuda
@@ -76,3 +75,39 @@ def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tenso
         bias_correction2,
         stream
     )
+
+def adam_bf16(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tensor, m_fp16: torch.Tensor,
+             v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
+             weight_decay: float, step: int) -> None:
+    assert CHECK_INPUT(param_fp32), "param_fp32 must be contiguous and on cuda"
+    assert CHECK_INPUT(param_fp16), "param_fp16 must be contiguous and on cuda"
+    assert CHECK_INPUT(g_fp16), "g_fp16 must be contiguous and on cuda"
+    assert CHECK_INPUT(m_fp16), "m_fp32 must be contiguous and on cuda"
+    assert CHECK_INPUT(v_fp32), "v_fp32 must be contiguous and on cuda"
+    assert param_fp32.dtype == torch.float32, "param_fp32 must be float32 tensor"
+    assert param_fp16.dtype == torch.float16, "param_fp16 must be float16 tensor"
+    assert g_fp16.dtype == torch.float16, "g_fp16 must be float16 tensor"
+    assert m_fp16.dtype == torch.float16, "m_fp16 must be float16 tensor"
+    assert v_fp32.dtype == torch.float32, "v_fp32 must be float32 tensor"
+    assert param_fp32.numel() == param_fp16.numel(), "param_fp32 and param_fp16 must have the same number of elements"
+    assert param_fp32.numel() == g_fp16.numel(), "param_fp32 and g_fp16 must have the same number of elements"
+    assert param_fp32.numel() == m_fp16.numel(), "param_fp32 and m_fp32 must have the same number of elements"
+    assert param_fp32.numel() == v_fp32.numel(), "param_fp32 and v_fp32 must have the same number of elements"
+    bias_correction1 = 1 - beta1 ** step
+    bias_correction2 = 1 - beta2 ** step
+    stream = torch.cuda.current_stream().cuda_stream
+    C.adam_launcher(
+        param_fp32.numel(),
+        param_fp32.data_ptr(),
+        param_fp16.data_ptr(),
+        g_fp16.data_ptr(),
+        m_fp16.data_ptr(),
+        v_fp32.data_ptr(),
+        beta1, beta2,
+        eps, lr,
+        scale,
+        weight_decay,
+        bias_correction1,
+        bias_correction2,
+        stream
+    )
\ No newline at end of file
diff --git a/bmtrain/optim/adam.py b/bmtrain/optim/adam.py
index 6630ffcf..c9482f6a 100644
--- a/bmtrain/optim/adam.py
+++ b/bmtrain/optim/adam.py
@@ -82,7 +82,7 @@ def step(self, closure=None, scale=1):
                         # Exponential moving average of squared gradient values
                         state['exp_avg_sq'] = torch.zeros(p.size(), dtype=torch.float32, device=p.device)# on device
                   
-                        if p.dtype == torch.half or p.dtype == torch.bfloat16:
+                        if p.dtype != torch.float32:
                             state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, device=p.device)   # on device
                             state['_param_fp32'].copy_(p)
 
@@ -95,7 +95,7 @@ def step(self, closure=None, scale=1):
                         grad = p.grad
                         
                     if p.dtype == torch.half:
-                        C.f_adam(
+                            F.adam(
                             state["_param_fp32"],    # fp32
                             p,                      # fp16
                             grad,                 # fp16
@@ -109,7 +109,7 @@ def step(self, closure=None, scale=1):
                             state['step']
                         )
                     elif p.dtype == torch.bfloat16:
-                        C.f_adam_bf16(
+                            F.adam_bf16(
                             state["_param_fp32"],    # fp32
                             p,                      # bf16
                             grad,                 # bf16
diff --git a/bmtrain/optim/adam_offload.py b/bmtrain/optim/adam_offload.py
index e33219bf..5e73f918 100644
--- a/bmtrain/optim/adam_offload.py
+++ b/bmtrain/optim/adam_offload.py
@@ -54,8 +54,8 @@ def step(self, closure=None, scale=1):
                 if p.grad is not None and p.requires_grad:
                     if p.grad.is_sparse:
                         raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
-                    if p.dtype not in [torch.float16, torch.float32]:
-                        raise RuntimeError('Adam only supports fp32 or fp16 gradients')
+                    if p.dtype not in [torch.float16, torch.float32, torch.bfloat16]:
+                        raise RuntimeError('Adam only supports fp32, fp16 and bf16 gradients')
 
                     state = self.state[p]
                     # Lazy state initialization
@@ -73,6 +73,15 @@ def step(self, closure=None, scale=1):
                             # placeholder
                             state["_param_fp16"] = torch.empty(p.size(), dtype=torch.float16, pin_memory=True)  # on host
                             state["_grad_fp16"] = torch.empty(p.size(), dtype=torch.float16, pin_memory=True)   # on host
+                            
+                        elif p.dtype == torch.bfloat16:
+                            state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, device="cpu")     # on host
+                            state['_param_fp32'].copy_(p)
+
+                            # placeholder
+                            state["_param_bf16"] = torch.empty(p.size(), dtype=torch.bfloat16, pin_memory=True)  # on host
+                            state["_grad_bf16"] = torch.empty(p.size(), dtype=torch.bfloat16, pin_memory=True)   # on host
+                                                        
                         else:
                             state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, pin_memory=True)     # on host
                             state['_param_fp32'].copy_(p)
@@ -89,6 +98,8 @@ def step(self, closure=None, scale=1):
         for param, state, event, _, _, _, _, _ in update_params:
             if param.dtype == torch.half:
                 state["_grad_fp16"].copy_(param.grad, non_blocking=True)
+            elif param.dtype == torch.bfloat16:
+                state ["_grad_bf16"].copy_(param.grad, non_blocking=True)
             else:
                 state["_grad_fp32"].copy_(param.grad, non_blocking=True)
             torch.cuda.current_stream().record_event(event)
@@ -119,6 +130,23 @@ def step(self, closure=None, scale=1):
                 )
                 # transfer parameters back to device asynchronously
                 param.copy_(state["_param_fp16"], non_blocking=True)
+            elif param.dtype == torch.bfloat16:
+                if ('maximize' in group) and (group['maximize'] is True):
+                    grad = -state["_grad_bf16"]
+                else:
+                    grad = state["_grad_bf16"]
+                F.adam_cpu_bf16(
+                    state["_param_fp32"].view(-1),
+                    state["_param_bf16"].view(-1),
+                    grad.view(-1),
+                    state["exp_avg"].view(-1),
+                    state["exp_avg_sq"].view(-1),
+                    beta1, beta2,
+                    eps,  0.0 if state["step"] <= self._hold_steps else lr,
+                    scale,
+                    weight_decay,
+                    state["step"]                      
+                )
             else:
                 state["_grad_fp32"].mul_(1.0 / scale)
                 if ('maximize' in group) and (group['maximize'] is True):
@@ -197,9 +225,12 @@ def load_state_dict(self, state_dict: dict) -> None:
                     # initialize placeholders
                     state[param]["_param_fp16"] = torch.empty(param.size(), dtype=torch.float16, pin_memory=True)  # on host
                     state[param]["_grad_fp16"] = torch.empty(param.size(), dtype=torch.float16, pin_memory=True)   # on host
+                elif param.dtype == torch.bfloat16:
+                    #initialize placeholders
+                    state[param]["_param_bf16"] = torch.empty(param.size(), dtype=torch.bfloat16, pin_memory=True)  # on host
+                    state[param]["_grad_bf16"] = torch.empty(param.size(), dtype=torch.bfloat16, pin_memory=True)   # on host
                 else:
-                    state[param]["_param_fp32"] = state[param]["_param_fp32"].pin_memory()
-
+                    state[param]["_param_fp32"] = state[param]["_param_fp32"].pin_memory()  # on host
                     # initialize placeholders
                     state[param]["_grad_fp32"] = torch.empty(param.size(), dtype=torch.float32, pin_memory=True)   # on host
             else:
@@ -254,5 +285,4 @@ def cut_states(state):
             
     #TODO zero_grad(set_to_none=True) makes optimizer crashed, maybe the reason of grad accu
     def zero_grad(self, set_to_none: bool = False):
-        super().zero_grad(set_to_none=set_to_none)
-
+        super().zero_grad(set_to_none=set_to_none)
\ No newline at end of file
diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 0ab55934..744e1fa8 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -20,7 +20,7 @@ __global__ void adam_fp32_accum(
     float bias_correction2
 ) {
     int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
-    if (col < n) {
+    if (col < n) {Supercomputing@ncu666
         float local_g = __half2float(g[col]);                                       // real_g * scale
         float local_m = beta1 * __half2float(m[col]) + (1 - beta1) * local_g;       // real_m * scale
         float local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g / scale;   // real_v * scale
@@ -60,4 +60,8 @@ void adam_launcher(
     dim3 block_size = dim3(threads, 1, 1);
     dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
     adam_fp32_accum<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, m_ptr, v_fp32_ptr, param_fp32_ptr, param_h_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
+}
+
+void adam_bf16_launcher(){
+
 }
\ No newline at end of file
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index b070c7a0..1e7d6f8f 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -329,4 +329,23 @@ void adam_cpu_launcher(
     }
 }
 
+void adam_cpu_bf16_launcher(
+    int64_t n,
+    std::uintptr_t param_fp32,
+    std::uintptr_t param_bf16,
+    std::uintptr_t g_bf16,
+    std::uintptr_t m_fp32,
+    std::uintptr_t v_fp32,
+    float beta1, float beta2,
+    float eps, float lr,
+    float scale,
+    float weight_decay,
+    float bias_correction1,
+    float bias_correction2
+) {
 
+    auto param_fp32_ptr = reinterpret_cast<float*>(param_fp32);
+    auto m_fp32_ptr = reinterpret_cast<float*>(m_fp32);
+    auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
+    auto param_bf16_ptr = reinterpret_cast<uint16_t*>(param_bf16);
+}

From a1e0e425c10f05f7510bcdc6d80fd6ac63c2dfa7 Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Wed, 2 Aug 2023 17:33:57 +0800
Subject: [PATCH 04/27] 	modified:   bmtrain/optim/_function.py 	modified:  
 csrc/bind.cpp 	modified:   csrc/cuda/adam_cuda.cu 	modified:  
 csrc/include/adam_cpu.hpp 	modified:   csrc/include/bind.hpp

---
 bmtrain/optim/_function.py |  2 +-
 csrc/bind.cpp              |  1 +
 csrc/cuda/adam_cuda.cu     | 29 ++++++++++++++++++++++++++---
 csrc/include/adam_cpu.hpp  |  3 +--
 csrc/include/bind.hpp      | 15 +++++++++++++++
 5 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/bmtrain/optim/_function.py b/bmtrain/optim/_function.py
index 9777b589..686a41d3 100644
--- a/bmtrain/optim/_function.py
+++ b/bmtrain/optim/_function.py
@@ -76,7 +76,7 @@ def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tenso
         stream
     )
 
-def adam_bf16(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tensor, m_fp16: torch.Tensor,
+def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tensor, m_fp16: torch.Tensor,
              v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
              weight_decay: float, step: int) -> None:
     assert CHECK_INPUT(param_fp32), "param_fp32 must be contiguous and on cuda"
diff --git a/csrc/bind.cpp b/csrc/bind.cpp
index 8324ba52..99d53970 100644
--- a/csrc/bind.cpp
+++ b/csrc/bind.cpp
@@ -3,6 +3,7 @@
 PYBIND11_MODULE(C, m) {
     m.def("has_nan_inf_launcher",&has_nan_inf_launcher,"has nan inf");
     m.def("adam_launcher", &adam_launcher, "adam function cpu");
+    m.def("adam_launcher", &adam_bf16_launcher, "adam function cpu");
     m.def("adam_cpu_launcher", &adam_cpu_launcher, "adam function cpu");
     m.def("cross_entropy_forward_launcher", &cross_entropy_forward_launcher, "cross entropy forward");
     m.def("cross_entropy_backward_launcher", &cross_entropy_backward_launcher, "cross entropy backward");
diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 744e1fa8..8e47ed12 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -20,7 +20,7 @@ __global__ void adam_fp32_accum(
     float bias_correction2
 ) {
     int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
-    if (col < n) {Supercomputing@ncu666
+    if (col < n) {
         float local_g = __half2float(g[col]);                                       // real_g * scale
         float local_m = beta1 * __half2float(m[col]) + (1 - beta1) * local_g;       // real_m * scale
         float local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g / scale;   // real_v * scale
@@ -62,6 +62,29 @@ void adam_launcher(
     adam_fp32_accum<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, m_ptr, v_fp32_ptr, param_fp32_ptr, param_h_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
 }
 
-void adam_bf16_launcher(){
-
+void adam_bf16_launcher(
+    int n,
+    std::uintptr_t param_fp32,
+    std::uintptr_t param_bf16,
+    std::uintptr_t g_bf16,
+    std::uintptr_t m_bf16,
+    std::uintptr_t v_fp32,
+    float beta1, float beta2,
+    float eps, float lr,
+    float scale,
+    float weight_decay,
+    float bias_correction1,
+    float bias_correction2,
+    uintptr_t stream
+) {
+    if (n <= 0) return;
+    auto g_ptr = reinterpret_cast<half*>(g_bf16);
+    auto m_ptr = reinterpret_cast<half*>(m_bf16);
+    auto param_h_ptr = reinterpret_cast<half*>(param_bf16);
+    auto param_fp32_ptr = reinterpret_cast<float*>(param_fp32);
+    auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
+    int32_t threads = 1024;
+    dim3 block_size = dim3(threads, 1, 1);
+    dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
+    adam_fp32_accum<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, m_ptr, v_fp32_ptr, param_fp32_ptr, param_h_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
 }
\ No newline at end of file
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index 1e7d6f8f..c7090300 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -101,7 +101,7 @@ inline uint16_t fp16_ieee_from_fp32_value(float f) {
             (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
           );
   }
-
+  
 inline float fp16_ieee_to_fp32_value(uint16_t h) {
 
   const uint32_t w = (uint32_t)h << 16;
@@ -343,7 +343,6 @@ void adam_cpu_bf16_launcher(
     float bias_correction1,
     float bias_correction2
 ) {
-
     auto param_fp32_ptr = reinterpret_cast<float*>(param_fp32);
     auto m_fp32_ptr = reinterpret_cast<float*>(m_fp32);
     auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
diff --git a/csrc/include/bind.hpp b/csrc/include/bind.hpp
index 0929de91..804ba0f2 100644
--- a/csrc/include/bind.hpp
+++ b/csrc/include/bind.hpp
@@ -52,4 +52,19 @@ void adam_launcher(
     float bias_correction1,
     float bias_correction2,
     uintptr_t stream
+);
+void adam_bf16_launcher(
+    int n,
+    std::uintptr_t param_fp32,
+    std::uintptr_t param_bf16,
+    std::uintptr_t g_bf16,
+    std::uintptr_t m_bf16,
+    std::uintptr_t v_fp32,
+    float beta1, float beta2,
+    float eps, float lr,
+    float scale,
+    float weight_decay,
+    float bias_correction1,
+    float bias_correction2,
+    uintptr_t stream
 );
\ No newline at end of file

From efe3cdc125185367a8ae67c826394a666d21fa6a Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Thu, 3 Aug 2023 11:53:49 +0800
Subject: [PATCH 05/27] modified: add bf16.h to csrc/cuda/cross_entropy.cu
 modified: add bf16.h to csrc/cuda/has_inf_nan.cu modified: add bf16 support
 to bmtrain/optim/_function.py modified: add bf16 support to
 csrc/cuda/adam_cuda.cu(Todo) modified: add bf16.h to csrc/cuda/reduce.cuh
 modified: add bf16 support to csrc/include/adam_cpu.hpp

---
 csrc/cuda/cross_entropy.cu | 1 +
 csrc/cuda/has_inf_nan.cu   | 1 +
 2 files changed, 2 insertions(+)

diff --git a/csrc/cuda/cross_entropy.cu b/csrc/cuda/cross_entropy.cu
index c0b742ac..c7644e2b 100644
--- a/csrc/cuda/cross_entropy.cu
+++ b/csrc/cuda/cross_entropy.cu
@@ -1,4 +1,5 @@
 #include <cuda_fp16.h>
+#include <cuda_bf16.h>
 #include "reduce.cuh"
 #include <cstdint>
 #include<cuda_runtime_api.h>
diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index b0e906ff..d53b09e4 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -1,4 +1,5 @@
 #include <cuda_fp16.h>
+#include <cuda_bf16.h>
 #include <cstdint>
 
 namespace{

From 6125f39deb19c8090324ec92d87f6825faabd156 Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Thu, 3 Aug 2023 11:58:04 +0800
Subject: [PATCH 06/27] 	modified:   bmtrain/optim/_function.py 	modified:  
 csrc/cuda/adam_cuda.cu 	modified:   csrc/cuda/reduce.cuh 
 modified:   csrc/include/adam_cpu.hpp

---
 bmtrain/optim/_function.py | 67 ++++++++++++++++++++++++++++++--------
 csrc/cuda/adam_cuda.cu     |  1 +
 csrc/cuda/reduce.cuh       |  1 +
 csrc/include/adam_cpu.hpp  |  1 +
 4 files changed, 56 insertions(+), 14 deletions(-)

diff --git a/bmtrain/optim/_function.py b/bmtrain/optim/_function.py
index 686a41d3..2968b859 100644
--- a/bmtrain/optim/_function.py
+++ b/bmtrain/optim/_function.py
@@ -76,32 +76,71 @@ def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tenso
         stream
     )
 
-def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tensor, m_fp16: torch.Tensor,
+def adam_cpu_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.Tensor, m_fp32: torch.Tensor,
+             v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
+             weight_decay: float, step: int) -> None:
+    assert param_fp32.is_contiguous(), "param_fp32 must be contiguous"
+    assert param_bf16.is_contiguous(), "param_bf16 must be contiguous"
+    assert g_bf16.is_contiguous(), "g_bf16 must be contiguous"
+    assert m_fp32.is_contiguous(), "m_fp32 must be contiguous"
+    assert v_fp32.is_contiguous(), "v_fp32 must be contiguous"
+    assert param_fp32.dtype == torch.float32, "param_fp32 must be float32 tensor"
+    assert param_bf16.dtype == torch.bfloat16, "param_bf16 must be bfloat16 tensor"
+    assert g_bf16.dtype == torch.bfloat16, "g_bf16 must be bfloat16 tensor"
+    assert m_fp32.dtype == torch.float32, "m_fp32 must be float32 tensor"
+    assert v_fp32.dtype == torch.float32, "v_fp32 must be float32 tensor"
+    assert param_fp32.device == torch.device("cpu"), "param_fp32 must be a cpu tensor"
+    assert param_bf16.device == torch.device("cpu"), "param_bf16 must be a cpu tensor"
+    assert g_bf16.device == torch.device("cpu"), "g_bf16 must be a cpu tensor"
+    assert m_fp32.device == torch.device("cpu"), "m_fp32 must be a cpu tensor"
+    assert v_fp32.device == torch.device("cpu"), "v_fp32 must be a cpu tensor"
+    assert param_fp32.numel() == param_bf16.numel(), "param_fp32 and param_bf16 must have the same number of elements"
+    assert param_fp32.numel() == g_bf16.numel(), "param_fp32 and g_bf16 must have the same number of elements"
+    assert param_fp32.numel() == m_fp32.numel(), "param_fp32 and m_fp32 must have the same number of elements"
+    assert param_fp32.numel() == v_fp32.numel(), "param_fp32 and v_fp32 must have the same number of elements"
+    bias_correction1 = 1 - beta1 ** step
+    bias_correction2 = 1 - beta2 ** step
+    C.adam_cpu_bf16_launcher(
+                    param_fp32.numel(),
+                    param_fp32.data_ptr(),
+                    param_bf16.data_ptr(),
+                    g_bf16.data_ptr(),
+                    m_fp32.data_ptr(),
+                    v_fp32.data_ptr(),
+                    beta1, beta2,
+                    eps, lr,
+                    scale,
+                    weight_decay,
+                    bias_correction1,
+                    bias_correction2,
+                )
+    
+def adam_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.Tensor, m_bf16: torch.Tensor,
              v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
              weight_decay: float, step: int) -> None:
     assert CHECK_INPUT(param_fp32), "param_fp32 must be contiguous and on cuda"
-    assert CHECK_INPUT(param_fp16), "param_fp16 must be contiguous and on cuda"
-    assert CHECK_INPUT(g_fp16), "g_fp16 must be contiguous and on cuda"
-    assert CHECK_INPUT(m_fp16), "m_fp32 must be contiguous and on cuda"
+    assert CHECK_INPUT(param_bf16), "param_bf16 must be contiguous and on cuda"
+    assert CHECK_INPUT(g_bf16), "g_bf16 must be contiguous and on cuda"
+    assert CHECK_INPUT(m_bf16), "m_bf16 must be contiguous and on cuda"
     assert CHECK_INPUT(v_fp32), "v_fp32 must be contiguous and on cuda"
     assert param_fp32.dtype == torch.float32, "param_fp32 must be float32 tensor"
-    assert param_fp16.dtype == torch.float16, "param_fp16 must be float16 tensor"
-    assert g_fp16.dtype == torch.float16, "g_fp16 must be float16 tensor"
-    assert m_fp16.dtype == torch.float16, "m_fp16 must be float16 tensor"
+    assert param_bf16.dtype == torch.bfloat16, "param_fp16 must be float16 tensor"
+    assert g_bf16.dtype == torch.bfloat16, "g_bf16 must be bfloat16 tensor"
+    assert m_bf16.dtype == torch.bfloat16, "m_bf16 must be bfloat16 tensor"
     assert v_fp32.dtype == torch.float32, "v_fp32 must be float32 tensor"
-    assert param_fp32.numel() == param_fp16.numel(), "param_fp32 and param_fp16 must have the same number of elements"
-    assert param_fp32.numel() == g_fp16.numel(), "param_fp32 and g_fp16 must have the same number of elements"
-    assert param_fp32.numel() == m_fp16.numel(), "param_fp32 and m_fp32 must have the same number of elements"
+    assert param_fp32.numel() == param_bf16.numel(), "param_fp32 and param_bf16 must have the same number of elements"
+    assert param_fp32.numel() == g_bf16.numel(), "param_fp32 and g_fp16 must have the same number of elements"
+    assert param_fp32.numel() == m_bf16.numel(), "param_fp32 and m_bf16 must have the same number of elements"
     assert param_fp32.numel() == v_fp32.numel(), "param_fp32 and v_fp32 must have the same number of elements"
     bias_correction1 = 1 - beta1 ** step
     bias_correction2 = 1 - beta2 ** step
     stream = torch.cuda.current_stream().cuda_stream
-    C.adam_launcher(
+    C.adam_bf16_launcher(
         param_fp32.numel(),
         param_fp32.data_ptr(),
-        param_fp16.data_ptr(),
-        g_fp16.data_ptr(),
-        m_fp16.data_ptr(),
+        param_bf16.data_ptr(),
+        g_bf16.data_ptr(),
+        m_bf16.data_ptr(),
         v_fp32.data_ptr(),
         beta1, beta2,
         eps, lr,
diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 8e47ed12..a3aee56f 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -1,4 +1,5 @@
 #include <cuda_fp16.h>
+#include <cuda_bf16.h>
 #include <cstdint>
 
 namespace {
diff --git a/csrc/cuda/reduce.cuh b/csrc/cuda/reduce.cuh
index 095e8593..a0c1e97a 100644
--- a/csrc/cuda/reduce.cuh
+++ b/csrc/cuda/reduce.cuh
@@ -1,4 +1,5 @@
 #include <cuda_fp16.h>
+#include <cuda_bf16.h>
 
 namespace {
 const int WARP_SZ = 32;
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index c7090300..6068c64e 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -6,6 +6,7 @@
 #include <pybind11/pybind11.h>
 #include<iostream>
 #include<cuda_fp16.h>
+#include <cuda_bf16.h>
 #include <vector>
 #include <thread>
 #include <algorithm>

From e2fdcc5109f1e884add31284ada5e07da0b81058 Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Thu, 3 Aug 2023 15:45:01 +0800
Subject: [PATCH 07/27] 	modified: add adam_fp32_accum_bf16 function

---
 csrc/cuda/adam_cuda.cu | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index a3aee56f..387c6ee9 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -34,6 +34,38 @@ __global__ void adam_fp32_accum(
         m[col] = __float2half(local_m);
     }
 }
+
+__global__ void adam_fp32_accum_bf16(
+    int32_t n,
+    const nv_bfloat16 *g,        // (n)
+    nv_bfloat16 *m,        // (n)
+    float *v,        // (n)
+    float* param,   // (n)
+    nv_bfloat16* param_h,  // (n)
+    float beta1,
+    float beta2,
+    float eps,
+    float lr,
+    float scale,
+    float weight_decay,
+    float bias_correction1,
+    float bias_correction2
+) {
+    int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
+    if (col < n) {
+        nv_bfloat16 local_g = __bfloat162float(g[col]) / scale;                                       // real_g
+        nv_bfloat16 local_m = beta1 * __bfloat162float(m[col]) + (1 - beta1) * local_g;       // real_m
+        nv_bfloat16 local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g;   // real_v
+        float local_p = param[col];
+        local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v * scale / bias_correction2) + eps * scale) - lr * weight_decay * local_p;
+
+        param_h[col] = __float2bfloat16(local_p);
+        param[col] = local_p;
+        v[col] = local_v;
+        m[col] = __float2bfloat16(local_m);
+    }
+}
+
 }
 
 void adam_launcher(

From 00eee3936340c3a5064fd18eb219a77d679f6a1c Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Thu, 3 Aug 2023 15:47:43 +0800
Subject: [PATCH 08/27] modified: add adam_fp32_accum_bf16 function

---
 csrc/cuda/adam_cuda.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 387c6ee9..05669f2b 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -119,5 +119,5 @@ void adam_bf16_launcher(
     int32_t threads = 1024;
     dim3 block_size = dim3(threads, 1, 1);
     dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
-    adam_fp32_accum<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, m_ptr, v_fp32_ptr, param_fp32_ptr, param_h_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
+    adam_fp32_accum_bf16<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, m_ptr, v_fp32_ptr, param_fp32_ptr, param_h_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
 }
\ No newline at end of file

From 2eb6d726a7926d27d91a1513312dc431cfd885d9 Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Thu, 3 Aug 2023 15:49:46 +0800
Subject: [PATCH 09/27] modified: add adam_fp32_accum_bf16 function

---
 csrc/cuda/adam_cuda.cu | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 05669f2b..ba6210df 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -111,9 +111,9 @@ void adam_bf16_launcher(
     uintptr_t stream
 ) {
     if (n <= 0) return;
-    auto g_ptr = reinterpret_cast<half*>(g_bf16);
-    auto m_ptr = reinterpret_cast<half*>(m_bf16);
-    auto param_h_ptr = reinterpret_cast<half*>(param_bf16);
+    auto g_ptr = reinterpret_cast<nv_bfloat16*>(g_bf16);
+    auto m_ptr = reinterpret_cast<nv_bfloat16*>(m_bf16);
+    auto param_h_ptr = reinterpret_cast<nv_bfloat16*>(param_bf16);
     auto param_fp32_ptr = reinterpret_cast<float*>(param_fp32);
     auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
     int32_t threads = 1024;

From 6dfd7392f87c8fe7687898923fa2e08ca224690d Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Mon, 7 Aug 2023 10:44:09 +0800
Subject: [PATCH 10/27] 	modified:   bmtrain/loss/_function.py 	modified:  
 bmtrain/optim/optim_manager.py 	modified:   csrc/bind.cpp 
 modified:   csrc/cuda/has_inf_nan.cu 	modified:   csrc/include/adam_cpu.hpp 
 new file:   csrc/include/test.cpp

---
 bmtrain/loss/_function.py      |  9 +++-
 bmtrain/optim/optim_manager.py |  3 +-
 csrc/bind.cpp                  |  2 +-
 csrc/cuda/has_inf_nan.cu       | 51 ++++++++++++++++++
 csrc/include/adam_cpu.hpp      | 66 ++++++++++++++++++++++-
 csrc/include/test.cpp          | 97 ++++++++++++++++++++++++++++++++++
 6 files changed, 223 insertions(+), 5 deletions(-)
 create mode 100644 csrc/include/test.cpp

diff --git a/bmtrain/loss/_function.py b/bmtrain/loss/_function.py
index 658ef242..9b9edfeb 100644
--- a/bmtrain/loss/_function.py
+++ b/bmtrain/loss/_function.py
@@ -11,7 +11,14 @@ def has_inf_nan(g_fp16: torch.Tensor, out: torch.Tensor) -> None:
     stream = torch.cuda.current_stream().cuda_stream
     C.has_nan_inf_launcher(g_fp16.numel(), g_fp16.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
 
-
+def has_inf_nan_bf16(g_bf16: torch.Tensor, out: torch.Tensor) -> None:
+    assert g_bf16.dtype == torch.bfloat16, "g_bf16 must be a bfloat16 tensor"
+    assert out.dtype == torch.uint8, "out must be a uint8 tensor"
+    assert CHECK_INPUT(g_bf16), "g_bf16 must be contiguous and on cuda"
+    assert CHECK_INPUT(out), "out must be contiguous and on cuda"
+    mid = torch.zeros(1024, device=out.device, dtype=out.dtype)
+    stream = torch.cuda.current_stream().cuda_stream
+    C.has_nan_inf_bf16_launcher(g_bf16.numel(), g_bf16.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
 
 def cross_entropy_forward(m: int, n: int, input: torch.Tensor, target: torch.Tensor,
                             softmax: torch.Tensor, output: torch.Tensor, ignore_index: int) -> None:
diff --git a/bmtrain/optim/optim_manager.py b/bmtrain/optim/optim_manager.py
index 78ad15f8..9a908773 100644
--- a/bmtrain/optim/optim_manager.py
+++ b/bmtrain/optim/optim_manager.py
@@ -13,7 +13,8 @@ def check_overflow(param_groups):
         for p in group['params']:
             if p.grad is not None and p.dtype == torch.half: # TODO support other types
                 has_inf_nan(p.grad, has_inf_or_nan)
-
+            elif p.grad is not None and p.dtype == torch.bfloat16:
+                has_inf_nan_bf16(p.grad, has_inf_or_nan)
     if "comm" in config:
         nccl.allReduce(has_inf_or_nan.storage(), has_inf_or_nan.storage(), "max", config["comm"])
 
diff --git a/csrc/bind.cpp b/csrc/bind.cpp
index 99d53970..a55ebd41 100644
--- a/csrc/bind.cpp
+++ b/csrc/bind.cpp
@@ -3,7 +3,7 @@
 PYBIND11_MODULE(C, m) {
     m.def("has_nan_inf_launcher",&has_nan_inf_launcher,"has nan inf");
     m.def("adam_launcher", &adam_launcher, "adam function cpu");
-    m.def("adam_launcher", &adam_bf16_launcher, "adam function cpu");
+    m.def("adam_bf16_launcher", &adam_bf16_launcher, "adam function cpu");
     m.def("adam_cpu_launcher", &adam_cpu_launcher, "adam function cpu");
     m.def("cross_entropy_forward_launcher", &cross_entropy_forward_launcher, "cross entropy forward");
     m.def("cross_entropy_backward_launcher", &cross_entropy_backward_launcher, "cross entropy backward");
diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index d53b09e4..1c4c6140 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -12,6 +12,14 @@ __inline__ __device__ bool isnan_(half v) {
 #endif
 }
 
+__inline__ __device__ bool isnan_(bfloat16 v) {
+#if __CUDA_ARCH__ >= 800
+    return __hisnan(v);
+#else
+    return !__heq(v, v);
+#endif
+}
+
 __inline__ __device__ int8_t warpReduceAny(int8_t x) {
     for (int offset = warpSize/2; offset > 0; offset /= 2) 
         x |= __shfl_down_sync(0xFFFFFFFF, x, offset);
@@ -65,6 +73,29 @@ __global__ void bmt_has_nan_inf_2(
     }
 }
 
+// grid <min(ceil(n/1024), 1024)>,        thread<1024>
+__global__ void bmt_has_nan_inf_3(
+    int32_t n,
+    const bfloat16* inp,        // (n,) 
+    uint8_t* mid            // (1024,)
+) {
+    int32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
+    int32_t span = blockDim.x * gridDim.x;
+
+    int8_t r = 0;
+    for (int i = gid; i < n; i += span) {
+        half v = inp[i];
+        if (__hisinf(v) || isnan_(v)) {
+            r = 1;
+            break;
+        }
+    }
+    r = blockReduceAny(r);
+    if (threadIdx.x == 0) {
+        mid[blockIdx.x] = r;
+    }
+}
+
 }
 
 void has_nan_inf_launcher(
@@ -85,4 +116,24 @@ void has_nan_inf_launcher(
     
     bmt_has_nan_inf_1<<<clamp_grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, mid_ptr);
     bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(mid_ptr, out_ptr);
+}
+
+void has_nan_inf_bf16_launcher(
+    int32_t n,
+    std::uintptr_t g_bf16,
+    std::uintptr_t mid,
+    std::uintptr_t out,
+    std::uintptr_t stream
+) {
+    if (n <= 0) return;
+    auto g_ptr = reinterpret_cast<bfloat16*>(g_bf16);
+    auto mid_ptr = reinterpret_cast<uint8_t*>(mid);
+    auto out_ptr = reinterpret_cast<uint8_t*>(out);
+    int32_t threads = 1024;
+    dim3 block_size = dim3(threads, 1, 1);
+    dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
+    dim3 clamp_grid_size = dim3(min((n + threads - 1) / threads, 1024), 1, 1);
+    
+    bmt_has_nan_inf_3<<<clamp_grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, mid_ptr);
+    bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(mid_ptr, out_ptr);
 }
\ No newline at end of file
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index 6068c64e..85f1d0f3 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -71,7 +71,7 @@ inline void parallel_for(int64_t begin, int64_t end, int64_t grain_size, const F
 }
 
 
-
+// fp32 -> fp16
 inline uint16_t fp16_ieee_from_fp32_value(float f) {
     // const float scale_to_inf = 0x1.0p+112f;
     // const float scale_to_zero = 0x1.0p-110f;
@@ -102,7 +102,28 @@ inline uint16_t fp16_ieee_from_fp32_value(float f) {
             (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
           );
   }
-  
+
+// fp32 -> bf16
+inline float bf16_ieee_from_fp32_value(float f){
+    uint32_t w = fp32_to_bits(f);
+    uint32_t sign = w & UINT32_C(0x80000000);
+    uint32_t exp = w & UINT32_C(0x7F800000);
+    uint32_t mantissa = w & UINT32_C(0x007FFFFF);
+    uint16_t bf16 = static_cast<uint16_t>(sign >> 16) | static_cast<uint16_t>(exp >> 16) | static_cast<uint16_t>(mantissa >> 16);
+    return fp32_from_bits(static_cast<uint32_t>(bf16) << 16);
+} 
+
+// bf16 -> fp32
+inline float bf16_ieee_to_fp32_value(u_int16_t h){
+    uint32_t w = static_cast<uint32_t>(h) << 16;
+    uint32_t sign = w & UINT32_C(0x80000000);
+    uint32_t exp = w & UINT32_C(0x7F800000);
+    uint32_t mantissa = w & UINT32_C(0x007FFFFF);
+    uint32_t fp32 = sign << 16 | exp << 16 | mantissa << 16;
+    return fp32_from_bits(fp32);
+}
+
+// fp16 -> fp32 
 inline float fp16_ieee_to_fp32_value(uint16_t h) {
 
   const uint32_t w = (uint32_t)h << 16;
@@ -126,6 +147,7 @@ inline float fp16_ieee_to_fp32_value(uint16_t h) {
   return  fp32_from_bits(result);
 }
 
+
 void adam_cpu_0(
     int64_t n,
     float* param_fp32_ptr,
@@ -161,6 +183,46 @@ void adam_cpu_0(
         });
 }
 
+void adam_cpu_bf16_0{
+    int64_t n,
+    float* param_fp32_ptr,
+    uint16_t* param_bf16_ptr,
+    uint16_t* g_bf16_ptr,
+    float* m_fp32_ptr,
+    float* v_fp32_ptr,
+    float beta1, float beta2,
+    float eps, float lr,
+    float scale,
+    float weight_decay,
+    float bias_correction1,
+    float bias_correction2
+}{
+    int64_t span = 1;
+    parallel_for(0, n, 0, [&](int64_t start, int64_t end) {
+    for (int64_t j = start; j < end; j += span) {
+        for (int64_t i = j; i < end; i++) {
+            float g = bf16_ieee_to_fp32_value(g_fp16_ptr[i]) / scale;
+            float m = m_fp32_ptr[i];
+            float v = v_fp32_ptr[i];
+            float p = param_fp32_ptr[i];
+            m = beta1 * m + (1 - beta1) * g;
+            v = beta2 * v + (1 - beta2) * g * g;
+            p = p - lr * m  / bias_correction1 / (sqrtf(v / bias_correction2) + eps) - lr * weight_decay * p;
+            param_fp32_ptr[i] = p;
+            param_bf16_ptr[i] = bf16_ieee_from_fp32_value(p);
+            m_fp32_ptr[i] = m;
+            v_fp32_ptr[i] = v;
+                    }
+            break; // must break here
+        }
+        });
+}
+
+static void __attribute__ ((__ta))
+
+
+
+
 static void __attribute__ ((__target__ ("avx,fma,f16c"))) adam_cpu_1(
     int64_t n,
     float* param_fp32_ptr,
diff --git a/csrc/include/test.cpp b/csrc/include/test.cpp
new file mode 100644
index 00000000..2a6e9e5d
--- /dev/null
+++ b/csrc/include/test.cpp
@@ -0,0 +1,97 @@
+#include <iostream>
+#include <iomanip>
+#include <cstring>
+#include <cmath>
+#include <cstdint>
+#include<sched.h>
+
+inline float fp32_from_bits(uint32_t w) {
+    union {
+        uint32_t as_bits;
+        float as_value;
+    } fp32 = {w};
+    return fp32.as_value;
+}
+
+inline uint32_t fp32_to_bits(float f) {
+    union {
+        float as_value;
+        uint32_t as_bits;
+    } fp32 = {f};
+    return fp32.as_bits;
+}
+
+// fp32 -> fp16
+inline uint16_t fp16_ieee_from_fp32_value(float f) {
+    // const float scale_to_inf = 0x1.0p+112f;
+    // const float scale_to_zero = 0x1.0p-110f;
+    uint32_t scale_to_inf_bits = (uint32_t) 239 << 23;
+    uint32_t scale_to_zero_bits = (uint32_t) 17 << 23;
+    float scale_to_inf_val, scale_to_zero_val;
+    std::memcpy(&scale_to_inf_val, &scale_to_inf_bits, sizeof(scale_to_inf_val));
+    std::memcpy(&scale_to_zero_val, &scale_to_zero_bits, sizeof(scale_to_zero_val));
+    const float scale_to_inf = scale_to_inf_val;
+    const float scale_to_zero = scale_to_zero_val;
+
+    float base = (fabsf(f) * scale_to_inf) * scale_to_zero;
+
+          const uint32_t w = (uint32_t)fp32_to_bits(f);
+          const uint32_t shl1_w = w + w;
+          const uint32_t sign = w & UINT32_C(0x80000000);
+          uint32_t bias = shl1_w & UINT32_C(0xFF000000);
+          if (bias < UINT32_C(0x71000000)) {
+                  bias = UINT32_C(0x71000000);
+          }
+
+          base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
+          const uint32_t bits = (uint32_t)fp32_to_bits(base);
+          const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
+          const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
+          const uint32_t nonsign = exp_bits + mantissa_bits;
+          return static_cast<uint16_t>(
+            (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
+          );
+}
+
+int main() {
+    float f = 2.71828;
+    uint16_t bf16_val = fp16_ieee_from_fp32_value(f);
+
+    std::cout << "Input: " << std::setprecision(7) << f << std::endl;
+    std::cout << "Output: " << bf16_val << std::endl;
+
+    return 0;
+}
+
+
+// // fp32 -> fp16
+// inline uint16_t fp16_ieee_from_fp32_value(float f) {
+//     // const float scale_to_inf = 0x1.0p+112f;
+//     // const float scale_to_zero = 0x1.0p-110f;
+//     uint32_t scale_to_inf_bits = (uint32_t) 239 << 23;
+//     uint32_t scale_to_zero_bits = (uint32_t) 17 << 23;
+//     float scale_to_inf_val, scale_to_zero_val;
+//     std::memcpy(&scale_to_inf_val, &scale_to_inf_bits, sizeof(scale_to_inf_val));
+//     std::memcpy(&scale_to_zero_val, &scale_to_zero_bits, sizeof(scale_to_zero_val));
+//     const float scale_to_inf = scale_to_inf_val;
+//     const float scale_to_zero = scale_to_zero_val;
+
+//     float base = (fabsf(f) * scale_to_inf) * scale_to_zero;
+
+//           const uint32_t w = (uint32_t)fp32_to_bits(f);
+//           const uint32_t shl1_w = w + w;
+//           const uint32_t sign = w & UINT32_C(0x80000000);
+//           uint32_t bias = shl1_w & UINT32_C(0xFF000000);
+//           if (bias < UINT32_C(0x71000000)) {
+//                   bias = UINT32_C(0x71000000);
+//           }
+
+//           base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
+//           const uint32_t bits = (uint32_t)fp32_to_bits(base);
+//           const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
+//           const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
+//           const uint32_t nonsign = exp_bits + mantissa_bits;
+//           return static_cast<uint16_t>(
+//             (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
+//           );
+//   }
\ No newline at end of file

From 8a6f68600f89ba3e0c78c48a9063772d01f30eae Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Mon, 7 Aug 2023 14:08:43 +0800
Subject: [PATCH 11/27] 	modified:   add bf16 to is_nan_inf() 	modified:  
 bind bf16 functions 	modified:   revise the logic of bf16 accum 
 modified:   add bf16 to is_nan_inf() of CUDA version 	modified:   revise the
 logic of bf16 accum 	modified:   bind bf16 functions 	deleted:   
 csrc/include/test.cpp

---
 bmtrain/optim/optim_manager.py |  6 +--
 csrc/bind.cpp                  |  1 +
 csrc/cuda/adam_cuda.cu         | 18 ++++---
 csrc/cuda/has_inf_nan.cu       | 10 ++--
 csrc/include/adam_cpu.hpp      | 20 +++----
 csrc/include/bind.hpp          |  3 +-
 csrc/include/test.cpp          | 97 ----------------------------------
 7 files changed, 28 insertions(+), 127 deletions(-)
 delete mode 100644 csrc/include/test.cpp

diff --git a/bmtrain/optim/optim_manager.py b/bmtrain/optim/optim_manager.py
index 9a908773..a21b3a12 100644
--- a/bmtrain/optim/optim_manager.py
+++ b/bmtrain/optim/optim_manager.py
@@ -1,6 +1,6 @@
 from typing import Optional, Union, List, Dict, Tuple
 import torch
-from ..loss._function import has_inf_nan
+from ..loss._function import has_inf_nan, has_inf_nan_bf16
 from ..utils import print_rank
 from ..lr_scheduler.warmup import WarmupLRScheduler
 from .. import nccl
@@ -11,10 +11,10 @@ def check_overflow(param_groups):
     has_inf_or_nan = torch.zeros(1, dtype=torch.uint8, device="cuda")[0]
     for group in param_groups:
         for p in group['params']:
-            if p.grad is not None and p.dtype == torch.half: # TODO support other types
+            if p.grad is not None and p.dtype == torch.half: 
                 has_inf_nan(p.grad, has_inf_or_nan)
             elif p.grad is not None and p.dtype == torch.bfloat16:
-                has_inf_nan_bf16(p.grad, has_inf_or_nan)
+                has_inf_nan_bf16(p.grad, has_inf_or_nan) # TODO support other types
     if "comm" in config:
         nccl.allReduce(has_inf_or_nan.storage(), has_inf_or_nan.storage(), "max", config["comm"])
 
diff --git a/csrc/bind.cpp b/csrc/bind.cpp
index a55ebd41..315d14ea 100644
--- a/csrc/bind.cpp
+++ b/csrc/bind.cpp
@@ -2,6 +2,7 @@
 
 PYBIND11_MODULE(C, m) {
     m.def("has_nan_inf_launcher",&has_nan_inf_launcher,"has nan inf");
+    m.def("has_nan_inf_bf16_launcher",&has_nan_inf_bf16_launcher,"has nan inf bf16")
     m.def("adam_launcher", &adam_launcher, "adam function cpu");
     m.def("adam_bf16_launcher", &adam_bf16_launcher, "adam function cpu");
     m.def("adam_cpu_launcher", &adam_cpu_launcher, "adam function cpu");
diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index ba6210df..1e9faa9c 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -52,18 +52,20 @@ __global__ void adam_fp32_accum_bf16(
     float bias_correction2
 ) {
     int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
+
     if (col < n) {
-        nv_bfloat16 local_g = __bfloat162float(g[col]) / scale;                                       // real_g
-        nv_bfloat16 local_m = beta1 * __bfloat162float(m[col]) + (1 - beta1) * local_g;       // real_m
-        nv_bfloat16 local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g;   // real_v
+        float local_g = __nv_bfloat162float(g[col]) / scale; // real_g
+        float local_m = beta1 * __nv_bfloat162float(m[col]) + (1 - beta1) * local_g; // real_m
+        float local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g; // real_v
         float local_p = param[col];
-        local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v * scale / bias_correction2) + eps * scale) - lr * weight_decay * local_p;
-
+        local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v / bias_correction2 / scale) + eps) - lr * weight_decay * local_p; 
+    
         param_h[col] = __float2bfloat16(local_p);
         param[col] = local_p;
         v[col] = local_v;
-        m[col] = __float2bfloat16(local_m);
+        m[col] = __float2bfloat16(local_m); 
     }
+       
 }
 
 }
@@ -100,7 +102,7 @@ void adam_bf16_launcher(
     std::uintptr_t param_fp32,
     std::uintptr_t param_bf16,
     std::uintptr_t g_bf16,
-    std::uintptr_t m_bf16,
+    std::uintptr_t m_fp32,
     std::uintptr_t v_fp32,
     float beta1, float beta2,
     float eps, float lr,
@@ -112,7 +114,7 @@ void adam_bf16_launcher(
 ) {
     if (n <= 0) return;
     auto g_ptr = reinterpret_cast<nv_bfloat16*>(g_bf16);
-    auto m_ptr = reinterpret_cast<nv_bfloat16*>(m_bf16);
+    auto m_ptr = reinterpret_cast<nv_bfloat16*>(m_fp32);
     auto param_h_ptr = reinterpret_cast<nv_bfloat16*>(param_bf16);
     auto param_fp32_ptr = reinterpret_cast<float*>(param_fp32);
     auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index 1c4c6140..dd3cf275 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -12,13 +12,15 @@ __inline__ __device__ bool isnan_(half v) {
 #endif
 }
 
-__inline__ __device__ bool isnan_(bfloat16 v) {
 #if __CUDA_ARCH__ >= 800
+__inline__ __device__ bool isnan_(nv_bfloat16 v) {
     return __hisnan(v);
+}
 #else
+__inline__ __device__ bool isnan_(nv_bfloat16 v) {
     return !__heq(v, v);
-#endif
 }
+#endif
 
 __inline__ __device__ int8_t warpReduceAny(int8_t x) {
     for (int offset = warpSize/2; offset > 0; offset /= 2) 
@@ -76,7 +78,7 @@ __global__ void bmt_has_nan_inf_2(
 // grid <min(ceil(n/1024), 1024)>,        thread<1024>
 __global__ void bmt_has_nan_inf_3(
     int32_t n,
-    const bfloat16* inp,        // (n,) 
+    const nv_bfloat16* inp,        // (n,) 
     uint8_t* mid            // (1024,)
 ) {
     int32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
@@ -84,7 +86,7 @@ __global__ void bmt_has_nan_inf_3(
 
     int8_t r = 0;
     for (int i = gid; i < n; i += span) {
-        half v = inp[i];
+        nv_bfloat16 v = inp[i];
         if (__hisinf(v) || isnan_(v)) {
             r = 1;
             break;
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index 85f1d0f3..a3890d7b 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -4,8 +4,8 @@
 #include <cstdint>
 #include<sched.h>
 #include <pybind11/pybind11.h>
-#include<iostream>
-#include<cuda_fp16.h>
+#include <iostream>
+#include <cuda_fp16.h>
 #include <cuda_bf16.h>
 #include <vector>
 #include <thread>
@@ -70,7 +70,6 @@ inline void parallel_for(int64_t begin, int64_t end, int64_t grain_size, const F
     }
 }
 
-
 // fp32 -> fp16
 inline uint16_t fp16_ieee_from_fp32_value(float f) {
     // const float scale_to_inf = 0x1.0p+112f;
@@ -183,7 +182,7 @@ void adam_cpu_0(
         });
 }
 
-void adam_cpu_bf16_0{
+void adam_cpu_bf16_0(
     int64_t n,
     float* param_fp32_ptr,
     uint16_t* param_bf16_ptr,
@@ -196,12 +195,12 @@ void adam_cpu_bf16_0{
     float weight_decay,
     float bias_correction1,
     float bias_correction2
-}{
+){
     int64_t span = 1;
     parallel_for(0, n, 0, [&](int64_t start, int64_t end) {
     for (int64_t j = start; j < end; j += span) {
         for (int64_t i = j; i < end; i++) {
-            float g = bf16_ieee_to_fp32_value(g_fp16_ptr[i]) / scale;
+            float g = bf16_ieee_to_fp32_value(g_bf16_ptr[i]) / scale;
             float m = m_fp32_ptr[i];
             float v = v_fp32_ptr[i];
             float p = param_fp32_ptr[i];
@@ -218,11 +217,6 @@ void adam_cpu_bf16_0{
         });
 }
 
-static void __attribute__ ((__ta))
-
-
-
-
 static void __attribute__ ((__target__ ("avx,fma,f16c"))) adam_cpu_1(
     int64_t n,
     float* param_fp32_ptr,
@@ -359,9 +353,6 @@ static void __attribute__ ((__target__ ("avx512f"))) adam_cpu_2(
         });
 }
 
-
-
-
 void adam_cpu_launcher(
     int64_t n,
     std::uintptr_t param_fp32,
@@ -410,4 +401,5 @@ void adam_cpu_bf16_launcher(
     auto m_fp32_ptr = reinterpret_cast<float*>(m_fp32);
     auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
     auto param_bf16_ptr = reinterpret_cast<uint16_t*>(param_bf16);
+    adam_cpu_bf16_0(n, param_fp32_ptr, param_bf16_ptr, g_bf16_ptr, m_fp32_ptr, v_fp32_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
 }
diff --git a/csrc/include/bind.hpp b/csrc/include/bind.hpp
index 804ba0f2..08d1ee69 100644
--- a/csrc/include/bind.hpp
+++ b/csrc/include/bind.hpp
@@ -3,6 +3,7 @@
 #include "adam_cpu.hpp"
 
 void has_nan_inf_launcher(int32_t n,std::uintptr_t g_fp16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
+void has_nan_inf_bf16_launcher(int32_t n,std:uintptr_t g_bf16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
 
 void cross_entropy_backward_launcher(
     int32_t m, int32_t n,
@@ -58,7 +59,7 @@ void adam_bf16_launcher(
     std::uintptr_t param_fp32,
     std::uintptr_t param_bf16,
     std::uintptr_t g_bf16,
-    std::uintptr_t m_bf16,
+    std::uintptr_t m_fp32,
     std::uintptr_t v_fp32,
     float beta1, float beta2,
     float eps, float lr,
diff --git a/csrc/include/test.cpp b/csrc/include/test.cpp
deleted file mode 100644
index 2a6e9e5d..00000000
--- a/csrc/include/test.cpp
+++ /dev/null
@@ -1,97 +0,0 @@
-#include <iostream>
-#include <iomanip>
-#include <cstring>
-#include <cmath>
-#include <cstdint>
-#include<sched.h>
-
-inline float fp32_from_bits(uint32_t w) {
-    union {
-        uint32_t as_bits;
-        float as_value;
-    } fp32 = {w};
-    return fp32.as_value;
-}
-
-inline uint32_t fp32_to_bits(float f) {
-    union {
-        float as_value;
-        uint32_t as_bits;
-    } fp32 = {f};
-    return fp32.as_bits;
-}
-
-// fp32 -> fp16
-inline uint16_t fp16_ieee_from_fp32_value(float f) {
-    // const float scale_to_inf = 0x1.0p+112f;
-    // const float scale_to_zero = 0x1.0p-110f;
-    uint32_t scale_to_inf_bits = (uint32_t) 239 << 23;
-    uint32_t scale_to_zero_bits = (uint32_t) 17 << 23;
-    float scale_to_inf_val, scale_to_zero_val;
-    std::memcpy(&scale_to_inf_val, &scale_to_inf_bits, sizeof(scale_to_inf_val));
-    std::memcpy(&scale_to_zero_val, &scale_to_zero_bits, sizeof(scale_to_zero_val));
-    const float scale_to_inf = scale_to_inf_val;
-    const float scale_to_zero = scale_to_zero_val;
-
-    float base = (fabsf(f) * scale_to_inf) * scale_to_zero;
-
-          const uint32_t w = (uint32_t)fp32_to_bits(f);
-          const uint32_t shl1_w = w + w;
-          const uint32_t sign = w & UINT32_C(0x80000000);
-          uint32_t bias = shl1_w & UINT32_C(0xFF000000);
-          if (bias < UINT32_C(0x71000000)) {
-                  bias = UINT32_C(0x71000000);
-          }
-
-          base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
-          const uint32_t bits = (uint32_t)fp32_to_bits(base);
-          const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
-          const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
-          const uint32_t nonsign = exp_bits + mantissa_bits;
-          return static_cast<uint16_t>(
-            (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
-          );
-}
-
-int main() {
-    float f = 2.71828;
-    uint16_t bf16_val = fp16_ieee_from_fp32_value(f);
-
-    std::cout << "Input: " << std::setprecision(7) << f << std::endl;
-    std::cout << "Output: " << bf16_val << std::endl;
-
-    return 0;
-}
-
-
-// // fp32 -> fp16
-// inline uint16_t fp16_ieee_from_fp32_value(float f) {
-//     // const float scale_to_inf = 0x1.0p+112f;
-//     // const float scale_to_zero = 0x1.0p-110f;
-//     uint32_t scale_to_inf_bits = (uint32_t) 239 << 23;
-//     uint32_t scale_to_zero_bits = (uint32_t) 17 << 23;
-//     float scale_to_inf_val, scale_to_zero_val;
-//     std::memcpy(&scale_to_inf_val, &scale_to_inf_bits, sizeof(scale_to_inf_val));
-//     std::memcpy(&scale_to_zero_val, &scale_to_zero_bits, sizeof(scale_to_zero_val));
-//     const float scale_to_inf = scale_to_inf_val;
-//     const float scale_to_zero = scale_to_zero_val;
-
-//     float base = (fabsf(f) * scale_to_inf) * scale_to_zero;
-
-//           const uint32_t w = (uint32_t)fp32_to_bits(f);
-//           const uint32_t shl1_w = w + w;
-//           const uint32_t sign = w & UINT32_C(0x80000000);
-//           uint32_t bias = shl1_w & UINT32_C(0xFF000000);
-//           if (bias < UINT32_C(0x71000000)) {
-//                   bias = UINT32_C(0x71000000);
-//           }
-
-//           base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
-//           const uint32_t bits = (uint32_t)fp32_to_bits(base);
-//           const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
-//           const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
-//           const uint32_t nonsign = exp_bits + mantissa_bits;
-//           return static_cast<uint16_t>(
-//             (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
-//           );
-//   }
\ No newline at end of file

From 0b22ed1190eca813c34bc70f98104ca065fdf489 Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Mon, 7 Aug 2023 16:12:59 +0800
Subject: [PATCH 12/27] =?UTF-8?q?FIX=EF=BC=9A=20=20=20csrc/bind.cpp=20modi?=
 =?UTF-8?q?fied=EF=BC=9A=20revise=20logic=20of=20has=5Finf=5Fnan.cu=20FIX?=
 =?UTF-8?q?=EF=BC=9Acsrc/include/adam=5Fcpu.hpp=20FIX=EF=BC=9Acsrc/include?=
 =?UTF-8?q?/bind.hpp?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 csrc/bind.cpp             |  2 +-
 csrc/cuda/has_inf_nan.cu  | 39 ++++++++++++++++++---------------------
 csrc/include/adam_cpu.hpp |  1 +
 csrc/include/bind.hpp     |  2 +-
 4 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/csrc/bind.cpp b/csrc/bind.cpp
index 315d14ea..0802bf5b 100644
--- a/csrc/bind.cpp
+++ b/csrc/bind.cpp
@@ -2,7 +2,7 @@
 
 PYBIND11_MODULE(C, m) {
     m.def("has_nan_inf_launcher",&has_nan_inf_launcher,"has nan inf");
-    m.def("has_nan_inf_bf16_launcher",&has_nan_inf_bf16_launcher,"has nan inf bf16")
+    m.def("has_nan_inf_bf16_launcher",&has_nan_inf_bf16_launcher,"has nan inf bf16");
     m.def("adam_launcher", &adam_launcher, "adam function cpu");
     m.def("adam_bf16_launcher", &adam_bf16_launcher, "adam function cpu");
     m.def("adam_cpu_launcher", &adam_cpu_launcher, "adam function cpu");
diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index dd3cf275..d0fab2f9 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -2,26 +2,24 @@
 #include <cuda_bf16.h>
 #include <cstdint>
 
-namespace{
 __inline__ __device__ bool isnan_(half v) {
-#if __CUDA_ARCH__ >= 700 || __CUDA_ARCH__ == 600
-    return __hisnan(v);
-#else
+    #if __CUDA_ARCH__ >= 700 || __CUDA_ARCH__ == 600
+        return __hisnan(v);
+    #else
+        return !__heq(v, v);
+    #endif
+    }
     
-    return !__heq(v, v);
-#endif
-}
-
-#if __CUDA_ARCH__ >= 800
 __inline__ __device__ bool isnan_(nv_bfloat16 v) {
-    return __hisnan(v);
-}
-#else
-__inline__ __device__ bool isnan_(nv_bfloat16 v) {
-    return !__heq(v, v);
-}
-#endif
-
+    #if __CUDA_ARCH__ >=800
+        float tmp = (float)v;
+        return __hisnan(__float2half(tmp));
+    #else
+        float tmp = (float)v; 
+        return !__heq(__float2half(tmp), __float2half(tmp)); 
+    #endif
+    }
+    
 __inline__ __device__ int8_t warpReduceAny(int8_t x) {
     for (int offset = warpSize/2; offset > 0; offset /= 2) 
         x |= __shfl_down_sync(0xFFFFFFFF, x, offset);
@@ -87,7 +85,8 @@ __global__ void bmt_has_nan_inf_3(
     int8_t r = 0;
     for (int i = gid; i < n; i += span) {
         nv_bfloat16 v = inp[i];
-        if (__hisinf(v) || isnan_(v)) {
+        float tmp = (float)v; 
+        if (isinf(tmp) || isnan(tmp)) { 
             r = 1;
             break;
         }
@@ -98,8 +97,6 @@ __global__ void bmt_has_nan_inf_3(
     }
 }
 
-}
-
 void has_nan_inf_launcher(
     int32_t n,
     std::uintptr_t g_fp16,
@@ -128,7 +125,7 @@ void has_nan_inf_bf16_launcher(
     std::uintptr_t stream
 ) {
     if (n <= 0) return;
-    auto g_ptr = reinterpret_cast<bfloat16*>(g_bf16);
+    auto g_ptr = reinterpret_cast<nv_bfloat16*>(g_bf16);
     auto mid_ptr = reinterpret_cast<uint8_t*>(mid);
     auto out_ptr = reinterpret_cast<uint8_t*>(out);
     int32_t threads = 1024;
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index a3890d7b..c519b895 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -401,5 +401,6 @@ void adam_cpu_bf16_launcher(
     auto m_fp32_ptr = reinterpret_cast<float*>(m_fp32);
     auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
     auto param_bf16_ptr = reinterpret_cast<uint16_t*>(param_bf16);
+    auto g_bf16_ptr  = reinterpret_cast<uint16_t*>(g_bf16);
     adam_cpu_bf16_0(n, param_fp32_ptr, param_bf16_ptr, g_bf16_ptr, m_fp32_ptr, v_fp32_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
 }
diff --git a/csrc/include/bind.hpp b/csrc/include/bind.hpp
index 08d1ee69..e419340b 100644
--- a/csrc/include/bind.hpp
+++ b/csrc/include/bind.hpp
@@ -3,7 +3,7 @@
 #include "adam_cpu.hpp"
 
 void has_nan_inf_launcher(int32_t n,std::uintptr_t g_fp16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
-void has_nan_inf_bf16_launcher(int32_t n,std:uintptr_t g_bf16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
+void has_nan_inf_bf16_launcher(int32_t n,std::uintptr_t g_bf16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
 
 void cross_entropy_backward_launcher(
     int32_t m, int32_t n,

From f8885cf3bab2e0dff0a65a3a7e0c3283ed89be54 Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Mon, 7 Aug 2023 16:59:10 +0800
Subject: [PATCH 13/27] 	modified:   tests/test_has_inf_nan.py 	new file:  
 tests/test_has_inf_nan_bf16.py

---
 tests/test_has_inf_nan.py      | 43 ++++++++++++++++++----------------
 tests/test_has_inf_nan_bf16.py | 35 +++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 20 deletions(-)
 create mode 100644 tests/test_has_inf_nan_bf16.py

diff --git a/tests/test_has_inf_nan.py b/tests/test_has_inf_nan.py
index b1b9b4a9..96aa602a 100644
--- a/tests/test_has_inf_nan.py
+++ b/tests/test_has_inf_nan.py
@@ -1,32 +1,35 @@
 from utils import *
-
 import torch
-import bmtrain.optim._cuda as G
+import bmtrain.loss._function as G
 import random
 
 def check(x, v):
     out = torch.zeros(1, dtype=torch.uint8, device="cuda")[0]
-    G.f_has_inf_nan(x, out)
+    G.has_inf_nan(x, out)
     assert_eq(out.item(), v)
 
 def test_main():
-    for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
-        x = torch.rand((i,)).half().cuda()
-        check(x, 0)
-        p = random.randint(0, i-1)
-        x[p] = x[p] / 0
-        check(x, 1)
-        x[p] = 2
-        check(x, 0)
-        p = random.randint(0, i-1)
-        x[p] = 0
-        x[p] = x[p] / 0
-        check(x, 1)
-        p = random.randint(0, i-1)
-        x[p] = x[p] / 0
-        p = random.randint(0, i-1)
-        x[p] = x[p] / 0
-        check(x, 1)
+    try:
+        for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
+            x = torch.rand((i,)).half().cuda()
+            check(x, 0)
+            p = random.randint(0, i-1)
+            x[p] = x[p] / 0
+            check(x, 1)
+            x[p] = 2
+            check(x, 0)
+            p = random.randint(0, i-1)
+            x[p] = 0
+            x[p] = x[p] / 0
+            check(x, 1)
+            p = random.randint(0, i-1)
+            x[p] = x[p] / 0
+            p = random.randint(0, i-1)
+            x[p] = x[p] / 0
+            check(x, 1)
+        print("That's right")
+    except AssertionError:
+        print("That's fail")
 
 if __name__ == "__main__":
     test_main()
\ No newline at end of file
diff --git a/tests/test_has_inf_nan_bf16.py b/tests/test_has_inf_nan_bf16.py
new file mode 100644
index 00000000..d7c08f19
--- /dev/null
+++ b/tests/test_has_inf_nan_bf16.py
@@ -0,0 +1,35 @@
+from utils import *
+import torch
+import bmtrain.loss._function as G
+import random
+
+def check(x, v):
+    out = torch.zeros(1, dtype=torch.uint8, device="cuda")[0]
+    G.has_inf_nan_bf16(x, out)
+    assert_eq(out.item(), v)
+
+def test_main():
+    try:
+        for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
+            x = torch.rand((i,)).bfloat16().cuda()
+            check(x, 0)
+            p = random.randint(0, i-1)
+            x[p] = x[p] / 0
+            check(x, 1)
+            x[p] = 2
+            check(x, 0)
+            p = random.randint(0, i-1)
+            x[p] = 0
+            x[p] = x[p] / 0
+            check(x, 1)
+            p = random.randint(0, i-1)
+            x[p] = x[p] / 0
+            p = random.randint(0, i-1)
+            x[p] = x[p] / 0
+            check(x, 1)
+        print("That's right")
+    except AssertionError:
+        print("That's fail")
+
+if __name__ == "__main__":
+    test_main()
\ No newline at end of file

From 9a9d526c057335e40f9c3b93ee663060f3cc4bfa Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Mon, 7 Aug 2023 19:10:19 +0800
Subject: [PATCH 14/27] 	modified:   bmtrain/optim/_function.py 	modified:  
 csrc/bind.cpp 	new file:   tests/test_optim_bf16.py 	new file:  
 tests/test_optim_fp16.py

---
 bmtrain/optim/_function.py | 10 +++---
 csrc/bind.cpp              |  1 +
 tests/test_optim_bf16.py   | 65 +++++++++++++++++++++++++++++++++++++
 tests/test_optim_fp16.py   | 66 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 137 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_optim_bf16.py
 create mode 100644 tests/test_optim_fp16.py

diff --git a/bmtrain/optim/_function.py b/bmtrain/optim/_function.py
index 2968b859..26cab409 100644
--- a/bmtrain/optim/_function.py
+++ b/bmtrain/optim/_function.py
@@ -115,22 +115,22 @@ def adam_cpu_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: to
                     bias_correction2,
                 )
     
-def adam_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.Tensor, m_bf16: torch.Tensor,
+def adam_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.Tensor, m_fp32: torch.Tensor,
              v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
              weight_decay: float, step: int) -> None:
     assert CHECK_INPUT(param_fp32), "param_fp32 must be contiguous and on cuda"
     assert CHECK_INPUT(param_bf16), "param_bf16 must be contiguous and on cuda"
     assert CHECK_INPUT(g_bf16), "g_bf16 must be contiguous and on cuda"
-    assert CHECK_INPUT(m_bf16), "m_bf16 must be contiguous and on cuda"
+    assert CHECK_INPUT(m_fp32), "m_fp32 must be contiguous and on cuda"
     assert CHECK_INPUT(v_fp32), "v_fp32 must be contiguous and on cuda"
     assert param_fp32.dtype == torch.float32, "param_fp32 must be float32 tensor"
     assert param_bf16.dtype == torch.bfloat16, "param_fp16 must be float16 tensor"
     assert g_bf16.dtype == torch.bfloat16, "g_bf16 must be bfloat16 tensor"
-    assert m_bf16.dtype == torch.bfloat16, "m_bf16 must be bfloat16 tensor"
+    assert m_fp32.dtype == torch.float32, "m_fp32 must be bfloat16 tensor"
     assert v_fp32.dtype == torch.float32, "v_fp32 must be float32 tensor"
     assert param_fp32.numel() == param_bf16.numel(), "param_fp32 and param_bf16 must have the same number of elements"
     assert param_fp32.numel() == g_bf16.numel(), "param_fp32 and g_fp16 must have the same number of elements"
-    assert param_fp32.numel() == m_bf16.numel(), "param_fp32 and m_bf16 must have the same number of elements"
+    assert param_fp32.numel() == m_fp32.numel(), "param_fp32 and m_m_fp32 must have the same number of elements"
     assert param_fp32.numel() == v_fp32.numel(), "param_fp32 and v_fp32 must have the same number of elements"
     bias_correction1 = 1 - beta1 ** step
     bias_correction2 = 1 - beta2 ** step
@@ -140,7 +140,7 @@ def adam_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.
         param_fp32.data_ptr(),
         param_bf16.data_ptr(),
         g_bf16.data_ptr(),
-        m_bf16.data_ptr(),
+        m_fp32.data_ptr(),
         v_fp32.data_ptr(),
         beta1, beta2,
         eps, lr,
diff --git a/csrc/bind.cpp b/csrc/bind.cpp
index 0802bf5b..b4bc6515 100644
--- a/csrc/bind.cpp
+++ b/csrc/bind.cpp
@@ -6,6 +6,7 @@ PYBIND11_MODULE(C, m) {
     m.def("adam_launcher", &adam_launcher, "adam function cpu");
     m.def("adam_bf16_launcher", &adam_bf16_launcher, "adam function cpu");
     m.def("adam_cpu_launcher", &adam_cpu_launcher, "adam function cpu");
+    m.def("adam_cpu_bf16_launcher", &adam_cpu_bf16_launcher, "adam function cpu");
     m.def("cross_entropy_forward_launcher", &cross_entropy_forward_launcher, "cross entropy forward");
     m.def("cross_entropy_backward_launcher", &cross_entropy_backward_launcher, "cross entropy backward");
     m.def("cross_entropy_forward_inplace_launcher", &cross_entropy_forward_inplace_launcher, "cross entropy forward inplace");
diff --git a/tests/test_optim_bf16.py b/tests/test_optim_bf16.py
new file mode 100644
index 00000000..d3cc7bca
--- /dev/null
+++ b/tests/test_optim_bf16.py
@@ -0,0 +1,65 @@
+import torch
+import bmtrain as bmt
+
+class TestModule(torch.nn.Module):
+    def __init__(self):
+        super(TestModule, self).__init__()
+        self.fc1 = torch.nn.Linear(128, 128)
+        self.fc2 = torch.nn.Linear(128, 128)
+        self.fc3 = torch.nn.Linear(128, 128)
+        self.fc4 = torch.nn.Linear(128, 128)
+        self.fc5 = torch.nn.Linear(128, 128)
+        self.param = torch.nn.Parameter(torch.empty(1237))
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.fc3(x)
+        x = self.fc4(x)
+        x = self.fc5(x)
+        return x
+
+def main():
+    model1 = TestModule()
+    model2 = TestModule()
+    model3 = TestModule()
+
+    state_dict = model1.state_dict()
+    for kw in state_dict.keys():
+        state_dict[kw] = torch.randn_like(state_dict[kw])
+
+    model1.load_state_dict(state_dict)
+    model2.load_state_dict(state_dict)
+    model3.load_state_dict(state_dict)
+
+    model1 = model1.cuda().to(dtype=torch.bfloat16)
+    model2 = model2.cuda().to(dtype=torch.bfloat16)
+    model3 = model3.cuda().to(dtype=torch.bfloat16)
+
+    opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)
+    opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), weight_decay=1e-3)
+    opt3 = torch.optim.Adam(model3.parameters(), weight_decay=1e-3)
+
+    for _ in range(100):
+        opt1.zero_grad()
+        opt2.zero_grad()
+        opt3.zero_grad()
+
+        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
+            grad_bf16 = torch.randn_like(p1).to(dtype=torch.bfloat16) 
+            p1.grad = grad_bf16
+            p2.grad = grad_bf16
+            p3.grad = grad_bf16
+
+        opt1.step()
+        opt2.step()
+        opt3.step()
+
+        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
+            diff1 = torch.abs(p1 - p2).max() 
+            diff2 = torch.abs(p1 - p3).max()
+            diff3 = torch.abs(p2 - p3).max()
+            print(diff1, diff2, diff3)
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/test_optim_fp16.py b/tests/test_optim_fp16.py
new file mode 100644
index 00000000..f3f05551
--- /dev/null
+++ b/tests/test_optim_fp16.py
@@ -0,0 +1,66 @@
+import torch
+import bmtrain as bmt
+
+class TestModule(torch.nn.Module):
+    def __init__(self):
+        super(TestModule, self).__init__()
+        self.fc1 = torch.nn.Linear(128, 128)
+        self.fc2 = torch.nn.Linear(128, 128)
+        self.fc3 = torch.nn.Linear(128, 128)
+        self.fc4 = torch.nn.Linear(128, 128)
+        self.fc5 = torch.nn.Linear(128, 128)
+        self.param = torch.nn.Parameter(torch.empty(1237))
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = self.fc3(x)
+        x = self.fc4(x)
+        x = self.fc5(x)
+        return x
+
+def main():
+    # FIXME: this test script is not working
+    model1 = TestModule()
+    model2 = TestModule()
+    model3 = TestModule()
+
+    state_dict = model1.state_dict()
+    for kw in state_dict.keys():
+        state_dict[kw] = torch.randn_like(state_dict[kw])
+    
+    model1.load_state_dict(state_dict)
+    model2.load_state_dict(state_dict)
+    model3.load_state_dict(state_dict)
+
+    model1 = model1.cuda().half()
+    model2 = model2.cuda().half()
+    model3 = model3.cuda()
+    
+    opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)
+    opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), weight_decay=1e-3)
+    opt3 = torch.optim.Adam(model3.parameters(), weight_decay=1e-3)
+
+    for _ in range(100):
+        opt1.zero_grad()
+        opt2.zero_grad()
+        opt3.zero_grad()
+
+        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
+            grad = torch.randn_like(p1).half()
+            p1.grad = grad
+            p2.grad = grad
+            p3.grad = grad.float()
+        
+        opt1.step()
+        opt2.step()
+        opt3.step()
+
+        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
+            diff1 = torch.abs(p1.float() - p2.float()).max()
+            diff2 = torch.abs(p1.float() - p3).max()
+            diff3 = torch.abs(p2.float() - p3).max()
+            print(diff1, diff2, diff3)
+
+if __name__ == "__main__":
+    main()

From 77c358510a46b8e04332947b8f79f1e85ebc65cd Mon Sep 17 00:00:00 2001
From: Congrui Yin <88324880+JerryYin777@users.noreply.github.com>
Date: Tue, 8 Aug 2023 00:39:25 +0800
Subject: [PATCH 15/27] add pybind11 in Update other_requirements.txt

---
 other_requirements.txt | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/other_requirements.txt b/other_requirements.txt
index 6654b1ac..232fee4a 100644
--- a/other_requirements.txt
+++ b/other_requirements.txt
@@ -3,4 +3,5 @@ cpm_kernels>=1.0.11
 jieba
 tensorboard
 setuptools_rust
-transformers
\ No newline at end of file
+transformers
+pybind11

From 5cc3611923d40615a220ec6db5cc038ec067cce6 Mon Sep 17 00:00:00 2001
From: Congrui Yin <88324880+JerryYin777@users.noreply.github.com>
Date: Tue, 8 Aug 2023 00:42:17 +0800
Subject: [PATCH 16/27] Update adam_cuda.cu

---
 csrc/cuda/adam_cuda.cu | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 1e9faa9c..15e2f805 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -54,8 +54,8 @@ __global__ void adam_fp32_accum_bf16(
     int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
 
     if (col < n) {
-        float local_g = __nv_bfloat162float(g[col]) / scale; // real_g
-        float local_m = beta1 * __nv_bfloat162float(m[col]) + (1 - beta1) * local_g; // real_m
+        float local_g = __bfloat162float(g[col]) / scale; // real_g
+        float local_m = beta1 * __bfloat162float(m[col]) + (1 - beta1) * local_g; // real_m
         float local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g; // real_v
         float local_p = param[col];
         local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v / bias_correction2 / scale) + eps) - lr * weight_decay * local_p; 
@@ -122,4 +122,4 @@ void adam_bf16_launcher(
     dim3 block_size = dim3(threads, 1, 1);
     dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
     adam_fp32_accum_bf16<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, m_ptr, v_fp32_ptr, param_fp32_ptr, param_h_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
-}
\ No newline at end of file
+}

From 2b414b8cbd1ea393bf84c4d714bfbd9067634f3f Mon Sep 17 00:00:00 2001
From: Congrui Yin <88324880+JerryYin777@users.noreply.github.com>
Date: Mon, 7 Aug 2023 19:47:18 -0700
Subject: [PATCH 17/27] Update test_optim_bf16.py

---
 tests/test_optim_bf16.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/test_optim_bf16.py b/tests/test_optim_bf16.py
index d3cc7bca..2a156e46 100644
--- a/tests/test_optim_bf16.py
+++ b/tests/test_optim_bf16.py
@@ -34,7 +34,7 @@ def main():
 
     model1 = model1.cuda().to(dtype=torch.bfloat16)
     model2 = model2.cuda().to(dtype=torch.bfloat16)
-    model3 = model3.cuda().to(dtype=torch.bfloat16)
+    model3 = model3.cuda()
 
     opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)
     opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), weight_decay=1e-3)
@@ -46,10 +46,10 @@ def main():
         opt3.zero_grad()
 
         for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
-            grad_bf16 = torch.randn_like(p1).to(dtype=torch.bfloat16) 
-            p1.grad = grad_bf16
-            p2.grad = grad_bf16
-            p3.grad = grad_bf16
+            grad = torch.randn_like(p1)
+            p1.grad = grad.to(dtype=torch.bfloat16)
+            p2.grad = grad.to(dtype=torch.bfloat16)
+            p3.grad = grad.float()
 
         opt1.step()
         opt2.step()

From c5f7e49b9765683d442c7dfe2e907104f0fa676e Mon Sep 17 00:00:00 2001
From: JerryYin777 <943321359@qq.com>
Date: Wed, 9 Aug 2023 09:22:45 +0800
Subject: [PATCH 18/27] FIX

---
 csrc/cuda/has_inf_nan.cu | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index d0fab2f9..90b42a01 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -16,7 +16,8 @@ __inline__ __device__ bool isnan_(nv_bfloat16 v) {
         return __hisnan(__float2half(tmp));
     #else
         float tmp = (float)v; 
-        return !__heq(__float2half(tmp), __float2half(tmp)); 
+        return !__heq(__float2half(tmp), __float2half(tmp));
+    printf("you can not use bf16 if __CUDA_ARCH__ < 800\n");
     #endif
     }
     

From 55839bed3552d69fdab7b8cb6c52fbdb6c011202 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Thu, 10 Aug 2023 17:51:00 +0800
Subject: [PATCH 19/27] refactor has_inf_nan_bf16

---
 csrc/cuda/has_inf_nan.cu       | 47 ++++++++++++---------------
 other_requirements.txt         |  3 +-
 setup.py                       |  7 ++--
 tests/test_all.py              |  4 +++
 tests/test_has_inf_nan.py      | 41 +++++++++++-------------
 tests/test_has_inf_nan_bf16.py | 41 +++++++++++-------------
 tests/test_optim.py            | 58 ----------------------------------
 tests/test_optim_bf16.py       |  4 +++
 tests/test_optim_fp16.py       |  4 +++
 9 files changed, 76 insertions(+), 133 deletions(-)
 delete mode 100644 tests/test_optim.py

diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index 90b42a01..9e73859e 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -1,25 +1,17 @@
-#include <cuda_fp16.h>
-#include <cuda_bf16.h>
 #include <cstdint>
+#include <cuda_fp16.h>
+#if __CUDA_ARCH__ >= 800
+    #include <cuda_bf16.h>
+#endif
 
+namespace{
 __inline__ __device__ bool isnan_(half v) {
     #if __CUDA_ARCH__ >= 700 || __CUDA_ARCH__ == 600
         return __hisnan(v);
     #else
         return !__heq(v, v);
     #endif
-    }
-    
-__inline__ __device__ bool isnan_(nv_bfloat16 v) {
-    #if __CUDA_ARCH__ >=800
-        float tmp = (float)v;
-        return __hisnan(__float2half(tmp));
-    #else
-        float tmp = (float)v; 
-        return !__heq(__float2half(tmp), __float2half(tmp));
-    printf("you can not use bf16 if __CUDA_ARCH__ < 800\n");
-    #endif
-    }
+}
     
 __inline__ __device__ int8_t warpReduceAny(int8_t x) {
     for (int offset = warpSize/2; offset > 0; offset /= 2) 
@@ -40,7 +32,7 @@ __inline__ __device__ float blockReduceAny(int8_t x) {
 }
 
 // grid <min(ceil(n/1024), 1024)>,        thread<1024>
-__global__ void bmt_has_nan_inf_1(
+__global__ void bmt_has_nan_inf_fp16(
     int32_t n,
     const half* inp,        // (n,) 
     uint8_t* mid            // (1024,)
@@ -63,7 +55,7 @@ __global__ void bmt_has_nan_inf_1(
 }
 
 // grid <1>,        thread<1024>
-__global__ void bmt_has_nan_inf_2(
+__global__ void bmt_has_nan_inf_reduce(
     const uint8_t* mid,    // (1024,) 
     uint8_t* out
 ) {
@@ -75,19 +67,20 @@ __global__ void bmt_has_nan_inf_2(
 }
 
 // grid <min(ceil(n/1024), 1024)>,        thread<1024>
-__global__ void bmt_has_nan_inf_3(
+__global__ void bmt_has_nan_inf_bf16(
     int32_t n,
-    const nv_bfloat16* inp,        // (n,) 
+    const uintptr_t inp,        // (n,) 
     uint8_t* mid            // (1024,)
 ) {
+#if __CUDA_ARCH__ >= 800
+    const __nv_bfloat16* bf_inp = reinterpret_cast<const __nv_bfloat16*>(inp);
     int32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
     int32_t span = blockDim.x * gridDim.x;
 
     int8_t r = 0;
     for (int i = gid; i < n; i += span) {
-        nv_bfloat16 v = inp[i];
-        float tmp = (float)v; 
-        if (isinf(tmp) || isnan(tmp)) { 
+        __nv_bfloat16 v = bf_inp[i];
+        if (__hisinf(v) || __hisnan(v)) { 
             r = 1;
             break;
         }
@@ -96,6 +89,9 @@ __global__ void bmt_has_nan_inf_3(
     if (threadIdx.x == 0) {
         mid[blockIdx.x] = r;
     }
+#endif
+}
+
 }
 
 void has_nan_inf_launcher(
@@ -114,8 +110,8 @@ void has_nan_inf_launcher(
     dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
     dim3 clamp_grid_size = dim3(min((n + threads - 1) / threads, 1024), 1, 1);
     
-    bmt_has_nan_inf_1<<<clamp_grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, mid_ptr);
-    bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(mid_ptr, out_ptr);
+    bmt_has_nan_inf_fp16<<<clamp_grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, mid_ptr);
+    bmt_has_nan_inf_reduce<<<1, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(mid_ptr, out_ptr);
 }
 
 void has_nan_inf_bf16_launcher(
@@ -126,7 +122,6 @@ void has_nan_inf_bf16_launcher(
     std::uintptr_t stream
 ) {
     if (n <= 0) return;
-    auto g_ptr = reinterpret_cast<nv_bfloat16*>(g_bf16);
     auto mid_ptr = reinterpret_cast<uint8_t*>(mid);
     auto out_ptr = reinterpret_cast<uint8_t*>(out);
     int32_t threads = 1024;
@@ -134,6 +129,6 @@ void has_nan_inf_bf16_launcher(
     dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
     dim3 clamp_grid_size = dim3(min((n + threads - 1) / threads, 1024), 1, 1);
     
-    bmt_has_nan_inf_3<<<clamp_grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, mid_ptr);
-    bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(mid_ptr, out_ptr);
+    bmt_has_nan_inf_bf16<<<clamp_grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_bf16, mid_ptr);
+    bmt_has_nan_inf_reduce<<<1, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(mid_ptr, out_ptr);
 }
\ No newline at end of file
diff --git a/other_requirements.txt b/other_requirements.txt
index 232fee4a..6654b1ac 100644
--- a/other_requirements.txt
+++ b/other_requirements.txt
@@ -3,5 +3,4 @@ cpm_kernels>=1.0.11
 jieba
 tensorboard
 setuptools_rust
-transformers
-pybind11
+transformers
\ No newline at end of file
diff --git a/setup.py b/setup.py
index a163ba54..d2898204 100644
--- a/setup.py
+++ b/setup.py
@@ -1,4 +1,4 @@
-import os
+import os, shutil
 from setuptools.command.build_ext import build_ext
 from setuptools import  setup, find_packages, Extension
 import setuptools
@@ -79,8 +79,9 @@ def build_extension(self, ext):
                 build_args += [f"-j{self.parallel}"]
 
         build_temp = os.path.join(self.build_temp, ext.name)
-        if not os.path.exists(build_temp):
-            os.makedirs(build_temp)
+        if os.path.exists(build_temp):
+            shutil.rmtree(build_temp)
+        os.makedirs(build_temp)
 
         cmake_args += ["-DPython_ROOT_DIR=" + os.path.dirname(os.path.dirname(sys.executable))]
         subprocess.check_call(["cmake", ext.sourcedir] + cmake_args, cwd=build_temp)
diff --git a/tests/test_all.py b/tests/test_all.py
index b614d3eb..24531042 100644
--- a/tests/test_all.py
+++ b/tests/test_all.py
@@ -12,9 +12,13 @@
     ("requires_grad", 1),
     ("requires_grad_multi_gpu", 2),
     ("has_inf_nan", 1),
+    ("has_inf_nan_bf16", 1),
     ("dropout", 1),
     ("loss_func", 1),
 
+    ("optim_fp16", 1),
+    ("optim_bf16", 1),
+
     ("multi_return", 2),
     ("middle_hidden", 4),
     ("other_hidden", 4),
diff --git a/tests/test_has_inf_nan.py b/tests/test_has_inf_nan.py
index 96aa602a..a61a70b9 100644
--- a/tests/test_has_inf_nan.py
+++ b/tests/test_has_inf_nan.py
@@ -9,27 +9,24 @@ def check(x, v):
     assert_eq(out.item(), v)
 
 def test_main():
-    try:
-        for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
-            x = torch.rand((i,)).half().cuda()
-            check(x, 0)
-            p = random.randint(0, i-1)
-            x[p] = x[p] / 0
-            check(x, 1)
-            x[p] = 2
-            check(x, 0)
-            p = random.randint(0, i-1)
-            x[p] = 0
-            x[p] = x[p] / 0
-            check(x, 1)
-            p = random.randint(0, i-1)
-            x[p] = x[p] / 0
-            p = random.randint(0, i-1)
-            x[p] = x[p] / 0
-            check(x, 1)
-        print("That's right")
-    except AssertionError:
-        print("That's fail")
+    for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
+        x = torch.rand((i,)).half().cuda()
+        check(x, 0)
+        p = random.randint(0, i-1)
+        x[p] = x[p] / 0
+        check(x, 1)
+        x[p] = 2
+        check(x, 0)
+        p = random.randint(0, i-1)
+        x[p] = 0
+        x[p] = x[p] / 0
+        check(x, 1)
+        p = random.randint(0, i-1)
+        x[p] = x[p] / 0
+        p = random.randint(0, i-1)
+        x[p] = x[p] / 0
+        check(x, 1)
+    print("That's right")
 
 if __name__ == "__main__":
-    test_main()
\ No newline at end of file
+    test_main()
diff --git a/tests/test_has_inf_nan_bf16.py b/tests/test_has_inf_nan_bf16.py
index d7c08f19..e510341b 100644
--- a/tests/test_has_inf_nan_bf16.py
+++ b/tests/test_has_inf_nan_bf16.py
@@ -9,27 +9,24 @@ def check(x, v):
     assert_eq(out.item(), v)
 
 def test_main():
-    try:
-        for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
-            x = torch.rand((i,)).bfloat16().cuda()
-            check(x, 0)
-            p = random.randint(0, i-1)
-            x[p] = x[p] / 0
-            check(x, 1)
-            x[p] = 2
-            check(x, 0)
-            p = random.randint(0, i-1)
-            x[p] = 0
-            x[p] = x[p] / 0
-            check(x, 1)
-            p = random.randint(0, i-1)
-            x[p] = x[p] / 0
-            p = random.randint(0, i-1)
-            x[p] = x[p] / 0
-            check(x, 1)
-        print("That's right")
-    except AssertionError:
-        print("That's fail")
+    for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
+        x = torch.rand((i,)).bfloat16().cuda()
+        check(x, 0)
+        p = random.randint(0, i-1)
+        x[p] = x[p] / 0
+        check(x, 1)
+        x[p] = 2
+        check(x, 0)
+        p = random.randint(0, i-1)
+        x[p] = 0
+        x[p] = x[p] / 0
+        check(x, 1)
+        p = random.randint(0, i-1)
+        x[p] = x[p] / 0
+        p = random.randint(0, i-1)
+        x[p] = x[p] / 0
+        check(x, 1)
+    print("That's right")
 
 if __name__ == "__main__":
-    test_main()
\ No newline at end of file
+    test_main()
diff --git a/tests/test_optim.py b/tests/test_optim.py
deleted file mode 100644
index 81356ede..00000000
--- a/tests/test_optim.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import torch
-import bmtrain as bmt
-
-class TestModule(torch.nn.Module):
-    def __init__(self):
-        super(TestModule, self).__init__()
-        self.fc1 = torch.nn.Linear(128, 128)
-        self.fc2 = torch.nn.Linear(128, 128)
-        self.fc3 = torch.nn.Linear(128, 128)
-        self.fc4 = torch.nn.Linear(128, 128)
-        self.fc5 = torch.nn.Linear(128, 128)
-        self.param = torch.nn.Parameter(torch.empty(1237))
-
-def main():
-    # FIXME: this test script is not working
-    model1 = TestModule()
-    model2 = TestModule()
-    model3 = TestModule()
-
-    state_dict = model1.state_dict()
-    for kw in state_dict.keys():
-        state_dict[kw] = torch.randn_like(state_dict[kw])
-    
-    model1.load_state_dict(state_dict)
-    model2.load_state_dict(state_dict)
-    model3.load_state_dict(state_dict)
-
-    model1 = model1.cuda()
-    model2 = model2.cuda()
-    model3 = model3.cuda()
-    
-    opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)
-    opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), weight_decay=1e-3)
-    opt3 = torch.optim.Adam(model3.parameters(), weight_decay=1e-3)
-
-    for _ in range(100):
-        opt1.zero_grad()
-        opt2.zero_grad()
-        opt3.zero_grad()
-
-        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
-            grad = torch.randn_like(p1)
-            p1.grad = grad
-            p2.grad = grad
-            p3.grad = grad.float()
-        
-        opt1.step()
-        opt2.step()
-        opt3.step()
-
-        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
-            diff1 = torch.abs(p1 - p2).max()
-            diff2 = torch.abs(p1 - p3).max()
-            diff3 = torch.abs(p2 - p3).max()
-            print(diff1, diff2, diff3)
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/test_optim_bf16.py b/tests/test_optim_bf16.py
index 2a156e46..521f42aa 100644
--- a/tests/test_optim_bf16.py
+++ b/tests/test_optim_bf16.py
@@ -1,3 +1,4 @@
+from utils import *
 import torch
 import bmtrain as bmt
 
@@ -60,6 +61,9 @@ def main():
             diff2 = torch.abs(p1 - p3).max()
             diff3 = torch.abs(p2 - p3).max()
             print(diff1, diff2, diff3)
+            assert_lt(diff1, 1)
+            assert_lt(diff2, 1)
+            assert_lt(diff3, 1)
 
 if __name__ == "__main__":
     main()
diff --git a/tests/test_optim_fp16.py b/tests/test_optim_fp16.py
index f3f05551..e995318b 100644
--- a/tests/test_optim_fp16.py
+++ b/tests/test_optim_fp16.py
@@ -1,3 +1,4 @@
+from utils import *
 import torch
 import bmtrain as bmt
 
@@ -61,6 +62,9 @@ def main():
             diff2 = torch.abs(p1.float() - p3).max()
             diff3 = torch.abs(p2.float() - p3).max()
             print(diff1, diff2, diff3)
+            assert_lt(diff1, 1)
+            assert_lt(diff2, 1)
+            assert_lt(diff3, 1)
 
 if __name__ == "__main__":
     main()

From 4d44dcef37be34729ac431c71ee104f4ee39ed11 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Thu, 10 Aug 2023 18:32:10 +0800
Subject: [PATCH 20/27] refactor has_inf_nan_bf16

---
 bmtrain/loss/_function.py                   | 21 +++----
 bmtrain/optim/optim_manager.py              |  9 ++-
 csrc/cuda/reduce.cuh                        |  3 -
 tests/test_all.py                           |  4 +-
 tests/test_has_inf_nan.py                   |  7 ++-
 tests/test_has_inf_nan_bf16.py              | 32 ----------
 tests/{test_optim_bf16.py => test_optim.py} | 13 ++--
 tests/test_optim_fp16.py                    | 70 ---------------------
 8 files changed, 24 insertions(+), 135 deletions(-)
 delete mode 100644 tests/test_has_inf_nan_bf16.py
 rename tests/{test_optim_bf16.py => test_optim.py} (89%)
 delete mode 100644 tests/test_optim_fp16.py

diff --git a/bmtrain/loss/_function.py b/bmtrain/loss/_function.py
index 9b9edfeb..e3c76ea9 100644
--- a/bmtrain/loss/_function.py
+++ b/bmtrain/loss/_function.py
@@ -2,23 +2,18 @@
 from .. import C 
 import torch
 CHECK_INPUT = lambda x: x.is_contiguous() and x.is_cuda
-def has_inf_nan(g_fp16: torch.Tensor, out: torch.Tensor) -> None:
-    assert g_fp16.dtype == torch.float16, "g_fp16 must be a half tensor"
+def has_inf_nan(g_half: torch.Tensor, out: torch.Tensor) -> None:
     assert out.dtype == torch.uint8, "out must be a uint8 tensor"
-    assert CHECK_INPUT(g_fp16), "g_fp16 must be contiguous and on cuda"
+    assert CHECK_INPUT(g_half), "g_fp16 must be contiguous and on cuda"
     assert CHECK_INPUT(out), "out must be contiguous and on cuda"
     mid = torch.zeros(1024, device=out.device, dtype=out.dtype)
     stream = torch.cuda.current_stream().cuda_stream
-    C.has_nan_inf_launcher(g_fp16.numel(), g_fp16.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
-
-def has_inf_nan_bf16(g_bf16: torch.Tensor, out: torch.Tensor) -> None:
-    assert g_bf16.dtype == torch.bfloat16, "g_bf16 must be a bfloat16 tensor"
-    assert out.dtype == torch.uint8, "out must be a uint8 tensor"
-    assert CHECK_INPUT(g_bf16), "g_bf16 must be contiguous and on cuda"
-    assert CHECK_INPUT(out), "out must be contiguous and on cuda"
-    mid = torch.zeros(1024, device=out.device, dtype=out.dtype)
-    stream = torch.cuda.current_stream().cuda_stream
-    C.has_nan_inf_bf16_launcher(g_bf16.numel(), g_bf16.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
+    if g_half.dtype == torch.float16:
+        C.has_nan_inf_launcher(g_half.numel(), g_half.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
+    elif g_half.dtype == torch.bfloat16:
+        C.has_nan_inf_bf16_launcher(g_half.numel(), g_half.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
+    else:
+        raise ValueError(f"has_inf_nan not supported for dtype {g_half.dtype}")
 
 def cross_entropy_forward(m: int, n: int, input: torch.Tensor, target: torch.Tensor,
                             softmax: torch.Tensor, output: torch.Tensor, ignore_index: int) -> None:
diff --git a/bmtrain/optim/optim_manager.py b/bmtrain/optim/optim_manager.py
index a21b3a12..9b7a3120 100644
--- a/bmtrain/optim/optim_manager.py
+++ b/bmtrain/optim/optim_manager.py
@@ -1,6 +1,6 @@
 from typing import Optional, Union, List, Dict, Tuple
 import torch
-from ..loss._function import has_inf_nan, has_inf_nan_bf16
+from ..loss._function import has_inf_nan
 from ..utils import print_rank
 from ..lr_scheduler.warmup import WarmupLRScheduler
 from .. import nccl
@@ -11,10 +11,9 @@ def check_overflow(param_groups):
     has_inf_or_nan = torch.zeros(1, dtype=torch.uint8, device="cuda")[0]
     for group in param_groups:
         for p in group['params']:
-            if p.grad is not None and p.dtype == torch.half: 
-                has_inf_nan(p.grad, has_inf_or_nan)
-            elif p.grad is not None and p.dtype == torch.bfloat16:
-                has_inf_nan_bf16(p.grad, has_inf_or_nan) # TODO support other types
+            if p.grad is not None:
+                if p.dtype != torch.float:
+                    has_inf_nan(p.grad, has_inf_or_nan)
     if "comm" in config:
         nccl.allReduce(has_inf_or_nan.storage(), has_inf_or_nan.storage(), "max", config["comm"])
 
diff --git a/csrc/cuda/reduce.cuh b/csrc/cuda/reduce.cuh
index a0c1e97a..a9c4c15b 100644
--- a/csrc/cuda/reduce.cuh
+++ b/csrc/cuda/reduce.cuh
@@ -1,6 +1,3 @@
-#include <cuda_fp16.h>
-#include <cuda_bf16.h>
-
 namespace {
 const int WARP_SZ = 32;
 
diff --git a/tests/test_all.py b/tests/test_all.py
index 24531042..a3b7feef 100644
--- a/tests/test_all.py
+++ b/tests/test_all.py
@@ -12,12 +12,10 @@
     ("requires_grad", 1),
     ("requires_grad_multi_gpu", 2),
     ("has_inf_nan", 1),
-    ("has_inf_nan_bf16", 1),
     ("dropout", 1),
     ("loss_func", 1),
 
-    ("optim_fp16", 1),
-    ("optim_bf16", 1),
+    ("optim", 1),
 
     ("multi_return", 2),
     ("middle_hidden", 4),
diff --git a/tests/test_has_inf_nan.py b/tests/test_has_inf_nan.py
index a61a70b9..5432aed9 100644
--- a/tests/test_has_inf_nan.py
+++ b/tests/test_has_inf_nan.py
@@ -8,9 +8,9 @@ def check(x, v):
     G.has_inf_nan(x, out)
     assert_eq(out.item(), v)
 
-def test_main():
+def test_main(dtype):
     for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
-        x = torch.rand((i,)).half().cuda()
+        x = torch.rand((i,)).to(dtype).cuda()
         check(x, 0)
         p = random.randint(0, i-1)
         x[p] = x[p] / 0
@@ -29,4 +29,5 @@ def test_main():
     print("That's right")
 
 if __name__ == "__main__":
-    test_main()
+    test_main(torch.float16)
+    test_main(torch.bfloat16)
diff --git a/tests/test_has_inf_nan_bf16.py b/tests/test_has_inf_nan_bf16.py
deleted file mode 100644
index e510341b..00000000
--- a/tests/test_has_inf_nan_bf16.py
+++ /dev/null
@@ -1,32 +0,0 @@
-from utils import *
-import torch
-import bmtrain.loss._function as G
-import random
-
-def check(x, v):
-    out = torch.zeros(1, dtype=torch.uint8, device="cuda")[0]
-    G.has_inf_nan_bf16(x, out)
-    assert_eq(out.item(), v)
-
-def test_main():
-    for i in list(range(1, 100)) + [1000]*10 + [10000]*10 + [100000]*10 + [1000000]*10:
-        x = torch.rand((i,)).bfloat16().cuda()
-        check(x, 0)
-        p = random.randint(0, i-1)
-        x[p] = x[p] / 0
-        check(x, 1)
-        x[p] = 2
-        check(x, 0)
-        p = random.randint(0, i-1)
-        x[p] = 0
-        x[p] = x[p] / 0
-        check(x, 1)
-        p = random.randint(0, i-1)
-        x[p] = x[p] / 0
-        p = random.randint(0, i-1)
-        x[p] = x[p] / 0
-        check(x, 1)
-    print("That's right")
-
-if __name__ == "__main__":
-    test_main()
diff --git a/tests/test_optim_bf16.py b/tests/test_optim.py
similarity index 89%
rename from tests/test_optim_bf16.py
rename to tests/test_optim.py
index 521f42aa..6ee5e0fc 100644
--- a/tests/test_optim_bf16.py
+++ b/tests/test_optim.py
@@ -20,7 +20,7 @@ def forward(self, x):
         x = self.fc5(x)
         return x
 
-def main():
+def main(dtype):
     model1 = TestModule()
     model2 = TestModule()
     model3 = TestModule()
@@ -33,8 +33,8 @@ def main():
     model2.load_state_dict(state_dict)
     model3.load_state_dict(state_dict)
 
-    model1 = model1.cuda().to(dtype=torch.bfloat16)
-    model2 = model2.cuda().to(dtype=torch.bfloat16)
+    model1 = model1.cuda().to(dtype)
+    model2 = model2.cuda().to(dtype)
     model3 = model3.cuda()
 
     opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)
@@ -48,8 +48,8 @@ def main():
 
         for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
             grad = torch.randn_like(p1)
-            p1.grad = grad.to(dtype=torch.bfloat16)
-            p2.grad = grad.to(dtype=torch.bfloat16)
+            p1.grad = grad.to(dtype)
+            p2.grad = grad.to(dtype)
             p3.grad = grad.float()
 
         opt1.step()
@@ -66,4 +66,5 @@ def main():
             assert_lt(diff3, 1)
 
 if __name__ == "__main__":
-    main()
+    main(torch.float16)
+    main(torch.bfloat16)
diff --git a/tests/test_optim_fp16.py b/tests/test_optim_fp16.py
deleted file mode 100644
index e995318b..00000000
--- a/tests/test_optim_fp16.py
+++ /dev/null
@@ -1,70 +0,0 @@
-from utils import *
-import torch
-import bmtrain as bmt
-
-class TestModule(torch.nn.Module):
-    def __init__(self):
-        super(TestModule, self).__init__()
-        self.fc1 = torch.nn.Linear(128, 128)
-        self.fc2 = torch.nn.Linear(128, 128)
-        self.fc3 = torch.nn.Linear(128, 128)
-        self.fc4 = torch.nn.Linear(128, 128)
-        self.fc5 = torch.nn.Linear(128, 128)
-        self.param = torch.nn.Parameter(torch.empty(1237))
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.fc2(x)
-        x = self.fc3(x)
-        x = self.fc4(x)
-        x = self.fc5(x)
-        return x
-
-def main():
-    # FIXME: this test script is not working
-    model1 = TestModule()
-    model2 = TestModule()
-    model3 = TestModule()
-
-    state_dict = model1.state_dict()
-    for kw in state_dict.keys():
-        state_dict[kw] = torch.randn_like(state_dict[kw])
-    
-    model1.load_state_dict(state_dict)
-    model2.load_state_dict(state_dict)
-    model3.load_state_dict(state_dict)
-
-    model1 = model1.cuda().half()
-    model2 = model2.cuda().half()
-    model3 = model3.cuda()
-    
-    opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)
-    opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), weight_decay=1e-3)
-    opt3 = torch.optim.Adam(model3.parameters(), weight_decay=1e-3)
-
-    for _ in range(100):
-        opt1.zero_grad()
-        opt2.zero_grad()
-        opt3.zero_grad()
-
-        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
-            grad = torch.randn_like(p1).half()
-            p1.grad = grad
-            p2.grad = grad
-            p3.grad = grad.float()
-        
-        opt1.step()
-        opt2.step()
-        opt3.step()
-
-        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
-            diff1 = torch.abs(p1.float() - p2.float()).max()
-            diff2 = torch.abs(p1.float() - p3).max()
-            diff3 = torch.abs(p2.float() - p3).max()
-            print(diff1, diff2, diff3)
-            assert_lt(diff1, 1)
-            assert_lt(diff2, 1)
-            assert_lt(diff3, 1)
-
-if __name__ == "__main__":
-    main()

From d008d7fc9917a82bcf9782de03c3a6ce948b5a3f Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Thu, 10 Aug 2023 19:42:09 +0800
Subject: [PATCH 21/27] refactor adam_offload

---
 bmtrain/optim/_function.py    |  33 +++++------
 bmtrain/optim/adam_offload.py | 102 ++++++++++++----------------------
 csrc/cuda/cross_entropy.cu    |   1 -
 csrc/include/adam_cpu.hpp     |  91 +++++++++++++-----------------
 tests/test_optim.py           |  14 ++---
 5 files changed, 98 insertions(+), 143 deletions(-)

diff --git a/bmtrain/optim/_function.py b/bmtrain/optim/_function.py
index 26cab409..37228a35 100644
--- a/bmtrain/optim/_function.py
+++ b/bmtrain/optim/_function.py
@@ -10,8 +10,8 @@ def adam_cpu(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.T
     assert m_fp32.is_contiguous(), "m_fp32 must be contiguous"
     assert v_fp32.is_contiguous(), "v_fp32 must be contiguous"
     assert param_fp32.dtype == torch.float32, "param_fp32 must be float32 tensor"
-    assert param_fp16.dtype == torch.float16, "param_fp16 must be float16 tensor"
-    assert g_fp16.dtype == torch.float16, "g_fp16 must be float16 tensor"
+    assert param_fp16.dtype == torch.float16 or param_fp16.dtype == torch.bfloat16, "param_fp16 must be float16/bfloat16 tensor"
+    assert g_fp16.dtype == torch.float16 or g_fp16.dtype == torch.bfloat16, "g_fp16 must be float16/bfloat16 tensor"
     assert m_fp32.dtype == torch.float32, "m_fp32 must be float32 tensor"
     assert v_fp32.dtype == torch.float32, "v_fp32 must be float32 tensor"
     assert param_fp32.device == torch.device("cpu"), "param_fp32 must be a cpu tensor"
@@ -25,20 +25,21 @@ def adam_cpu(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.T
     assert param_fp32.numel() == v_fp32.numel(), "param_fp32 and v_fp32 must have the same number of elements"
     bias_correction1 = 1 - beta1 ** step
     bias_correction2 = 1 - beta2 ** step
-    C.adam_cpu_launcher(
-                    param_fp32.numel(),
-                    param_fp32.data_ptr(),
-                    param_fp16.data_ptr(),
-                    g_fp16.data_ptr(),
-                    m_fp32.data_ptr(),
-                    v_fp32.data_ptr(),
-                    beta1, beta2,
-                    eps, lr,
-                    scale,
-                    weight_decay,
-                    bias_correction1,
-                    bias_correction2,
-                )
+    launcher = C.adam_cpu_launcher if g_fp16.dtype == torch.float16 else C.adam_cpu_bf16_launcher
+    launcher(
+        param_fp32.numel(),
+        param_fp32.data_ptr(),
+        param_fp16.data_ptr(),
+        g_fp16.data_ptr(),
+        m_fp32.data_ptr(),
+        v_fp32.data_ptr(),
+        beta1, beta2,
+        eps, lr,
+        scale,
+        weight_decay,
+        bias_correction1,
+        bias_correction2,
+    )
 
 def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tensor, m_fp16: torch.Tensor,
              v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
diff --git a/bmtrain/optim/adam_offload.py b/bmtrain/optim/adam_offload.py
index 5e73f918..5408f12c 100644
--- a/bmtrain/optim/adam_offload.py
+++ b/bmtrain/optim/adam_offload.py
@@ -54,7 +54,7 @@ def step(self, closure=None, scale=1):
                 if p.grad is not None and p.requires_grad:
                     if p.grad.is_sparse:
                         raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
-                    if p.dtype not in [torch.float16, torch.float32, torch.bfloat16]:
+                    if p.dtype not in [torch.float32, torch.float16, torch.bfloat16]:
                         raise RuntimeError('Adam only supports fp32, fp16 and bf16 gradients')
 
                     state = self.state[p]
@@ -66,28 +66,19 @@ def step(self, closure=None, scale=1):
                         # Exponential moving average of squared gradient values
                         state['exp_avg_sq'] = torch.zeros(p.size(), dtype=torch.float32, device="cpu")      # on host
 
-                        if p.dtype == torch.half:
-                            state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, device="cpu")     # on host
-                            state['_param_fp32'].copy_(p)
-
-                            # placeholder
-                            state["_param_fp16"] = torch.empty(p.size(), dtype=torch.float16, pin_memory=True)  # on host
-                            state["_grad_fp16"] = torch.empty(p.size(), dtype=torch.float16, pin_memory=True)   # on host
-                            
-                        elif p.dtype == torch.bfloat16:
-                            state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, device="cpu")     # on host
+                        if p.dtype == torch.float32:
+                            state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, pin_memory=True)     # on host
                             state['_param_fp32'].copy_(p)
 
                             # placeholder
-                            state["_param_bf16"] = torch.empty(p.size(), dtype=torch.bfloat16, pin_memory=True)  # on host
-                            state["_grad_bf16"] = torch.empty(p.size(), dtype=torch.bfloat16, pin_memory=True)   # on host
-                                                        
+                            state["_grad_fp32"] = torch.empty(p.size(), dtype=torch.float32, pin_memory=True)   # on host
                         else:
-                            state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, pin_memory=True)     # on host
+                            state['_param_fp32'] = torch.empty(p.size(), dtype=torch.float32, device="cpu")     # on host
                             state['_param_fp32'].copy_(p)
 
                             # placeholder
-                            state["_grad_fp32"] = torch.empty(p.size(), dtype=torch.float32, pin_memory=True)   # on host
+                            state["_param_fp16"] = torch.empty(p.size(), dtype=p.dtype, pin_memory=True)  # on host
+                            state["_grad_fp16"] = torch.empty(p.size(), dtype=p.dtype, pin_memory=True)   # on host
 
                     if p not in self._events:
                         self._events[p] = torch.cuda.Event()
@@ -96,12 +87,10 @@ def step(self, closure=None, scale=1):
 
         # transfer parameters to host asynchronously
         for param, state, event, _, _, _, _, _ in update_params:
-            if param.dtype == torch.half:
-                state["_grad_fp16"].copy_(param.grad, non_blocking=True)
-            elif param.dtype == torch.bfloat16:
-                state ["_grad_bf16"].copy_(param.grad, non_blocking=True)
-            else:
+            if param.dtype == torch.float32:
                 state["_grad_fp32"].copy_(param.grad, non_blocking=True)
+            else:
+                state["_grad_fp16"].copy_(param.grad, non_blocking=True)
             torch.cuda.current_stream().record_event(event)
 
         for param, state, event, beta1, beta2, eps, lr, weight_decay in update_params:
@@ -111,43 +100,7 @@ def step(self, closure=None, scale=1):
             state["step"] += 1
 
             # update parameters
-            if param.dtype == torch.half:
-                if ('maximize' in group) and (group['maximize'] is True):
-                    grad = -state["_grad_fp16"]
-                else:
-                    grad = state["_grad_fp16"]
-                F.adam_cpu(
-                    state["_param_fp32"].view(-1),
-                    state["_param_fp16"].view(-1),
-                    grad.view(-1),
-                    state["exp_avg"].view(-1),
-                    state["exp_avg_sq"].view(-1),
-                    beta1, beta2,
-                    eps,  0.0 if state["step"] <= self._hold_steps else lr,
-                    scale,
-                    weight_decay,
-                    state["step"]
-                )
-                # transfer parameters back to device asynchronously
-                param.copy_(state["_param_fp16"], non_blocking=True)
-            elif param.dtype == torch.bfloat16:
-                if ('maximize' in group) and (group['maximize'] is True):
-                    grad = -state["_grad_bf16"]
-                else:
-                    grad = state["_grad_bf16"]
-                F.adam_cpu_bf16(
-                    state["_param_fp32"].view(-1),
-                    state["_param_bf16"].view(-1),
-                    grad.view(-1),
-                    state["exp_avg"].view(-1),
-                    state["exp_avg_sq"].view(-1),
-                    beta1, beta2,
-                    eps,  0.0 if state["step"] <= self._hold_steps else lr,
-                    scale,
-                    weight_decay,
-                    state["step"]                      
-                )
-            else:
+            if param.dtype == torch.float32:
                 state["_grad_fp32"].mul_(1.0 / scale)
                 if ('maximize' in group) and (group['maximize'] is True):
                     grad = -state["_grad_fp32"]
@@ -174,6 +127,25 @@ def step(self, closure=None, scale=1):
                 )
                 # transfer parameters back to device asynchronously
                 param.copy_(state["_param_fp32"], non_blocking=True)
+            else:
+                if ('maximize' in group) and (group['maximize'] is True):
+                    grad = -state["_grad_fp16"]
+                else:
+                    grad = state["_grad_fp16"]
+                F.adam_cpu(
+                    state["_param_fp32"].view(-1),
+                    state["_param_fp16"].view(-1),
+                    grad.view(-1),
+                    state["exp_avg"].view(-1),
+                    state["exp_avg_sq"].view(-1),
+                    beta1, beta2,
+                    eps,  0.0 if state["step"] <= self._hold_steps else lr,
+                    scale,
+                    weight_decay,
+                    state["step"]
+                )
+                # transfer parameters back to device asynchronously
+                param.copy_(state["_param_fp16"], non_blocking=True)
 
         return loss
 
@@ -221,18 +193,14 @@ def load_state_dict(self, state_dict: dict) -> None:
                         v[name] = v[name].to("cpu").to(dtype)
 
                 state[param] = v
-                if param.dtype == torch.half:
-                    # initialize placeholders
-                    state[param]["_param_fp16"] = torch.empty(param.size(), dtype=torch.float16, pin_memory=True)  # on host
-                    state[param]["_grad_fp16"] = torch.empty(param.size(), dtype=torch.float16, pin_memory=True)   # on host
-                elif param.dtype == torch.bfloat16:
-                    #initialize placeholders
-                    state[param]["_param_bf16"] = torch.empty(param.size(), dtype=torch.bfloat16, pin_memory=True)  # on host
-                    state[param]["_grad_bf16"] = torch.empty(param.size(), dtype=torch.bfloat16, pin_memory=True)   # on host
-                else:
+                if param.dtype == torch.float32:
                     state[param]["_param_fp32"] = state[param]["_param_fp32"].pin_memory()  # on host
                     # initialize placeholders
                     state[param]["_grad_fp32"] = torch.empty(param.size(), dtype=torch.float32, pin_memory=True)   # on host
+                else:
+                    # initialize placeholders
+                    state[param]["_param_fp16"] = torch.empty(param.size(), dtype=torch.float16, pin_memory=True)  # on host
+                    state[param]["_grad_fp16"] = torch.empty(param.size(), dtype=torch.float16, pin_memory=True)   # on host
             else:
                 state[k] = v
 
diff --git a/csrc/cuda/cross_entropy.cu b/csrc/cuda/cross_entropy.cu
index c7644e2b..c0b742ac 100644
--- a/csrc/cuda/cross_entropy.cu
+++ b/csrc/cuda/cross_entropy.cu
@@ -1,5 +1,4 @@
 #include <cuda_fp16.h>
-#include <cuda_bf16.h>
 #include "reduce.cuh"
 #include <cstdint>
 #include<cuda_runtime_api.h>
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index 480ac488..83e93ff3 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -2,11 +2,9 @@
 #include <immintrin.h>
 #include <cmath>
 #include <cstdint>
-#include<sched.h>
+#include <sched.h>
 #include <pybind11/pybind11.h>
 #include <iostream>
-#include <cuda_fp16.h>
-#include <cuda_bf16.h>
 #include <vector>
 #include <thread>
 #include <algorithm>
@@ -84,68 +82,56 @@ inline uint16_t fp16_ieee_from_fp32_value(float f) {
 
     float base = (fabsf(f) * scale_to_inf) * scale_to_zero;
 
-          const uint32_t w = (uint32_t)fp32_to_bits(f);
-          const uint32_t shl1_w = w + w;
-          const uint32_t sign = w & UINT32_C(0x80000000);
-          uint32_t bias = shl1_w & UINT32_C(0xFF000000);
-          if (bias < UINT32_C(0x71000000)) {
-                  bias = UINT32_C(0x71000000);
-          }
-
-          base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
-          const uint32_t bits = (uint32_t)fp32_to_bits(base);
-          const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
-          const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
-          const uint32_t nonsign = exp_bits + mantissa_bits;
-          return static_cast<uint16_t>(
-            (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
-          );
-  }
-
-// fp32 -> bf16
-inline float bf16_ieee_from_fp32_value(float f){
-    uint32_t w = fp32_to_bits(f);
-    uint32_t sign = w & UINT32_C(0x80000000);
-    uint32_t exp = w & UINT32_C(0x7F800000);
-    uint32_t mantissa = w & UINT32_C(0x007FFFFF);
-    uint16_t bf16 = static_cast<uint16_t>(sign >> 16) | static_cast<uint16_t>(exp >> 16) | static_cast<uint16_t>(mantissa >> 16);
-    return fp32_from_bits(static_cast<uint32_t>(bf16) << 16);
-} 
+    const uint32_t w = (uint32_t)fp32_to_bits(f);
+    const uint32_t shl1_w = w + w;
+    const uint32_t sign = w & UINT32_C(0x80000000);
+    uint32_t bias = shl1_w & UINT32_C(0xFF000000);
+    if (bias < UINT32_C(0x71000000)) {
+        bias = UINT32_C(0x71000000);
+    }
 
-// bf16 -> fp32
-inline float bf16_ieee_to_fp32_value(u_int16_t h){
-    uint32_t w = static_cast<uint32_t>(h) << 16;
-    uint32_t sign = w & UINT32_C(0x80000000);
-    uint32_t exp = w & UINT32_C(0x7F800000);
-    uint32_t mantissa = w & UINT32_C(0x007FFFFF);
-    uint32_t fp32 = sign << 16 | exp << 16 | mantissa << 16;
-    return fp32_from_bits(fp32);
+    base = fp32_from_bits((bias >> 1) + UINT32_C(0x07800000)) + base;
+    const uint32_t bits = (uint32_t)fp32_to_bits(base);
+    const uint32_t exp_bits = (bits >> 13) & UINT32_C(0x00007C00);
+    const uint32_t mantissa_bits = bits & UINT32_C(0x00000FFF);
+    const uint32_t nonsign = exp_bits + mantissa_bits;
+    return static_cast<uint16_t>(
+        (sign >> 16) | (shl1_w > UINT32_C(0xFF000000) ? UINT16_C(0x7E00) : nonsign)
+    );
 }
 
 // fp16 -> fp32 
 inline float fp16_ieee_to_fp32_value(uint16_t h) {
-
   const uint32_t w = (uint32_t)h << 16;
   const uint32_t sign = w & UINT32_C(0x80000000);
   const uint32_t two_w = w + w;
 
   const uint32_t exp_offset = UINT32_C(0xE0) << 23;
   const float exp_scale = 0x1.0p-112f;
-  const float normalized_value =
-      fp32_from_bits((two_w >> 4) + exp_offset) * exp_scale;
+  const float normalized_value = fp32_from_bits((two_w >> 4) + exp_offset) * exp_scale;
 
   const uint32_t magic_mask = UINT32_C(126) << 23;
   const float magic_bias = 0.5f;
-  const float denormalized_value =
-      fp32_from_bits((two_w >> 17) | magic_mask) - magic_bias;
+  const float denormalized_value = fp32_from_bits((two_w >> 17) | magic_mask) - magic_bias;
 
   const uint32_t denormalized_cutoff = UINT32_C(1) << 27;
   const uint32_t result =
-      sign | (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value)
-                                          : fp32_to_bits(normalized_value));
-  return  fp32_from_bits(result);
+        sign | (two_w < denormalized_cutoff ? fp32_to_bits(denormalized_value)
+                                            : fp32_to_bits(normalized_value));
+  return fp32_from_bits(result);
 }
 
+// fp32 -> bf16
+inline uint16_t bf16_from_fp32_value(float f){
+  return *reinterpret_cast<uint32_t*>(&f) >> 16;
+} 
+
+// bf16 -> fp32
+inline float bf16_to_fp32_value(uint16_t h){
+    uint32_t src = h;
+    src <<= 16;
+    return *reinterpret_cast<float*>(&src);
+}
 
 void adam_cpu_0(
     int64_t n,
@@ -179,7 +165,7 @@ void adam_cpu_0(
                     }
             break; // must break here
         }
-        });
+    });
 }
 
 void adam_cpu_bf16_0(
@@ -200,7 +186,7 @@ void adam_cpu_bf16_0(
     parallel_for(0, n, 0, [&](int64_t start, int64_t end) {
     for (int64_t j = start; j < end; j += span) {
         for (int64_t i = j; i < end; i++) {
-            float g = bf16_ieee_to_fp32_value(g_bf16_ptr[i]) / scale;
+            float g = bf16_to_fp32_value(g_bf16_ptr[i]) / scale;
             float m = m_fp32_ptr[i];
             float v = v_fp32_ptr[i];
             float p = param_fp32_ptr[i];
@@ -208,13 +194,13 @@ void adam_cpu_bf16_0(
             v = beta2 * v + (1 - beta2) * g * g;
             p = p - lr * m  / bias_correction1 / (sqrtf(v / bias_correction2) + eps) - lr * weight_decay * p;
             param_fp32_ptr[i] = p;
-            param_bf16_ptr[i] = bf16_ieee_from_fp32_value(p);
+            param_bf16_ptr[i] = bf16_from_fp32_value(p);
             m_fp32_ptr[i] = m;
             v_fp32_ptr[i] = v;
                     }
             break; // must break here
         }
-        });
+    });
 }
 
 static void __attribute__ ((__target__ ("avx,fma,f16c"))) adam_cpu_1(
@@ -280,7 +266,8 @@ static void __attribute__ ((__target__ ("avx,fma,f16c"))) adam_cpu_1(
                 _mm256_storeu_ps(&m_fp32_ptr[j], m);
                 _mm256_storeu_ps(&v_fp32_ptr[j], v);
             }
-    }});
+        }
+    });
 }
 
 static void __attribute__ ((__target__ ("avx512f"))) adam_cpu_2(
@@ -350,7 +337,7 @@ static void __attribute__ ((__target__ ("avx512f"))) adam_cpu_2(
                 _mm512_storeu_ps(&v_fp32_ptr[j], v);
             }
         }
-        });
+    });
 }
 
 void adam_cpu_launcher(
diff --git a/tests/test_optim.py b/tests/test_optim.py
index 6ee5e0fc..9b0b6cfa 100644
--- a/tests/test_optim.py
+++ b/tests/test_optim.py
@@ -37,9 +37,9 @@ def main(dtype):
     model2 = model2.cuda().to(dtype)
     model3 = model3.cuda()
 
-    opt1 = bmt.optim.AdamOptimizer(model1.parameters(), weight_decay=1e-3)
-    opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), weight_decay=1e-3)
-    opt3 = torch.optim.Adam(model3.parameters(), weight_decay=1e-3)
+    opt1 = bmt.optim.AdamOptimizer(model1.parameters(), lr=1)
+    opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), lr=1)
+    opt3 = torch.optim.Adam(model3.parameters(), lr=1)
 
     for _ in range(100):
         opt1.zero_grad()
@@ -57,10 +57,10 @@ def main(dtype):
         opt3.step()
 
         for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
-            diff1 = torch.abs(p1 - p2).max() 
-            diff2 = torch.abs(p1 - p3).max()
-            diff3 = torch.abs(p2 - p3).max()
-            print(diff1, diff2, diff3)
+            diff1 = torch.abs(p1 - p2).max().item() 
+            diff2 = torch.abs(p1 - p3).max().item()
+            diff3 = torch.abs(p2 - p3).max().item()
+            print(f"{diff1:4.6f}, {diff2:4.6f}, {diff3:4.6f}")
             assert_lt(diff1, 1)
             assert_lt(diff2, 1)
             assert_lt(diff3, 1)

From 145d90f716d0e3fab5276f6faef0c6b5a7072c26 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 11 Aug 2023 00:31:40 +0800
Subject: [PATCH 22/27] refactor adam

---
 bmtrain/optim/_function.py    | 41 +---------------------
 bmtrain/optim/adam.py         | 66 ++++++++++++++---------------------
 bmtrain/optim/adam_offload.py |  9 ++---
 csrc/include/adam_cpu.hpp     | 56 ++++++++++++++---------------
 tests/test_optim.py           | 27 +++++++++++---
 5 files changed, 83 insertions(+), 116 deletions(-)

diff --git a/bmtrain/optim/_function.py b/bmtrain/optim/_function.py
index 37228a35..b5ad60b1 100644
--- a/bmtrain/optim/_function.py
+++ b/bmtrain/optim/_function.py
@@ -41,7 +41,7 @@ def adam_cpu(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.T
         bias_correction2,
     )
 
-def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tensor, m_fp16: torch.Tensor,
+def adam_fp16(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tensor, m_fp16: torch.Tensor,
              v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
              weight_decay: float, step: int) -> None:
     assert CHECK_INPUT(param_fp32), "param_fp32 must be contiguous and on cuda"
@@ -76,45 +76,6 @@ def adam(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.Tenso
         bias_correction2,
         stream
     )
-
-def adam_cpu_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.Tensor, m_fp32: torch.Tensor,
-             v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
-             weight_decay: float, step: int) -> None:
-    assert param_fp32.is_contiguous(), "param_fp32 must be contiguous"
-    assert param_bf16.is_contiguous(), "param_bf16 must be contiguous"
-    assert g_bf16.is_contiguous(), "g_bf16 must be contiguous"
-    assert m_fp32.is_contiguous(), "m_fp32 must be contiguous"
-    assert v_fp32.is_contiguous(), "v_fp32 must be contiguous"
-    assert param_fp32.dtype == torch.float32, "param_fp32 must be float32 tensor"
-    assert param_bf16.dtype == torch.bfloat16, "param_bf16 must be bfloat16 tensor"
-    assert g_bf16.dtype == torch.bfloat16, "g_bf16 must be bfloat16 tensor"
-    assert m_fp32.dtype == torch.float32, "m_fp32 must be float32 tensor"
-    assert v_fp32.dtype == torch.float32, "v_fp32 must be float32 tensor"
-    assert param_fp32.device == torch.device("cpu"), "param_fp32 must be a cpu tensor"
-    assert param_bf16.device == torch.device("cpu"), "param_bf16 must be a cpu tensor"
-    assert g_bf16.device == torch.device("cpu"), "g_bf16 must be a cpu tensor"
-    assert m_fp32.device == torch.device("cpu"), "m_fp32 must be a cpu tensor"
-    assert v_fp32.device == torch.device("cpu"), "v_fp32 must be a cpu tensor"
-    assert param_fp32.numel() == param_bf16.numel(), "param_fp32 and param_bf16 must have the same number of elements"
-    assert param_fp32.numel() == g_bf16.numel(), "param_fp32 and g_bf16 must have the same number of elements"
-    assert param_fp32.numel() == m_fp32.numel(), "param_fp32 and m_fp32 must have the same number of elements"
-    assert param_fp32.numel() == v_fp32.numel(), "param_fp32 and v_fp32 must have the same number of elements"
-    bias_correction1 = 1 - beta1 ** step
-    bias_correction2 = 1 - beta2 ** step
-    C.adam_cpu_bf16_launcher(
-                    param_fp32.numel(),
-                    param_fp32.data_ptr(),
-                    param_bf16.data_ptr(),
-                    g_bf16.data_ptr(),
-                    m_fp32.data_ptr(),
-                    v_fp32.data_ptr(),
-                    beta1, beta2,
-                    eps, lr,
-                    scale,
-                    weight_decay,
-                    bias_correction1,
-                    bias_correction2,
-                )
     
 def adam_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.Tensor, m_fp32: torch.Tensor,
              v_fp32: torch.Tensor, beta1: float, beta2: float, eps: float, lr: float, scale: float,
diff --git a/bmtrain/optim/adam.py b/bmtrain/optim/adam.py
index c9482f6a..a3138980 100644
--- a/bmtrain/optim/adam.py
+++ b/bmtrain/optim/adam.py
@@ -40,10 +40,7 @@ def _on_justify_scale(self, old_scale, new_scale):
                 if p in self.state:
                     state = self.state[p]
                     if len(state) > 0:
-                        #if p belongs to bf16, do not justify scale  
-                        if p.dtype == torch.bfloat16:
-                            continue
-                        else:                 
+                        if p.dtype == torch.float16:
                             state['exp_avg'] *= delta
                             state['exp_avg_sq'] *= delta
 
@@ -75,10 +72,10 @@ def step(self, closure=None, scale=1):
                     if len(state) == 0:
                         state['step'] = 0
                         # Exponential moving average of gradient values
-                        if p.dtype == torch.bfloat16:
-                            state['exp_avg'] = torch.zeros(p.size(), dtype=torch.float32, device=p.device) # on device
+                        if p.dtype == torch.float16:
+                            state['exp_avg'] = torch.zeros(p.size(), dtype=torch.float16, device=p.device) # on device
                         else:
-                            state['exp_avg'] = torch.zeros(p.size(), dtype=p.dtype, device=p.device) # on device
+                            state['exp_avg'] = torch.zeros(p.size(), dtype=torch.float32, device=p.device) # on device
                         # Exponential moving average of squared gradient values
                         state['exp_avg_sq'] = torch.zeros(p.size(), dtype=torch.float32, device=p.device)# on device
                   
@@ -87,42 +84,12 @@ def step(self, closure=None, scale=1):
                             state['_param_fp32'].copy_(p)
 
                     # update the steps for each param group update
-                    state['step'] += 1
-
                     if ('maximize' in group) and (group['maximize'] is True):
                         grad = -p.grad
                     else:
                         grad = p.grad
                         
-                    if p.dtype == torch.half:
-                            F.adam(
-                            state["_param_fp32"],    # fp32
-                            p,                      # fp16
-                            grad,                 # fp16
-                            state['exp_avg'],       # fp16: m
-                            state["exp_avg_sq"],    # fp32: v
-                            group['betas'][0], group['betas'][1],
-                            group['eps'],
-                            0.0 if state["step"] <= self._hold_steps else group['lr'],
-                            scale,
-                            group['weight_decay'],
-                            state['step']
-                        )
-                    elif p.dtype == torch.bfloat16:
-                            F.adam_bf16(
-                            state["_param_fp32"],    # fp32
-                            p,                      # bf16
-                            grad,                 # bf16
-                            state['exp_avg'],       # fp32: m
-                            state["exp_avg_sq"],    # fp32: v
-                            group['betas'][0], group['betas'][1],
-                            group['eps'],
-                            0.0 if state["step"] <= self._hold_steps else group['lr'],
-                            scale,
-                            group['weight_decay'],
-                            state['step']
-                        )
-                    else:
+                    if p.dtype == torch.float32:
                         other_kwargs = {}
                         if 'maximize' in inspect.signature(torch.optim._functional.adam).parameters:
                             other_kwargs['maximize'] = False
@@ -137,11 +104,30 @@ def step(self, closure=None, scale=1):
                             amsgrad=False,
                             beta1=group['betas'][0],
                             beta2=group['betas'][1],
-                            lr=0.0 if state["step"] <= self._hold_steps else group['lr'],
+                            lr=0.0 if state["step"] < self._hold_steps else group['lr'],
                             weight_decay=group['weight_decay'],
                             eps=group['eps'],
                             **other_kwargs
                         )
+                        state['step'] += 1
+                    else:
+                        f = F.adam_fp16 if p.dtype == torch.float16 else F.adam_bf16
+                        state['step'] += 1
+                        f(
+                            state["_param_fp32"],    # fp32
+                            p,                      # fp16
+                            grad,                 # fp16
+                            state['exp_avg'],       # fp16: m
+                            state["exp_avg_sq"],    # fp32: v
+                            group['betas'][0], group['betas'][1],
+                            group['eps'],
+                            0.0 if state["step"] < self._hold_steps else group['lr'],
+                            scale,
+                            group['weight_decay'],
+                            state['step']
+                        )
+
+
 
         return loss
 
@@ -184,7 +170,7 @@ def load_state_dict(self, state_dict: dict) -> None:
                     v["_param_fp32"] = torch.empty(param.size(), dtype=torch.float32, device=param.device)
                     v["_param_fp32"].copy_(param)
 
-                for name, dtype in [("exp_avg", torch.float32 if param.dtype == torch.bfloat16 else param.dtype), ("exp_avg_sq", torch.float32), ("_param_fp32", torch.float32)]:
+                for name, dtype in [("exp_avg", torch.float16 if param.dtype == torch.float16 else torch.float32), ("exp_avg_sq", torch.float32), ("_param_fp32", torch.float32)]:
                     if name in v:
                         v[name] = v[name].to(param.device).to(dtype)
 
diff --git a/bmtrain/optim/adam_offload.py b/bmtrain/optim/adam_offload.py
index 5408f12c..5b34a287 100644
--- a/bmtrain/optim/adam_offload.py
+++ b/bmtrain/optim/adam_offload.py
@@ -97,8 +97,6 @@ def step(self, closure=None, scale=1):
             # wait for transfer to host
             event.synchronize()
 
-            state["step"] += 1
-
             # update parameters
             if param.dtype == torch.float32:
                 state["_grad_fp32"].mul_(1.0 / scale)
@@ -120,14 +118,16 @@ def step(self, closure=None, scale=1):
                     amsgrad=False,
                     beta1=beta1,
                     beta2=beta2,
-                    lr=0.0 if state["step"] <= self._hold_steps else lr,
+                    lr=0.0 if state["step"] < self._hold_steps else lr,
                     weight_decay=weight_decay,
                     eps=eps,
                     **other_kwargs
                 )
                 # transfer parameters back to device asynchronously
                 param.copy_(state["_param_fp32"], non_blocking=True)
+                state["step"] += 1
             else:
+                state["step"] += 1
                 if ('maximize' in group) and (group['maximize'] is True):
                     grad = -state["_grad_fp16"]
                 else:
@@ -139,7 +139,7 @@ def step(self, closure=None, scale=1):
                     state["exp_avg"].view(-1),
                     state["exp_avg_sq"].view(-1),
                     beta1, beta2,
-                    eps,  0.0 if state["step"] <= self._hold_steps else lr,
+                    eps,  0.0 if state["step"] < self._hold_steps else lr,
                     scale,
                     weight_decay,
                     state["step"]
@@ -147,6 +147,7 @@ def step(self, closure=None, scale=1):
                 # transfer parameters back to device asynchronously
                 param.copy_(state["_param_fp16"], non_blocking=True)
 
+
         return loss
 
     def load_state_dict(self, state_dict: dict) -> None:
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index 83e93ff3..6ff901a5 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -149,20 +149,20 @@ void adam_cpu_0(
 ){
     int64_t span = 1;
     parallel_for(0, n, 0, [&](int64_t start, int64_t end) {
-    for (int64_t j = start; j < end; j += span) {
-        for (int64_t i = j; i < end; i++) {
-            float g = fp16_ieee_to_fp32_value(g_fp16_ptr[i]) / scale;
-            float m = m_fp32_ptr[i];
-            float v = v_fp32_ptr[i];
-            float p = param_fp32_ptr[i];
-            m = beta1 * m + (1 - beta1) * g;
-            v = beta2 * v + (1 - beta2) * g * g;
-            p = p - lr * m  / bias_correction1 / (sqrtf(v / bias_correction2) + eps) - lr * weight_decay * p;
-            param_fp32_ptr[i] = p;
-            param_fp16_ptr[i] = fp16_ieee_from_fp32_value(p);
-            m_fp32_ptr[i] = m;
-            v_fp32_ptr[i] = v;
-                    }
+        for (int64_t j = start; j < end; j += span) {
+            for (int64_t i = j; i < end; i++) {
+                float g = fp16_ieee_to_fp32_value(g_fp16_ptr[i]) / scale;
+                float m = m_fp32_ptr[i];
+                float v = v_fp32_ptr[i];
+                float p = param_fp32_ptr[i];
+                m = beta1 * m + (1 - beta1) * g;
+                v = beta2 * v + (1 - beta2) * g * g;
+                p = p - lr * m  / bias_correction1 / (sqrtf(v / bias_correction2) + eps) - lr * weight_decay * p;
+                param_fp32_ptr[i] = p;
+                param_fp16_ptr[i] = fp16_ieee_from_fp32_value(p);
+                m_fp32_ptr[i] = m;
+                v_fp32_ptr[i] = v;
+            }
             break; // must break here
         }
     });
@@ -184,20 +184,20 @@ void adam_cpu_bf16_0(
 ){
     int64_t span = 1;
     parallel_for(0, n, 0, [&](int64_t start, int64_t end) {
-    for (int64_t j = start; j < end; j += span) {
-        for (int64_t i = j; i < end; i++) {
-            float g = bf16_to_fp32_value(g_bf16_ptr[i]) / scale;
-            float m = m_fp32_ptr[i];
-            float v = v_fp32_ptr[i];
-            float p = param_fp32_ptr[i];
-            m = beta1 * m + (1 - beta1) * g;
-            v = beta2 * v + (1 - beta2) * g * g;
-            p = p - lr * m  / bias_correction1 / (sqrtf(v / bias_correction2) + eps) - lr * weight_decay * p;
-            param_fp32_ptr[i] = p;
-            param_bf16_ptr[i] = bf16_from_fp32_value(p);
-            m_fp32_ptr[i] = m;
-            v_fp32_ptr[i] = v;
-                    }
+        for (int64_t j = start; j < end; j += span) {
+            for (int64_t i = j; i < end; i++) {
+                float g = bf16_to_fp32_value(g_bf16_ptr[i]) / scale;
+                float m = m_fp32_ptr[i];
+                float v = v_fp32_ptr[i];
+                float p = param_fp32_ptr[i];
+                m = beta1 * m + (1 - beta1) * g;
+                v = beta2 * v + (1 - beta2) * g * g;
+                p = p - lr * m  / bias_correction1 / (sqrtf(v / bias_correction2) + eps) - lr * weight_decay * p;
+                param_fp32_ptr[i] = p;
+                param_bf16_ptr[i] = bf16_from_fp32_value(p);
+                m_fp32_ptr[i] = m;
+                v_fp32_ptr[i] = v;
+            }
             break; // must break here
         }
     });
diff --git a/tests/test_optim.py b/tests/test_optim.py
index 9b0b6cfa..38c83af6 100644
--- a/tests/test_optim.py
+++ b/tests/test_optim.py
@@ -5,7 +5,7 @@
 class TestModule(torch.nn.Module):
     def __init__(self):
         super(TestModule, self).__init__()
-        self.fc1 = torch.nn.Linear(128, 128)
+        self.fc1 = torch.nn.Linear(128, 128, bias=False)
         self.fc2 = torch.nn.Linear(128, 128)
         self.fc3 = torch.nn.Linear(128, 128)
         self.fc4 = torch.nn.Linear(128, 128)
@@ -24,6 +24,8 @@ def main(dtype):
     model1 = TestModule()
     model2 = TestModule()
     model3 = TestModule()
+    model4 = TestModule()
+    model5 = TestModule()
 
     state_dict = model1.state_dict()
     for kw in state_dict.keys():
@@ -32,38 +34,55 @@ def main(dtype):
     model1.load_state_dict(state_dict)
     model2.load_state_dict(state_dict)
     model3.load_state_dict(state_dict)
+    model4.load_state_dict(state_dict)
+    model5.load_state_dict(state_dict)
 
     model1 = model1.cuda().to(dtype)
     model2 = model2.cuda().to(dtype)
     model3 = model3.cuda()
+    model4 = model4.cuda()
+    model5 = model5.cuda()
 
     opt1 = bmt.optim.AdamOptimizer(model1.parameters(), lr=1)
     opt2 = bmt.optim.AdamOffloadOptimizer(model2.parameters(), lr=1)
     opt3 = torch.optim.Adam(model3.parameters(), lr=1)
+    opt4 = bmt.optim.AdamOptimizer(model4.parameters(), lr=1)
+    opt5 = bmt.optim.AdamOffloadOptimizer(model5.parameters(), lr=1)
 
     for _ in range(100):
         opt1.zero_grad()
         opt2.zero_grad()
         opt3.zero_grad()
+        opt4.zero_grad()
+        opt5.zero_grad()
 
-        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
+        for p1, p2, p3, p4, p5 in zip(model1.parameters(), model2.parameters(), model3.parameters(), model4.parameters(), model5.parameters()):
             grad = torch.randn_like(p1)
             p1.grad = grad.to(dtype)
             p2.grad = grad.to(dtype)
             p3.grad = grad.float()
+            p4.grad = grad.float()
+            p5.grad = grad.float()
 
         opt1.step()
         opt2.step()
         opt3.step()
+        opt4.step()
+        opt5.step()
+        torch.cuda.synchronize()
 
-        for p1, p2, p3 in zip(model1.parameters(), model2.parameters(), model3.parameters()):
+        for p1, p2, p3, p4, p5 in zip(model1.parameters(), model2.parameters(), model3.parameters(), model4.parameters(), model5.parameters()):
             diff1 = torch.abs(p1 - p2).max().item() 
             diff2 = torch.abs(p1 - p3).max().item()
             diff3 = torch.abs(p2 - p3).max().item()
-            print(f"{diff1:4.6f}, {diff2:4.6f}, {diff3:4.6f}")
+            diff4 = torch.abs(p3 - p4).max().item()
+            diff5 = torch.abs(p3 - p5).max().item()
+            print(f"{diff1:.6f}, {diff2:.6f}, {diff3:.6f}, {diff4:.6f}, {diff5:.6f}")
             assert_lt(diff1, 1)
             assert_lt(diff2, 1)
             assert_lt(diff3, 1)
+            assert_eq(diff4, 0)
+            assert_lt(diff5, 0.00001)
 
 if __name__ == "__main__":
     main(torch.float16)

From da14e7e7eed788be83f074f0de7b1114b483f91a Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 11 Aug 2023 09:17:23 +0800
Subject: [PATCH 23/27] fix adam_cuda

---
 csrc/cuda/adam_cuda.cu   | 34 ++++++++++++++++++----------------
 csrc/cuda/has_inf_nan.cu |  2 +-
 tests/test_optim.py      | 21 +++++++++++----------
 3 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 15e2f805..9e701881 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -1,6 +1,8 @@
+#include <cstdint>
 #include <cuda_fp16.h>
+#if __CUDA_ARCH__ >= 800
 #include <cuda_bf16.h>
-#include <cstdint>
+#endif
 
 namespace {
 // blocks <n // 1024>,      threads<min(n, 1024)>
@@ -9,8 +11,8 @@ __global__ void adam_fp32_accum(
     const half *g,        // (n)
     half *m,        // (n)
     float *v,        // (n)
-    float* param,   // (n)
-    half* param_h,  // (n)
+    float *param,   // (n)
+    half *param_h,  // (n)
     float beta1,
     float beta2,
     float eps,
@@ -37,11 +39,11 @@ __global__ void adam_fp32_accum(
 
 __global__ void adam_fp32_accum_bf16(
     int32_t n,
-    const nv_bfloat16 *g,        // (n)
-    nv_bfloat16 *m,        // (n)
+    const std::uintptr_t g_ptr,        // (n)
+    float *m,        // (n)
     float *v,        // (n)
-    float* param,   // (n)
-    nv_bfloat16* param_h,  // (n)
+    float *param,   // (n)
+    std::uintptr_t param_h_ptr,  // (n)
     float beta1,
     float beta2,
     float eps,
@@ -51,21 +53,23 @@ __global__ void adam_fp32_accum_bf16(
     float bias_correction1,
     float bias_correction2
 ) {
+#if __CUDA_ARCH__ >= 800
+    const __nv_bfloat16* g = reinterpret_cast<const __nv_bfloat16*>(g_ptr);
+    __nv_bfloat16* param_h = reinterpret_cast<__nv_bfloat16*>(param_h_ptr);
     int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
-
     if (col < n) {
         float local_g = __bfloat162float(g[col]) / scale; // real_g
-        float local_m = beta1 * __bfloat162float(m[col]) + (1 - beta1) * local_g; // real_m
+        float local_m = beta1 * m[col] + (1 - beta1) * local_g; // real_m
         float local_v = beta2 * v[col] + (1 - beta2) * local_g * local_g; // real_v
         float local_p = param[col];
-        local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v / bias_correction2 / scale) + eps) - lr * weight_decay * local_p; 
+        local_p = local_p - lr * local_m / bias_correction1 / (sqrtf(local_v / bias_correction2) + eps) - lr * weight_decay * local_p; 
     
         param_h[col] = __float2bfloat16(local_p);
         param[col] = local_p;
         v[col] = local_v;
-        m[col] = __float2bfloat16(local_m); 
+        m[col] = local_m; 
     }
-       
+#endif
 }
 
 }
@@ -113,13 +117,11 @@ void adam_bf16_launcher(
     uintptr_t stream
 ) {
     if (n <= 0) return;
-    auto g_ptr = reinterpret_cast<nv_bfloat16*>(g_bf16);
-    auto m_ptr = reinterpret_cast<nv_bfloat16*>(m_fp32);
-    auto param_h_ptr = reinterpret_cast<nv_bfloat16*>(param_bf16);
+    auto m_ptr = reinterpret_cast<float*>(m_fp32);
     auto param_fp32_ptr = reinterpret_cast<float*>(param_fp32);
     auto v_fp32_ptr = reinterpret_cast<float*>(v_fp32);
     int32_t threads = 1024;
     dim3 block_size = dim3(threads, 1, 1);
     dim3 grid_size = dim3((n + threads - 1) / threads, 1, 1);
-    adam_fp32_accum_bf16<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_ptr, m_ptr, v_fp32_ptr, param_fp32_ptr, param_h_ptr, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
+    adam_fp32_accum_bf16<<<grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_bf16, m_ptr, v_fp32_ptr, param_fp32_ptr, param_bf16, beta1, beta2, eps, lr, scale, weight_decay, bias_correction1, bias_correction2);
 }
diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index 9e73859e..b4c3113d 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -1,7 +1,7 @@
 #include <cstdint>
 #include <cuda_fp16.h>
 #if __CUDA_ARCH__ >= 800
-    #include <cuda_bf16.h>
+#include <cuda_bf16.h>
 #endif
 
 namespace{
diff --git a/tests/test_optim.py b/tests/test_optim.py
index 38c83af6..2888501e 100644
--- a/tests/test_optim.py
+++ b/tests/test_optim.py
@@ -49,12 +49,15 @@ def main(dtype):
     opt4 = bmt.optim.AdamOptimizer(model4.parameters(), lr=1)
     opt5 = bmt.optim.AdamOffloadOptimizer(model5.parameters(), lr=1)
 
+    optim_manager = bmt.optim.OptimManager(loss_scale=4)
+    optim_manager.add_optimizer(opt1)
+    optim_manager.add_optimizer(opt2)
+    optim_manager.add_optimizer(opt3)
+    optim_manager.add_optimizer(opt4)
+    optim_manager.add_optimizer(opt5)
+
     for _ in range(100):
-        opt1.zero_grad()
-        opt2.zero_grad()
-        opt3.zero_grad()
-        opt4.zero_grad()
-        opt5.zero_grad()
+        optim_manager.zero_grad()
 
         for p1, p2, p3, p4, p5 in zip(model1.parameters(), model2.parameters(), model3.parameters(), model4.parameters(), model5.parameters()):
             grad = torch.randn_like(p1)
@@ -64,11 +67,7 @@ def main(dtype):
             p4.grad = grad.float()
             p5.grad = grad.float()
 
-        opt1.step()
-        opt2.step()
-        opt3.step()
-        opt4.step()
-        opt5.step()
+        optim_manager.step()
         torch.cuda.synchronize()
 
         for p1, p2, p3, p4, p5 in zip(model1.parameters(), model2.parameters(), model3.parameters(), model4.parameters(), model5.parameters()):
@@ -85,5 +84,7 @@ def main(dtype):
             assert_lt(diff5, 0.00001)
 
 if __name__ == "__main__":
+    bmt.init_distributed()
     main(torch.float16)
+    print("==============================================================================")
     main(torch.bfloat16)

From 21d52186822696f8d501d810f8d7e14bf5a0c613 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 11 Aug 2023 10:22:14 +0800
Subject: [PATCH 24/27] test nccl

---
 .gitignore                  | 3 ++-
 tests/test_nccl_backward.py | 7 ++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index 0222862f..2e8c0dcd 100644
--- a/.gitignore
+++ b/.gitignore
@@ -150,4 +150,5 @@ log
 .vscode
 
 !bmtrain/dist
-tests/test_log.txt
\ No newline at end of file
+tests/test_log.txt
+tests/*.opt
\ No newline at end of file
diff --git a/tests/test_nccl_backward.py b/tests/test_nccl_backward.py
index 3dcd0560..5e7b22d8 100644
--- a/tests/test_nccl_backward.py
+++ b/tests/test_nccl_backward.py
@@ -3,8 +3,8 @@
 import bmtrain as bmt
 import torch
 
-def test_main():
-    x = torch.full((1,), bmt.rank() + 1, dtype=torch.half, device="cuda").requires_grad_(True)
+def test_main(dtype):
+    x = torch.full((1,), bmt.rank() + 1, dtype=dtype, device="cuda").requires_grad_(True)
     y = bmt.distributed.all_reduce(x, "prod").view(-1)
     loss = (y * y).sum() / 2
     loss.backward()
@@ -17,4 +17,5 @@ def test_main():
 if __name__ == "__main__":
     bmt.init_distributed()
 
-    test_main()
\ No newline at end of file
+    test_main(torch.half)
+    test_main(torch.bfloat16)
\ No newline at end of file

From d29b22aeb38a742d2e38757114235275e8700d6b Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 11 Aug 2023 11:01:38 +0800
Subject: [PATCH 25/27] fix optim state test

---
 tests/test_optim_state.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_optim_state.py b/tests/test_optim_state.py
index df697f49..9de61c26 100644
--- a/tests/test_optim_state.py
+++ b/tests/test_optim_state.py
@@ -9,7 +9,7 @@ def __init__(self):
         self.fc1 = bmt.BMTrainModelWrapper(torch.nn.Linear(768, 3072))
         self.fc2 = bmt.BMTrainModelWrapper(torch.nn.Linear(3072, 1024))
         self.fc3 = bmt.BMTrainModelWrapper(torch.nn.Linear(1024, 768))
-        self.param = bmt.DistributedParameter(torch.empty(1237))
+        self.param = bmt.DistributedParameter(torch.zeros(1237))
         self.fc4 = bmt.BMTrainModelWrapper(torch.nn.Linear(768, 300))
         self.fc5 = bmt.BMTrainModelWrapper(torch.nn.Linear(300, 768))
         self.dropout = torch.nn.Dropout(0.0)

From 4aeb638074871e53f652fb236835c8abe39f4304 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 11 Aug 2023 15:52:31 +0800
Subject: [PATCH 26/27] fix cuda version if; refactor cross_entropy

---
 bmtrain/loss/_function.py     | 63 +++++++---------------
 bmtrain/loss/cross_entropy.py | 37 +------------
 bmtrain/optim/_function.py    | 11 +++-
 csrc/bind.cpp                 | 15 +++---
 csrc/cuda/adam_cuda.cu        |  9 ++--
 csrc/cuda/bfloat16.cuh        |  5 ++
 csrc/cuda/cross_entropy.cu    | 99 ++++++++++++++++++-----------------
 csrc/cuda/has_inf_nan.cu      | 21 ++++++--
 csrc/include/adam_cpu.hpp     |  2 +-
 csrc/include/bind.hpp         | 28 +++++-----
 tests/test_has_inf_nan.py     |  6 ++-
 tests/test_loss_func.py       | 40 ++++++--------
 tests/test_optim.py           |  5 +-
 13 files changed, 156 insertions(+), 185 deletions(-)
 create mode 100644 csrc/cuda/bfloat16.cuh

diff --git a/bmtrain/loss/_function.py b/bmtrain/loss/_function.py
index e3c76ea9..3a93df2e 100644
--- a/bmtrain/loss/_function.py
+++ b/bmtrain/loss/_function.py
@@ -9,8 +9,11 @@ def has_inf_nan(g_half: torch.Tensor, out: torch.Tensor) -> None:
     mid = torch.zeros(1024, device=out.device, dtype=out.dtype)
     stream = torch.cuda.current_stream().cuda_stream
     if g_half.dtype == torch.float16:
-        C.has_nan_inf_launcher(g_half.numel(), g_half.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
+        C.has_nan_inf_fp16_launcher(g_half.numel(), g_half.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
     elif g_half.dtype == torch.bfloat16:
+        # print(C.is_bf16_supported())
+        # if not C.is_bf16_supported():
+        #     raise NotImplementedError(f"bfloat16 is not supported on current GPU")
         C.has_nan_inf_bf16_launcher(g_half.numel(), g_half.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
     else:
         raise ValueError(f"has_inf_nan not supported for dtype {g_half.dtype}")
@@ -21,9 +24,7 @@ def cross_entropy_forward(m: int, n: int, input: torch.Tensor, target: torch.Ten
     CHECK_INPUT(target)
     CHECK_INPUT(softmax)
     CHECK_INPUT(output)
-    assert input.dtype == torch.float16, "input must be a half tensor"
     assert target.dtype == torch.int32, "target must be an int tensor"
-    assert softmax.dtype == torch.float16, "softmax must be a half tensor"
     assert output.dtype == torch.float32, "output must be a float tensor"
     assert input.numel() == softmax.numel(), "input and softmax must have the same number of elements"
     assert target.numel() == output.numel(), "target and output must have the same number of elements"
@@ -32,43 +33,14 @@ def cross_entropy_forward(m: int, n: int, input: torch.Tensor, target: torch.Ten
     softmax_ptr = softmax.data_ptr()
     output_ptr = output.data_ptr()
     cuda_stream = torch.cuda.current_stream().cuda_stream
-    C.cross_entropy_forward_launcher(m, n, input_ptr, target_ptr, softmax_ptr, output_ptr, ignore_index, cuda_stream)
-
-def cross_entropy_backward(m: int, n: int, grad_output: torch.Tensor, target: torch.Tensor,
-                             softmax: torch.Tensor, grad_input: torch.Tensor, ignore_index: int) -> None:
-    CHECK_INPUT(grad_output)
-    CHECK_INPUT(target)
-    CHECK_INPUT(softmax)
-    CHECK_INPUT(grad_input)
-    assert grad_output.dtype == torch.float32, "grad_output must be a float tensor"
-    assert target.dtype == torch.int32, "target must be an int tensor"
-    assert softmax.dtype == torch.float16, "softmax must be a half tensor"
-    assert grad_input.dtype == torch.float16, "grad_input must be a half tensor"
-    assert grad_input.numel() == softmax.numel(), "grad_input and softmax must have the same number of elements"
-    assert target.numel() == grad_output.numel(), "target and grad_output must have the same number of elements"
-    grad_output_ptr = grad_output.data_ptr()
-    target_ptr = target.data_ptr()
-    softmax_ptr = softmax.data_ptr()
-    grad_input_ptr = grad_input.data_ptr()
-    cuda_stream = torch.cuda.current_stream().cuda_stream
-    C.cross_entropy_backward_launcher(m, n, grad_output_ptr, target_ptr, softmax_ptr, grad_input_ptr, ignore_index, cuda_stream)
-
-def cross_entropy_forward_inplace(m: int, n: int, x: torch.Tensor, target: torch.Tensor,
-                                    output: torch.Tensor, ignore_index: int) -> None:
-    CHECK_INPUT(x)
-    CHECK_INPUT(target)
-    CHECK_INPUT(output)
-    assert x.dtype == torch.float16, "x must be a half tensor"
-    assert target.dtype == torch.int32, "target must be an int tensor"
-    assert output.dtype == torch.float32, "output must be a float tensor"
-    assert target.numel() == output.numel(), "target and output must have the same number of elements"
-    cuda_stream = torch.cuda.current_stream().cuda_stream
-    x_ptr = x.data_ptr()
-    output_ptr = output.data_ptr()
-    target_ptr = target.data_ptr()
-    output_ptr = output.data_ptr()
-    
-    C.cross_entropy_forward_inplace_launcher(m, n, x_ptr, target_ptr, output_ptr, ignore_index, cuda_stream)
+    if input.dtype == torch.float16:
+        C.cross_entropy_forward_fp16_launcher(m, n, input_ptr, target_ptr, softmax_ptr, output_ptr, ignore_index, cuda_stream)
+    elif input.dtype == torch.bfloat16:
+        if not C.is_bf16_supported():
+            raise NotImplementedError(f"bfloat16 is not supported on current GPU")
+        C.cross_entropy_forward_bf16_launcher(m, n, input_ptr, target_ptr, softmax_ptr, output_ptr, ignore_index, cuda_stream)
+    else:
+        raise ValueError(f"cross_entropy_forward not supported for dtype {input.dtype}")
 
 def cross_entropy_backward_inplace(m: int, n: int, grad_output: torch.Tensor, target: torch.Tensor,
                                      x: torch.Tensor, ignore_index: int) -> None:
@@ -77,12 +49,17 @@ def cross_entropy_backward_inplace(m: int, n: int, grad_output: torch.Tensor, ta
     CHECK_INPUT(x)
     assert grad_output.dtype == torch.float32, "grad_output must be a float tensor"
     assert target.dtype == torch.int32, "target must be an int tensor"
-    assert x.dtype == torch.float16, "x must be a half tensor"
     assert target.numel() == grad_output.numel(), "target and grad_output must have the same number of elements"
     cuda_stream = torch.cuda.current_stream().cuda_stream
     grad_output_ptr = grad_output.data_ptr()
     target_ptr = target.data_ptr()
     x_ptr = x.data_ptr()
 
-    C.cross_entropy_backward_inplace_launcher(m, n, grad_output_ptr, target_ptr, x_ptr, ignore_index, cuda_stream)
-
+    if x.dtype == torch.float16:
+        C.cross_entropy_backward_inplace_fp16_launcher(m, n, grad_output_ptr, target_ptr, x_ptr, ignore_index, cuda_stream)
+    elif x.dtype == torch.bfloat16:
+        if not C.is_bf16_supported():
+            raise NotImplementedError(f"bfloat16 is not supported on current GPU")
+        C.cross_entropy_backward_inplace_bf16_launcher(m, n, grad_output_ptr, target_ptr, x_ptr, ignore_index, cuda_stream)
+    else:
+        raise ValueError(f"cross_entropy_backward not supported for dtype {input.dtype}")
diff --git a/bmtrain/loss/cross_entropy.py b/bmtrain/loss/cross_entropy.py
index 160ef421..663a9e8d 100644
--- a/bmtrain/loss/cross_entropy.py
+++ b/bmtrain/loss/cross_entropy.py
@@ -32,36 +32,6 @@ def backward(ctx, grad_output : torch.Tensor):
         )
         return (softmax, None, None)
 
-class OpFusedCrossEntropyInplace(torch.autograd.Function):
-    """
-    CrossEntropy dim = 1
-    """
-    @staticmethod
-    def forward(ctx, x : torch.Tensor, target : torch.Tensor, ignore_index: int):
-        assert x.ndim == 2
-        out = torch.empty(x.size(0), device=x.device, dtype=torch.float)
-        F.cross_entropy_forward_inplace(
-            x.size(0), x.size(1),
-            x, target,
-            out,
-            ignore_index,
-        ) # x is inplace modify to softmax result
-        ctx.ignore_index = ignore_index
-        ctx.save_for_backward(x, target)
-        return out # float tensor
-        
-    @staticmethod
-    def backward(ctx, grad_output : torch.Tensor):
-        grad_output = grad_output.contiguous()
-        softmax, target = ctx.saved_tensors
-        F.cross_entropy_backward_inplace(
-            softmax.size(0), softmax.size(1),
-            grad_output, target,
-            softmax,
-            ctx.ignore_index,
-        ) # softmax is inplace modify to grad_input
-        return (softmax, None, None)
-
 class FusedCrossEntropy(torch.nn.Module):
     r"""This criterion computes the cross entropy loss between input and target.
 
@@ -175,20 +145,15 @@ def __init__(self,
                  ignore_index: int = -100,
                  reduction: str = 'mean',
                  label_smoothing: float = 0.0, # TODO not supported yet
-                 inplace: bool = False,
                 ) -> None:
         super().__init__()
         self.weight = weight
         self.ignore_index = ignore_index
         self.reduction = reduction
         self.label_smoothing = label_smoothing
-        self.inplace = inplace
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        if self.inplace:
-            ret = OpFusedCrossEntropyInplace.apply(input, target.int(), self.ignore_index) # return float tensor
-        else:
-            ret = OpFusedCrossEntropy.apply(input, target.int(), self.ignore_index) # return float tensor
+        ret = OpFusedCrossEntropy.apply(input, target.int(), self.ignore_index) # return float tensor
 
         if self.weight is not None:
             if self.weight.dim() != 1 or self.weight.size(0) != input.size(1):
diff --git a/bmtrain/optim/_function.py b/bmtrain/optim/_function.py
index b5ad60b1..f04f9ca0 100644
--- a/bmtrain/optim/_function.py
+++ b/bmtrain/optim/_function.py
@@ -25,7 +25,12 @@ def adam_cpu(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.T
     assert param_fp32.numel() == v_fp32.numel(), "param_fp32 and v_fp32 must have the same number of elements"
     bias_correction1 = 1 - beta1 ** step
     bias_correction2 = 1 - beta2 ** step
-    launcher = C.adam_cpu_launcher if g_fp16.dtype == torch.float16 else C.adam_cpu_bf16_launcher
+    if g_fp16.dtype == torch.float16:
+        launcher = C.adam_cpu_fp16_launcher
+    elif g_fp16.dtype == torch.bfloat16:
+        if not C.is_bf16_supported():
+            raise NotImplementedError(f"bfloat16 is not supported on current GPU")
+        launcher = C.adam_cpu_bf16_launcher
     launcher(
         param_fp32.numel(),
         param_fp32.data_ptr(),
@@ -61,7 +66,7 @@ def adam_fp16(param_fp32: torch.Tensor, param_fp16: torch.Tensor, g_fp16: torch.
     bias_correction1 = 1 - beta1 ** step
     bias_correction2 = 1 - beta2 ** step
     stream = torch.cuda.current_stream().cuda_stream
-    C.adam_launcher(
+    C.adam_fp16_launcher(
         param_fp32.numel(),
         param_fp32.data_ptr(),
         param_fp16.data_ptr(),
@@ -97,6 +102,8 @@ def adam_bf16(param_fp32: torch.Tensor, param_bf16: torch.Tensor, g_bf16: torch.
     bias_correction1 = 1 - beta1 ** step
     bias_correction2 = 1 - beta2 ** step
     stream = torch.cuda.current_stream().cuda_stream
+    if not C.is_bf16_supported():
+        raise NotImplementedError(f"bfloat16 is not supported on current GPU")
     C.adam_bf16_launcher(
         param_fp32.numel(),
         param_fp32.data_ptr(),
diff --git a/csrc/bind.cpp b/csrc/bind.cpp
index b4bc6515..73f79a61 100644
--- a/csrc/bind.cpp
+++ b/csrc/bind.cpp
@@ -1,16 +1,17 @@
 #include "include/bind.hpp"
 
 PYBIND11_MODULE(C, m) {
-    m.def("has_nan_inf_launcher",&has_nan_inf_launcher,"has nan inf");
+    m.def("is_bf16_supported",&is_bf16_supported,"whether bf16 supported");
+    m.def("has_nan_inf_fp16_launcher",&has_nan_inf_fp16_launcher,"has nan inf");
     m.def("has_nan_inf_bf16_launcher",&has_nan_inf_bf16_launcher,"has nan inf bf16");
-    m.def("adam_launcher", &adam_launcher, "adam function cpu");
+    m.def("adam_fp16_launcher", &adam_fp16_launcher, "adam function cpu");
     m.def("adam_bf16_launcher", &adam_bf16_launcher, "adam function cpu");
-    m.def("adam_cpu_launcher", &adam_cpu_launcher, "adam function cpu");
+    m.def("adam_cpu_fp16_launcher", &adam_cpu_fp16_launcher, "adam function cpu");
     m.def("adam_cpu_bf16_launcher", &adam_cpu_bf16_launcher, "adam function cpu");
-    m.def("cross_entropy_forward_launcher", &cross_entropy_forward_launcher, "cross entropy forward");
-    m.def("cross_entropy_backward_launcher", &cross_entropy_backward_launcher, "cross entropy backward");
-    m.def("cross_entropy_forward_inplace_launcher", &cross_entropy_forward_inplace_launcher, "cross entropy forward inplace");
-    m.def("cross_entropy_backward_inplace_launcher", &cross_entropy_backward_inplace_launcher, "cross entropy backward inplace");
+    m.def("cross_entropy_forward_fp16_launcher", &cross_entropy_forward_fp16_launcher, "cross entropy forward");
+    m.def("cross_entropy_forward_bf16_launcher", &cross_entropy_forward_bf16_launcher, "cross entropy forward");
+    m.def("cross_entropy_backward_inplace_fp16_launcher", &cross_entropy_backward_inplace_fp16_launcher, "cross entropy backward inplace");
+    m.def("cross_entropy_backward_inplace_bf16_launcher", &cross_entropy_backward_inplace_bf16_launcher, "cross entropy backward inplace");
     m.def("ncclGetUniqueId", &pyNCCLGetUniqueID, "nccl get unique ID");
     m.def("ncclCommInitRank", &pyNCCLCommInitRank, "nccl init rank");
     m.def("ncclCommDestroy", &pyNCCLCommDestroy, "nccl delete rank");
diff --git a/csrc/cuda/adam_cuda.cu b/csrc/cuda/adam_cuda.cu
index 9e701881..0510ac12 100644
--- a/csrc/cuda/adam_cuda.cu
+++ b/csrc/cuda/adam_cuda.cu
@@ -1,8 +1,7 @@
 #include <cstdint>
+#include <cuda.h>
 #include <cuda_fp16.h>
-#if __CUDA_ARCH__ >= 800
-#include <cuda_bf16.h>
-#endif
+#include "bfloat16.cuh"
 
 namespace {
 // blocks <n // 1024>,      threads<min(n, 1024)>
@@ -53,7 +52,7 @@ __global__ void adam_fp32_accum_bf16(
     float bias_correction1,
     float bias_correction2
 ) {
-#if __CUDA_ARCH__ >= 800
+#ifdef BF16_SUPPORT
     const __nv_bfloat16* g = reinterpret_cast<const __nv_bfloat16*>(g_ptr);
     __nv_bfloat16* param_h = reinterpret_cast<__nv_bfloat16*>(param_h_ptr);
     int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
@@ -74,7 +73,7 @@ __global__ void adam_fp32_accum_bf16(
 
 }
 
-void adam_launcher(
+void adam_fp16_launcher(
     int n,
     std::uintptr_t param_fp32,
     std::uintptr_t param_fp16,
diff --git a/csrc/cuda/bfloat16.cuh b/csrc/cuda/bfloat16.cuh
new file mode 100644
index 00000000..564d8bec
--- /dev/null
+++ b/csrc/cuda/bfloat16.cuh
@@ -0,0 +1,5 @@
+#include <cuda.h>
+#if defined(__CUDACC__) && CUDA_VERSION >= 11000
+#include <cuda_bf16.h>
+#define BF16_SUPPORT
+#endif
\ No newline at end of file
diff --git a/csrc/cuda/cross_entropy.cu b/csrc/cuda/cross_entropy.cu
index c0b742ac..bdd5a08e 100644
--- a/csrc/cuda/cross_entropy.cu
+++ b/csrc/cuda/cross_entropy.cu
@@ -1,11 +1,12 @@
-#include <cuda_fp16.h>
 #include "reduce.cuh"
 #include <cstdint>
-#include<cuda_runtime_api.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include "bfloat16.cuh"
 
 namespace {
 // blocks <m>,      threads<1024>
-__global__ void cross_entropy_forward(
+__global__ void cross_entropy_forward_fp16(
     int64_t n,
     const half *input,      // (m, n)
     const int32_t *target,  // (m)
@@ -42,12 +43,11 @@ __global__ void cross_entropy_forward(
 }
 
 // blocks <m>,      threads<1024>
-__global__ void cross_entropy_backward(
+__global__ void cross_entropy_backward_inplace_fp16(
     int64_t n,
     const float *grad_output,   // (m)
     const int32_t *target,      // (m)
-    const half *softmax,        // (m, n)
-    half *grad_input,           // (m, n)
+    half *x,                    // (m, n)
     int32_t ignore_index
 ) {
     int64_t base_idx = blockIdx.x * n;
@@ -56,83 +56,99 @@ __global__ void cross_entropy_backward(
     if (t == ignore_index) {
         half v = __float2half(0.);
         for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
-            grad_input[base_idx + i] = v;
+            x[base_idx + i] = v;
         }
     }
     else {
         half v = __float2half(grad_output[blockIdx.x]);
+        __syncthreads();
         for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
-            grad_input[base_idx + i] = i==t ? __hsub(__hmul(softmax[base_idx + i], v), v) : __hmul(softmax[base_idx + i], v);
+            x[base_idx + i] = i==t ? __hsub(__hmul(x[base_idx + i], v), v) : __hmul(x[base_idx + i], v);
         }
     }
 }
 
 // blocks <m>,      threads<1024>
-__global__ void cross_entropy_forward_inplace(
+__global__ void cross_entropy_forward_bf16(
     int64_t n,
-    half *x,                // (m, n)
+    std::uintptr_t input_ptr,      // (m, n)
     const int32_t *target,  // (m)
+    std::uintptr_t softmax_ptr,          // (m, n)
     float *output,          // (m)
     int32_t ignore_index
 ) {
+#ifdef BF16_SUPPORT
+    __nv_bfloat16* input = reinterpret_cast<__nv_bfloat16*>(input_ptr);
+    __nv_bfloat16* softmax = reinterpret_cast<__nv_bfloat16*>(softmax_ptr);
     int64_t base_idx = blockIdx.x * n;
 
     float local_max = -INFINITY;
     for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
-        local_max = fmaxf(__half2float(x[base_idx + i]), local_max);
+        local_max = fmaxf(__bfloat162float(input[base_idx + i]), local_max);
     }
+
     local_max = fmaxf(block_allreduce_max(local_max), -1e6);
     
     float local_sum = 0;
     for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
-        local_sum += expf(__half2float(x[base_idx + i]) - local_max);
+        local_sum += expf(__bfloat162float(input[base_idx + i]) - local_max);
     }
     local_sum = block_allreduce_sum(local_sum) + 1e-10; // avoid nan
+    
+    for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
+        softmax[base_idx + i] = __float2bfloat16( expf(__bfloat162float(input[base_idx + i]) - local_max) / local_sum );
+    }
 
     if (threadIdx.x == 0) {
         if (target[blockIdx.x] != ignore_index) {
-            output[blockIdx.x] = -__half2float(x[base_idx + target[blockIdx.x]]) + local_max + logf(local_sum);
+            output[blockIdx.x] = -__bfloat162float(input[base_idx + target[blockIdx.x]]) + local_max + logf(local_sum);
         } else {
             output[blockIdx.x] = 0;
         }
     }
-
-    __syncthreads();
-    
-    for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
-        x[base_idx + i] = __float2half( expf(__half2float(x[base_idx + i]) - local_max) / local_sum );
-    }
+#endif
 }
 
 // blocks <m>,      threads<1024>
-__global__ void cross_entropy_backward_inplace(
+__global__ void cross_entropy_backward_inplace_bf16(
     int64_t n,
     const float *grad_output,   // (m)
     const int32_t *target,      // (m)
-    half *x,                    // (m, n)
+    std::uintptr_t x_ptr,                    // (m, n)
     int32_t ignore_index
 ) {
+#ifdef BF16_SUPPORT
+    __nv_bfloat16* x = reinterpret_cast<__nv_bfloat16*>(x_ptr);
     int64_t base_idx = blockIdx.x * n;
 
     int32_t t = target[blockIdx.x];
     if (t == ignore_index) {
-        half v = __float2half(0.);
+        __nv_bfloat16 v = __float2bfloat16(0.);
         for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
             x[base_idx + i] = v;
         }
     }
     else {
-        half v = __float2half(grad_output[blockIdx.x]);
+        #if __CUDA_ARCH__ >= 800
+        __nv_bfloat16 v = __float2bfloat16(grad_output[blockIdx.x]);
         __syncthreads();
         for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
             x[base_idx + i] = i==t ? __hsub(__hmul(x[base_idx + i], v), v) : __hmul(x[base_idx + i], v);
         }
+        #else
+        float v = grad_output[blockIdx.x];
+        __syncthreads();
+        for (int64_t i = threadIdx.x; i < n; i += blockDim.x) {
+            x[base_idx + i] = __float2bfloat16(i==t ? (__bfloat162float(x[base_idx + i])*v)-v : __bfloat162float(x[base_idx + i])*v);
+        }
+        #endif
     }
+#endif
 }
 
 }
 
-void cross_entropy_forward_launcher(
+void cross_entropy_forward_fp16_launcher(
     int32_t m, int32_t n,
     std::uintptr_t input,
     std::uintptr_t target,
@@ -146,48 +162,40 @@ void cross_entropy_forward_launcher(
     auto softmax_ptr = reinterpret_cast<half*>(softmax);
     auto output_ptr = reinterpret_cast<float*>(output);
     int32_t threads = 1024;
-    cross_entropy_forward<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, input_ptr, target_ptr, softmax_ptr, output_ptr, ignore_index);
+    cross_entropy_forward_fp16<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, input_ptr, target_ptr, softmax_ptr, output_ptr, ignore_index);
 }
 
-void cross_entropy_backward_launcher(
+void cross_entropy_backward_inplace_fp16_launcher(
     int32_t m, int32_t n,
     std::uintptr_t grad_output,
     std::uintptr_t target,
-    std::uintptr_t softmax,
-    std::uintptr_t grad_input,
+    std::uintptr_t x,
     int32_t ignore_index,
     std::uintptr_t stream
 ) {
-    // auto output_ptr = grad_output.data_ptr<float>();
     auto output_ptr = reinterpret_cast<float*>(grad_output);
-    // auto target_ptr = target.data_ptr<int32_t>();
     auto target_ptr = reinterpret_cast<int32_t*>(target);
-    auto softmax_ptr = reinterpret_cast<half*>(softmax);
-    auto input_ptr = reinterpret_cast<half*>(grad_input);
+    auto x_ptr = reinterpret_cast<half*>(x);
     int32_t threads = 1024;
-    cross_entropy_backward<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, output_ptr, target_ptr, softmax_ptr, input_ptr, ignore_index);
+    cross_entropy_backward_inplace_fp16<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, output_ptr, target_ptr, x_ptr, ignore_index);
 }
 
-void cross_entropy_forward_inplace_launcher(
+void cross_entropy_forward_bf16_launcher(
     int32_t m, int32_t n,
-    std::uintptr_t x,
+    std::uintptr_t input,
     std::uintptr_t target,
+    std::uintptr_t softmax,
     std::uintptr_t output,
     int32_t ignore_index,
     std::uintptr_t stream
 ) {
-    // auto x_ptr = reinterpret_cast<half*>(x.data_ptr<at::Half>());
-    auto x_ptr = reinterpret_cast<half*>(x);
-    // auto target_ptr = target.data_ptr<int32_t>();
     auto target_ptr = reinterpret_cast<int32_t*>(target);
-    // auto output_ptr = output.data_ptr<float>();
     auto output_ptr = reinterpret_cast<float*>(output);
     int32_t threads = 1024;
-    // auto stream = at::cuda::getCurrentCUDAStream();
-    cross_entropy_forward_inplace<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, x_ptr, target_ptr, output_ptr, ignore_index);
+    cross_entropy_forward_bf16<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, input, target_ptr, softmax, output_ptr, ignore_index);
 }
 
-void cross_entropy_backward_inplace_launcher(
+void cross_entropy_backward_inplace_bf16_launcher(
     int32_t m, int32_t n,
     std::uintptr_t grad_output,
     std::uintptr_t target,
@@ -195,13 +203,8 @@ void cross_entropy_backward_inplace_launcher(
     int32_t ignore_index,
     std::uintptr_t stream
 ) {
-    // auto output_ptr = grad_output.data_ptr<float>();
     auto output_ptr = reinterpret_cast<float*>(grad_output);
-    // auto target_ptr = target.data_ptr<int32_t>();
     auto target_ptr = reinterpret_cast<int32_t*>(target);
-    // auto x_ptr = reinterpret_cast<half*>(x.data_ptr<at::Half>());
-    auto x_ptr = reinterpret_cast<half*>(x);
     int32_t threads = 1024;
-    // auto stream = at::cuda::getCurrentCUDAStream();
-    cross_entropy_backward_inplace<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, output_ptr, target_ptr, x_ptr, ignore_index);
+    cross_entropy_backward_inplace_bf16<<<m, threads, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, output_ptr, target_ptr, x, ignore_index);
 }
\ No newline at end of file
diff --git a/csrc/cuda/has_inf_nan.cu b/csrc/cuda/has_inf_nan.cu
index b4c3113d..32bc5a5f 100644
--- a/csrc/cuda/has_inf_nan.cu
+++ b/csrc/cuda/has_inf_nan.cu
@@ -1,8 +1,8 @@
 #include <cstdint>
+#include <cstdio>
+#include <cuda.h>
 #include <cuda_fp16.h>
-#if __CUDA_ARCH__ >= 800
-#include <cuda_bf16.h>
-#endif
+#include "bfloat16.cuh"
 
 namespace{
 __inline__ __device__ bool isnan_(half v) {
@@ -72,7 +72,7 @@ __global__ void bmt_has_nan_inf_bf16(
     const uintptr_t inp,        // (n,) 
     uint8_t* mid            // (1024,)
 ) {
-#if __CUDA_ARCH__ >= 800
+#ifdef BF16_SUPPORT
     const __nv_bfloat16* bf_inp = reinterpret_cast<const __nv_bfloat16*>(inp);
     int32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
     int32_t span = blockDim.x * gridDim.x;
@@ -80,7 +80,11 @@ __global__ void bmt_has_nan_inf_bf16(
     int8_t r = 0;
     for (int i = gid; i < n; i += span) {
         __nv_bfloat16 v = bf_inp[i];
+        #if __CUDA_ARCH__ >= 800
         if (__hisinf(v) || __hisnan(v)) { 
+        #else
+        if (isinf(__bfloat162float(v)) || isnan(__bfloat162float(v))) {
+        #endif
             r = 1;
             break;
         }
@@ -94,7 +98,7 @@ __global__ void bmt_has_nan_inf_bf16(
 
 }
 
-void has_nan_inf_launcher(
+void has_nan_inf_fp16_launcher(
     int32_t n,
     std::uintptr_t g_fp16,
     std::uintptr_t mid,
@@ -131,4 +135,11 @@ void has_nan_inf_bf16_launcher(
     
     bmt_has_nan_inf_bf16<<<clamp_grid_size, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(n, g_bf16, mid_ptr);
     bmt_has_nan_inf_reduce<<<1, block_size, 0, reinterpret_cast<cudaStream_t>(stream)>>>(mid_ptr, out_ptr);
+}
+
+int is_bf16_supported() {
+#ifdef BF16_SUPPORT
+    return 1;
+#endif
+    return 0;
 }
\ No newline at end of file
diff --git a/csrc/include/adam_cpu.hpp b/csrc/include/adam_cpu.hpp
index 6ff901a5..1e497bb3 100644
--- a/csrc/include/adam_cpu.hpp
+++ b/csrc/include/adam_cpu.hpp
@@ -340,7 +340,7 @@ static void __attribute__ ((__target__ ("avx512f"))) adam_cpu_2(
     });
 }
 
-void adam_cpu_launcher(
+void adam_cpu_fp16_launcher(
     int64_t n,
     std::uintptr_t param_fp32,
     std::uintptr_t param_fp16,
diff --git a/csrc/include/bind.hpp b/csrc/include/bind.hpp
index e419340b..94d6af95 100644
--- a/csrc/include/bind.hpp
+++ b/csrc/include/bind.hpp
@@ -1,20 +1,22 @@
-#include<pybind11/pybind11.h>
+#include <pybind11/pybind11.h>
 #include "nccl.hpp"
 #include "adam_cpu.hpp"
 
-void has_nan_inf_launcher(int32_t n,std::uintptr_t g_fp16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
+int is_bf16_supported();
+
+void has_nan_inf_fp16_launcher(int32_t n,std::uintptr_t g_fp16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
 void has_nan_inf_bf16_launcher(int32_t n,std::uintptr_t g_bf16,std::uintptr_t mid,std::uintptr_t out,std::uintptr_t stream);
 
-void cross_entropy_backward_launcher(
+void cross_entropy_forward_fp16_launcher(
     int32_t m, int32_t n,
-    std::uintptr_t grad_output,
+    std::uintptr_t input,
     std::uintptr_t target,
     std::uintptr_t softmax,
-    std::uintptr_t grad_input,
+    std::uintptr_t output,
     int32_t ignore_index,
     std::uintptr_t stream
 );
-void cross_entropy_backward_inplace_launcher(
+void cross_entropy_backward_inplace_fp16_launcher(
     int32_t m, int32_t n,
     std::uintptr_t grad_output,
     std::uintptr_t target,
@@ -22,24 +24,24 @@ void cross_entropy_backward_inplace_launcher(
     int32_t ignore_index,
     std::uintptr_t stream
 );
- void cross_entropy_forward_inplace_launcher(
+void cross_entropy_forward_bf16_launcher(
     int32_t m, int32_t n,
-    std::uintptr_t x,
+    std::uintptr_t input,
     std::uintptr_t target,
+    std::uintptr_t softmax,
     std::uintptr_t output,
     int32_t ignore_index,
     std::uintptr_t stream
 );
-void cross_entropy_forward_launcher(
+void cross_entropy_backward_inplace_bf16_launcher(
     int32_t m, int32_t n,
-    std::uintptr_t input,
+    std::uintptr_t grad_output,
     std::uintptr_t target,
-    std::uintptr_t softmax,
-    std::uintptr_t output,
+    std::uintptr_t x,
     int32_t ignore_index,
     std::uintptr_t stream
 );
-void adam_launcher(
+void adam_fp16_launcher(
     int n,
     std::uintptr_t param_fp32,
     std::uintptr_t param_fp16,
diff --git a/tests/test_has_inf_nan.py b/tests/test_has_inf_nan.py
index 5432aed9..7b7c16ee 100644
--- a/tests/test_has_inf_nan.py
+++ b/tests/test_has_inf_nan.py
@@ -30,4 +30,8 @@ def test_main(dtype):
 
 if __name__ == "__main__":
     test_main(torch.float16)
-    test_main(torch.bfloat16)
+    print("==============================================================================")
+    try:
+        test_main(torch.bfloat16)
+    except NotImplementedError: 
+        pass
diff --git a/tests/test_loss_func.py b/tests/test_loss_func.py
index a76be5f4..a448b6d1 100644
--- a/tests/test_loss_func.py
+++ b/tests/test_loss_func.py
@@ -27,59 +27,53 @@ def check(x, tgt, loss_func1, loss_func2, bigmodel=None):
     loss_2, grad_2 = run(x, tgt, loss_func2, bigmodel=bigmodel, use_float=True)
     assert_eq(grad_1.isnan().sum(), 0)
     assert_eq(grad_2.isnan().sum(), 0)
+    print(f"{(loss_1 - loss_2).abs().item():.6f} {(grad_1 - grad_2).abs().max().item():.6f}")
     assert_lt((loss_1 - loss_2).abs().item(), 1e-5)
-    assert_lt((grad_1 - grad_2).abs().max().item(), 1e-2)
+    assert_lt((grad_1 - grad_2).abs().max().item(), 1e-1)
 
-def test_simple():
+def test_simple(dtype):
     loss_func1 = bmt.loss.FusedCrossEntropy()
     loss_func2 = torch.nn.CrossEntropyLoss()
 
     N = 32 * 512
     for i in range(1, 10):
         C = i * 10
-        x = torch.randn(N, C).cuda().half()
+        x = torch.randn(N, C).cuda().to(dtype)
         tgt = torch.randint(0, C, (N,)).cuda().long()
         check(x, tgt, loss_func1, loss_func2)
     for i in range(1, 10):
         C = i * 100
-        x = torch.randn(N, C).cuda().half()
+        x = torch.randn(N, C).cuda().to(dtype)
         tgt = torch.randint(0, C, (N,)).cuda().long()
         check(x, tgt, loss_func1, loss_func2)
     for i in range(1, 31):
         C = i * 1000
-        x = torch.randn(N, C).cuda().half()
+        x = torch.randn(N, C).cuda().to(dtype)
         tgt = torch.randint(0, C, (N,)).cuda().long()
         check(x, tgt, loss_func1, loss_func2)
 
-def test_other():
+def test_other(dtype):
     N = 32 * 512
     for i in range(1, 11):
         C = i * 10
         weight = [i+1 for i in range(C)]
         random.shuffle(weight)
         weight = torch.tensor(weight, device="cuda")
-        loss_func1 = bmt.loss.FusedCrossEntropy(weight=weight.clone().half())
+        loss_func1 = bmt.loss.FusedCrossEntropy(weight=weight.clone().to(dtype))
         loss_func2 = torch.nn.CrossEntropyLoss(weight=weight.clone().float())
 
-        x = torch.randn(N, C).cuda().half()
+        x = torch.randn(N, C).cuda().to(dtype)
         tgt = torch.randint(0, C, (N,)).cuda().long()
         mask = torch.randint(0, 2, (N,)).cuda().bool()
         tgt[mask] = -100
         check(x, tgt, loss_func1, loss_func2)
 
-def test_inplace():
-    loss_func1 = bmt.loss.FusedCrossEntropy(inplace=True)
-    loss_func2 = torch.nn.CrossEntropyLoss()
-    N = 32 * 512
-
-    for i in range(1, 11):
-        C = i * 10
-        bigmodel = torch.nn.Linear(5, C).cuda().half()
-        x = torch.randn(N, 5).cuda().half()
-        tgt = torch.randint(0, C, (N,)).cuda().long()
-        check(x, tgt, loss_func1, loss_func2, bigmodel=bigmodel)
-
 if __name__ == "__main__":
-    test_other()
-    test_inplace()
-    test_simple()
\ No newline at end of file
+    test_other(torch.float16)
+    test_simple(torch.float16)
+    print("==============================================================================")
+    try:
+        test_other(torch.bfloat16)
+        test_simple(torch.bfloat16)
+    except NotImplementedError: 
+        pass
\ No newline at end of file
diff --git a/tests/test_optim.py b/tests/test_optim.py
index 2888501e..680fe41a 100644
--- a/tests/test_optim.py
+++ b/tests/test_optim.py
@@ -87,4 +87,7 @@ def main(dtype):
     bmt.init_distributed()
     main(torch.float16)
     print("==============================================================================")
-    main(torch.bfloat16)
+    try:
+        main(torch.bfloat16)
+    except NotImplementedError: 
+        pass

From 53d8f8b2af309335d88e963d6c024cd8e91e5748 Mon Sep 17 00:00:00 2001
From: Achazwl <acha131441373@gmail.com>
Date: Fri, 11 Aug 2023 16:11:36 +0800
Subject: [PATCH 27/27] fix bf16 not support info

---
 bmtrain/loss/_function.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/bmtrain/loss/_function.py b/bmtrain/loss/_function.py
index 3a93df2e..e2b67bb8 100644
--- a/bmtrain/loss/_function.py
+++ b/bmtrain/loss/_function.py
@@ -11,9 +11,8 @@ def has_inf_nan(g_half: torch.Tensor, out: torch.Tensor) -> None:
     if g_half.dtype == torch.float16:
         C.has_nan_inf_fp16_launcher(g_half.numel(), g_half.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
     elif g_half.dtype == torch.bfloat16:
-        # print(C.is_bf16_supported())
-        # if not C.is_bf16_supported():
-        #     raise NotImplementedError(f"bfloat16 is not supported on current GPU")
+        if not C.is_bf16_supported():
+            raise NotImplementedError(f"bfloat16 is not supported on current GPU")
         C.has_nan_inf_bf16_launcher(g_half.numel(), g_half.data_ptr(), mid.data_ptr(), out.data_ptr(), stream)
     else:
         raise ValueError(f"has_inf_nan not supported for dtype {g_half.dtype}")
