GPT输入层：Token ID、Embedding与RoPE位置编码

上一篇文本预处理已经解决了一个问题：人类语言如何被切成token，并映射成整数id。到这里为止，模型还没有真正开始“理解”文本，因为输入仍然只是一个整数矩阵：

inputs:  (B, T)
targets: (B, T)

其中：

B 是 batch size，一次喂给模型多少条序列。
T 是 sequence length，每条序列包含多少个token。
每个元素都是一个整数，表示词表中的某个token id。

GPT输入层要做的事情，就是把这个二维整数矩阵变成三维连续向量：

(B, T) -> (B, T, C)

其中 C = n_embd，也就是每个token的隐藏向量维度。下面结合 /Volumes/PSSD/工具/nanochat 中的源码来看这一层到底怎么实现。

从Tokenizer到模型配置

在 scripts/base_train.py 中，训练脚本会先加载Tokenizer。原因很直接：模型的输入词表大小必须和Tokenizer的词表大小一致。

tokenizer = get_tokenizer()
token_bytes = get_token_bytes(device=device)
vocab_size = tokenizer.get_vocab_size()
print0(f"Vocab size: {vocab_size:,}")

拿到 vocab_size 后，nanochat再构造GPT配置：

def build_model_meta(depth):
    base_dim = depth * args.aspect_ratio
    model_dim = ((base_dim + args.head_dim - 1) // args.head_dim) * args.head_dim
    num_heads = model_dim // args.head_dim
    config = GPTConfig(
        sequence_len=args.max_seq_len,
        vocab_size=vocab_size,
        n_layer=depth,
        n_head=num_heads,
        n_kv_head=num_heads,
        n_embd=model_dim,
        window_pattern=args.window_pattern,
    )

这里有两个关键点。

第一，vocab_size 来自Tokenizer。Tokenizer有多少个token，模型的embedding表至少就要有多少行。

第二，n_embd 不是随便定的。nanochat通过 depth * aspect_ratio 得到基础宽度，再向上取整到 head_dim 的倍数。这样可以保证：

n_embd = n_head * head_dim

否则后面做多头注意力时，通道维度就无法整齐切成多个head。

nanochat/nanochat/gpt.py 里的默认配置如下：

@dataclass
class GPTConfig:
    sequence_len: int = 2048
    vocab_size: int = 32768
    n_layer: int = 12
    n_head: int = 6
    n_kv_head: int = 6
    n_embd: int = 768
    window_pattern: str = "SSSL"

如果使用默认值，那么每个token最终会被映射成一个768维向量。

Token Embedding是什么

Embedding层本质上是一张可以训练的查表矩阵。

假设：

vocab_size = 32768
n_embd = 768

那么token embedding表的形状就是：

(32768, 768)

第 i 行就是第 i 个token对应的向量。输入token id为 123，模型就取出embedding表的第123行。

从数学上看，可以把token id理解成一个one-hot向量，再乘以embedding矩阵：

x_i = onehot(i) W_E

但工程上不会真的构造巨大的one-hot矩阵，而是直接查表：

embedding = nn.Embedding(vocab_size, n_embd)
x = embedding(idx)

如果：

idx.shape = (B, T)
embedding.weight.shape = (vocab_size, C)

那么输出就是：

x.shape = (B, T, C)

这就是GPT主干网络真正接收的数据形状。

nanochat里的wte

nanochat在 GPT.__init__ 中定义token embedding层：

def __init__(self, config, pad_vocab_size_to=64):
    padded_vocab_size = ((config.vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to) * pad_vocab_size_to
    if padded_vocab_size != config.vocab_size:
        print0(f"Padding vocab_size from {config.vocab_size} to {padded_vocab_size} for efficiency")
    self.transformer = nn.ModuleDict({
        "wte": nn.Embedding(padded_vocab_size, config.n_embd),
        "h": nn.ModuleList([Block(config, layer_idx) for layer_idx in range(config.n_layer)]),
    })
    self.lm_head = Linear(config.n_embd, padded_vocab_size, bias=False)

这里的 wte 就是 Word Token Embedding。

注意它用的不是原始 config.vocab_size，而是 padded_vocab_size。也就是说，如果真实词表大小不是64的倍数，nanochat会把词表行数向上补齐到64的倍数。

比如：

vocab_size = 32769
pad_vocab_size_to = 64
padded_vocab_size = 32832

这不会改变Tokenizer真实能产生的token id范围，只是让embedding表和 lm_head 的矩阵尺寸更适合GPU矩阵计算。最后输出logits时，nanochat会再切回真实词表大小：

logits = self.lm_head(x)
logits = logits[..., :self.config.vocab_size]

也就是说，padding出来的那些词表行只是工程优化，不代表模型真的多了可生成的正常token。

Embedding初始化

Embedding不是固定表，而是模型参数，会在训练中更新。

nanochat在 init_weights() 中初始化它：

torch.nn.init.normal_(self.transformer.wte.weight, mean=0.0, std=0.8)
torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)

wte 和 lm_head 是两张不同的矩阵。很多GPT实现会共享输入embedding和输出embedding的权重，也就是所谓 weight tying；但nanochat没有共享。源码开头也明确写了：

untied weights for token embedding and lm_head

因此这里有两张表：

wte: token id -> hidden vector
lm_head: hidden vector -> vocab logits

前者负责把离散token带入模型，后者负责把模型最后的隐藏状态映射回词表得分。

forward中的输入流动

真正前向传播时，idx 就是DataLoader给出的token id矩阵：

def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
    B, T = idx.size()

此时：

idx.shape = (B, T)

然后进入embedding层：

x = self.transformer.wte(idx)
x = x.to(COMPUTE_DTYPE)
x = norm(x)

形状变化是：

idx: (B, T)
x:   (B, T, n_embd)

这里还做了两件事。

第一，x.to(COMPUTE_DTYPE) 把激活值转换成nanochat统一使用的计算精度。nanochat没有依赖 torch.amp.autocast，而是通过 COMPUTE_DTYPE 显式控制前向计算精度。A100/H100这类支持bf16的GPU上，通常会使用 bfloat16。

第二，norm(x) 会对embedding后的向量做RMSNorm：

def norm(x):
    return F.rms_norm(x, (x.size(-1),))

这一步可以理解成进入Transformer主干前，先把每个token向量的尺度稳定下来。

GPT需要位置信息

到目前为止，wte 只知道“这个位置是什么token”，但并不知道“这个token在第几个位置”。

假设有两个输入：

我 爱 你
你 爱 我

如果只看token embedding，那么每个token自己的向量是固定的：

embedding("我") 不会因为它出现在第1位或第3位而自动变化

而Attention本身主要依赖向量之间的点积。没有位置信息时，模型很难区分“谁在前，谁在后”。所以Transformer必须引入位置编码。

早期Transformer会把位置向量直接加到token embedding上：

x = token_embedding + position_embedding

但是nanochat没有使用这种可学习的绝对位置embedding。源码开头直接说明：

rotary embeddings (and no positional embeddings)

也就是说，nanochat不在输入层把位置向量加到 x 上，而是在Attention内部，把位置信息注入到 Q 和 K 中。

RoPE缓存什么时候建立

RoPE需要大量 cos 和 sin 值。它们只和位置、head维度、base有关，不依赖具体输入文本，所以可以提前算好。

nanochat在 GPT.__init__ 中预计算RoPE缓存：

self.rotary_seq_len = config.sequence_len * 10
head_dim = config.n_embd // config.n_head
cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim)
self.register_buffer("cos", cos, persistent=False)
self.register_buffer("sin", sin, persistent=False)

这里的 rotary_seq_len = config.sequence_len * 10 是一个工程取舍。比如训练序列长度是2048，那么RoPE表会预先算到20480。

这样做的好处是推理时如果上下文略长，不需要频繁重新分配缓存。坏处是如果真的超过这个长度，前向传播里会直接断言失败：

assert T <= self.cos.size(1)

register_buffer(..., persistent=False) 表示 cos 和 sin 跟着模型移动设备，但不保存进checkpoint。原因是它们不是训练参数，也可以由配置重新计算出来。

RoPE的预计算代码

核心函数是：

def _precompute_rotary_embeddings(self, seq_len, head_dim, base=100000, device=None):
    if device is None:
        device = self.transformer.wte.weight.device
    channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
    inv_freq = 1.0 / (base ** (channel_range / head_dim))
    t = torch.arange(seq_len, dtype=torch.float32, device=device)
    freqs = torch.outer(t, inv_freq)
    cos, sin = freqs.cos(), freqs.sin()
    cos, sin = cos.to(COMPUTE_DTYPE), sin.to(COMPUTE_DTYPE)
    cos, sin = cos[None, :, None, :], sin[None, :, None, :]
    return cos, sin

这段代码可以拆成四步。

第一步，取偶数通道：

channel_range = torch.arange(0, head_dim, 2)

如果 head_dim = 128，那么这里得到：

[0, 2, 4, ..., 126]

长度是 head_dim / 2。

第二步，计算每个通道对应的频率：

inv_freq = 1.0 / (base ** (channel_range / head_dim))

越靠前的通道频率越高，越靠后的通道频率越低。直观上，高频通道更敏感于短距离变化，低频通道更适合表示长距离变化。

第三步，列出所有位置：

t = torch.arange(seq_len)

如果 seq_len = 20480，那么 t 就是：

[0, 1, 2, ..., 20479]

第四步，做外积：

freqs = torch.outer(t, inv_freq)

形状变成：

freqs.shape = (seq_len, head_dim / 2)

其中第 m 行、第 i 列表示：

第 m 个位置，在第 i 个旋转平面上应该旋转多少角度

最后取 cos 和 sin，并调整成方便广播的形状：

cos, sin = cos[None, :, None, :], sin[None, :, None, :]

最终：

cos.shape = (1, seq_len, 1, head_dim / 2)
sin.shape = (1, seq_len, 1, head_dim / 2)

这个形状很重要，因为Attention里的Q/K形状是：

q.shape = (B, T, n_head, head_dim)
k.shape = (B, T, n_kv_head, head_dim)

cos/sin 的第0维和第2维都是1，可以自动广播到不同batch和不同head。

forward中如何截取RoPE

前向传播时，不需要每次都用完整RoPE表，只要取当前输入长度对应的那一段：

T0 = 0 if kv_cache is None else kv_cache.get_pos()
cos_sin = self.cos[:, T0:T0+T], self.sin[:, T0:T0+T]

训练时通常没有KV Cache：

T0 = 0

所以输入长度为 T，就取位置 [0, T) 的RoPE。

推理时如果使用KV Cache，当前这次forward可能只输入一个新token，但这个token并不是第0个位置，而是整个上下文中的第 T0 个位置。因此需要：

从RoPE缓存中取 [T0, T0 + T)

这就是为什么RoPE截取要考虑 kv_cache.get_pos()。

RoPE在哪里真正作用

Embedding阶段只得到 x：

x.shape = (B, T, C)

RoPE不是直接加到 x 上，而是在Attention内部，先生成Q/K/V：

q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)

然后只对 q 和 k 做旋转：

cos, sin = cos_sin
q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)

注意：v 不做RoPE。

原因是Attention的匹配关系来自：

QK^T

位置应该影响“当前query应该看哪些key”，所以注入到Q和K里。V是被加权汇总的内容本体，通常不需要旋转。

apply_rotary_emb怎么旋转

nanochat的RoPE应用函数如下：

def apply_rotary_emb(x, cos, sin):
    assert x.ndim == 4
    d = x.shape[3] // 2
    x1, x2 = x[..., :d], x[..., d:]
    y1 = x1 * cos + x2 * sin
    y2 = x1 * (-sin) + x2 * cos
    return torch.cat([y1, y2], 3)

输入 x 的形状是：

(B, T, H, head_dim)

先把最后一维切成两半：

x1: (B, T, H, head_dim / 2)
x2: (B, T, H, head_dim / 2)

然后把 x1 和 x2 看成一组组二维平面上的坐标。第一个半区的第 i 个通道，和第二个半区的第 i 个通道组成一个二维向量：

(x1_i, x2_i)

RoPE对这个二维向量做旋转：

\begin{bmatrix} y_1 \\ y_2 \end{bmatrix} = \begin{bmatrix} \cos \theta & \sin \theta \\ -\sin \theta & \cos \theta \end{bmatrix} \begin{bmatrix} x_1 \\ x_2 \end{bmatrix}

对应代码就是：

y1 = x1 * cos + x2 * sin
y2 = x1 * (-sin) + x2 * cos

最后再拼回原来的 head_dim：

return torch.cat([y1, y2], 3)

所以RoPE前后形状不变：

(B, T, H, head_dim) -> (B, T, H, head_dim)

它改变的是向量内容，不改变张量尺寸。

一个小例子：形状如何流动

假设：

B = 2
T = 8
vocab_size = 32768
n_embd = 768
n_head = 6
head_dim = 128

输入token id：

idx.shape = (2, 8)

查embedding表：

wte.weight.shape = (32768, 768)
x.shape = (2, 8, 768)

进入Attention后：

q.shape = (2, 8, 6, 128)
k.shape = (2, 8, 6, 128)
v.shape = (2, 8, 6, 128)

RoPE缓存切片：

cos.shape = (1, 8, 1, 64)
sin.shape = (1, 8, 1, 64)

应用RoPE时，Q/K的最后一维被切成两个64维：

q1, q2: (2, 8, 6, 64)

旋转后再拼回：

q_rot.shape = (2, 8, 6, 128)
k_rot.shape = (2, 8, 6, 128)

所以这一段完整流程是：

token ids
  -> wte查表
  -> RMSNorm
  -> Q/K/V线性映射
  -> 对Q/K注入RoPE
  -> 进入Attention打分

为什么不是在Embedding里加位置向量

传统绝对位置编码通常是：

x = token_embedding + position_embedding

这种做法简单，但位置被混进了残差流的每一层，模型需要自己从绝对位置向量中学出相对距离关系。

RoPE的思路不一样：它直接改造Attention中的点积。

如果第 m 个位置的query旋转 mθ，第 n 个位置的key旋转 nθ，那么它们做内积时，会自然出现和相对距离有关的项：

q_m^T k_n \quad \Rightarrow \quad \text{包含 } (n - m) \text{ 的信息}

因此RoPE更适合表达“当前位置和被关注位置之间相隔多远”。这也是它在现代GPT类模型中非常常见的原因。

和文本预处理的衔接

把上一篇和这一篇连起来，GPT输入阶段的完整路径就是：

原始文本
  -> BPE Tokenizer
  -> token ids
  -> DataLoader打包成 inputs / targets
  -> inputs: (B, T)
  -> wte查表
  -> x: (B, T, C)
  -> Attention内部对Q/K应用RoPE

到这里，模型已经完成了两件基础工作：

把离散token变成连续向量。
给Attention注入顺序和相对位置信息。

下一步就可以正式进入Transformer Block，继续看一个block内部如何通过Attention、MLP、残差连接和归一化不断更新这些token向量。

翻检卷帙

GPT输入层：Token ID、Embedding与RoPE位置编码

GPT输入层：Token ID、Embedding与RoPE位置编码

从Tokenizer到模型配置

Token Embedding是什么

nanochat里的wte

Embedding初始化

forward中的输入流动

GPT需要位置信息

RoPE缓存什么时候建立

RoPE的预计算代码

forward中如何截取RoPE

RoPE在哪里真正作用

apply_rotary_emb怎么旋转

一个小例子：形状如何流动

为什么不是在Embedding里加位置向量

和文本预处理的衔接

尺素

对话助手