如何在HuggingFace的GPT2模型中替换自定义全连接网络？

我正在学习从零开始训练因果语言模型的教程，当前加载标准GPT2的代码如下：

from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)
model = GPT2LMHeadModel(config)

我希望保留GPT2的其他结构，仅替换其中的标准全连接网络，用来尝试不同层数、激活函数的变体。虽然找到了GPT2的源码，但结构复杂，搞不清怎么替换，也不知道自定义网络的输入输出尺寸要求。

更新
比如我想用下面这个自定义全连接网络：

class FC_model(nn.Module):
    def __init__(self):
        super(FC_model, self).__init__()

        self.fc1 = nn.Linear(768,256)
        self.fc2 = nn.Linear(256,256)
        self.fc3 = nn.Linear(256,50000)

    def forward(self, x):
        x = torch.sin(self.fc1(x)) + torch.rand(1)
        x = torch.sin(self.fc2(x))
        x = self.fc3(x)
        return x

两种替换场景及实现方法 #

场景1：替换最后的语言模型预测头（LMHead） #

你的自定义FC正好匹配GPT2最后一层lm_head的输入输出要求：输入是模型的隐藏层维度（GPT2-small默认768），输出是词汇表大小（示例中为50000）。直接替换这个模块即可：

import torch
import torch.nn as nn
from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

# 改造自定义FC，用参数适配模型配置，避免硬编码
class FC_model(nn.Module):
    def __init__(self, hidden_size, vocab_size):
        super(FC_model, self).__init__()
        self.fc1 = nn.Linear(hidden_size, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, vocab_size)

    def forward(self, x):
        # 随机数和输入张量同设备，避免设备不匹配报错
        x = torch.sin(self.fc1(x)) + torch.rand(1, device=x.device)
        x = torch.sin(self.fc2(x))
        x = self.fc3(x)
        return x

# 加载tokenizer与配置
tokenizer = AutoTokenizer.from_pretrained("gpt2")
context_length = 128
config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)

# 初始化原模型
model = GPT2LMHeadModel(config)

# 替换lm_head
model.lm_head = FC_model(config.hidden_size, config.vocab_size)

# 测试前向传播
inputs = tokenizer("Hello, my name is", return_tensors="pt")
outputs = model(**inputs)
print(outputs.logits.shape)  # 输出应为 [batch_size, seq_len, vocab_size]

场景2：替换Transformer层内的MLP模块 #

如果要替换的是GPT2每个Transformer块中的全连接网络（残差结构里的MLP），需满足输入输出维度一致（都等于模型的hidden_size），才能适配残差连接：

import torch
import torch.nn as nn
from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig

# 自定义Transformer层内的MLP，输入输出维度均为hidden_size
class CustomMLP(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.fc1 = nn.Linear(hidden_size, 2048)  # 可自定义中间层维度
        self.fc2 = nn.Linear(2048, hidden_size)

    def forward(self, x):
        x = torch.sin(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载配置与模型
tokenizer = AutoTokenizer.from_pretrained("gpt2")
context_length = 128
config = AutoConfig.from_pretrained(
    "gpt2",
    vocab_size=len(tokenizer),
    n_ctx=context_length,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
)
model = GPT2LMHeadModel(config)

# 遍历所有Transformer块，替换内部MLP
for layer in model.transformer.h:
    layer.mlp = CustomMLP(config.hidden_size)

# 测试前向传播
inputs = tokenizer("Hello, my name is", return_tensors="pt")
outputs = model(**inputs)
print(outputs.logits.shape)

核心注意事项 #

• 维度匹配：替换LMHead时输入=模型hidden_size、输出=模型vocab_size；替换Transformer层MLP时输入输出必须等于hidden_size
• 设备对齐：自定义模块的参数、前向传播中的张量要和原模型在同一设备（CPU/GPU）
• 保留原有结构：仅替换目标模块，其余部分复用GPT2的注意力机制等核心逻辑

内容的提问来源于stack exchange，提问作者Penguin