使用HuggingFace微调SmolLM3模型时遇张量尺寸不匹配RuntimeError

问题现象 #

• 令牌配置警告：tokenizer的PAD/BOS/EOS令牌与模型配置、生成配置不一致，系统对齐后令牌为{'bos_token_id': None, 'pad_token_id': None}
• 运行时错误：RuntimeError: The size of tensor a (4) must match the size of tensor b (8) at non-singleton dimension 0

报错栈信息 #

The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': None}.
  0%|          | 0/1000 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/tmp/script.py", line 199, in <module>
    trainer.train()
  File "/root/.cache/uv/environments-v2/script-912247c0edd68a55/lib/python3.12/site-packages/transformers/trainer.py", line 2328, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/uv/environments-v2/script-912247c0edd68a55/lib/python3.12/site-packages/transformers/trainer.py", line 2672, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/uv/environments-v2/script-912247c0edd68a55/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 1189, in training_step
    return super().training_step(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/uv/environments-v2/script-912247c0edd68a55/lib/python3.12/site-packages/transformers/trainer.py", line 4009, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/.cache/uv/environments-v2/script-912247c0edd68a55/lib/python3.12/site-packages/trl/trainer/sft_trainer.py", line 1123, in compute_loss
    entropy = torch.sum(per_token_entropy * attention_mask) / attention_mask.sum()
                        ~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (4) must match the size of tensor b (8) at non-singleton dimension 0
  0%|          | 0/1000 [00:01<?, ?it/s]

触发错误的代码 #

from trl import SFTTrainer, SFTConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset
from typing import Optional, Dict, Any

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM3-3B-Base")
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM3-3B-Base")

# Load dataset
dataset = load_dataset("HuggingFaceTB/smoltalk2", "SFT", split="s1k_1.1_think")

SYSTEM_PROMPT = "You are a helpful assistant."

CHAT_TEMPLATE = (
    "<|im_start|>system\n{system_prompt}<|im_end|>\n"
    "<|im_start|>user\n{instruction}{input_section}<|im_end|>\n"
    "<|im_start|>assistant\n"
)

CHAT_CONFIG = {
        "custom_instructions": "",
        "enable_thinking": True,
        # Tool lists are accepted but unused by this simple SFT preprocessing.
        "python_tools": [],
        "xml_tools": [],
}

def get_system_prompt(chat_config: Dict[str, Any]) -> str:
        """Compose the short system prompt from `CHAT_CONFIG`.

        - If `custom_instructions` is provided, it is used instead of the default
            `SYSTEM_PROMPT`.
        - If `enable_thinking` is True we append a short instruction that allows
            the model to 'think' or include chain-of-thought style tokens. This is a
            lightweight mechanism used only to alter the system prompt text.
        """
        prompt = chat_config.get("custom_instructions") or SYSTEM_PROMPT
        if chat_config.get("enable_thinking"):
                # A minimal, explicit hint that encourages internal 'thinking' — keep it
                # short so tokenization and label alignment remain straightforward.
                prompt += "\nYou may think through steps before answering."
        return prompt

def build_prompt(instruction: str, input_text: Optional[str], system_prompt: Optional[str] = None) -> str:
    """Construct the prompt (system + user) but stop before the assistant response.

    The returned string ends with the assistant start marker so the response can be
    concatenated and tokenized as a single sequence.
    """
    input_section = f"\n{input_text}" if input_text and len(input_text) > 0 else ""
    # If a system_prompt is explicitly passed use it; otherwise build one from
    # the `CHAT_CONFIG` so callers can control `custom_instructions` and
    # `enable_thinking` settings without changing this file.
    system_prompt = system_prompt or get_system_prompt(CHAT_CONFIG)
    return CHAT_TEMPLATE.format(instruction=instruction, input_section=input_section, system_prompt=system_prompt)

def tokenize_for_sft(example: Dict[str, Any], max_length: int = 512) -> Dict[str, Any]:
    """Tokenize a single example for SFT.

    Expects the dataset example to contain one of these response fields: `response`, `output`,
    `answer`, or `text`. For the prompt it prefers `instruction` and optional `input`.

    Returns a dict with `input_ids`, `attention_mask`, and `labels` where label tokens
    corresponding to the prompt are set to -100.
    """
    # Find prompt fields
    instruction = example.get("instruction") or example.get("prompt") or example.get("query") or ""
    input_text = example.get("input") or example.get("context") or ""

    # Find response field
    response = (
        example.get("response")
        or example.get("output")
        or example.get("answer")
        or example.get("text")
        or ""
    )

    prompt = build_prompt(instruction, input_text)
    assistant_end = "<|im_end|>\n"
    full_text = prompt + response + assistant_end + (tokenizer.eos_token or "")

    # Ensure pad_token_id is set (some tokenizers/models don't have it)
    if tokenizer.pad_token_id is None:
        # Prefer eos_token_id as pad_token_id fallback
        tokenizer.pad_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
        model.config.pad_token_id = tokenizer.pad_token_id

    # Tokenize prompt and full_text separately so we can find prompt length
    tokenized_full = tokenizer(
        full_text,
        truncation=True,
        max_length=max_length,
        padding="max_length",
        return_tensors=None
    )

    tokenized_prompt = tokenizer(
        prompt,
        truncation=True,
        max_length=max_length,
        padding=False,
        return_tensors=None,
    )

    input_ids = tokenized_full["input_ids"]
    attention_mask = tokenized_full["attention_mask"]

    # Build labels: copy input_ids then mask prompt tokens with -100
    labels = list(input_ids)
    prompt_len = len(tokenized_prompt["input_ids"])
    # If prompt_len >= max_length then there's no room for the response; labels become all -100
    for i in range(min(prompt_len, len(labels))):
        labels[i] = -100

    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

# Map preprocessing over the training dataset and keep the processed dataset for trainer
# Note: we remove other columns so the dataset only contains tensors needed by the trainer.
max_seq_length = 512
processed_train = dataset.map(
    lambda ex: tokenize_for_sft(ex, max_length=max_seq_length),
    remove_columns=[c for c in dataset.column_names],
)


# Configure training

config = SFTConfig(
    output_dir="./smollm3-jobs-sft",
    # Balanced GPU Memory
    per_device_train_batch_size=4,
    gradient_accumulation_steps = 4,
    learning_rate=1e-4,
    max_steps=1000,
    # Short sequences
    max_length = 1024,
    #max_length=4096,
    #packing=True,
    #assistant_only_loss=True,
    logging_steps=50,
    save_steps=200,
    push_to_hub=False,
    hub_model_id="hubsnippetai/smollm3-jobs-sft"
)
# Train
trainer = SFTTrainer(
    model=model,
    train_dataset=processed_train,
    args=config,
)
trainer.train()

解决方案 #

1. 修复令牌配置问题 #

SmolLM3的tokenizer默认未设置pad_token和bos_token，需手动指定并同步到模型配置：

# 加载模型和tokenizer后立即执行
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM3-3B-Base")
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM3-3B-Base")

# 设置pad_token为eos_token，bos_token也用eos_token替代（SmolLM3无单独bos_token）
tokenizer.pad_token = tokenizer.eos_token
tokenizer.bos_token = tokenizer.eos_token

# 同步模型配置
model.config.pad_token_id = tokenizer.pad_token_id
model.config.bos_token_id = tokenizer.bos_token_id
model.config.eos_token_id = tokenizer.eos_token_id

2. 统一序列长度参数 #

报错核心原因是预处理时的max_length（512）与SFTConfig中的max_length（1024）不一致，导致训练时张量维度冲突。修改为统一值：

# 预处理部分
max_seq_length = 512
processed_train = dataset.map(
    lambda ex: tokenize_for_sft(ex, max_length=max_seq_length),
    remove_columns=[c for c in dataset.column_names],
)

# SFTConfig部分
config = SFTConfig(
    output_dir="./smollm3-jobs-sft",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=1e-4,
    max_steps=1000,
    max_length=512,  # 和预处理时保持一致
    logging_steps=50,
    save_steps=200,
    push_to_hub=False,
    hub_model_id="hubsnippetai/smollm3-jobs-sft"
)

3. 优化标签掩码逻辑（可选） #

原代码的标签掩码逻辑存在冗余，可简化为更直观的写法，避免索引越界：

def tokenize_for_sft(example: Dict[str, Any], max_length: int = 512) -> Dict[str, Any]:
    # ... 其他代码保持不变 ...

    input_ids = tokenized_full["input_ids"]
    attention_mask = tokenized_full["attention_mask"]

    # 初始化标签为全-100，仅保留响应部分的标签
    labels = [-100] * len(input_ids)
    prompt_len = len(tokenized_prompt["input_ids"])
    
    # 确保prompt_len不超过序列长度
    if prompt_len < len(input_ids):
        labels[prompt_len:] = input_ids[prompt_len:]

    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

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