modified: add adam_fp32_accum_bf16 function

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(