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 的垃圾回收机制会自动释放内存，因此无需手动释放
}

// ... (剩余代码省略) ...