char-fork/llama.cpp
git clone https://git.t4t.associates/char-fork/llama.cpp
e29acf74f
char/aliceai-t5-moe
1#version 450 2 3#include "generic_head.glsl" 4#include "types.glsl" 5 6#extension GL_EXT_control_flow_attributes : enable 7 8layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; 9 10layout (binding = 0) buffer X {A_TYPE x[];}; 11layout (binding = 1) readonly buffer G {A_TYPE grad[];}; 12layout (binding = 2) buffer GM {A_TYPE gradm[];}; 13layout (binding = 3) buffer GV {A_TYPE gradv[];}; 14layout (binding = 4) readonly buffer P {float params[7];}; 15 16void main() { 17 const uint i = gl_GlobalInvocationID.z * 262144 + gl_GlobalInvocationID.y * 512 + gl_GlobalInvocationID.x; 18 19 if (i >= p.KX) { 20 return; 21 } 22 23 const float alpha = params[0]; 24 const float beta1 = params[1]; 25 const float beta2 = params[2]; 26 const float eps = params[3]; 27 const float wd = params[4]; 28 const float beta1h = params[5]; 29 const float beta2h = params[6]; 30 31 const float gi = grad[i]; 32 const float gmi = gradm[i]*beta1 + gi*(1.0f - beta1); 33 const float gvi = gradv[i]*beta2 + gi*gi*(1.0f - beta2); 34 35 gradm[i] = gmi; 36 gradv[i] = gvi; 37 38 const float mh = gmi*beta1h; 39 const float vh = sqrt(gvi*beta2h) + eps; 40 41 x[i] = x[i]*(1.0f - alpha*wd) - alpha*mh/vh; 42}