# llm.c 筆記 (0) 開發方式

TLDR. 先用 pytorch 簡單做出該層，並產生測資 ln.bin
再跑C實作，C實作讀取 ln.bin後計算並驗證，再輸出驗證結果。

import torch

eps = 1e-5

class LayerNorm:

    @staticmethod
    def forward(x, w, b):
        B, T, C = x.size()
        mean = x.sum(-1, keepdim=True) / C # B,T,1
        xshift = x - mean # B,T,C
        var = (xshift**2).sum(-1, keepdim=True) / C # B,T,1
        rstd = (var + eps) ** -0.5 # B,T,1
        norm = xshift * rstd # B,T,C
        out = norm * w + b # B,T,C

        cache = (x, w, mean, rstd)
        return out, cache

    @staticmethod
    def backward(dout, cache):
        x, w, mean, rstd = cache
        # recompute the norm (save memory at the cost of compute)
        norm = (x - mean) * rstd
        # gradients for weights, bias
        db = dout.sum((0, 1))
        dw = (dout * norm).sum((0, 1))
        # gradients for input
        dnorm = dout * w
        dx = dnorm - dnorm.mean(-1, keepdim=True) - norm * (dnorm * norm).mean(-1, keepdim=True)
        dx *= rstd
        return dx, dw, db

# create a small dummy example and check w.r.t PyTorch backward
B = 2
T = 3
C = 4
x = torch.randn(B, T, C, requires_grad=True)
w = torch.randn(C, requires_grad=True)
b = torch.randn(C, requires_grad=True)
out, cache = LayerNorm.forward(x, w, b)

dout = torch.randn(B, T, C)
dx, dw, db = LayerNorm.backward(dout, cache)

# compare to PyTorch autograd
fakeloss = (out * dout).sum()
fakeloss.backward()
print("dx error:", (x.grad - dx).abs().max().item())
print("dw error:", (w.grad - dw).abs().max().item())
print("db error:", (b.grad - db).abs().max().item())

# for reference checking in C also
x, w, mean, rstd = cache

def write(tensor, handle):
    handle.write(tensor.detach().numpy().astype("float32").tobytes())

# Write to file
with open('ln.bin', 'wb') as file:
    write(x, file) # (B, T, C)
    write(w, file) # (C, )
    write(b, file) # (C, )
    write(out, file) # (B, T, C)
    write(mean, file) # (B, T)
    write(rstd, file) # (B, T)
    write(dout, file) # (B, T, C)
    write(dx, file) # (B, T, C)
    write(dw, file) # (C, )
    write(db, file) # (C, )

// must run `python layernorm.py` first to generate the reference data
// then compile for example as `gcc layernorm.c -o layernorm -lm`
// and then run as `./layernorm` to see the output

#include <stdio.h>
#include <stdlib.h>
#include <math.h>

void layernorm_forward(float* out, float* mean, float* rstd,
                       float* inp, float* weight, float* bias,
                       int B, int T, int C) {
    float eps = 1e-5f;
    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            // seek to the input position inp[b,t,:]
            float* x = inp + b * T * C + t * C;
            // calculate the mean
            float m = 0.0f;
            for (int i = 0; i < C; i++) {
                m += x[i];
            }
            m = m/C;
            // calculate the variance (without any bias correction)
            float v = 0.0f;
            for (int i = 0; i < C; i++) {
                float xshift = x[i] - m;
                v += xshift * xshift;
            }
            v = v/C;
            // calculate the rstd
            float s = 1.0f / sqrtf(v + eps);
            // seek to the output position in out[b,t,:]
            float* out_bt = out + b * T * C + t * C;
            for (int i = 0; i < C; i++) {
                float n = (s * (x[i] - m)); // normalized output
                float o = n * weight[i] + bias[i]; // scale and shift it
                out_bt[i] = o; // write
            }
            // cache the mean and rstd for the backward pass later
            mean[b * T + t] = m;
            rstd[b * T + t] = s;
        }
    }
}

void layernorm_backward(float* dinp, float* dweight, float* dbias,
                        float* dout, float* inp, float* weight, float* mean, float* rstd,
                        int B, int T, int C) {
    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            float* dout_bt = dout + b * T * C + t * C;
            float* inp_bt = inp + b * T * C + t * C;
            float* dinp_bt = dinp + b * T * C + t * C;
            float mean_bt = mean[b * T + t];
            float rstd_bt = rstd[b * T + t];

            // first: two reduce operations
            float dnorm_mean = 0.0f;
            float dnorm_norm_mean = 0.0f;
            for (int i = 0; i < C; i++) {
                float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
                float dnorm_i = weight[i] * dout_bt[i];
                dnorm_mean += dnorm_i;
                dnorm_norm_mean += dnorm_i * norm_bti;
            }
            dnorm_mean = dnorm_mean / C;
            dnorm_norm_mean = dnorm_norm_mean / C;

            // now iterate again and accumulate all the gradients
            for (int i = 0; i < C; i++) {
                float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
                float dnorm_i = weight[i] * dout_bt[i];
                // gradient contribution to bias
                dbias[i] += dout_bt[i];
                // gradient contribution to weight
                dweight[i] += norm_bti * dout_bt[i];
                // gradient contribution to input
                float dval = 0.0f;
                dval += dnorm_i; // term 1
                dval -= dnorm_mean; // term 2
                dval -= norm_bti * dnorm_norm_mean; // term 3
                dval *= rstd_bt; // final scale
                dinp_bt[i] += dval;
            }
        }
    }
}

// poor man's tensor checker
int check_tensor(float *a, float *b, int n, char* label) {
    int ok = 1;
    printf("%s\n", label);
    for (int i = 0; i < n; i++) {
        if (fabs(a[i] - b[i]) <= 1e-5) {
            printf("OK ");
        } else {
            printf("NOT OK ");
            ok = 0;
        }
        printf("%f %f\n", a[i], b[i]);
    }
    return ok;
}

int main() {

    int B = 2; // batch
    int T = 3; // time / sequence length
    int C = 4; // number of channels

    float* x = (float*) malloc(B * T * C * sizeof(float));
    float* w = (float*) malloc(C * sizeof(float));
    float* b = (float*) malloc(C * sizeof(float));
    float* out = (float*) malloc(B * T * C * sizeof(float));
    float* mean = (float*) malloc(B * T * sizeof(float));
    float* rstd = (float*) malloc(B * T * sizeof(float));
    float* dout = (float*) malloc(B * T * C * sizeof(float));
    float* dx = (float*) malloc(B * T * C * sizeof(float));
    float* dw = (float*) malloc(C * sizeof(float));
    float* db = (float*) malloc(C * sizeof(float));

    // read reference information from Python
    FILE *file = fopen("ln.bin", "rb");
    if (file == NULL) {
        printf("Error opening file\n");
        return 1;
    }
    fread(x, sizeof(float), B * T * C, file);
    fread(w, sizeof(float), C, file);
    fread(b, sizeof(float), C, file);
    fread(out, sizeof(float), B * T * C, file);
    fread(mean, sizeof(float), B * T, file);
    fread(rstd, sizeof(float), B * T, file);
    fread(dout, sizeof(float), B * T * C, file);
    fread(dx, sizeof(float), B * T * C, file);
    fread(dw, sizeof(float), C, file);
    fread(db, sizeof(float), C, file);
    fclose(file);

    // now let's calculate everything ourselves

    // forward pass
    float* c_out = (float*) malloc(B * T * C * sizeof(float));
    float* c_mean = (float*) malloc(B * T * sizeof(float));
    float* c_rstd = (float*) malloc(B * T * sizeof(float));
    layernorm_forward(c_out, c_mean, c_rstd, x, w, b, B, T, C);

    // check correctness of forward pass
    check_tensor(out, c_out, B*T*C, "out");
    check_tensor(mean, c_mean, B*T, "mean");
    check_tensor(rstd, c_rstd, B*T, "rstd");

    // backward pass (note calloc inits grads to zero)
    float* c_dx = (float*) calloc(B * T * C, sizeof(float));
    float* c_dw = (float*) calloc(B * T, sizeof(float));
    float* c_db = (float*) calloc(B * T, sizeof(float));
    layernorm_backward(c_dx, c_dw, c_db, dout, x, w, c_mean, c_rstd, B, T, C);

    // check correctness of backward pass
    check_tensor(c_dx, dx, B*T*C, "dx");
    check_tensor(c_dw, dw, C, "dw");
    check_tensor(c_db, db, C, "db");

    free(x);
    free(w);
    free(b);
    free(out);
    free(mean);
    free(rstd);
    free(dout);
    free(dx);
    free(dw);
    free(db);
    return 0;
}

LayerNorm 快速教程。
讓我們來看看如何處理 LayerNorm，作為模型中的一個例子層。
我們從 PyTorch 的 LayerNorm 文檔開始。
LayerNorm 當然是源自 Ba 等人在 2016 年的原始論文，
並被納入 Vaswani 等人的著名論文《Attention is All You Need》中的 Transformer。
GPT-2 採用了與 Transformer 相同的架構，
但 LayerNorm 的位置被著名地移至現在稱為預標準化版本的位置。
也就是說，Transformer 的殘差路徑保持清潔，LayerNorms 現在是 Transformer 
每個塊的第一層。這顯著改善了訓練穩定性。

當查看 PyTorch 的 LayerNorm 時，
首先要注意的是，你很可能找不到方程式的實際實現。
這是因為它深藏在代碼的第 30 層，背後是一個難以理解的動態分配器，
在一些可能自動生成的 CUDA 代碼中
（有興趣的可以查看 layer_norm.cpp 和 layer_norm_kernel.cu）。PyTorch 真的非常關心效率，
這是可以理解的。
但對於我們的目的來說，
我們首先要使用較簡單的 PyTorch 操作手動實現 LayerNorm。
這會比直接使用 LayerNorm 模塊的效率低得多，
但在演算法上具有教學意義。這是使用較簡單的 PyTorch 操作直接實現 LayerNorm 數學公式的方法：

import torch
eps = 1e-5

class LayerNorm:

    @staticmethod
    def forward(x, w, b):
        # x 是輸入激活值，形狀為 B,T,C
        # w 是權重，形狀為 C
        # b 是偏差，形狀為 C
        B, T, C = x.size()
        # 計算平均值
        mean = x.sum(-1, keepdim=True) / C  # B,T,1
        # 計算方差
        xshift = x - mean  # B,T,C
        var = (xshift**2).sum(-1, keepdim=True) / C  # B,T,1
        # 計算逆標準差：**0.5 是平方根, **-0.5 是 1/平方根
        rstd = (var + eps) ** -0.5  # B,T,1
        # 正常化輸入激活值
        norm = xshift * rstd  # B,T,C
        # 在最後對正常化的激活值進行縮放和移位
        out = norm * w + b  # B,T,C

        # 返回輸出和 cache，cache 中包含之後反向傳播過程中需要的變數
        cache = (x, w, mean, rstd)
        return out, cache


這裡我們用一些隨機數來進行這一層的前向傳播：


B = 2  # 這裡有一些玩具數字
T = 3
C = 4
x = torch.randn(B, T, C, requires_grad=True)
w = torch.randn(C, requires_grad=True)
b = torch.randn(C, requires_grad=True)
out, cache = LayerNorm.forward(x, w, b)

我們得到的輸出是形狀為 B,T,C 的張量 out，其中每個 C 維的“纖維”（我們稱之為激活）都被正常化，然後縮放，最後還被這一層的權重和偏差移位。重要的是，我們還返回了一個變數 cache，它是一個包含輸入激活值 x、權重 w、平均值 mean 和倒數標準差 rstd 的元組。這些都是我們在反向傳播過程中需要的變數。

PyTorch 當然可以為我們進行這一層的反向傳播：


dout = torch.randn(B, T, C)
fakeloss = (out * dout).sum()
fakeloss.backward()
你可以看到，我們創建了一個 fakeloss，
它簡單地將所有 layernorm 的輸出的（隨機）加權組合起來。
這就是將所有 B,T,C 數字投影到一個單一的標量值（損失）上，
這樣我們就有了我們“計算圖”的一個單一輸出。
通常這將是模型的損失，但在這裡我們只是進行一個假的損失。
然後我們在這個標量上調用 backward()，PyTorch 將為我們計算所有輸入到這個圖的梯度——即輸入激活值 x、權重 w 和偏差 b。
如果你對 autograd 不太熟悉，我建議你觀看我的 micrograd 視頻，
在那裡我們構建了一個微小的 autograd 引擎。
所以 PyTorch autograd 的神奇之處在於，
在我們調用 .backward 之後，它將用損失相對於該張量的梯度填充所有具有 requires_grad=True 的張量的 .grad 屬性。
這些梯度告訴我們 x、w、b 中所有輸入數字的損失斜率。因此，x.grad、w.grad 和 b.grad 的形狀與 x、w 和 b 的形狀完全相同。

但我們不想使用 PyTorch Autograd。
我們希望手動進行反向傳播。因此，我們拿出筆和紙，
寫出 LayerNorm 的表達式。前向傳播有以下數學形式：

其中：

是逐元素乘法，
是平均值，
是方差，以及
是一個小常數，以避免除以零。根據微積分的微分規則，
我們現在想要導出梯度。對於這部分，我的視頻《成為反向傳播忍者》
可能非常有幫助，因為我詳細地處理了一個類似的層——批量歸一化層。
當你處理微分時，你會注意到表達式在分析上簡化了，
你可以移動條款並簡化表達式。所以你不必手動反向每一行在前向傳播中。
特別是，我們得到：


    @staticmethod
    def backward(dout, cache):
        x, w, mean, rstd = cache
        # 重新計算 norm（以計算代價節省記憶體）
        norm = (x - mean) * rstd
        # 權重、偏差的梯度
        db = dout.sum((0, 1))
        dw = (dout * norm).sum((0, 1))
        # 輸入的梯度
        dnorm = dout * w
        dx = dnorm - dnorm.mean(-1, keepdim=True) - norm * (dnorm * norm).mean(-1, keepdim=True)
        dx *= rstd
        return dx, dw, db

因此，根據儲存在 dout 中的每個個別輸出數字的梯度和前向傳播的緩存，
我們現在可以通過這一層反向傳播到輸入中，
以繼續反向傳播的鏈式規則。
所以現在我們可以進行我們自己的反向傳播並看到它們匹配（誤差很小）：


dx, dw, db = LayerNorm.backward(dout, cache)
print("dx error:", (x.grad - dx).abs().max().item())
print("dw error:", (w.grad - dw).abs().max().item())
print("db error:", (b.grad - db).abs().max().item())

注意一件事。在反向傳播中，我们重新计算了变量 norm。
我们已经在前向传播中计算过这个变量，但之后我们把它丢掉了！
我们其实可以把这也作为缓存的一部分，以节省这次重新计算？
实际上我们完全可以这样做，你当然会得到完全相同的结果。
我们保存到我们缓存中的东西完全取决于我们。
我们甚至不必保存平均值和 rstd，我们可以在反向传播中重新计算它们。
区别在于平均值和 rstd 都非常小，只有 B,T 的形状，
而 norm 的形状是 B,T,C。所以这只是内存和计算之间的一种权衡。
通过不保持 norm 在缓存中，我们节省了内存，
但我们在反向传播中用了一点计算来换取。这在所有层中都很常见，
你会看到深度学习框架中各种层的不同实现可能都有不同的“检查点设置”。
是的，令人困惑的是，这被称为检查点，并且与将模型权重保存到磁盘无关。
这是关于在前向传播中保存中间变量以节省在反向传播中的计算。

好了，这就是使用 PyTorch 张量的版本。
现在我们必须将这转移到 C 语言并去掉 Tensor 抽象。
在给你完整的前向传播实现之前，简单谈谈 Tensors。
什么是 Tensors？它们是 1) 一个称为 Storage 的 1D 内存块，
用来存储原始数据，以及 2) 
一个在该存储上保持其形状的 View。 
PyTorch 内部机制在这里可能会有帮助。例如，如果我们有 3D 张量：


torch.manual_seed(42)
B, T, C = 2, 3, 4
a = torch.randn(B, T, C)
print(a)

tensor([[[ 1.9269,  1.4873,  0.9007, -2.1055],
         [ 0.6784, -1.2345, -0.0431, -1.6047],
         [ 0.3559, -0.6866, -0.4934,  0.2415]],

        [[-1.1109,  0.0915, -2.3169, -0.2168],
         [-0.3097, -0.3957,  0.8034, -0.6216],
         [-0.5920, -0.0631, -0.8286,  0.3309]]])

这是一个 2x3x4 的张量，
但它的底层内存只是一个大小为 234=24 的单一 1D 数组。View 只是这个 1D 数组上的一个形状。
所以现在当我们索引这个 PyTorch 张量，
例如 a[1,2,3]，PyTorch 
计算到 1D 数组的偏移量为 134 + 24 + 3 = 23，
并返回该偏移量的值。一般的公式是，
如果你想检索任何元素 b,t,c，
你计算到 Storage 的偏移量为 bTC + tC + c。例如：


b,t,c = 1,2,3
print(a[b,t,c])
print(a.view(-1)[b*T*C + t*C + c])
这两个都打印 0.3309。所以这样，我们就知道如何访问所有个别元素，
并如何偏移所有指针。特别注意的是，通道维度是最内层维度。
所以当我们增加偏移量时，我们正在遍历通道维度。
这对我们 C 实现的内存布局很重要。相当于在 C 中的前向传播变成：


#include <stdio.h>
#include <stdlib.h>
#include <math.h>

void layernorm_forward(float* out, float* mean, float* rstd,
                       float* inp, float* weight, float* bias,
                       int B, int T, int C) {
    float eps = 1e-5f;
    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            // 定位到输入位置 inp[b,t,:]
            float* x = inp + b * T * C + t * C;
            // 计算平均值
            float m = 0.0f;
            for (int i = 0; i < C; i++) {
                m += x[i];
            }
            m = m/C;
            // 计算方差（不进行任何偏差修正）
            float v = 0.0f;
            for (int i = 0; i < C; i++) {
                float xshift = x[i] - m;
                v += xshift * xshift;
            }
            v = v/C;
            // 计算 rstd
            float s = 1.0f / sqrtf(v + eps);
            // 定位到输出位置 out[b,t,:]
            float* out_bt = out + b * T * C + t * C;
            for (int i = 0; i < C; i++) {
                float n = (s * (x[i] - m)); // 正规化输出
                float o = n * weight[i] + bias[i]; // 缩放并移位
                out_bt[i] = o; // 写入
            }
            // 缓存平均值和 rstd 以备后续反向传播使用
            mean[b * T + t] = m;
            rstd[b * T + t] = s;
        }
    }
}
你会看到我如何偏移指针到 inp[b,t]，然后你知道接下来的 C 元素是该位置（批次，时间）的通道。反向传播：

void layernorm_backward(float* dinp, float* dweight, float* dbias,
                        float* dout, float* inp, float* weight, float* mean, float* rstd,
                        int B, int T, int C) {
    for (int b = 0; b < B; b++) {
        for (int t = 0; t < T; t++) {
            float* dout_bt = dout + b * T * C + t * C;
            float* inp_bt = inp + b * T * C + t * C;
            float* dinp_bt = dinp + b * T * C + t * C;
            float mean_bt = mean[b * T + t];
            float rstd_bt = rstd[b * T + t];

            // 首先：两个 reduce 操作
            float dnorm_mean = 0.0f;
            float dnorm_norm_mean = 0.0f;
            for (int i = 0; i < C; i++) {
                float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
                float dnorm_i = weight[i] * dout_bt[i];
                dnorm_mean += dnorm_i;
                dnorm_norm_mean += dnorm_i * norm_bti;
            }
            dnorm_mean = dnorm_mean / C;
            dnorm_norm_mean = dnorm_norm_mean / C;

            // 现在再次迭代并累积所有梯度
            for (int i = 0; i < C; i++) {
                float norm_bti = (inp_bt[i] - mean_bt) * rstd_bt;
                float dnorm_i = weight[i] * dout_bt[i];
                // 对偏差的梯度贡献
                dbias[i] += dout_bt[i];
                // 对权重的梯度贡献
                dweight[i] += norm_bti * dout_bt[i];
                // 对输入的梯度贡献
                float dval = 0.0f;
                dval += dnorm_i; // 第 1 项
                dval -= dnorm_mean; // 第 2 项
                dval -= norm_bti * dnorm_norm_mean; // 第 3 项
                dval *= rstd_bt; // 最后的缩放
                dinp_bt[i] += dval;
            }
        }
    }
}

还有一个细节需要注意，我们总是 += 进入梯度。我们从不使用 =，
也不使用 *=。这在风格上很重要，
因为如果你有一个变量在图中多次使用，反向传播的梯度总是累加。
在这个仓库中这不重要，因为我们没有复杂的分支，但这是正确的。
所以在训练中我们总是先做 zero_grad 把所有梯度设置为零，
然后在反向传播中累积它们。

关于训练和推理之间的差异的另一点说明。
你们中的一些人可能已经看过我之前的项目 llama2.c，
它在纯 C 中推理 Llama 2 架构。
与 GPT-2 不同，Llama 2 用更简单的 RMSNorm 替换了 LayerNorm。
你可以在 llama2.c 中看到 RMSNorm 的实现，这里复制粘贴：


void rmsnorm(float* o, float* x, float* weight, int size) {
    // 计算平方和
    float ss = 0.0f;
    for (int j = 0; j < size; j++) {
        ss += x[j] * x[j];
    }
    ss /= size;
    ss += 1e-5f;
    ss = 1.0f / sqrtf(ss);
    // 正规化并缩放
    for (int j = 0; j < size; j++) {
        o[j] = weight[j] * (ss * x[j]);
    }
}

这与我们上面的 LayerNorm 有何不同？

首先，从算法上看，你会注意到 RMSNorm 不追踪或减去平均值，
它只通过 norm 正规化。注意：norm，不是标准差，
因为我们没有减去平均值。这是对层的一种简化，现在已经变得非常流行，
因为它的效果一样好，如果不是更好的话。
此外，RMSNorm 没有偏差，它只有一个权重用于正规化后的缩放。
一般来说，GPT-2 在各处使用了太多的偏差，
事实证明你可以去掉这些——从所有线性层和 LayerNorms 中去掉。
网络可以“模拟”偏差，如果需要的话，
例如通过分配一个通道维度为常数（数据无关），
然后任何乘
以这个常数维度的权重实际上就像偏差一样工作。这大大简化了很多代码。

第二，推理代码没有批次维度 B，即假设批次大小为 1。
原则上，你也可以进行批量推理，
特别是如果你希望托管一个你期望许多同时查询的 LLM。
但如果你只是在本地运行一个 LLM，你可能只想有一个“生成流”，所以没有并行支持多个流的批次大小。为了保持简单，llama2.c 没有批次化，
因此你不会看到像 for (int b = 0; b < B; b++) 这样的循环。

第三，这个推理代码在这个单独的层内没有时间维度 T。在训练期间，
我们可以在每个层内循环时间并计算所有时间步的 layernorm。但在推理时，我们必须一次生成一个令牌，将在时间 t 预测的令牌输入到下一个时间步骤 t+1 的 Transformer 前向传播中。所以这就是为什么你不会看到像 for (int t = 0; t < T; t++) 这样的循环在单独层内部。这种时间循环确实存在，但它在 Transformer 前向传播的外部。

你会看到我们没有跟踪任何中间计算、内存或缓存。这是因为在推理中，之后不会有 .backward 传递。我们只需要计算输出，不需要保留任何中间变量。因此，推理的内存消耗显著低于训练。我们可以负担得起只丢弃激活，并只保留“激活前沿”的内存。同样，没有必要为这个 RMSNorm 实现反向功能，因为没有反向传递。
由于所有这些差异，训练在算法和计算上都显著更复杂和复杂，这也是我为什么先写推理（llama2.c）再实现训练（llm.c，这里）的部分原因。最后，我在同一目录下附上两个辅助文件，包含完整的代码。首先：


python layernorm.py

用来写出来自 PyTorch 的参考数据。然后编译并运行 C 版本：


gcc layernorm.c -o layernorm -lm
./layernorm
你会看到一切都匹配得很好。

这只是 LayerNorm 的教程。
我们对所有其他层也进行了完全相同的过程。
大多数其他层实际上比 LayerNorm 更容易。希望这有帮助！