如何创建兼容HF Trainer的自定义Hugging Face解码器模型？

一、Tokenizer集成（复用现有Tokenizer，适配自定义模型） #

分词时机明确：必须在数据加载器阶段完成分词，模型内部不做分词处理。这样能保证数据预处理效率，适配DataLoader和Trainer的批量处理逻辑，同时避免模型内重复预处理的资源浪费。

from transformers import AutoTokenizer, PreTrainedTokenizerFast

# 复用现有优秀Tokenizer（以T5为例，也可替换为BERT、GPT等Tokenizer）
tokenizer = AutoTokenizer.from_pretrained("t5-small")

# 若需自定义Tokenizer（基于现有词表扩展）
# tokenizer = PreTrainedTokenizerFast(
#     tokenizer_file="custom_tokenizer.json",
#     bos_token="<s>", eos_token="</s>", pad_token="<pad>"
# )

核心要求：确保Tokenizer的vocab_size、特殊token（bos/eos/pad）与后续模型的embedding层参数完全对齐。

二、自定义编码器-解码器模型（遵循HF标准API） #

参考T5的实现风格，继承HF的PreTrainedModel，确保适配Trainer、GPU加速及PyTorch生态：

import torch
import torch.nn as nn
from transformers import PreTrainedModel, PretrainedConfig

# 自定义模型配置类（对应HF的Config规范）
class CustomSeq2SeqConfig(PretrainedConfig):
    model_type = "custom_seq2seq"
    def __init__(self, vocab_size=32128, d_model=512, num_layers=6, num_heads=8, dropout=0.1, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.dropout = dropout

# 编码器模块（参考T5 Encoder结构，复用PyTorch原生Transformer层）
class CustomEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        self.encoder_layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=config.d_model, nhead=config.num_heads,
                dim_feedforward=config.d_model*4, dropout=config.dropout
            ) for _ in range(config.num_layers)
        ])
    
    def forward(self, input_ids, attention_mask=None):
        embeds = self.embedding(input_ids)
        embeds = self.layer_norm(embeds)
        embeds = self.dropout(embeds)
        # 转换attention_mask为Transformer所需的key_padding_mask格式
        key_padding_mask = ~attention_mask.bool() if attention_mask is not None else None
        for layer in self.encoder_layers:
            embeds = layer(embeds, src_key_padding_mask=key_padding_mask)
        return embeds

# 解码器模块（参考T5 Decoder结构，集成LM Head）
class CustomDecoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.layer_norm = nn.LayerNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        self.decoder_layers = nn.ModuleList([
            nn.TransformerDecoderLayer(
                d_model=config.d_model, nhead=config.num_heads,
                dim_feedforward=config.d_model*4, dropout=config.dropout
            ) for _ in range(config.num_layers)
        ])
        self.lm_head = nn.Linear(config.d_model, config.vocab_size)
    
    def forward(self, input_ids, encoder_hidden_states, attention_mask=None, decoder_attention_mask=None):
        embeds = self.embedding(input_ids)
        embeds = self.layer_norm(embeds)
        embeds = self.dropout(embeds)
        # 生成自注意力掩码（防止前瞻）
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(input_ids.size(1)).to(input_ids.device)
        # 处理padding掩码
        tgt_key_padding_mask = ~decoder_attention_mask.bool() if decoder_attention_mask is not None else None
        memory_key_padding_mask = ~attention_mask.bool() if attention_mask is not None else None
        
        for layer in self.decoder_layers:
            embeds = layer(
                embeds, encoder_hidden_states,
                tgt_mask=tgt_mask, tgt_key_padding_mask=tgt_key_padding_mask,
                memory_key_padding_mask=memory_key_padding_mask
            )
        logits = self.lm_head(embeds)
        return logits

# 完整编码器-解码器模型（符合HF API规范）
class CustomSeq2SeqModel(PreTrainedModel):
    config_class = CustomSeq2SeqConfig
    
    def __init__(self, config):
        super().__init__(config)
        self.encoder = CustomEncoder(config)
        self.decoder = CustomDecoder(config)
        self.init_weights() # 调用HF的权重初始化逻辑
    
    def forward(self, input_ids, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None, labels=None):
        # 自动处理Teacher Forcing：无decoder_input_ids时用labels移位
        if decoder_input_ids is None and labels is not None:
            decoder_input_ids = labels[:, :-1].contiguous()
            decoder_attention_mask = decoder_attention_mask[:, :-1].contiguous() if decoder_attention_mask is not None else None
        
        encoder_hidden = self.encoder(input_ids, attention_mask)
        logits = self.decoder(decoder_input_ids, encoder_hidden, attention_mask, decoder_attention_mask)
        
        # 适配HF Trainer的损失计算逻辑
        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels[:, 1:].contiguous().view(-1))
        
        return {"loss": loss, "logits": logits} if loss is not None else {"logits": logits}

三、数据预处理（数据加载器阶段完成分词） #

from datasets import Dataset
from transformers import DataCollatorForSeq2Seq

# 模拟训练数据
data = {
    "source": ["这是测试输入1", "这是测试输入2", "这是测试输入3"],
    "target": ["对应测试输出1", "对应测试输出2", "对应测试输出3"]
}
dataset = Dataset.from_dict(data)

# 批量分词函数（在数据加载前执行）
def preprocess_function(examples):
    inputs = tokenizer(examples["source"], max_length=128, truncation=True, padding="max_length")
    # 单独处理目标序列（避免与源序列共享tokenizer配置）
    with tokenizer.as_target_tokenizer():
        targets = tokenizer(examples["target"], max_length=128, truncation=True, padding="max_length")
    inputs["labels"] = targets["input_ids"]
    inputs["decoder_attention_mask"] = targets["attention_mask"]
    return inputs

# 应用分词处理
tokenized_dataset = dataset.map(preprocess_function, batched=True)

# 数据整理器（自动处理批量数据的padding、设备迁移）
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=CustomSeq2SeqModel)

四、模型训练与验证 #

方式1：自定义PyTorch训练循环 #

# 初始化设备、模型、优化器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
config = CustomSeq2SeqConfig(vocab_size=tokenizer.vocab_size)
model = CustomSeq2SeqModel(config).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

# 构建DataLoader
from torch.utils.data import DataLoader
train_dataloader = DataLoader(tokenized_dataset, batch_size=2, shuffle=True, collate_fn=data_collator)

# 训练循环
model.train()
for epoch in range(3):
    total_loss = 0.0
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs["loss"]
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    print(f"Epoch {epoch+1} | 平均损失: {total_loss/len(train_dataloader):.4f}")

方式2：使用HF Trainer #

from transformers import Trainer, TrainingArguments

# 训练参数配置
training_args = TrainingArguments(
    output_dir="./custom_seq2seq_checkpoints",
    evaluation_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=2,
    num_train_epochs=3,
    weight_decay=0.01,
    fp16=torch.cuda.is_available(), # 自动启用混合精度训练
    logging_steps=10,
)

# 初始化Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

# 启动训练
trainer.train()

关于参考T5实现风格的说明 #

完全可以参考T5的设计逻辑：

• 结构上：采用编码器-解码器架构，共享词表，底层复用Transformer层，与T5的核心结构一致；
• API规范上：继承PreTrainedModel，实现支持labels参数的forward方法，自动计算损失，适配HF生态；
• 数据处理上：使用DataCollatorForSeq2Seq处理批量数据，分词在预处理阶段完成，符合T5的最佳实践。

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