From 321063284ff9462b77f8f56a81383574d4db46e7 Mon Sep 17 00:00:00 2001 From: ycr0776 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 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 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<<(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(param_fp32); + auto m_fp32_ptr = reinterpret_cast(m_fp32); + auto v_fp32_ptr = reinterpret_cast(v_fp32); + auto param_bf16_ptr = reinterpret_cast(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<<(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(g_bf16); + auto m_ptr = reinterpret_cast(m_bf16); + auto param_h_ptr = reinterpret_cast(param_bf16); + auto param_fp32_ptr = reinterpret_cast(param_fp32); + auto v_fp32_ptr = reinterpret_cast(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<<(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(param_fp32); auto m_fp32_ptr = reinterpret_cast(m_fp32); auto v_fp32_ptr = reinterpret_cast(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 +#include #include "reduce.cuh" #include #include 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 +#include #include 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 +#include #include 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 +#include 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 #include #include +#include #include #include #include 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<<(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<<(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(g_bf16); - auto m_ptr = reinterpret_cast(m_bf16); - auto param_h_ptr = reinterpret_cast(param_bf16); + auto g_ptr = reinterpret_cast(g_bf16); + auto m_ptr = reinterpret_cast(m_bf16); + auto param_h_ptr = reinterpret_cast(param_bf16); auto param_fp32_ptr = reinterpret_cast(param_fp32); auto v_fp32_ptr = reinterpret_cast(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 , 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<<(stream)>>>(n, g_ptr, mid_ptr); bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast(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(g_bf16); + auto mid_ptr = reinterpret_cast(mid); + auto out_ptr = reinterpret_cast(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<<(stream)>>>(n, g_ptr, mid_ptr); + bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast(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(sign >> 16) | static_cast(exp >> 16) | static_cast(mantissa >> 16); + return fp32_from_bits(static_cast(bf16) << 16); +} + +// bf16 -> fp32 +inline float bf16_ieee_to_fp32_value(u_int16_t h){ + uint32_t w = static_cast(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 +#include +#include +#include +#include +#include + +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( + (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( +// (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(g_bf16); - auto m_ptr = reinterpret_cast(m_bf16); + auto m_ptr = reinterpret_cast(m_fp32); auto param_h_ptr = reinterpret_cast(param_bf16); auto param_fp32_ptr = reinterpret_cast(param_fp32); auto v_fp32_ptr = reinterpret_cast(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 , 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 #include #include -#include -#include +#include +#include #include #include #include @@ -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(m_fp32); auto v_fp32_ptr = reinterpret_cast(v_fp32); auto param_bf16_ptr = reinterpret_cast(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 -#include -#include -#include -#include -#include - -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( - (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( -// (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 #include -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(g_bf16); + auto g_ptr = reinterpret_cast(g_bf16); auto mid_ptr = reinterpret_cast(mid); auto out_ptr = reinterpret_cast(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(m_fp32); auto v_fp32_ptr = reinterpret_cast(v_fp32); auto param_bf16_ptr = reinterpret_cast(param_bf16); + auto g_bf16_ptr = reinterpret_cast(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<<(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 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 -#include #include +#include +#if __CUDA_ARCH__ >= 800 + #include +#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 , 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 , 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(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<<(stream)>>>(n, g_ptr, mid_ptr); - bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast(stream)>>>(mid_ptr, out_ptr); + bmt_has_nan_inf_fp16<<(stream)>>>(n, g_ptr, mid_ptr); + bmt_has_nan_inf_reduce<<<1, block_size, 0, reinterpret_cast(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(g_bf16); auto mid_ptr = reinterpret_cast(mid); auto out_ptr = reinterpret_cast(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<<(stream)>>>(n, g_ptr, mid_ptr); - bmt_has_nan_inf_2<<<1, block_size, 0, reinterpret_cast(stream)>>>(mid_ptr, out_ptr); + bmt_has_nan_inf_bf16<<(stream)>>>(n, g_bf16, mid_ptr); + bmt_has_nan_inf_reduce<<<1, block_size, 0, reinterpret_cast(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 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 -#include - 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 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 -#include #include "reduce.cuh" #include #include 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 #include #include -#include +#include #include #include -#include -#include #include #include #include @@ -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( - (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(sign >> 16) | static_cast(exp >> 16) | static_cast(mantissa >> 16); - return fp32_from_bits(static_cast(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(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( + (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(&f) >> 16; +} + +// bf16 -> fp32 +inline float bf16_to_fp32_value(uint16_t h){ + uint32_t src = h; + src <<= 16; + return *reinterpret_cast(&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 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 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 #include +#if __CUDA_ARCH__ >= 800 #include -#include +#endif namespace { // blocks , threads @@ -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(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(g_bf16); - auto m_ptr = reinterpret_cast(m_fp32); - auto param_h_ptr = reinterpret_cast(param_bf16); + auto m_ptr = reinterpret_cast(m_fp32); auto param_fp32_ptr = reinterpret_cast(param_fp32); auto v_fp32_ptr = reinterpret_cast(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<<(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<<(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 #include #if __CUDA_ARCH__ >= 800 - #include +#include #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 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 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 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 +#include #include -#if __CUDA_ARCH__ >= 800 -#include -#endif +#include "bfloat16.cuh" namespace { // blocks , threads @@ -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(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 +#if defined(__CUDACC__) && CUDA_VERSION >= 11000 +#include +#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 #include "reduce.cuh" #include -#include +#include +#include +#include "bfloat16.cuh" namespace { // blocks , 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 , 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 , 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 , 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(softmax); auto output_ptr = reinterpret_cast(output); int32_t threads = 1024; - cross_entropy_forward<<(stream)>>>(n, input_ptr, target_ptr, softmax_ptr, output_ptr, ignore_index); + cross_entropy_forward_fp16<<(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(); auto output_ptr = reinterpret_cast(grad_output); - // auto target_ptr = target.data_ptr(); auto target_ptr = reinterpret_cast(target); - auto softmax_ptr = reinterpret_cast(softmax); - auto input_ptr = reinterpret_cast(grad_input); + auto x_ptr = reinterpret_cast(x); int32_t threads = 1024; - cross_entropy_backward<<(stream)>>>(n, output_ptr, target_ptr, softmax_ptr, input_ptr, ignore_index); + cross_entropy_backward_inplace_fp16<<(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(x.data_ptr()); - auto x_ptr = reinterpret_cast(x); - // auto target_ptr = target.data_ptr(); auto target_ptr = reinterpret_cast(target); - // auto output_ptr = output.data_ptr(); auto output_ptr = reinterpret_cast(output); int32_t threads = 1024; - // auto stream = at::cuda::getCurrentCUDAStream(); - cross_entropy_forward_inplace<<(stream)>>>(n, x_ptr, target_ptr, output_ptr, ignore_index); + cross_entropy_forward_bf16<<(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(); auto output_ptr = reinterpret_cast(grad_output); - // auto target_ptr = target.data_ptr(); auto target_ptr = reinterpret_cast(target); - // auto x_ptr = reinterpret_cast(x.data_ptr()); - auto x_ptr = reinterpret_cast(x); int32_t threads = 1024; - // auto stream = at::cuda::getCurrentCUDAStream(); - cross_entropy_backward_inplace<<(stream)>>>(n, output_ptr, target_ptr, x_ptr, ignore_index); + cross_entropy_backward_inplace_bf16<<(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 +#include +#include #include -#if __CUDA_ARCH__ >= 800 -#include -#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(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<<(stream)>>>(n, g_bf16, mid_ptr); bmt_has_nan_inf_reduce<<<1, block_size, 0, reinterpret_cast(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 +#include #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 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}")