装载 GPT-2 · from_pretrained 与 Conv1D 转置 · nanoGPT 源码解读

逐行读 from_pretrained——怎么把 OpenAI 公开的 GPT-2 权重搬进 nanoGPT。重点是那个烧脑的坑：GPT-2 用的 Conv1D 模块把权重存成了转置，所以四类权重要 .t() 一下才能对上 nn.Linear。

装载 GPT-2 · from_pretrained 与 Conv1D 转置

nanoGPT 不只能从零训，还能直接装载 OpenAI 公开的 GPT-2 权重，拿来推理或微调。这活儿由 from_pretrained 干（model.py 约 206–261 行）。它是个典型的「跨框架搬权重」的脏活——大部分是对名字、对形状，但藏着一个很经典的坑。

第一步：按型号定好配置

@classmethod
def from_pretrained(cls, model_type, override_args=None):
    assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'}
    from transformers import GPT2LMHeadModel

    config_args = {
        'gpt2':         dict(n_layer=12, n_head=12, n_embd=768),  # 124M
        'gpt2-medium':  dict(n_layer=24, n_head=16, n_embd=1024), # 350M
        'gpt2-large':   dict(n_layer=36, n_head=20, n_embd=1280), # 774M
        'gpt2-xl':      dict(n_layer=48, n_head=25, n_embd=1600), # 1558M
    }[model_type]
    config_args['vocab_size'] = 50257  # GPT-2 固定
    config_args['block_size'] = 1024   # GPT-2 固定
    config_args['bias'] = True         # GPT-2 固定
    config = GPTConfig(**config_args)
    model = GPT(config)

四个 GPT-2 尺寸，层数/头数/维度各不同——但正好印证全景那章说的：结构一模一样，只是把 GPTConfig 的几个数字调大。三个 GPT-2 铁定的值（词表 50257、上下文 1024、带 bias）直接写死。然后先造一个同规格、随机初始化的 nanoGPT 模型 model——它就是即将被 GPT-2 权重「灌满」的空壳。

第二步：两边的 state_dict 对齐

    sd = model.state_dict()
    sd_keys = [k for k in sd.keys() if not k.endswith('.attn.bias')]  # 丢掉因果掩码 buffer

    model_hf = GPT2LMHeadModel.from_pretrained(model_type)  # 从 HuggingFace 下载 GPT-2
    sd_hf = model_hf.state_dict()
    sd_keys_hf = [k for k in sd_hf.keys() if not k.endswith('.attn.masked_bias')]
    sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')]

sd 是 nanoGPT 空壳的参数表，sd_hf 是从 HuggingFace 下载的 GPT-2 参数表。
两边都剔掉掩码类的 buffer（.attn.bias、.attn.masked_bias）——它们是因果掩码那样的固定张量，不是要学习的参数，不用搬。
剔完之后，断言两边 key 数量相等(assert len(sd_keys_hf) == len(sd_keys))，确保没漏没多。

第三步：那个经典的坑——Conv1D 转置

搬权重本该是「同名的直接 copy」，但 GPT-2 有个历史包袱：

    transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight']
    # OpenAI 的 checkpoint 用的是 "Conv1D" 模块,而我们用的是普通 Linear,
    # 所以这几个权重导入时要转置一下
    for k in sd_keys_hf:
        if any(k.endswith(w) for w in transposed):
            # 需要转置的:形状应当正好互为倒序
            assert sd_hf[k].shape[::-1] == sd[k].shape
            with torch.no_grad():
                sd[k].copy_(sd_hf[k].t())
        else:
            # 其余的:形状一致,直接 copy
            assert sd_hf[k].shape == sd[k].shape
            with torch.no_grad():
                sd[k].copy_(sd_hf[k])
    return model

坑在这：OpenAI 最初的 GPT-2 用的是从 TensorFlow 传下来的 Conv1D 模块（名字叫卷积，其实就是个线性层），它存权重矩阵时是转置的——行列和 PyTorch 的 nn.Linear 正好反过来。

nanoGPT 用的是标准 nn.Linear。于是那四类权重（注意力的 c_attn、c_proj，MLP 的 c_fc、c_proj）搬进来时，必须先 .t() 转置，行列摆正了才对得上；其余参数（嵌入表、LayerNorm 等）形状一致，直接 copy。

那两行 assert 是贴心的保险丝：

需要转置的：断言 sd_hf[k].shape[::-1] == sd[k].shape（形状互为倒序）；
直接复制的：断言 sd_hf[k].shape == sd[k].shape（形状全等）。

一旦哪个权重形状对不上，assert 立刻炸出来，而不是默默 copy 出一个错乱的模型。

这一段是**「跨框架加载权重」的活教材**：两边的参数命名、形状约定、哪些是可学习参数哪些是 buffer、甚至某个模块的历史实现细节（Conv1D 存成转置），全得一一对上。真实工程里，「把模型 A 的权重装进模型 B」十有八九都要处理这类又琐碎又致命的不一致——错一个转置，模型就废，还不报错。nanoGPT 用一张 transposed 名单 + 两行 assert，把这件脏活做得干净又稳。

[1] Radford et al. “Language Models are Unsupervised Multitask Learners (GPT-2).” 2019. —— GPT-2 权重的出处。
[2] Wolf et al. “Transformers: State-of-the-Art Natural Language Processing.” EMNLP 2020. —— HuggingFace transformers（GPT2LMHeadModel 的来源）。
源码：Andrej Karpathy, nanoGPT， model.py 的 from_pretrained. https://github.com/karpathy/nanoGPT

装载 GPT-2 · from_pretrained 与 Conv1D 转置