rmsnorm backward simple baseline kernel

[ngc92]

baseline kernel and cpu versions

[azret]

Thank you for this! As a result of this PR I have found a bug in my own RMSNorm impl. where I forgot to multiply by the weight in the backward pass. Spent days chasing it.

PS: I've ran this rmsnorm_backward_cpu through my test cases that verify against torch and all is passing.

Thanks again!

[azret]

    void rms_backward(
        float* _In, float* _δ_Out, float* _δ_In,
        float* _δ_Gamma, float* _Gamma,
        int N, float EPSILON) {
        float rms = 0;
        for (int i = 0; i < N; i++) {
            rms += _In[i] * _In[i];
        }
        rms = std.sqrtf(rms / N + EPSILON);
        float rrms = 1.0f / rms;
        float rnrms = 1.0f / (N * rms * rms);
        float δ = 0;
        for (int i = 0; i < N; i++) {
            δ += _δ_Out[i] * _In[i] * rrms * **_Gamma[i]**; // This is what I missed!!
        }
        for (int i = 0; i < N; i++) {
            _δ_In[i] += _δ_Out[i] * _Gamma[i] * rrms - δ * _In[i] * rnrms;
            _δ_Gamma[i] += _δ_Out[i] * _In[i] * rrms;
        }
    }

[ngc92]

great to hear that this passes an independent set of tests. I did some testing, but wasn't yet fully sure that my implementation was correct.