diff --git a/ggml-metal.m b/ggml-metal.m index b47a98e21..372d3e696 100644 --- a/ggml-metal.m +++ b/ggml-metal.m @@ -1035,10 +1035,11 @@ void ggml_metal_graph_compute( const int n_dims = ((int32_t *) dst->op_params)[1]; const int mode = ((int32_t *) dst->op_params)[2]; - float freq_base; - float freq_scale; + float freq_base, freq_scale, ntk_factor, ext_factor; memcpy(&freq_base, (int32_t *) dst->op_params + 4, sizeof(float)); memcpy(&freq_scale, (int32_t *) dst->op_params + 5, sizeof(float)); + memcpy(&ntk_factor, (int32_t *) dst->op_params + 6, sizeof(float)); + memcpy(&ext_factor, (int32_t *) dst->op_params + 7, sizeof(float)); [encoder setComputePipelineState:ctx->pipeline_rope]; [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0]; @@ -1064,6 +1065,8 @@ void ggml_metal_graph_compute( [encoder setBytes:&mode length:sizeof( int) atIndex:20]; [encoder setBytes:&freq_base length:sizeof(float) atIndex:21]; [encoder setBytes:&freq_scale length:sizeof(float) atIndex:22]; + [encoder setBytes:&ntk_factor length:sizeof(float) atIndex:23]; + [encoder setBytes:&ext_factor length:sizeof(float) atIndex:24]; [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)]; } break; diff --git a/ggml-metal.metal b/ggml-metal.metal index 8d26b5ec2..347fd17ac 100644 --- a/ggml-metal.metal +++ b/ggml-metal.metal @@ -597,6 +597,55 @@ kernel void kernel_alibi_f32( } } +static float rope_ntkv2_ramp(const float low, const float high, const int i0) { + const float y = (i0 / 2 - low) / min(0.001f, high - low); + return 1.0f - min(1.0f, max(0.0f, y)); +} + +// NTKv2 algorithm based on LlamaPartNTKScaledRotaryEmbedding.py from https://github.com/jquesnelle/scaled-rope +// MIT licensed. Copyright (c) 2023 Jeffrey Quesnelle and Bowen Peng. +static float rope_ntkv2( + const float theta_base, + const float theta_linear, + const float theta_ntk, + const float corr_factors[4], + const int64_t i0, + const float ntk_factor, + const float ext_factor) { + float ramp_mix; + float theta; + + ramp_mix = rope_ntkv2_ramp(corr_factors[0], corr_factors[1], i0) * ntk_factor; + theta = theta_linear * (1 - ramp_mix) + theta_ntk * ramp_mix; + + ramp_mix = rope_ntkv2_ramp(corr_factors[2], corr_factors[3], i0) * ext_factor; + theta = theta * (1 - ramp_mix) + theta_base * ramp_mix; + return theta; +} + +// Interpolation constants found experimentally for LLaMA (might not be totally optimal though) +// Do not change unless there is a good reason for doing so! +constant float BETA_0 = 1.75f; +constant float BETA_1 = 1.25f; +constant float GAMMA_0 = 16.0f; +constant float GAMMA_1 = 2.0f; + +constant float max_pos_emb = 2048; + +// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get +// `corr_fac(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))` +static float rope_ntkv2_corr_factor(const int n_dims, const float n_rot, const float base) { + return n_dims * log(max_pos_emb / (n_rot * 2 * M_PI_F)) / (2 * log(base)); +} + +static void rope_ntkv2_corr_factors(int n_dims, const float freq_base, float factors[4]) { + // start and end correction factors + factors[0] = max(0.0f, floor(rope_ntkv2_corr_factor(n_dims, BETA_0, freq_base))); + factors[1] = min(n_dims - 1.0f, ceil(rope_ntkv2_corr_factor(n_dims, BETA_1, freq_base))); + factors[2] = max(0.0f, floor(rope_ntkv2_corr_factor(n_dims, GAMMA_0, freq_base))); + factors[3] = min(n_dims - 1.0f, ceil(rope_ntkv2_corr_factor(n_dims, GAMMA_1, freq_base))); +} + kernel void kernel_rope( device const void * src0, device float * dst, @@ -621,24 +670,33 @@ kernel void kernel_rope( constant int & mode, constant float & freq_base, constant float & freq_scale, + constant float & ntk_factor, + constant float & ext_factor, uint3 tpig[[thread_position_in_grid]]) { const int64_t i3 = tpig[2]; const int64_t i2 = tpig[1]; const int64_t i1 = tpig[0]; - const bool is_neox = mode & 2; const float theta_scale = pow(freq_base, -2.0f/n_dims); + const float theta_ntk_scale = pow(freq_base * pow(freq_scale, (n_dims / (n_dims - 2.0f))), -2.0f/n_dims); + float corr_factors[4]; + rope_ntkv2_corr_factors(n_dims, freq_base, corr_factors); - const int64_t p = ((mode & 1) == 0 ? n_past + i2 : i2); + float theta_base = (mode & 1) == 0 ? n_past + i2 : i2; + float theta_ntk = theta_base; - float theta = freq_scale * (float)p; + const bool is_neox = mode & 2; if (!is_neox) { for (int64_t i0 = 0; i0 < ne0; i0 += 2) { + const float theta_linear = freq_scale * theta_base; + const float theta = rope_ntkv2(theta_base, theta_linear, theta_ntk, corr_factors, + i0, ntk_factor, ext_factor); const float cos_theta = cos(theta); const float sin_theta = sin(theta); - theta *= theta_scale; + theta_base *= theta_scale; + theta_ntk *= theta_ntk_scale; device const float * const src = (device float *)((device char *) src0 + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); device float * dst_data = (device float *)((device char *) dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); @@ -650,6 +708,7 @@ kernel void kernel_rope( dst_data[1] = x0*sin_theta + x1*cos_theta; } } else { + theta_base *= freq_scale; // TODO: implement } }