From bb34fd50a1fec74e62beb6e23d51f0142c7d0ab6 Mon Sep 17 00:00:00 2001 From: Tim Dettmers Date: Wed, 20 Oct 2021 18:37:44 -0700 Subject: Initial plumbing for skip_zeros. --- csrc/ops.cuh | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) (limited to 'csrc/ops.cuh') diff --git a/csrc/ops.cuh b/csrc/ops.cuh index e6033cb..465b4a4 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -49,7 +49,7 @@ template void dequantizeBlockwise(float *code, unsigned char *A, flo template void optimizer32bit(T* g, T* p, float* state1, float* state2, float *unorm, float max_unorm, float param_norm, float beta1, float beta2, float eps, float weight_decay, - int step, float lr, const float gnorm_scale, int n); + int step, float lr, const float gnorm_scale, bool skip_zeros, int n); template void optimizerStatic8bit(T* p, T* g, unsigned char* state1, unsigned char* state2, float *unorm, float max_unorm, float param_norm, @@ -62,7 +62,8 @@ template void optimizerStatic8bit(T* p, T* g, unsigne template void optimizerStatic8bitBlockwise(T* p, T* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float eps, int step, float lr, - float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, int n); + float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, float weight_decay, const float gnorm_scale, + bool skip_zeros, int n); template void percentileClipping(T * g, float *gnorm_vec, int step, const int n); -- cgit v1.2.3