训练阶段扩展Tokenizer后，如何加载Huggingface T5预训练模型？

错误原因 #

你在训练阶段扩展了Tokenizer的词汇表（从32128扩展到32138），但加载的模型checkpoint中，encoder和decoder的embedding层仍保持原始词汇表尺寸，而当前初始化的模型使用了扩展后的词汇表尺寸，导致两者形状不匹配，触发加载错误。

解决方案分两种场景处理： #

场景1：加载自己训练好的适配扩展Tokenizer的模型 #

如果你的模型是在扩展Tokenizer后完成训练并保存的，问题出在加载时误用了原始预训练模型的config，而非自己模型的config。修改代码如下：

# 确保model_name_or_path是你自己保存的模型目录
config = T5Config.from_pretrained(
    model_name_or_path,
    cache_dir=model_args.cache_dir,
    revision=model_args.model_revision,
    use_auth_token=True if model_args.use_auth_token else None,
)
config.train_task_adapters = adapter_args.train_task_adapters

# 加载和模型配套的扩展Tokenizer（默认保存在同一目录）
tokenizer = AutoTokenizer.from_pretrained(
    model_name_or_path,
    cache_dir=model_args.cache_dir,
    use_fast=model_args.use_fast_tokenizer,
    revision=model_args.model_revision,
    use_auth_token=True if model_args.use_auth_token else None,
)

# 加载模型
model = T5ForConditionalGeneration.from_pretrained(
    model_name_or_path,
    from_tf=bool(".ckpt" in model_name_or_path),
    config=config,
    cache_dir=model_args.cache_dir,
    revision=model_args.model_revision,
    use_auth_token=True if model_args.use_auth_token else None,
    adapter_config=adapter_config
)

自己保存的模型目录中包含更新后的config（已同步扩展后的vocab_size），加载后模型与checkpoint的embedding尺寸会自动匹配。

场景2：加载原始预训练T5模型，适配自定义扩展的Tokenizer #

如果需要在原始预训练模型基础上适配扩展后的Tokenizer，需要手动调整embedding层尺寸并初始化新增token的embedding：

步骤1：加载扩展后的Tokenizer #

# 替换为你的扩展Tokenizer保存路径
tokenizer = AutoTokenizer.from_pretrained("your_extended_tokenizer_path")

步骤2：更新模型config的词汇量 #

config = T5Config.from_pretrained(
    model_name_or_path,
    cache_dir=model_args.cache_dir,
    revision=model_args.model_revision,
    use_auth_token=True if model_args.use_auth_token else None,
)
# 将config的vocab_size设置为扩展后Tokenizer的词汇量
config.vocab_size = len(tokenizer)
config.train_task_adapters = adapter_args.train_task_adapters

步骤3：初始化模型并处理embedding权重 #

有两种实现方式：

方式A：快速处理（自动随机初始化新增embedding）

利用from_pretrained的ignore_mismatched_sizes参数，自动跳过尺寸不匹配的参数并随机初始化新增部分：

model = T5ForConditionalGeneration.from_pretrained(
    model_name_or_path,
    from_tf=bool(".ckpt" in model_name_or_path),
    config=config,
    cache_dir=model_args.cache_dir,
    revision=model_args.model_revision,
    use_auth_token=True if model_args.use_auth_token else None,
    adapter_config=adapter_config,
    ignore_mismatched_sizes=True  # 新增该参数
)

这种方式操作简单，但新增token的embedding是随机初始化的，建议后续进行微调优化。

方式B：手动初始化（推荐，贴合原有模型分布）

手动复制原有embedding权重，并用现有embedding的均值初始化新增token的embedding，更贴合模型原有分布：

import torch

# 基于更新后的config初始化模型
model = T5ForConditionalGeneration(config)

# 加载原始预训练模型的权重文件
checkpoint = torch.load(f"{model_name_or_path}/pytorch_model.bin")

# 处理encoder和decoder的embedding权重
for emb_key in ["encoder.model_embeddings.weight", "decoder.model_embeddings.weight"]:
    old_emb = checkpoint[emb_key]
    new_emb = model.state_dict()[emb_key]
    
    # 复制原有token的embedding权重
    new_emb[:old_emb.shape[0], :] = old_emb
    
    # 用现有embedding的均值初始化新增token的embedding
    mean_emb = old_emb.mean(dim=0, keepdim=True)
    new_emb[old_emb.shape[0]:, :] = mean_emb
    
    # 更新checkpoint中的权重
    checkpoint[emb_key] = new_emb

# 加载处理后的权重
model.load_state_dict(checkpoint, strict=False)

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