import torch
from torch import nn
x = torch.tensor([[1.0], [2.0], [3.0]])  # 已知的输入
y = torch.tensor([[3.0], [5.0], [7.0]])  # 对应的答案
model = nn.Linear(1, 1)  # 创建模型，内部随机初始化 w 和 b
pred = model(x)  # Linear 内部用当前的 w、b 计算 w * x + b
print("训练前的预测：", pred)

答案：[3, 5, 7]
预测：[2, 4, 6]
loss = ((2-3)² + (4-5)² + (6-7)²) / 3 = 1

h = w1 * x + b1
y = w2 * h + b2
  = (w2 * w1) * x + (w2 * b1 + b2)

# Sequential 按顺序执行，上一层的输出交给下一层
model = nn.Sequential(
    nn.Linear(1, 64),   # 一个输入，算出 64 个中间值
    nn.Tanh(),         # 对每个值做非线性变换
    nn.Linear(64, 64),  # 对这些值加权求和，加上偏置，得到 64 个输出
    nn.Tanh(),
    nn.Linear(64, 1),   # 最后得到一个预测值
)

from __future__ import annotations
import math
import random
import torch
from torch import nn
def make_dataset(n: int = 2048) -> tuple[torch.Tensor, torch.Tensor]:
    # x 原本是一维向量，unsqueeze(1) 把它变成 (n, 1)。
    # 神经网络通常按二维 batch 输入理解数据：每一行是一个样本，每一列是一个特征。
    # Java 类比：原来像 double[n]，unsqueeze 后更像 double[n][1]。
    x = torch.linspace(-2 * math.pi, 2 * math.pi, n).unsqueeze(1)
    # randn_like(x) 生成和 x 形状相同的随机噪声。
    # 这里加噪声是为了让任务更像真实数据：真实数据通常不会完美落在一条函数曲线上。
    noise = 0.05 * torch.randn_like(x)
    # 这里造一个有规律但不完全干净的函数，让模型学 sin/cos 的组合。
    # x 是输入，y 是标准答案。训练时模型只看到 x，loss 会拿预测值和 y 比较。
    y = torch.sin(x) + 0.3 * torch.cos(3 * x) + noise
    return x, y
class TinyMLP(nn.Module):
    def __init__(self) -> None:
        # nn.Module 是所有 PyTorch 模型/层的基类。
        # super().__init__() 会初始化父类里负责登记参数、子模块等基础设施。
        super().__init__()
        # nn.Sequential 会按顺序执行这些层：
        # Linear 做仿射变换，Tanh 提供非线性，否则模型只能学直线。
        # Linear(1, 64)：每个样本输入 1 个数字，输出 64 维隐藏表示。
        # Linear(64, 1)：最后把 64 维隐藏表示压回 1 个预测值。
        self.net = nn.Sequential(
            nn.Linear(1, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, 1),
        )
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 在 PyTorch 里，调用 model(x) 时，实际会转到 forward(x)。
        # 这有点像 Java 里某个框架约定你实现 handle/request 方法，然后框架负责调用。
        return self.net(x)
def main() -> None:
    # 固定随机种子，方便你多次运行时看到接近的结果。
    torch.manual_seed(42)
    random.seed(42)
    x, y = make_dataset()
    # 简单切分：前 80% 做训练集，后 20% 做验证集。
    split = int(0.8 * len(x))
    x_train, y_train = x[:split], y[:split]
    x_val, y_val = x[split:], y[split:]
    model = TinyMLP()
    # optimizer 负责根据梯度更新模型参数；lr 是每次更新的步子大小。
    # model.parameters() 来自 nn.Module，会递归收集 self.net 里 Linear 的 weight/bias。
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)
    # MSELoss 衡量预测值和真实值的平方误差，适合这个回归任务。
    loss_fn = nn.MSELoss()
    print("shape check")
    print(f"x_train: {tuple(x_train.shape)}")
    print(f"y_train: {tuple(y_train.shape)}")
    print(f"prediction: {tuple(model(x_train[:8]).shape)}")
    print()
    batch_size = 128
    for step in range(1, 801):
        # 真实训练通常不会每一步都用完整训练集，而是抽一个 mini-batch。
        # 好处：计算更快，也让每步梯度带一点随机性，常常更容易训练。
        # 随机抽一批样本。idx 的 shape 是 (batch_size,)。
        idx = torch.randint(0, len(x_train), (batch_size,))
        # 前向传播：输入 x，得到预测 pred。
        # x_train[idx] 的 shape 是 (batch_size, 1)，对应一批样本。
        pred = model(x_train[idx])
        loss = loss_fn(pred, y_train[idx])
        # PyTorch 默认会累积梯度，所以每一步训练前要先清空旧梯度。
        optimizer.zero_grad()
        # 反向传播：从 loss 出发，计算每个参数的梯度。
        loss.backward()
        # 根据梯度真正更新参数。
        optimizer.step()
        if step == 1 or step % 100 == 0:
            # 验证时不需要梯度；no_grad 会省内存，也避免误把验证计算放进计算图。
            # 注意：验证集不参与 optimizer.step，只用于观察模型有没有泛化。
            with torch.no_grad():
                val_loss = loss_fn(model(x_val), y_val)
            print(
                f"step {step:04d} | "
                f"train_loss={loss.item():.5f} | "
                f"val_loss={val_loss.item():.5f}"
            )
    with torch.no_grad():
        # 拿几个没放进 batch 的点，看模型现在会输出什么。
        sample_x = torch.tensor([[-3.0], [0.0], [3.0]])
        sample_y = model(sample_x)
    print()
    print("sample predictions")
    for value, pred in zip(sample_x.squeeze().tolist(), sample_y.squeeze().tolist()):
        print(f"x={value:+.1f} -> y_hat={pred:+.4f}")
if __name__ == "__main__":
    main()

from __future__ import annotations
import math
import torch
from torch import nn
import torch.nn.functional as F
TEXT = (
    "agents plan actions observe results and update context. "
    "language models predict tokens from previous tokens. "
    "attention lets each token read useful earlier tokens. "
) * 40
class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, num_heads: int, block_size: int) -> None:
        super().__init__()
        # 多头 attention 会把 d_model 平均分给每个 head。
        # d_model 是每个 token 的向量维度；block_size 才是最大上下文长度。
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        # qkv 一次性算出 query/key/value，减少三次线性层的样板代码。
        # 输入和输出 shape 都围绕 (batch_size, seq_len, d_model) 展开。
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        # register_buffer 注册的是“不是参数、但要跟着模型移动/保存”的 tensor。
        # mask 不需要训练，所以不用 nn.Parameter。
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: (batch_size, seq_len, d_model)
        batch_size, seq_len, d_model = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        # 多头 attention 标准 shape: (batch_size, num_heads, seq_len, head_dim)。
        # 多头不是多跑几个模型，而是把 d_model 切成几份，让不同 head 学不同关系。
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        # q @ k^T 得到 token 之间的相似度分数。
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim)
        # causal mask 保证当前位置不能看未来 token。
        scores = scores.masked_fill(self.mask[:seq_len, :seq_len] == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        out = weights @ v
        # 把 heads 维度拼回 d_model，恢复 (batch_size, seq_len, d_model)。
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model)
        return self.proj(out)
class Block(nn.Module):
    def __init__(self, d_model: int, num_heads: int, block_size: int) -> None:
        super().__init__()
        # 一个 GPT block 通常由两块组成：
        # attention 负责 token 之间交流，MLP 负责每个 token 自己的非线性加工。
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, num_heads, block_size)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            # 常见设计会先把 hidden 扩到 4 倍，再压回原维度。
            # 这不是改变序列长度，而是改变每个 token 向量内部的维度。
            nn.Linear(d_model, 4 * d_model),
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),
        )
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-LN Transformer block：先 LayerNorm，再 attention/MLP。
        # 残差连接 x + f(x) 保持 shape 不变，也让梯度更容易流动。
        # 第一行：每个 token 先从上下文里读信息。
        x = x + self.attn(self.ln1(x))
        # 第二行：每个 token 再独立过一段 MLP 做加工。
        x = x + self.mlp(self.ln2(x))
        return x
class TinyGPT(nn.Module):
    def __init__(self, vocab_size: int, block_size: int) -> None:
        super().__init__()
        d_model = 64
        self.block_size = block_size
        # token_embedding 负责“这个 token 是什么”。
        # 输入 idx 是整数 id；embedding 输出连续向量。神经网络只能处理数值张量。
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        # position_embedding 负责“这个 token 在第几个位置”。
        # 如果没有位置信息，attention 本身不天然知道顺序。
        self.position_embedding = nn.Embedding(block_size, d_model)
        self.blocks = nn.Sequential(
            Block(d_model, num_heads=4, block_size=block_size),
            Block(d_model, num_heads=4, block_size=block_size),
        )
        self.ln = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size)
    def forward(self, idx: torch.Tensor, targets: torch.Tensor | None = None):
        # idx: (batch_size, seq_len)，里面是 token ids。
        batch_size, seq_len = idx.shape
        # positions 是 0..seq_len-1，对应序列里的每个位置。
        # shape 是 (seq_len,)，加到 token_embedding 时会自动 broadcast 到 batch 维度。
        positions = torch.arange(seq_len, device=idx.device)
        # token 信息和位置信息相加，得到每个位置的初始表示。
        x = self.token_embedding(idx) + self.position_embedding(positions)
        x = self.blocks(x)
        x = self.ln(x)
        # logits: (batch_size, seq_len, vocab_size)，每个位置预测下一个 token。
        # 对某个位置来说，vocab_size 个数字就是“下一个 token 是词表中每个 id 的分数”。
        logits = self.lm_head(x)
        loss = None
        if targets is not None:
            # cross_entropy 要求 (N, classes) 和 (N,)，所以合并 batch_size/seq_len。
            loss = F.cross_entropy(logits.view(batch_size * seq_len, -1), targets.view(batch_size * seq_len))
        return logits, loss
    @torch.no_grad()
    def generate(self, idx: torch.Tensor, steps: int) -> torch.Tensor:
        # generate 是推理阶段：不再给 targets，也不算 loss，只反复预测下一个 token。
        for _ in range(steps):
            # 如果序列超过 block_size，只保留最后 block_size 个 token 作为上下文。
            context = idx[:, -self.block_size :]
            logits, _ = self(context)
            # 只取最后一个位置的 logits，因为生成时只需要预测“下一个 token”。
            probs = F.softmax(logits[:, -1, :], dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_id], dim=1)
        return idx
def make_data():
    # 字符级 tokenizer：简单但足够演示语言模型训练。
    # 后面如果换成 BPE tokenizer，这里 data 的构造方式会变，但模型仍然吃 token ids。
    chars = sorted(set(TEXT))
    stoi = {ch: i for i, ch in enumerate(chars)}
    itos = {i: ch for ch, i in stoi.items()}
    data = torch.tensor([stoi[ch] for ch in TEXT], dtype=torch.long)
    return data, stoi, itos
def get_batch(data: torch.Tensor, batch_size: int, block_size: int):
    # 每个样本是一段连续 token ids。x 是输入上下文，y 是每个位置的下一个 token。
    starts = torch.randint(0, len(data) - block_size - 1, (batch_size,))
    x = torch.stack([data[i : i + block_size] for i in starts])
    # y 是 x 的下一个字符序列，用来做 next-token prediction。
    y = torch.stack([data[i + 1 : i + block_size + 1] for i in starts])
    return x, y
def decode(ids: torch.Tensor, itos: dict[int, str]) -> str:
    return "".join(itos[i] for i in ids.tolist())
def main() -> None:
    torch.manual_seed(123)
    # block_size 是模型一次最多能看的上下文长度。
    block_size = 32
    data, stoi, itos = make_data()
    model = TinyGPT(vocab_size=len(stoi), block_size=block_size)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)
    x, y = get_batch(data, batch_size=4, block_size=block_size)
    logits, loss = model(x, y)
    assert loss is not None
    print("shape check")
    print(f"x: {tuple(x.shape)}")
    print(f"logits: {tuple(logits.shape)}")
    print(f"loss: {loss.item():.4f}")
    print()
    for step in range(1, 301):
        x, y = get_batch(data, batch_size=32, block_size=block_size)
        _, loss = model(x, y)
        assert loss is not None
        # 标准训练三步：清梯度 -> 反向传播 -> 更新参数。
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if step == 1 or step % 50 == 0:
            print(f"step {step:04d} | loss={loss.item():.4f}")
    start = torch.tensor([[stoi["a"]]], dtype=torch.long)
    sample = model.generate(start, steps=180)[0]
    print()
    print(decode(sample, itos))
if __name__ == "__main__":
    main()

字符：a  g  e  n  t  s
编号：2  8  6  14  19  18
"agents" → [2, 8, 6, 14, 19, 18]

文字 → 字符编号 → TinyGPT → 每个字符的分数 → 概率 → 选出下一个字符

x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))

看到 a     → 预测 g
看到 ag    → 预测 e
看到 age   → 预测 n
看到 agen  → 预测 t
看到 agent → 预测 s

输入 x：a  g  e  n  t
答案 y：g  e  n  t  s

agents plan actions observe results and update context.
language models predict tokens from previous tokens.
attention lets each token read useful earlier tokens.

step 0001 | loss=3.3545
step 0050 | loss=0.9013
step 0100 | loss=0.1608
step 0300 | loss=0.1020

anguage models predich tokens from previous tokens. attention lets each token read useful
 earlier tokens. agents plan actions observe results and update context. language models pre

>>> text = "A中"
>>> data = text.encode("utf-8")
>>> list(data)
[65, 228, 184, 173]
>>> data.decode("utf-8")
'A中'

"我爱AI！"
→ [230, 136, 145, 231, 136, 177, 65, 73, 239, 188, 129]

"Hello world!" → ["Hello", " world", "!"]

import regex
from collections import Counter
PATTERN = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
def train_bpe(input_path, vocab_size, special_tokens):
    # 以 bytes 读取语料，之后显式按 UTF-8 解码，避免默认编码影响结果。
    with open(input_path, "rb") as f:
        corpus = f.read()
    text = corpus.decode("utf-8")
    # special token 是文档边界：不把它的内部字符交给预分词或 BPE 统计。
    # 这里先按照 special 把原始文本分开成 spilt list
    spilt_list = []
    if len(special_tokens) == 0:
        spilt_list.append(text)
    else:
        special_token_pattern = "|".join([regex.escape(special) for special in special_tokens])
        spilt_list = regex.split(special_token_pattern, text)
    # 再对拆开来的每一项，按照作业里面的预分词器（实际上就是一段正则）进行预分词
    # 为什么要预分词，可以看这个：https://zhuanlan.zhihu.com/p/692508797
    all_pretokens = []
    for spilt in spilt_list:
        all_pretokens.extend(regex.findall(PATTERN, spilt))
    # 先统计下预分词后，每个词的频率，例如 "low" -> 5。
    pretoken_text_counts = Counter(all_pretokens)
    # 将每个 pre-token 从字符串转为 utf-8 的 bytes，例如 "low"-> b"low"
    # 然后再拆成独立的
    # token_sequence_counts 里面的内容：(b"l", b"o", b"w") -> 5
    token_sequence_counts = {}
    for key, value in pretoken_text_counts.items():
        raw_bytes = key.encode("utf-8")
        bytes_tokens = []
        # bytes 实际上是 int 的数组，所以这里取出来的是 int，还需要手动创建一个 bytes
        for byte_int in raw_bytes:
            bytes_tokens.append(bytes([byte_int]))
        token_sequence_counts[tuple(bytes_tokens)] = value
    # 初始词表包含全部 256 个可能的单 byte 值。
    # special token 作为一个完整 token 加入 vocab，不拆成内部的 byte token。
    merges = []
    vocab = {}
    for token_id in range(256):
        vocab[token_id] = bytes([token_id])
    for token_id, special_token in enumerate(special_tokens, start=256):
        vocab[token_id] = special_token.encode("utf-8")
    # 每轮只学习一种最高频 pair；新 token 会让 vocab 的大小增加 1。
    current_vocab_size = len(vocab)
    while current_vocab_size < vocab_size:
        # 基于「当前」token 序列统计相邻 pair 的加权次数。
        # 例如某序列出现 5 次，其中的每个相邻 pair 都贡献 5 次，而不是 1 次。
        pair_counts = Counter()
        for seq, count in token_sequence_counts.items():
            for pair in zip(seq, seq[1:]):
                pair_counts[pair] += count
        # 没有相邻 pair 时，说明无法继续学习新的 merge。
        if not pair_counts:
            break
        # 找最高频 pair；同频时由 find_max 按字典序选择更大的 pair。
        pair, count = find_max(pair_counts)
        # 两个 bytes token 拼接成一个新 bytes token，例如 b"l" + b"o" -> b"lo"。
        target = pair[0] + pair[1]
        vocab[current_vocab_size] = target
        # merges 必须保留创建顺序；编码时会按这个顺序应用规则。
        merges.append(pair)
        # 不能边遍历边修改旧表：读取旧 token_sequence_counts，
        # 将 merge 后的结果写入新表，全部完成后再整体替换。
        new_token_sequence_counts = {}
        for seq, count in token_sequence_counts.items():
            new_seq = []
            index = 0
            while index < len(seq):
                if index + 1 < len(seq) and (seq[index], seq[index + 1]) == pair:
                    # 命中时消费两个旧 token，写入一个新 token；因此不会重叠合并。
                    new_seq.append(target)
                    index += 2
                else:
                    # 未命中时保留当前 token，继续检查下一个位置。
                    new_seq.append(seq[index])
                    index += 1
            new_token_sequence_counts[tuple(new_seq)] = count
        token_sequence_counts = new_token_sequence_counts
        # 本轮新增了一个 vocab token，下一轮将基于更新后的序列重新统计 pair。
        current_vocab_size += 1
    return vocab, merges
def find_max(pair_counts):
    return max(pair_counts.items(), key=lambda item: (item[1], item[0]))

import regex
from collections.abc import Iterable, Iterator
PATTERN = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
class Tokenizer:
    def __init__(
        self,
        vocab: dict[int, bytes],
        merges: list[tuple[bytes, bytes]],
        special_tokens: list[str] | None = None
    ):
        if special_tokens is None:
            special_tokens = []
        self.vocab = vocab
        self.merges = merges
        self.special_tokens = special_tokens
        # special token 最好校验一下，如果不在 vocab 里面，我们就给他手动加一下然后分配一个 id
        for sp_token in special_tokens:
            sp_token_bytes = sp_token.encode("utf-8")
            if(sp_token_bytes not in self.vocab.values()):
                self.vocab[max(self.vocab.keys()) + 1] = sp_token_bytes

        # 一个字典推导式，实际上就是把 vocab 的 kv 调转一下，因为它是 int2bytes 的    
        self.vocab_b2i = {value: key for key, value in self.vocab.items()}
    def encode(self, text: str) -> list[int]:
        token_int_list = []

        # 和 train_bpe 一样，还是先按照 special tokens 切分一下
        if len(self.special_tokens) == 0:
            split_list = [text]
        else:
            special_token_pattern = "|".join(
                regex.escape(special)
                for special in sorted(self.special_tokens, key=len, reverse=True)
            )
            split_list = regex.split(f"({special_token_pattern})", text)

        # 切分后的 split_list: ["Hi", "<|endoftext|>", " there"]
        # 对齐进行遍历，对每一个单独的结果进行 encode
        for part in split_list:
            if(part == ""):
                continue

            # 如果是 special token，就直接从 vocab 里面找到 id，然后 append 到最终结果里面
            if(part in self.special_tokens):
                token_int_list.append(self.vocab_b2i[part.encode("utf-8")])
            else:
                # 否则就是普通的正常文本
                # 先预分词，Tokenizer 的 encode 也是不能跨 pre-token 的
                # pre_tokens_bytes_list 的内容：[[b"a", b"b", b"d"],[b"a", b"b", b"d"],....]
                pre_tokens = regex.findall(PATTERN, part)
                pre_tokens_bytes_list = []
                for pre_token in pre_tokens:
                    raw_bytes = pre_token.encode("utf-8")
                    pre_tokens_bytes = []
                    for pre_token_byte_int in raw_bytes:
                        pre_tokens_bytes.append(bytes([pre_token_byte_int]))
                    pre_tokens_bytes_list.append(pre_tokens_bytes)

                # 这里三层循环
                current_list_index = 0
                while current_list_index < len(pre_tokens_bytes_list):
                    # 1. 先依次取出预分词的结果，取出来的内容 pre_token_bytes：[b"a", b"b", b"d"]
                    pre_token_bytes = pre_tokens_bytes_list[current_list_index]
                    for merge in self.merges:
                        # 2. 遍历 merges，一条 merge 规则的输入是一条 token 序列，输出也是一条新的 token 序列
                        new_bytes = []
                        i = 0
                        while i < len(pre_token_bytes) - 1: 
                            # 3. 从左到右扫描最新 token 序列，根据 merge 关系，拼装新的 token 序列
                            if merge == (pre_token_bytes[i], pre_token_bytes[i+1]):
                                new_bytes.append(pre_token_bytes[i]+pre_token_bytes[i+1])
                                i += 2
                            else:
                                new_bytes.append(pre_token_bytes[i])
                                i += 1
                        # 从左到右扫描 完 token 序列后，边界要处理下
                        # 例如 pre_token_bytes = [b"a", b"b", b"c"] merge = (b"a", b"b")
                        # merge 完了之后还剩最后一个 b"c" 就退出了，要手动添加到最新的 token 序列里面
                        if(i == len(pre_token_bytes) - 1):
                            new_bytes.append(pre_token_bytes[i])
                        # 新的 token 序列需要赋值，进入下一个 merge 规则的循环
                        pre_token_bytes = new_bytes

                    # 2、3循环跑完后，pre_token_bytes 就是一个已经完整应用过所有 merges 的新的 token 序列了，此时就可以查词表转为 int 了
                    for pre_token in pre_token_bytes:
                        token_int_list.append(self.vocab_b2i[pre_token])

                    current_list_index += 1

        return token_int_list

    def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
        for part in iterable:
            for token_id in self.encode(part):
                yield token_id

    def decode(self, ids: list[int]) -> str:
        bytes_list = []
        for token_id in ids:
            bytes_list.append(self.vocab[token_id])
        return b"".join(bytes_list).decode("utf-8", errors = "replace")

import torch
import math
class Embedding(torch.nn.Module):
    def __init__(self, vocab_size: int, d_model: int, device=None, dtype=None):
        super().__init__()

        # vocab_size 词表大小
        # d_model 每个 token 的维度
        self.embedding = torch.nn.Parameter(
            torch.empty(
                (vocab_size, d_model),
                device = device,
                dtype = dtype
            )
        )

        # 用 assignment1 里面给出的初始化公式进行初始化
        # W ~ 𝒩(μ = 0, σ² = 1，截断到 [-3, 3]
        torch.nn.init.trunc_normal_(
            self.embedding,
            mean = 0,
            std = 1,
            a = -3,
            b = 3
        )

    # embedding 实际上就是一个查表操作
    # forward 的输入 tensor shape 为 (batch, seq_len)
    # 里面执行的逻辑就是对每个 batch 的每个 seq 元素，执行一下查表操作

    # 怎么理解这个 []？
    # 实际上 python 中 A[B] 等价于 A.__getitem__(B)
    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        return self.embedding[token_ids]

import
 torchimport mathclass Linear(torch.nn.Module):    
def
 
__init__
(
self, in_features, out_features, device=
None
, dtype=
None
):        
super
().__init__()        
self
.in_features = in_features        
self
.out_features = out_features                
# 用 troch empty 申请出一个 tensor，devie、dtype 直接透传        # in_features 输入维度        # out_features 输出维度        self.W = torch.nn.Parameter(            torch.empty(                (out_features, in_features),                device = device,                dtype = dtype            )        )                # 用 assignment1 里面给出的初始化公式进行初始化        # W ~ 𝒩(μ = 0, σ² = 2 / (d_in + d_out))，截断到 [-3σ, 3σ]                # 先计算标准差        std = (2 / (in_features+out_features)) ** 0.5        torch.nn.init.trunc_normal_(            self.W,            mean = 0,            std = std,            a = -3 * std,            b = 3 * std        )        # linear 实际上就是一个简单的线性变换，在 pytorch 中表达为 y=x@W.T，「.T」 表示转置    def forward(self, x: torch.Tensor) -> torch.Tensor:        return x @ self.W.T

import torch
import math
def silu(in_features: torch.Tensor):
    return in_features * torch.sigmoid(in_features)
class SwiGLU(torch.nn.Module):
    def __init__(self, d_model: int, d_ff: int, device=None, dtype=None):
        # d_ff：隐藏层的维度

        super().__init__()
        self.d_model = d_model
        self.d_ff = d_ff

        # w1、w3 负责升维
        self.w1 = Linear(d_model,d_ff,device,dtype)
        self.w3 = Linear(d_model,d_ff,device,dtype)

        # w2 负责降维
        self.w2 = Linear(d_ff,d_model,device,dtype)

    def forward(self, x: torch.Tensor):
        a = silu(self.w1(x)) 
        b = self.w3(x)
        # 注意这里是是逐元素相乘
        return self.w2(a * b)

import torch
import math
def softmax(x: torch.Tensor, dim: int):
    # 因为 exp 是指数函数，如果 x 过大可能会导致溢出，而 softmax 实际上只关心每个 x 的差值，不关心具体的绝对值，所以先全都减掉 x 的最大值
    x = x - x.max(dim = dim, keepdim = True).values
    exp_x = torch.exp(x)
    return exp_x / exp_x.sum(dim = dim,keepdim = True)
def scaled_dot_product_attention(
    Q: torch.Tensor,
    K: torch.Tensor,
    V: torch.Tensor,
    mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    参数：
        Q: [..., n, d_k]   （n 为 Query 序列长度）
        K: [..., m, d_k]   （m 为 Key 序列长度）
        V: [..., m, d_v]
        mask: [n, m] 布尔矩阵，True 为保留，False 为屏蔽
    返回：
        [..., n, d_v] 与输入 batch 维一致的张量
    """
    # 先获取 K 向量维度
    d_k = K.size(-1)

    # 计算 Q K 向量的分数，并除以 √d_k 进行缩放
    socres = (Q @ K.mT) / math.sqrt(d_k)

    """
    Masked Scores = [[10, -inf, -inf],
                     [ 5,  10, -inf],
                     [ 2,   8,  10]]
        以下权重只示意遮罩后的形状，并非上面分数的实际 softmax 结果：
        [[1.0, 0.0, 0.0],
                       [0.3, 0.7, 0.0],
                       [0.1, 0.4, 0.5]]
    """
    if mask is not None:
        # mask 为 flase的地方设为负无穷
        socres = socres.masked_fill(mask == False, float('-inf'))

    socres = softmax(socres, dim = -1)

    return socres @ V
class MultiheadSelfAttention(torch.nn.Module):
    def __init__(self, d_model: int, num_heads: int, device=None, dtype=None):
        super().__init__()

        # 要用 head 数量均分 d_model，所以需要校验能否整除
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # 先定义获取 QKV 向量的 Linear
        self.W_Q = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)
        self.W_K = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)
        self.W_V = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)

        # 输出经过 QKV 运算后需要再过一次 Linear
        self.out_proj = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)

    def forward(self, x: torch.Tensor):
        # 先计算出完整的 QKV，再按照 head 拆分
        q = self.W_Q(x)
        k = self.W_K(x)
        v = self.W_V(x)

        # 拆分后的多头 QKV 的 shape：(batch_size, num_heads, seq_len, head_dim)
        # 这里 view 拆分需要先拆 qkv 的最后维度，所以前两维 shape 要展示保持 batch_size, seq_len，通过 transpose 处理成最终的 shape
        batch_size, seq_len, d_model = x.shape
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 这里需要一个下三角矩阵
        """
        True  False False False
        True  True  False False
        True  True  True  False
        True  True  True  True
        """
        mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))

        # 计算注意力
        out = scaled_dot_product_attention(Q=q,K=k,V=v,mask=mask)

        # 把 heads 维度拼回 d_model，恢复 (batch_size, seq_len, d_model)。
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.out_proj(out)

import torch
import math
class RotaryPositionalEmbedding(torch.nn.Module):
    def __init__(self, theta: float, d_k: int, max_seq_len: int, device=None):
        # theta：RoPE 的频率基数
        # d_k：q、k 向量维度
        # max_seq_len：最大 token 长度
        super().__init__()
        self.d_k = d_k

        # 先计算每个位置需要旋转的角度
        # 公式：f_k = 1 / theta^((2k - 2) / d_k)
        # arange(0, d_k, 2) 产生 [0, 2, 4, ..., d_k-2]，对应公式中的2k-2(k从1开始)
        # 注意这里 f_k 的 shape 是 (d_k / 2,)
        f_k = 1 / (theta ** (torch.arange(0, d_k, 2, device=device).float() / d_k))
        # 这里要用外积，外积的结果是前行后列，angle 是一个矩阵
        # angle[n, k] = n × f_k = n / theta^((2k - 2) / d_k)
        # angle 的shape：(max_seq_len, d_k / 2)
        angle = torch.outer(torch.arange(max_seq_len, device=device).float(), f_k)

        # cos_table、sin_table shape 都是 (max_seq_len, d_k / 2)
        self.register_buffer("cos_table", angle.cos())
        self.register_buffer("sin_table", angle.sin())

    def forward(self, x: torch.Tensor, token_positions: torch.Tensor) -> torch.Tensor:
        # x：shape 为 (..., seq_len, d_k) 的 q 或 k 向量
        # token_positions：shape 为 (..., seq_len) 的整数位置下标；每个位置对应 x 中一个 q/k 向量所在的 token 位置。

        # 这里 cos 和 sin 的 shape 都是 token_positions.shape + (d_k // 2,)，得用 repeat_interleave 复制一下元素
        cos_value = self.cos_table[token_positions].repeat_interleave(2, dim=-1).to(x.dtype)
        sin_value = self.sin_table[token_positions].repeat_interleave(2, dim=-1).to(x.dtype)

        # 公式：rotated_q = q ⊙ cos_value + rotate_half(q) ⊙ sin_value
        # rotate_half：[x0, x1, x2, x3, ...] → [-x1, x0, -x3, x2, ...]

        """
        最后一维两两拆分，例如：
        原来 (3, 6)：
        [
            [ 1,  2,  3,  4,  5,  6],
            [ 7,  8,  9, 10, 11, 12],
            [13, 14, 15, 16, 17, 18],
        ]
        现在 (3, 3, 2)：
        [
            [ [1, 2],  [3, 4],  [5, 6] ],      ← 第 0 个 token，3 组，每组 2 个
            [ [7, 8],  [9,10],  [11,12] ],     ← 第 1 个 token
            [ [13,14], [15,16], [17,18] ],     ← 第 2 个 token
        ]
        """
        x_unflatten = x.unflatten(-1,(-1,2))

        """
        按照奇偶拆分开
        even = 每组的第 0 个元素：
        [
            [ 1,  3,  5],
            [ 7,  9, 11],
            [13, 15, 17],
        ]
        odd = 每组的第 1 个元素：
        [
            [ 2,  4,  6],
            [ 8, 10, 12],
            [14, 16, 18],
        ]
        """
        even, odd = x_unflatten.unbind(-1)

        """
        先奇偶互换，然后压缩最后两维
        stack 后的结果：
        [
            [ [-2,  1], [-4,  3], [ -6,  5] ],     ← 第 0 个 token
            [ [-8,  7], [-10, 9], [-12, 11] ],     ← 第 1 个 token
            [ [-14,13], [-16,15], [-18, 17] ],     ← 第 2 个 token
        ]
        flatten 后：
        [
            [ -2,  1,  -4,   3,  -6,   5],
            [ -8,  7, -10,   9, -12,  11],
            [-14, 13, -16,  15, -18,  17],
        ]

        """
        x_rotate_half = torch.stack((-odd, even), dim = -1).flatten(-2)

        rotated_x = x * cos_value + x_rotate_half * sin_value

        return rotated_x

class MultiheadSelfAttention(torch.nn.Module):
    def __init__(self, d_model: int, num_heads: int, theta=None, max_seq_len=None, device=None, dtype=None):
        super().__init__()

        # 要用 head 数量均分 d_model，所以需要校验能否整除
        assert d_model % num_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # 先定义获取 QKV 向量的 Linear
        self.W_Q = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)
        self.W_K = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)
        self.W_V = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)

        # 输出经过 QKV 运算后需要再过一次 Linear
        self.out_proj = Linear(in_features=d_model, out_features=d_model, device=device, dtype=dtype)

        # 如果传入了 theta 和 max_seq_len，则处理 rope 的逻辑
        if theta is not None and max_seq_len is not None:
            self.rope = RotaryPositionalEmbedding(theta=theta, d_k=self.head_dim, max_seq_len=max_seq_len, device=device)
        else:
            self.rope = None

    def forward(self, x: torch.Tensor, token_positions=None):
        # 先计算出完整的 QKV，再按照 head 拆分
        q = self.W_Q(x)
        k = self.W_K(x)
        v = self.W_V(x)

        # 拆分后的多头 QKV 的 shape：(batch_size, num_heads, seq_len, head_dim)
        # 这里 view 拆分需要先拆 qkv 的最后维度，所以前两维 shape 要展示保持 batch_size, seq_len，通过 transpose 处理成最终的 shape
        batch_size, seq_len, d_model = x.shape
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 如果 rope 不为 null 就对 q、k 进行 rope 计算
        if self.rope is not None:
            # token_positions shape：(batch_size, seq_len)
            if token_positions is None:
                token_positions = torch.arange(seq_len, device=x.device).expand(batch_size, seq_len)
            # 每条文本的位置沿 head 轴广播；一维位置可由所有文本共用。
            if token_positions.ndim == 2:
                token_positions = token_positions.unsqueeze(1)
            q=self.rope(q, token_positions)
            k=self.rope(k, token_positions)

        # 这里需要一个下三角矩阵
        """
        True  False False False
        True  True  False False
        True  True  True  False
        True  True  True  True
        """
        mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))

        # 计算注意力
        out = scaled_dot_product_attention(Q=q,K=k,V=v,mask=mask)

        # 把 heads 维度拼回 d_model，恢复 (batch_size, seq_len, d_model)。
        out = out.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.out_proj(out)

import torch
import math
class RMSNorm(torch.nn.Module):
    def __init__(self, d_model: int, eps: float = 1e-5, device=None, dtype=None):
        super().__init__()

        # d_model 模型 token 维度；weight 用于最后做缩放，维度要和 d_model 一致
        # eps 用于防止分母为 0

        self.weight = torch.nn.Parameter(
            torch.ones(
                d_model,
                device = device,
                dtype = dtype
            )
        )
        self.eps = eps
        pass
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 先转成 float32，降低平方时溢出的风险
        in_dtype = x.dtype
        x_float32 = x.to(torch.float32)

        # 应用公式
        # 注意 mean 只需要对最后维度求平均，keepdim 表示需要保留维度，这里 mean 后的 shape 就是 (batch, seq_len, 1)
        rms = torch.sqrt(
            torch.mean(
                torch.square(x_float32), 
                dim = -1,
                keepdim = True
            ) + self.eps
        )
        # 注意这里是 * weight 而不是 @，因为是逐元素相乘缩放
        # 这里计算需要注意一下 shape
        # x_floagt32(batch, seq_len, d_model)，rms(batch, seq_len, 1)
        # x_floagt32/rms 这里会用到 pytorch 里面的广播机制
        # 广播：两个 shape 不完全一样的 tensor 做逐元素运算时，PyTorch 会在尺寸为 1 的维度上，自动“重复使用”那个值，让 shape 对齐。
        result = (x_float32 / rms) * self.weight

        # 返回前转回原来的 dtype
        return result.to(in_dtype)

import torch
import math
class TransformerBlock(torch.nn.Module):
    def __init__(self, d_model: int, num_heads: int, d_ff: int, theta=None, max_seq_len=None, device=None, dtype=None):
        super().__init__()

        # 初始化 MHA
        self.attention = MultiheadSelfAttention(
            d_model=d_model, 
            num_heads=num_heads, 
            theta=theta, 
            max_seq_len=max_seq_len,
            device=device,
            dtype=dtype
        )

        # 初始化 FF
        self.ffn = SwiGLU(d_model=d_model, d_ff=d_ff, device=device, dtype=dtype)

        # 初始化两个 RMSNorm，分别用于 MHA、FF
        self.norm1 = RMSNorm(d_model=d_model, device=device, dtype=dtype)
        self.norm2 = RMSNorm(d_model=d_model, device=device, dtype=dtype)

    def forward(self, x: torch.Tensor, token_positions: torch.Tensor = None):
        # 先进行 MHA 操作
        x = x + self.attention(x=self.norm1(x), token_positions=token_positions)
        # 再进行 FF 操作
        x = x + self.ffn(x=self.norm2(x))
        return x
class TransformerLM(torch.nn.Module):
    def __init__(self, vocab_size: int, nums_layer: int, d_model: int, num_heads: int, d_ff: int, 
                 theta=None, max_seq_len=None, device=None, dtype=None):

        super().__init__()

        self.token_embedding = Embedding(vocab_size=vocab_size, d_model=d_model, device=device, dtype=dtype)
        self.layers = torch.nn.ModuleList([
            TransformerBlock(
                d_model=d_model,
                num_heads=num_heads,
                d_ff=d_ff,
                theta=theta,
                max_seq_len=max_seq_len,
                device=device,
                dtype=dtype
            )
            for _ in range(nums_layer)
        ])

        self.norm = RMSNorm(d_model=d_model,device=device,dtype=dtype)

        # 注意最后是把向量映射回词表
        self.linear = Linear(in_features=d_model, out_features=vocab_size, device=device, dtype=dtype)

    def forward(self, token_ids: torch.Tensor):
        batch_size, seq_len = token_ids.shape

        # 计算后 x.shape: (batch, size, d_model)
        x = self.token_embedding(token_ids)

        # 需要构建一个 token positions，传递到 transformer block 里面给 rope 使用
        # 比如
        # batch_size = 2
        # seq_len = 4
        # [
        #  [0, 1, 2, 3],
        #  [0, 1, 2, 3],
        # ]
        token_positions = torch.arange(seq_len, device=token_ids.device).unsqueeze(0).expand(batch_size, seq_len)

        # 计算后 x.shape: (batch, size, d_model)
        for layer in self.layers:
            x = layer(x, token_positions)

        # 计算后 x.shape: (batch, size, d_model)
        x = self.norm(x)

        # 这里要把向量通过 lm header 计算为词表分数
        # 计算后 x.shape: (batch, size, vocab_size)
        x = self.linear(x)

        return x

Problem (transformer_accounting): Transformer LM resource accounting (5 points)

(a) Consider a GPT-2 XL-sized model using our assignment architecture, which has the following configuration:

○ vocab_size: 50,257

○ context_length: 1,024

○ num_layers: 48

○ d_model: 1,600

○ num_heads: 25

○ d_ff: 4,288 (the nearest multiple of 64 to 8/3 × 1,600)

Suppose we constructed our model using this configuration. How many trainable parameters would our model have? Assuming each parameter is represented using single-precision floating point, how much memory is required to just load this model?

Deliverable: A one-to-two sentence response.

(b) Identify the matrix multiplies required to complete a forward pass of our GPT-2 XL-shaped model. How many FLOPs do these matrix multiplies require in total? Assume that our input sequence has context_length tokens.

Deliverable: A list of matrix multiplies (with descriptions), and the total number of FLOPs required.

(c) Based on your analysis above, which parts of the model require the most FLOPs?

Deliverable: A one-to-two sentence response.

(d) Repeat your analysis with GPT-2 small (12 layers, 768 d_model, 12 heads), GPT-2 medium (24 layers, 1024 d_model, 16 heads), and GPT-2 large (36 layers, 1280 d_model, 20 heads). As the model size increases, which parts of the Transformer LM take up proportionally more or less of the total FLOPs?

Deliverable: For each model, provide a breakdown of model components and its associated FLOPs (as a proportion of the total FLOPs required for a forward pass). In addition, provide a one-to-two sentence description of how varying the model size changes the proportional FLOPs of each component.

(e) Take GPT-2 XL and increase the context length to 16,384. How does the total FLOPs for one forward pass change? How does the relative contribution of FLOPs of the model components change?

Deliverable: A one-to-two sentence response