总目录›Axolotl›第 9 章
Chapter 09

自定义 prompt 模板

📌 commit ab1a0d8三种自定义层次

Axolotl 的 prompt 模板有三种自定义层次：(1) 直接用 tokenizer 的 chat_template；(2) 写一段 Jinja2；(3) 写一个 PromptStrategy 类。

9.1层次 1：tokenizer 自带 chat_template

大部分情况这就够了。每个模型在 HF 上的 tokenizer 都带 chat_template，Axolotl 直接复用：

datasets:
  - path: ./mydata.jsonl
    type: chat_template
    chat_template: tokenizer_default      # ★ 用 tokenizer 自带

chat_template 取值：

9.2层次 2：自定义 Jinja2

chat_template: "{%- for message in messages %}{{ '<|' + message.role + '|>\n' + message.content + '\n<|end|>\n' }}{%- endfor %}"

# 或从文件
chat_template_jinja: ./templates/my.jinja2

9.3层次 3：自定义 PromptStrategy 类

当 Jinja2 不够（比如要根据数据字段动态决定算不算 loss），写一个 Python 类：

# src/axolotl/prompt_strategies/my_prompt.py
from axolotl.prompt_strategies.user_defined import UserDefinedPromptStrategy

def load(tokenizer, cfg):
    return MyStrategy(tokenizer, cfg)

class MyStrategy(UserDefinedPromptStrategy):
    def tokenize_prompt(self, example):
        instr = example["instruction"]
        resp = example["response"]
        # 自定义拼接
        full = f"<|user|>{instr}<|end|>\n<|assistant|>{resp}<|end|>"
        ids = self.tokenizer(full, add_special_tokens=False)["input_ids"]
        # 算 loss 部分（assistant token 起点）
        prefix = f"<|user|>{instr}<|end|>\n<|assistant|>"
        prefix_len = len(self.tokenizer(prefix, add_special_tokens=False)["input_ids"])
        labels = [-100] * prefix_len + ids[prefix_len:]
        return {"input_ids": ids, "labels": labels, "attention_mask": [1]*len(ids)}

YAML 里：

datasets:
  - path: ./mydata.jsonl
    type: my_prompt        # ★ 文件名（不带 .py）

9.4system prompt 注入

datasets:
  - path: ./mydata.jsonl
    type: chat_template
    chat_template: llama3
    field_messages: messages
    system_prompt: "You are a helpful assistant who replies in JSON."

所有样本前面自动 prepend system 消息。也可以在数据每条单独写 system。

9.5tool use / function calling 模板

函数调用数据需要特殊模板：

# Hermes-2 / Qwen2.5 风格
chat_template: chatml
datasets:
  - path: ./tool_use.jsonl
    type: chat_template
    chat_template: chatml
    field_messages: messages
    message_field_role: role
    message_field_content: content
    roles:
      user: [human, user]
      assistant: [assistant, gpt]
      tool: [tool, function_response]    # ★ 工具调用响应

数据 jsonl 里可以有 tool_calls / function_response 等额外字段。

9.6调试 prompt 模板

训练前用 axolotl preprocess 跑一遍，输出 ./prepared/*。可以读出来反 tokenize 看：

from datasets import load_from_disk
from transformers import AutoTokenizer

ds = load_from_disk("./prepared/...")
tok = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
print(tok.decode(ds[0]["input_ids"]))      # 看真实拼接
print(ds[0]["labels"][:50])                  # 看 labels mask 位置

9.7常见踩坑

现象处理
训完模型不会 EOSchat_template 末尾没 EOS；显式 add_eos_token: true
多轮训成单轮roles 字段没列全所有 role 名
tokenize 后 labels 全 -100roles 里 assistant 拼错了
system_prompt 重复出现jsonl 已经带 system 又传了一份；二选一

9.8jinja_template_analyzer.py：HF chat_template 的解析

v2024-04 引入 chat_template 重构后，prompt_strategies/jinja_template_analyzer.py 这个工具是"为什么自动判断哪段是 user / assistant 不需要写 roles"的功臣。

它做的事：

  1. 读 tokenizer 的 chat_template（Jinja2 字符串）；
  2. 静态分析 Jinja2 AST，找从模板中产生的 string literal（如 "<|im_start|>user\n"）；
  3. 把这些字面量映射成"角色边界 token"，供 RoleBoundary（processing_strategies.py:30）用；
  4. 训练时 collator 用这些 boundary 标记 mask label：assistant 段算 loss，其他段 -100。

所以 chat_template 改了 → roles 不写也能正确 mask。这是 Axolotl 比 LLaMA-Factory 抽象更高的设计点。

9.9roles_to_train + train_on_eos 4 模式

processing_strategies.py 的 RoleBoundary 定义 4 种 EOS 训练模式（field train_on_eos）：

train_on_eos含义典型用
"none" EOS token 永远 mask 不想模型主动结束（continued pretrain）
"turn"（默认）每一 turn 的结束 token 都计 loss 多轮对话标准
"all" 所有 EOS 都计 loss（包括拼接中的） 极少用
"last" 仅整段对话的最后一个 EOS 计 loss 避免训成主动断话

配合 roles_to_train: ["assistant"]（指定训哪个角色的输出），就能精细控制 label mask。

9.10调试 chat_template 的 3 个工具

命令 / 方法看什么
axolotl preprocess my.yml 生成 ./prepared/ 后人工 spot check
axolotl train --debug 模式 训练 step 0 打印若干条 decoded sample
tok.apply_chat_template(messages, tokenize=False)不通过 Axolotl，直接调 HF 看 Jinja 输出

常见 bug 90% 一开始就能用上面这三个之一定位：要么 chat_template 选错（如 Qwen 模板配 Llama base），要么数据列名不对（field_messages: messages 但数据列叫 conversations）。

9.11这章你需要带走的