Skip to content

Instantly share code, notes, and snippets.

@douglarek

douglarek/train_gpt2.go

Created April 10, 2024 06:25
Show Gist options
  • Save douglarek/fbb8a373b52ef46deeac29d7b752b42b to your computer and use it in GitHub Desktop.
Save douglarek/fbb8a373b52ef46deeac29d7b752b42b to your computer and use it in GitHub Desktop.
Download ZIP
Raw
train_gpt2.go
package main
import (
"fmt"
"math"
"os"
"unsafe"
)
// ----------------------------------------------------------------------------
// 所有单个层的正向和反向传递
// encoder_forward 将输入编码为嵌入和位置编码的总和。
// out: 输出张量 (B, T, C)
// inp: 输入标记索引 (B, T)
// wte: 词嵌入权重 (V, C)
// wpe: 位置嵌入权重 (maxT, C)
// B: 批量大小
// T: 序列长度
// C: 通道数
func encoder_forward(out []float32, inp []int, wte []float32, wpe []float32, B, T, C int) {
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
// 在 out[b,t,:] 中找到输出位置
out_bt := out[b*T*C+t*C:]
// 获取 inp[b, t] 处标记的索引
ix := inp[b*T+t]
// 在 wte 中找到与标记对应的词嵌入位置
wte_ix := wte[ix*C:]
// 在 wpe 中找到与位置对应的嵌入位置
wpe_t := wpe[t*C:]
// 将两个向量相加并将结果存储在 out[b,t,:] 中
copy(out_bt, wte_ix) // 将 wte_ix 复制到 out_bt
for i := 0; i < C; i++ {
out_bt[i] += wpe_t[i]
}
}
}
}
// encoder_backward 计算词嵌入和位置嵌入梯度。
// dwte: 词嵌入梯度 (V, C)
// dwpe: 位置嵌入梯度 (maxT, C)
// dout: 输出梯度 (B, T, C)
// inp: 输入标记索引 (B, T)
// B: 批量大小
// T: 序列长度
// C: 通道数
func encoder_backward(dwte []float32, dwpe []float32, dout []float32, inp []int, B, T, C int) {
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
dout_bt := dout[b*T*C+t*C:]
ix := inp[b*T+t]
dwte_ix := dwte[ix*C:]
dwpe_t := dwpe[t*C:]
for i := 0; i < C; i++ {
d := dout_bt[i]
dwte_ix[i] += d
dwpe_t[i] += d
}
}
}
}
// layernorm_forward 对输入执行层归一化。
// out: 输出张量 (B, T, C)
// mean: 缓存的均值 (B, T)
// rstd: 缓存的倒数标准差 (B, T)
// inp: 输入张量 (B, T, C)
// weight: 权重 (C)
// bias: 偏差 (C)
// B: 批量大小
// T: 序列长度
// C: 通道数
func layernorm_forward(out []float32, mean []float32, rstd []float32, inp []float32, weight []float32, bias []float32, B, T, C int) {
eps := 1e-5
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
// 在 inp[b,t,:] 中找到输入位置
x := inp[b*T*C+t*C:]
// 计算均值
m := 0.0
for i := 0; i < C; i++ {
m += x[i]
}
m /= float64(C)
// 计算方差（无偏差校正）
v := 0.0
for i := 0; i < C; i++ {
xshift := x[i] - float32(m)
v += float64(xshift * xshift)
}
v /= float64(C)
// 计算倒数标准差
s := 1.0 / math.Sqrt(v+eps)
// 在 out[b,t,:] 中找到输出位置
out_bt := out[b*T*C+t*C:]
for i := 0; i < C; i++ {
n := float32(s) * (x[i] - float32(m)) // 归一化输出
o := n*weight[i] + bias[i] // 缩放和偏移
out_bt[i] = o // 写入
}
// 缓存均值和倒数标准差，以便稍后进行反向传递
mean[b*T+t] = float32(m)
rstd[b*T+t] = float32(s)
}
}
}
// layernorm_backward 计算层归一化的梯度。
// dinp: 输入梯度 (B, T, C)
// dweight: 权重梯度 (C)
// dbias: 偏差梯度 (C)
// dout: 输出梯度 (B, T, C)
// inp: 输入张量 (B, T, C)
// weight: 权重 (C)
// mean: 缓存的均值 (B, T)
// rstd: 缓存的倒数标准差 (B, T)
// B: 批量大小
// T: 序列长度
// C: 通道数
func layernorm_backward(dinp []float32, dweight []float32, dbias []float32, dout []float32, inp []float32, weight []float32, mean []float32, rstd []float32, B, T, C int) {
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
dout_bt := dout[b*T*C+t*C:]
inp_bt := inp[b*T*C+t*C:]
dinp_bt := dinp[b*T*C+t*C:]
mean_bt := mean[b*T+t]
rstd_bt := rstd[b*T+t]
// 首先：两个约简操作
dnorm_mean := 0.0
dnorm_norm_mean := 0.0
for i := 0; i < C; i++ {
norm_bti := (inp_bt[i] - mean_bt) * rstd_bt
dnorm_i := weight[i] * dout_bt[i]
dnorm_mean += float64(dnorm_i)
dnorm_norm_mean += float64(dnorm_i * norm_bti)
}
dnorm_mean /= float64(C)
dnorm_norm_mean /= float64(C)
// 现在再次迭代并累积所有梯度
for i := 0; i < C; i++ {
norm_bti := (inp_bt[i] - mean_bt) * rstd_bt
dnorm_i := weight[i] * dout_bt[i]
// 对偏差的梯度贡献
dbias[i] += dout_bt[i]
// 对权重的梯度贡献
dweight[i] += norm_bti * dout_bt[i]
// 对输入的梯度贡献
dval := 0.0
dval += float64(dnorm_i) // 第一项
dval -= dnorm_mean // 第二项
dval -= float64(norm_bti) * dnorm_norm_mean // 第三项
dval *= float64(rstd_bt) // 最终缩放
dinp_bt[i] += float32(dval)
}
}
}
}
// matmul_forward 执行矩阵乘法。
// out: 输出张量 (B, T, OC)
// inp: 输入张量 (B, T, C)
// weight: 权重 (OC, C)
// bias: 偏差 (OC)
// B: 批量大小
// T: 序列长度
// C: 输入通道数
// OC: 输出通道数
func matmul_forward(out []float32, inp []float32, weight []float32, bias []float32, B, T, C, OC int) {
// 大部分运行时间都花在这里和 matmul_backward 中
// OC 是“输出通道”的缩写
// inp 是 (B,T,C)，weight 是 (OC, C)，bias 是 (OC)
// out 将是 (B,T,OC)
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
out_bt := out[b*T*OC+t*OC:]
inp_bt := inp[b*T*C+t*C:]
for o := 0; o < OC; o++ {
val := 0.0
if bias != nil {
val = float64(bias[o])
}
wrow := weight[o*C:]
for i := 0; i < C; i++ {
val += float64(inp_bt[i] * wrow[i])
}
out_bt[o] = float32(val)
}
}
}
}
// matmul_backward 计算矩阵乘法的梯度。
// dinp: 输入梯度 (B, T, C)
// dweight: 权重梯度 (OC, C)
// dbias: 偏差梯度 (OC)
// dout: 输出梯度 (B, T, OC)
// inp: 输入张量 (B, T, C)
// weight: 权重 (OC, C)
// B: 批量大小
// T: 序列长度
// C: 输入通道数
// OC: 输出通道数
func matmul_backward(dinp []float32, dweight []float32, dbias []float32, dout []float32, inp []float32, weight []float32, B, T, C, OC int) {
// 大部分运行时间都花在这里和 matmul_forward 中
// 这个反向传递可以在一“轮”循环中完成
// 但这不能提供有效的并行化策略
// 首先反向传播到 inp，并行化 B,T
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
dout_bt := dout[b*T*OC+t*OC:]
dinp_bt := dinp[b*T*C+t*C:]
for o := 0; o < OC; o++ {
wrow := weight[o*C:]
d := dout_bt[o]
for i := 0; i < C; i++ {
dinp_bt[i] += wrow[i] * d
}
}
}
}
// 反向传播到 weight/bias，并行化输出通道 OC
for o := 0; o < OC; o++ {
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
dout_bt := dout[b*T*OC+t*OC:]
inp_bt := inp[b*T*C+t*C:]
dwrow := dweight[o*C:]
d := dout_bt[o]
if dbias != nil {
dbias[o] += d
}
for i := 0; i < C; i++ {
dwrow[i] += inp_bt[i] * d
}
}
}
}
}
// attention_forward 执行多头注意力。
// out: 输出张量 (B, T, C)
// preatt: 缓存的注意力分数 (B, NH, T, T)
// att: 缓存的注意力权重 (B, NH, T, T)
// inp: 输入张量 (B, T, 3C) Q, K, V
// B: 批量大小
// T: 序列长度
// C: 通道数
// NH: 头数
func attention_forward(out []float32, preatt []float32, att []float32, inp []float32, B, T, C, NH int) {
// 输入是 (B, T, 3C) Q, K, V
// preatt, att 是 (B, NH, T, T)
// 输出是 (B, T, C)
C3 := C * 3
hs := C / NH // 头大小
scale := 1.0 / math.Sqrt(float64(hs))
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
for h := 0; h < NH; h++ {
query_t := inp[b*T*C3+t*C3+h*hs:]
preatt_bth := preatt[b*NH*T*T+h*T*T+t*T:]
att_bth := att[b*NH*T*T+h*T*T+t*T:]
// 第一阶段：计算查询点积键和最大值
maxval := -10000.0 // TODO 更好的初始化
for t2 := 0; t2 <= t; t2++ {
key_t2 := inp[b*T*C3+t2*C3+h*hs+C:] // +C 因为它
// (query_t) 点积 (key_t2)
val := 0.0
for i := 0; i < hs; i++ {
val += float64(query_t[i] * key_t2[i])
}
val *= scale
if val > maxval {
maxval = val
}
preatt_bth[t2] = float32(val)
}
// 第二阶段：计算指数并跟踪总和
expsum := 0.0
for t2 := 0; t2 <= t; t2++ {
expv := math.Exp(float64(preatt_bth[t2]) - maxval)
expsum += expv
att_bth[t2] = float32(expv)
}
expsum_inv := 0.0
if expsum != 0.0 {
expsum_inv = 1.0 / expsum
}
// 第三阶段：归一化以获得 softmax
for t2 := 0; t2 < T; t2++ {
if t2 <= t {
att_bth[t2] *= float32(expsum_inv)
} else {
// 因果注意力掩码。这里不严格设置为零
// 仅为了调试和检查 PyTorch 而明确执行此操作
att_bth[t2] = 0.0
}
}
// 第四阶段：将加权值累积到注意力的输出中
out_bth := out[b*T*C+t*C+h*hs:]
for i := 0; i < hs; i++ {
out_bth[i] = 0.0
}
for t2 := 0; t2 <= t; t2++ {
value_t2 := inp[b*T*C3+t2*C3+h*hs+C*2:] // +C*2 因为它
att_btht2 := att_bth[t2]
for i := 0; i < hs; i++ {
out_bth[i] += att_btht2 * value_t2[i]
}
}
}
}
}
}
// attention_backward 计算多头注意力的梯度。
// dinp: 输入梯度 (B, T, 3C) Q, K, V
// dpreatt: 缓存的注意力分数梯度 (B, NH, T, T)
// datt: 缓存的注意力权重梯度 (B, NH, T, T)
// dout: 输出梯度 (B, T, C)
// inp: 输入张量 (B, T, 3C) Q, K, V
// att: 缓存的注意力权重 (B, NH, T, T)
// B: 批量大小
// T: 序列长度
// C: 通道数
// NH: 头数
func attention_backward(dinp []float32, dpreatt []float32, datt []float32, dout []float32, inp []float32, att []float32, B, T, C, NH int) {
// inp/dinp 是 (B, T, 3C) Q, K, V
// att/datt/dpreatt 是 (B, NH, T, T)
// dout 是 (B, T, C)
C3 := C * 3
hs := C / NH // 头大小
scale := 1.0 / math.Sqrt(float64(hs))
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
for h := 0; h < NH; h++ {
att_bth := att[b*NH*T*T+h*T*T+t*T:]
datt_bth := datt[b*NH*T*T+h*T*T+t*T:]
dpreatt_bth := dpreatt[b*NH*T*T+h*T*T+t*T:]
dquery_t := dinp[b*T*C3+t*C3+h*hs:]
query_t := inp[b*T*C3+t*C3+h*hs:]
// 反向传递 4，通过值累积
dout_bth := dout[b*T*C+t*C+h*hs:]
for t2 := 0; t2 <= t; t2++ {
value_t2 := inp[b*T*C3+t2*C3+h*hs+C*2:] // +C*2 因为它
dvalue_t2 := dinp[b*T*C3+t2*C3+h*hs+C*2:]
for i := 0; i < hs; i++ {
// 在正向传递中，这是：
// out_bth[i] += att_bth[t2] * value_t2[i];
// 所以现在我们有：
datt_bth[t2] += value_t2[i] * dout_bth[i]
dvalue_t2[i] += att_bth[t2] * dout_bth[i]
}
}
// 反向传递 2 和 3，softmax
// 注意 softmax（如 tanh）不需要输入 (preatt) 来反向传播
for t2 := 0; t2 <= t; t2++ {
for t3 := 0; t3 <= t; t3++ {
indicator := 0.0
if t2 == t3 {
indicator = 1.0
}
local_derivative := att_bth[t2] * (float32(indicator) - att_bth[t3])
dpreatt_bth[t3] += local_derivative * datt_bth[t2]
}
}
// 反向传递 1，查询 @ 键矩阵乘法
for t2 := 0; t2 <= t; t2++ {
key_t2 := inp[b*T*C3+t2*C3+h*hs+C:] // +C 因为它
dkey_t2 := dinp[b*T*C3+t2*C3+h*hs+C:] // +C 因为它
for i := 0; i < hs; i++ {
// 在正向传递中，这是：
// preatt_bth[t2] += (query_t[i] * key_t2[i]) * scale;
// 所以现在我们有：
dquery_t[i] += key_t2[i] * dpreatt_bth[t2] * float32(scale)
dkey_t2[i] += query_t[i] * dpreatt_bth[t2] * float32(scale)
}
}
}
}
}
}
// gelu_forward 对输入应用 GELU 激活函数。
// out: 输出张量 (N)
// inp: 输入张量 (N)
// N: 张量大小
func gelu_forward(out []float32, inp []float32, N int) {
s := math.Sqrt(2.0 / math.Pi)
for i := 0; i < N; i++ {
x := inp[i]
cube := 0.044715 * x * x * x
out[i] = 0.5 * x * (1.0 + float32(math.Tanh(s*(float64(x)+cube))))
}
}
// gelu_backward 计算 GELU 激活函数的梯度。
// dinp: 输入梯度 (N)
// inp: 输入张量 (N)
// dout: 输出梯度 (N)
// N: 张量大小
func gelu_backward(dinp []float32, inp []float32, dout []float32, N int) {
s := math.Sqrt(2.0 / math.Pi)
for i := 0; i < N; i++ {
x := inp[i]
cube := 0.044715 * x * x * x
tanh_arg := s * (float64(x) + cube)
tanh_out := math.Tanh(tanh_arg)
coshf_out := math.Cosh(tanh_arg)
sech_out := 1.0 / (coshf_out * coshf_out)
local_grad := 0.5*(1.0+tanh_out) + float64(x)*0.5*sech_out*s*(1.0+3.0*0.044715*float64(x)*float64(x))
dinp[i] += float32(local_grad) * dout[i]
}
}
// residual_forward 执行残差连接。
// out: 输出张量 (N)
// inp1: 第一个输入张量 (N)
// inp2: 第二个输入张量 (N)
// N: 张量大小
func residual_forward(out []float32, inp1 []float32, inp2 []float32, N int) {
copy(out, inp1) // 将 inp1 复制到 out
for i := 0; i < N; i++ {
out[i] += inp2[i]
}
}
// residual_backward 计算残差连接的梯度。
// dinp1: 第一个输入梯度 (N)
// dinp2: 第二个输入梯度 (N)
// dout: 输出梯度 (N)
// N: 张量大小
func residual_backward(dinp1 []float32, dinp2 []float32, dout []float32, N int) {
copy(dinp1, dout) // 将 dout 复制到 dinp1
copy(dinp2, dout) // 将 dout 复制到 dinp2
}
// softmax_forward 对输入应用 softmax 函数。
// probs: 输出概率 (B, T, V)
// logits: 输入 logits (B, T, V)
// B: 批量大小
// T: 序列长度
// V: 词汇表大小
func softmax_forward(probs []float32, logits []float32, B, T, V int) {
// 输出：probs 是概率的 (B,T,V)
// 输入：logits 是未归一化对数概率的 (B,T,V)
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
// probs <- softmax(logits)
logits_bt := logits[b*T*V+t*V:]
probs_bt := probs[b*T*V+t*V:]
maxval := -10000.0 // TODO 更好的初始化
for i := 0; i < V; i++ {
if logits_bt[i] > maxval {
maxval = logits_bt[i]
}
}
sum := 0.0
for i := 0; i < V; i++ {
probs_bt[i] = float32(math.Exp(float64(logits_bt[i]) - float64(maxval)))
sum += float64(probs_bt[i])
}
for i := 0; i < V; i++ {
probs_bt[i] /= float32(sum)
}
}
}
}
// crossentropy_forward 计算交叉熵损失。
// losses: 损失 (B, T)
// probs: 概率 (B, T, V)
// targets: 目标标记索引 (B, T)
// B: 批量大小
// T: 序列长度
// V: 词汇表大小
func crossentropy_forward(losses []float32, probs []float32, targets []int, B, T, V int) {
// 输出：losses 是每个位置的单个损失的 (B,T)
// 输入：probs 是概率的 (B,T,V)
// 输入：targets 是 (B,T) 的整数，给出 logits 中的正确索引
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
// loss = -log(probs[target])
probs_bt := probs[b*T*V+t*V:]
ix := targets[b*T+t]
losses[b*T+t] = -float32(math.Log(float64(probs_bt[ix])))
}
}
}
// crossentropy_softmax_backward 计算 softmax 和交叉熵的梯度。
// dlogits: logits 梯度 (B, T, V)
// dlosses: 损失梯度 (B, T)
// probs: 概率 (B, T, V)
// targets: 目标标记索引 (B, T)
// B: 批量大小
// T: 序列长度
// V: 词汇表大小
func crossentropy_softmax_backward(dlogits []float32, dlosses []float32, probs []float32, targets []int, B, T, V int) {
// 通过 softmax 和交叉熵反向传播
for b := 0; b < B; b++ {
for t := 0; t < T; t++ {
dlogits_bt := dlogits[b*T*V+t*V:]
probs_bt := probs[b*T*V+t*V:]
dloss := dlosses[b*T+t]
ix := targets[b*T+t]
for i := 0; i < V; i++ {
p := probs_bt[i]
indicator := 0.0
if i == ix {
indicator = 1.0
}
dlogits_bt[i] += (p - float32(indicator)) * dloss
}
}
}
}
// ----------------------------------------------------------------------------
// GPT-2 模型定义
// 模型参数
const NUM_PARAMETER_TENSORS = 16
type ParameterTensors struct {
wte []float32 // (V, C)
wpe []float32 // (maxT, C)
ln1w []float32 // (L, C)
ln1b []float32 // (L, C)
qkvw []float32 // (L, 3*C, C)
qkvb []float32 // (L, 3*C)
attprojw []float32 // (L, C, C)
attprojb []float32 // (L, C)
ln2w []float32 // (L, C)
ln2b []float32 // (L, C)
fcw []float32 // (L, 4*C, C)
fcb []float32 // (L, 4*C)
fcprojw []float32 // (L, C, 4*C)
fcprojb []float32 // (L, C)
lnfw []float32 // (C)
lnfb []float32 // (C)
}
// malloc_and_point_parameters 为参数分配内存并将各个张量指向正确的位置。
// params: 参数张量结构
// param_sizes: 每个参数张量的大小
// 返回值: 分配的内存块
func malloc_and_point_parameters(params *ParameterTensors, param_sizes []int) []float32 {
num_parameters := 0
for _, size := range param_sizes {
num_parameters += size
}
// 一次性分配所有参数的内存
params_memory := make([]float32, num_parameters)
// 分配所有张量
ptrs := []*[]float32{
&params.wte, &params.wpe, &params.ln1w, &params.ln1b, &params.qkvw, &params.qkvb,
&params.attprojw, &params.attprojb, &params.ln2w, &params.ln2b, &params.fcw, &params.fcb,
&params.fcprojw, &params.fcprojb, &params.lnfw, &params.lnfb,
}
params_memory_iterator := 0
for i, size := range param_sizes {
*ptrs[i] = params_memory[params_memory_iterator : params_memory_iterator+size]
params_memory_iterator += size
}
return params_memory
}
// 激活张量的数量
const NUM_ACTIVATION_TENSORS = 23
// 激活张量结构
type ActivationTensors struct {
encoded []float32 // (B, T, C)
ln1 []float32 // (L, B, T, C)
ln1_mean []float32 // (L, B, T)
ln1_rstd []float32 // (L, B, T)
qkv []float32 // (L, B, T, 3*C)
atty []float32 // (L, B, T, C)
preatt []float32 // (L, B, NH, T, T)
att []float32 // (L, B, NH, T, T)
attproj []float32 // (L, B, T, C)
residual2 []float32 // (L, B, T, C)
ln2 []float32 // (L, B, T, C)
ln2_mean []float32 // (L, B, T)
ln2_rstd []float32 // (L, B, T)
fch []float32 // (L, B, T, 4*C)
fch_gelu []float32 // (L, B, T, 4*C)
fcproj []float32 // (L, B, T, C)
residual3 []float32 // (L, B, T, C)
lnf []float32 // (B, T, C)
lnf_mean []float32 // (B, T)
lnf_rstd []float32 // (B, T)
logits []float32 // (B, T, V)
probs []float32 // (B, T, V)
losses []float32 // (B, T)
}
// malloc_and_point_activations 为激活分配内存并将各个张量指向正确的位置。
// acts: 激活张量结构
// act_sizes: 每个激活张量的大小
// 返回值: 分配的内存块
func malloc_and_point_activations(acts *ActivationTensors, act_sizes []int) []float32 {
num_activations := 0
for _, size := range act_sizes {
num_activations += size
}
// 一次性分配所有激活的内存
acts_memory := make([]float32, num_activations)
// 分配所有张量
ptrs := []*[]float32{
&acts.encoded, &acts.ln1, &acts.ln1_mean, &acts.ln1_rstd, &acts.qkv, &acts.atty,
&acts.preatt, &acts.att, &acts.attproj, &acts.residual2, &acts.ln2, &acts.ln2_mean,
&acts.ln2_rstd, &acts.fch, &acts.fch_gelu, &acts.fcproj, &acts.residual3, &acts.lnf,
&acts.lnf_mean, &acts.lnf_rstd, &acts.logits, &acts.probs, &acts.losses,
}
acts_memory_iterator := 0
for i, size := range act_sizes {
*ptrs[i] = acts_memory[acts_memory_iterator : acts_memory_iterator+size]
acts_memory_iterator += size
}
return acts_memory
}
// GPT2Config 定义 GPT-2 模型的配置。
type GPT2Config struct {
max_seq_len int // 最大序列长度，例如 1024
vocab_size int // 词汇表大小，例如 50257
num_layers int // 层数，例如 12
num_heads int // 注意力头数，例如 12
channels int // 通道数，例如 768
}
type GPT2 struct {
config GPT2Config
// 模型权重及其大小
params ParameterTensors
param_sizes [NUM_PARAMETER_TENSORS]int
params_memory []float32
num_parameters int // 模型参数数量
// 权重梯度
grads ParameterTensors
grads_memory []float32
// AdamW 优化器的缓冲区
m_memory []float32
v_memory []float32
// 模型激活及其大小
acts ActivationTensors
act_sizes [NUM_ACTIVATION_TENSORS]int
acts_memory []float32
// 激活梯度
grads_acts ActivationTensors
grads_acts_memory []float32
// 其他运行状态配置
batch_size int // 当前正向传递的批量大小 (B)
seq_len int // 当前正向传递的序列长度 (T)
inputs []int // 当前正向传递的输入标记
targets []int // 当前正向传递的目标标记
mean_loss float32 // 在具有目标的正向传递后，将填充平均损失
}
// gpt2_build_from_checkpoint 从检查点文件构建 GPT-2 模型。
// model: GPT-2 模型结构
// checkpoint_path: 检查点文件路径
func gpt2_build_from_checkpoint(model *GPT2, checkpoint_path string) {
// 从检查点文件读取模型
model_file, err := os.Open(checkpoint_path)
if err != nil {
fmt.Println("Error opening model file:", err)
os.Exit(1)
}
defer model_file.Close()
// 读取模型头
model_header := make([]int32, 256)
_, err = model_file.Read((*[1 << 30]byte)(unsafe.Pointer(&model_header[0]))[:256*4])
if err != nil {
fmt.Println("Error reading model header:", err)
os.Exit(1)
}
// 检查模型文件魔数和版本
if model_header[0] != 20240326 {
fmt.Println("Bad magic model file")
os.Exit(1)
}
if model_header[1] != 1 {
fmt.Println("Bad version in model file")
os.Exit(1)
}
// 读取超参数
maxT, V, L, NH, C := int(model_header[2]), int(model_header[3]), int(model_header[4]), int(model_header[5]), int(model_header[6])
model.config.max_seq_len = maxT
model.config.vocab_size = V
model.config.num_layers = L
model.config.num_heads = NH
model.config.channels = C
fmt.Println("[GPT-2]")
fmt.Println("max_seq_len:", maxT)
fmt.Println("vocab_size:", V)
fmt.Println("num_layers:", L)
fmt.Println("num_heads:", NH)
fmt.Println("channels:", C)
// 为所有参数分配空间并读取它们
model.param_sizes[0] = V * C
model.param_sizes[1] = maxT * C
model.param_sizes[2] = L * C
model.param_sizes[3] = L * C
model.param_sizes[4] = L * (3 * C) * C
model.param_sizes[5] = L * (3 * C)
model.param_sizes[6] = L * C * C
model.param_sizes[7] = L * C
model.param_sizes[8] = L * C
model.param_sizes[9] = L * C
model.param_sizes[10] = L * (4 * C) * C
model.param_sizes[11] = L * (4 * C)
model.param_sizes[12] = L * C * (4 * C)
model.param_sizes[13] = L * C
model.param_sizes[14] = C
model.param_sizes[15] = C
// 计算参数数量
num_parameters := 0
for _, size := range model.param_sizes {
num_parameters += size
}
fmt.Println("num_parameters:", num_parameters)
// 从文件读取所有参数
model.params_memory = malloc_and_point_parameters(&model.params, model.param_sizes)
_, err = model_file.Read((*[1 << 30]byte)(unsafe.Pointer(&model.params_memory[0]))[:num_parameters*4])
if err != nil {
fmt.Println("Error reading model parameters:", err)
os.Exit(1)
}
// 其他初始化
model.acts_memory = nil
model.grads_memory = nil
model.m_memory = nil
model.v_memory = nil
model.grads_acts_memory = nil
model.inputs = nil
model.targets = nil
model.batch_size = 0
model.seq_len = 0
model.mean_loss = -1.0 // -1.0 表示没有损失
}
// ----------------------------------------------------------------------------
// 模型正向和反向传播
// gpt2_forward 执行 GPT-2 模型的正向传递。
// model: GPT-2 模型结构
// inputs: 输入标记索引 (B, T)
// targets: 目标标记索引 (B, T) (可选，可以为 nil)
// B: 批量大小
// T: 序列长度
func gpt2_forward(model *GPT2, inputs []int, targets []int, B, T int) {
// 确保模型已初始化，否则报错
if model.params_memory == nil {
fmt.Println("Error: model was not initialized properly.")
os.Exit(1)
}
// 方便使用的参数
V := model.config.vocab_size
L := model.config.num_layers
NH := model.config.num_heads
C := model.config.channels
// 如果需要，为所有激活分配空间（在此处延迟完成）
if model.acts_memory == nil {
// 记录当前 B, T
model.batch_size = B
model.seq_len = T
// 分配空间
model.act_sizes[0] = B * T * C
model.act_sizes[1] = L * B * T * C
model.act_sizes[2] = L * B * T
model.act_sizes[3] = L * B * T
model.act_sizes[4] = L * B * T * 3 * C
model.act_sizes[5] = L * B * T * C
model.act_sizes[6] = L * B * NH * T * T
model.act_sizes[7] = L * B * NH * T * T
model.act_sizes[8] = L * B * T * C
model.act_sizes[9] = L * B * T * C
model.act_sizes[10] = L * B * T * C
model.act_sizes[11] = L * B * T
model.act_sizes[12] = L * B * T
model.act_sizes[13] = L * B * T * 4 * C
model.act_sizes[14] = L * B * T * 4 * C
model.act_sizes[15] = L * B * T * C
model.act_sizes[16] = L * B * T * C
model.act_sizes[17] = B * T * C
model.act_sizes[18] = B * T
model.act_sizes[19] = B * T
model.act_sizes[20] = B * T * V
model.act_sizes[21] = B * T * V
model.act_sizes[22] = B * T
num_activations := 0
for _, size := range model.act_sizes {
num_activations += size
}
fmt.Println("num_activations:", num_activations)
model.acts_memory = malloc_and_point_activations(&model.acts, model.act_sizes)
// 创建内存以缓存输入和目标
model.inputs = make([]int, B*T)
model.targets = make([]int, B*T) // 如果我们从未有过目标，则可能未使用，但它很小
} else {
// 验证 B, T 不大于先前分配的大小
// 原则上，我们可以重新分配更大的内存块，但现在我们只是报错
if B > model.batch_size || T > model.seq_len {
fmt.Println("Error: batch size or sequence length is inadequately large")
fmt.Printf("Model: B=%d T=%d, Desired: B=%d T=%d\n", model.batch_size, model.seq_len, B, T)
os.Exit(1)
}
}
// 缓存输入/目标
copy(model.inputs, inputs)
if targets != nil {
copy(model.targets, targets)
}
// 正向传递
params := model.params // 简洁起见
acts := model.acts
var residual []float32
encoder_forward(acts.encoded, inputs, params.wte, params.wpe, B, T, C) // 编码进入 residual[0]
for l := 0; l < L; l++ {
if l == 0 {
residual = acts.encoded
} else {
residual = acts.residual3[(l-1)*B*T*C:]
}
// 获取此层的权重指针
l_ln1w := params.ln1w[l*C:]
l_ln1b := params.ln1b[l*C:]
l_qkvw := params.qkvw[l*3*C*C:]
l_qkvb := params.qkvb[l*3*C:]
l_attprojw := params.attprojw[l*C*C:]
l_attprojb := params.attprojb[l*C:]
l_ln2w := params.ln2w[l*C:]
l_ln2b := params.ln2b[l*C:]
l_fcw := params.fcw[l*4*C*C:]
l_fcb := params.fcb[l*4*C:]
l_fcprojw := params.fcprojw[l*C*4*C:]
l_fcprojb := params.fcprojb[l*C:]
// 获取此层的激活指针
l_ln1 := acts.ln1[l*B*T*C:]
l_ln1_mean := acts.ln1_mean[l*B*T:]
l_ln1_rstd := acts.ln1_rstd[l*B*T:]
l_qkv := acts.qkv[l*B*T*3*C:]
l_atty := acts.atty[l*B*T*C:]
l_preatt := acts.preatt[l*B*NH*T*T:]
l_att := acts.att[l*B*NH*T*T:]
l_attproj := acts.attproj[l*B*T*C:]
l_residual2 := acts.residual2[l*B*T*C:]
l_ln2 := acts.ln2[l*B*T*C:]
l_ln2_mean := acts.ln2_mean[l*B*T:]
l_ln2_rstd := acts.ln2_rstd[l*B*T:]
l_fch := acts.fch[l*B*T*4*C:]
l_fch_gelu := acts.fch_gelu[l*B*T*4*C:]
l_fcproj := acts.fcproj[l*B*T*C:]
l_residual3 := acts.residual3[l*B*T*C:]
// 执行正向传递
layernorm_forward(l_ln1, l_ln1_mean, l_ln1_rstd, residual, l_ln1w, l_ln1b, B, T, C)
matmul_forward(l_qkv, l_ln1, l_qkvw, l_qkvb, B, T, C, 3*C)
attention_forward(l_atty, l_preatt, l_att, l_qkv, B, T, C, NH)
matmul_forward(l_attproj, l_atty, l_attprojw, l_attprojb, B, T, C, C)
residual_forward(l_residual2, residual, l_attproj, B*T*C)
layernorm_forward(l_ln2, l_ln2_mean, l_ln2_rstd, l_residual2, l_ln2w, l_ln2b, B, T, C)
matmul_forward(l_fch, l_ln2, l_fcw, l_fcb, B, T, C, 4*C)
gelu_forward(l_fch_gelu, l_fch, B*T*4*C)
matmul_forward(l_fcproj, l_fch_gelu, l_fcprojw, l_fcprojb, B, T, 4*C, C)
residual_forward(l_residual3, l_residual2, l_fcproj, B*T*C)
}
residual = acts.residual3[(L-1)*B*T*C:] // 最后一个残差在 residual3 中
layernorm_forward(acts.lnf, acts.lnf_mean, acts.lnf_rstd, residual, params.lnfw, params.lnfb, B, T, C)
matmul_forward(acts.logits, acts.lnf, params.wte, nil, B, T, C, V)
softmax_forward(acts.probs, acts.logits, B, T, V)
// 如果有目标，也正向传递交叉熵损失函数
if targets != nil {
crossentropy_forward(model.acts.losses, model.acts.probs, targets, B, T, V)
// 为了方便，也计算平均损失
mean_loss := 0.0
for _, loss := range model.acts.losses {
mean_loss += float64(loss)
}
mean_loss /= float64(B * T)
model.mean_loss = float32(mean_loss)
} else {
// 如果没有目标，则没有损失
model.mean_loss = -1.0
}
}
// gpt2_zero_grad 将模型的梯度清零。
// model: GPT-2 模型结构
func gpt2_zero_grad(model *GPT2) {
if model.grads_memory != nil {
for i := range model.grads_memory {
model.grads_memory[i] = 0
}
}
if model.grads_acts_memory != nil {
for i := range model.grads_acts_memory {
model.grads_acts_memory[i] = 0
}
}
}
// gpt2_backward 执行 GPT-2 模型的反向传播。
// model: GPT-2 模型结构
func gpt2_backward(model *GPT2) {
// 再次检查我们之前是否转发过，并带有目标
if model.mean_loss == -1.0 {
fmt.Println("Error: must forward with targets before backward")
os.Exit(1)
}
// 延迟分配权重和激活梯度的内存，如果需要
if model.grads_memory == nil {
model.grads_memory = malloc_and_point_parameters(&model.grads, model.param_sizes)
model.grads_acts_memory = malloc_and_point_activations(&model.grads_acts, model.act_sizes)
gpt2_zero_grad(model)
}
// 方便的快捷方式
B := model.batch_size
T := model.seq_len
V := model.config.vocab_size
L := model.config.num_layers
NH := model.config.num_heads
C := model.config.channels
// 反向传递
params := model.params // 简洁起见
grads := model.grads
acts := model.acts
grads_acts := model.grads_acts
// 我们通过用 1.0f/(B*T) 填充 dlosses 来启动链，以获得平均损失
dloss_mean := 1.0 / float32(B*T)
for i := range grads_acts.losses {
grads_acts.losses[i] = dloss_mean
}
crossentropy_softmax_backward(grads_acts.logits, grads_acts.losses, acts.probs, model.targets, B, T, V)
matmul_backward(grads_acts.lnf, grads.wte, nil, grads_acts.logits, acts.lnf, params.wte, B, T, C, V)
residual := acts.residual3[(L-1)*B*T*C:] // 最后一层的残差
dresidual := grads_acts.residual3[(L-1)*B*T*C:] // 写入最后一层的残差
layernorm_backward(dresidual, grads.lnfw, grads.lnfb, grads_acts.lnf, residual, params.lnfw, acts.lnf_mean, acts.lnf_rstd, B, T, C)
for l := L - 1; l >= 0; l-- {
if l == 0 {
residual = acts.encoded
dresidual = grads_acts.encoded
} else {
residual = acts.residual3[(l-1)*B*T*C:]
dresidual = grads_acts.residual3[(l-1)*B*T*C:]
}
// 获取此层的权重指针
l_ln1w := params.ln1w[l*C:]
l_qkvw := params.qkvw[l*3*C*C:]
l_attprojw := params.attprojw[l*C*C:]
l_ln2w := params.ln2w[l*C:]
l_fcw := params.fcw[l*4*C*C:]
l_fcprojw := params.fcprojw[l*C*4*C:]
// 获取此层的权重梯度指针
dl_ln1w := grads.ln1w[l*C:]
dl_ln1b := grads.ln1b[l*C:]
dl_qkvw := grads.qkvw[l*3*C*C:]
dl_qkvb := grads.qkvb[l*3*C:]
dl_attprojw := grads.attprojw[l*C*C:]
dl_attprojb := grads.attprojb[l*C:]
dl_ln2w := grads.ln2w[l*C:]
dl_ln2b := grads.ln2b[l*C:]
dl_fcw := grads.fcw[l*4*C*C:]
dl_fcb := grads.fcb[l*4*C:]
dl_fcprojw := grads.fcprojw[l*C*4*C:]
dl_fcprojb := grads.fcprojb[l*C:]
// 获取此层的激活指针
l_ln1 := acts.ln1[l*B*T*C:]
l_ln1_mean := acts.ln1_mean[l*B*T:]
l_ln1_rstd := acts.ln1_rstd[l*B*T:]
l_qkv := acts.qkv[l*B*T*3*C:]
l_atty := acts.atty[l*B*T*C:]
l_att := acts.att[l*B*NH*T*T:]
l_residual2 := acts.residual2[l*B*T*C:]
l_ln2 := acts.ln2[l*B*T*C:]
l_ln2_mean := acts.ln2_mean[l*B*T:]
l_ln2_rstd := acts.ln2_rstd[l*B*T:]
l_fch := acts.fch[l*B*T*4*C:]
l_fch_gelu := acts.fch_gelu[l*B*T*4*C:]
// 获取此层的激活梯度指针
dl_ln1 := grads_acts.ln1[l*B*T*C:]
dl_qkv := grads_acts.qkv[l*B*T*3*C:]
dl_atty := grads_acts.atty[l*B*T*C:]
dl_preatt := grads_acts.preatt[l*B*NH*T*T:]
dl_att := grads_acts.att[l*B*NH*T*T:]
dl_attproj := grads_acts.attproj[l*B*T*C:]
dl_residual2 := grads_acts.residual2[l*B*T*C:]
dl_ln2 := grads_acts.ln2[l*B*T*C:]
dl_fch := grads_acts.fch[l*B*T*4*C:]
dl_fch_gelu := grads_acts.fch_gelu[l*B*T*4*C:]
dl_fcproj := grads_acts.fcproj[l*B*T*C:]
dl_residual3 := grads_acts.residual3[l*B*T*C:]
// 反向传播此层
residual_backward(dl_residual2, dl_fcproj, dl_residual3, B*T*C)
matmul_backward(dl_fch_gelu, dl_fcprojw, dl_fcprojb, dl_fcproj, l_fch_gelu, l_fcprojw, B, T, 4*C, C)
gelu_backward(dl_fch, l_fch, dl_fch_gelu, B*T*4*C)
matmul_backward(dl_ln2, dl_fcw, dl_fcb, dl_fch, l_ln2, l_fcw, B, T, C, 4*C)
layernorm_backward(dl_residual2, dl_ln2w, dl_ln2b, dl_ln2, l_residual2, l_ln2w, l_ln2_mean, l_ln2_rstd, B, T, C)
residual_backward(dresidual, dl_attproj, dl_residual2, B*T*C)
matmul_backward(dl_atty, dl_attprojw, dl_attprojb, dl_attproj, l_atty, l_attprojw, B, T, C, C)
attention_backward(dl_qkv, dl_preatt, dl_att, dl_atty, l_qkv, l_att, B, T, C, NH)
matmul_backward(dl_ln1, dl_qkvw, dl_qkvb, dl_qkv, l_ln1, l_qkvw, B, T, C, 3*C)
layernorm_backward(dresidual, dl_ln1w, dl_ln1b, dl_ln1, residual, l_ln1w, l_ln1_mean, l_ln1_rstd, B, T, C)
}
encoder_backward(grads.wte, grads.wpe, grads_acts.encoded, model.inputs, B, T, C)
}
// ----------------------------------------------------------------------------
// 模型更新和释放
// gpt2_update 使用 AdamW 优化器更新模型参数。
// model: GPT-2 模型结构
// learning_rate: 学习率
// beta1: AdamW 参数 beta1
// beta2: AdamW 参数 beta2
// eps: AdamW 参数 epsilon
// weight_decay: 权重衰减
// t: 当前时间步
func gpt2_update(model *GPT2, learning_rate, beta1, beta2, eps, weight_decay float32, t int) {
// 参考：https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html
// 延迟分配 m_memory 和 v_memory 的内存
if model.m_memory == nil {
model.m_memory = make([]float32, model.num_parameters)
model.v_memory = make([]float32, model.num_parameters)
}
for i, param := range model.params_memory {
grad := model.grads_memory[i]
// 更新一阶矩（动量）
m := beta1*model.m_memory[i] + (1.0-beta1)*grad
// 更新二阶矩（RMSprop）
v := beta2*model.v_memory[i] + (1.0-beta2)*grad*grad
// 偏差校正两个矩
m_hat := m / (1.0 - float32(math.Pow(float64(beta1), float64(t))))
v_hat := v / (1.0 - float32(math.Pow(float64(beta2), float64(t))))
// 更新
model.m_memory[i] = m
model.v_memory[i] = v
model.params_memory[i] -= learning_rate * (m_hat/(float32(math.Sqrt(float64(v_hat)))+eps) + weight_decay*param)
}
}
// gpt2_free 释放 GPT-2 模型占用的内存。
// model: GPT-2 模型结构
func gpt2_free(model *GPT2) {
// Go 的垃圾回收机制会自动释放内存，因此无需手动释放
}
// ... (剩余代码省略) ...
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment