微调多语言推理器,使用 Hugging Face

作者:Edward Beeching, Quentin Gallouédec, and Lewis Tunstall

OpenAI o3 这样的大型推理模型会生成一个思维链,以提高其响应的准确性和质量。然而,这些模型中的大多数都在英语中进行推理,即使问题是用另一种语言提出的。

在本笔记本中,我们将展示如何对 OpenAI 的开放权重推理模型 OpenAI gpt-oss-20b 进行微调,以使其能够有效地进行多种语言的推理。我们将通过向模型的系统提示中添加一个新的 “推理语言” 选项,并使用 Hugging Face 的 TRL 库 在多语言推理数据集上应用 监督微调

我们将涵盖以下步骤:

  1. 设置: 安装所需的库。
  2. 准备数据集: 下载并格式化用于微调的数据集。
  3. 准备模型: 加载基础模型并配置其进行 LoRA 微调,这是一种内存高效的技术。
  4. 微调: 使用我们的多语言推理数据训练模型。
  5. 推理: 使用微调后的模型生成不同语言的推理响应。

最终结果是一个多语言推理模型,可以生成英语、西班牙语、法语、意大利语或德语的思维链。您甚至可以 混合语言 — 例如,用西班牙语提问,用德语请求推理,然后用西班牙语接收最终响应:

User:
    ¿Cuál es el capital de Australia?
Assistant reasoning:
    Okay, der Benutzer fragt nach der Hauptstadt Australiens. Ich erinnere mich, dass Canberra die Hauptstadt ist. Ich
    sollte das bestätigen. Lass mich sehen, ob es irgendwelche potenziellen Verwirrungen gibt. Der Benutzer könnte auch
    an der größten Stadt interessiert sein. Die größte Stadt ist Sydney, aber die Hauptstadt ist Canberra. Ich sollte
    das klarstellen. Vielleicht auch erwähnen, dass Canberra eine geplante Stadt ist und nicht die größte. Der Benutzer
    könnte auch nach der Geografie fragen. Vielleicht erwähne ich, dass Canberra im südwestlichen Teil der Australian
    Capital Territory liegt. Ich sollte die Antwort präzise und freundlich halten. Vielleicht auch erwähnen, dass
    Canberra oft mit Sydney verwechselt wird. Ich sollte sicherstellen, dass die Antwort klar und korrekt ist.
Assistant response:
    La capital de Australia es **Canberra**. Aunque es la ciudad más pequeña de las principales capitales del país, fue
    elegida en 1908 como la sede del gobierno federal para equilibrar la influencia entre las ciudades de Sydney y
    Melbourne. Canberra está ubicada en el Territorio de la Capital Australiana (ACT), en el este de Australia.

我们希望本教程能够使使用代表性不足的语言的 AI 开发人员能够提高 openai/gpt-oss-20b 在其母语中的可解释性。

注意: 本笔记本设计为在具有 80GB 内存的单个 H100 GPU 上运行。如果您使用的是较小的 GPU,可以在下面的超参数中减小批处理大小和序列长度。

设置

首先,让我们安装所有必需的库。首先安装 PyTorch:

%pip install torch --index-url https://download.pytorch.org/whl/cu128

接下来,安装其余依赖项:

%pip install "trl>=0.20.0" "peft>=0.17.0" "transformers>=4.55.0" trackio

最后,按如下方式登录您的 Hugging Face 帐户:

from huggingface_hub import notebook_login

notebook_login()

现在我们已经安装了所需的库,让我们来看看我们将用于微调的数据集。

准备数据集

我们将使用 Multilingual-Thinking,这是一个推理数据集,其思维链已被翻译成多种语言,如法语、西班牙语和德语。通过在此数据集上微调 openai/gpt-oss-20b,它将学会用这些语言生成推理步骤,从而使其推理过程能够被说这些语言的用户所理解。

让我们从 Hugging Face Hub 下载此数据集:

from datasets import load_dataset

dataset = load_dataset("HuggingFaceH4/Multilingual-Thinking", split="train")
dataset

这是一个包含 1,000 个示例的小型数据集,但这通常足以满足像 openai/gpt-oss-20b 这样经过广泛的训练后模型。让我们看看其中一个训练示例:

dataset[0]

gpt-oss 模型在 Harmony 响应格式上进行了训练,用于定义对话结构、生成推理输出和构建函数调用。该格式旨在模仿 OpenAI 响应 API,下表总结了数据集中使用的不同消息类型:

|||

developer 开发人员消息用于为模型提供自定义指令(我们通常称之为 system 角色)。
user 用户消息用于向模型提供输入。
assistant 模型输出,可以是工具调用或消息输出。输出也可能与特定的“通道”相关联,该通道标识消息的意图。
analysis 这些是模型用于其思维链的消息
final 标记为最终通道的消息是打算显示给最终用户的消息,代表模型的响应。
messages 组合上述内容以生成完整对话的消息列表。这是模型的输入。

如果您熟悉 OpenAI 的消息格式,您会发现它与之非常相似,但有一个重要区别:

assistant 回合包含两个特殊字段:一个 thinking 字段,其中包含模型的推理过程,以及一个 content 字段,其中包含对用户的最终响应。

为了微调模型,我们需要将这些消息转换为模型可以理解的格式。实际上,这是通过使用模型的 聊天模板 格式化每条消息,然后对生成的文本进行标记来实现的。TRL 库会自动执行此操作,但让我们逐步了解它是如何工作的。

为此,让我们首先加载分词器:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

然后我们可以使用分词器的 apply_chat_template() 方法来格式化消息:

messages = dataset[0]["messages"]
conversation = tokenizer.apply_chat_template(messages, tokenize=False)
print(conversation)

这个聊天模板相当复杂,所以让我们仔细看看!首先,我们可以看到有特殊标记 <|start|><|end|> 来指示每条消息的开始和结束。还有一个 <|return|> 标记来标记对话的结束。这些标记有助于模型理解对话的结构。

我们还可以看到有两种系统消息:

  • 一个默认的 system 消息,用于所有消息。在上面的示例中,这指的是文本 “您是 ChatGPT,一个由 OpenAI 训练的大型语言模型……”
  • 一个特殊的 developer 消息,其中包含自定义指令(由我们 messages 对象中的 system 角色定义)。这允许我们为模型提供关于其如何处理给定对话的附加上下文。在上面的示例中,这指的是文本 “您是一位性格活泼、精力充沛的人工智能聊天机器人。”

最后,我们可以看到助手响应包含在一系列_通道_中:

  • analysis 通道用于模型的推理过程,它可以在其中逐步思考用户的问题。在上面的示例中,这指的是法语文本 “D'accord, l'utilisateur demande les tendances Twitter……”
  • final 通道用于模型的最终响应给用户。在上面的示例中,这指的是文本 “Hey there! While I can't check Twitter……”

现在我们了解了数据集将如何准备,让我们继续准备模型进行训练。

准备模型

为了准备模型进行训练,让我们首先从 Hugging Face Hub 下载权重。我们将使用 🤗 Transformers 中的 AutoModelForCausalLM 类来加载模型:

import torch
from transformers import AutoModelForCausalLM, Mxfp4Config

quantization_config = Mxfp4Config(dequantize=True)
model_kwargs = dict(
    attn_implementation="eager",
    torch_dtype=torch.bfloat16,
    quantization_config=quantization_config,
    use_cache=False,
    device_map="auto",
)

model = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-20b", **model_kwargs)

这将加载具有训练所需配置的模型。attn_implementation 设置为 eager 以获得更好的性能,而 use_cache 设置为 False,因为我们将使用梯度检查点来微调模型。

如果您熟悉 🤗 Transformers,您可能会注意到我们正在使用 Mxfp4Config 进行量化。这是 OpenAI 模型的一种特定配置,它允许我们使用一种称为 MXFP4 的特殊 4 位浮点格式进行混合精度训练,该格式针对 AI 工作负载进行了优化。

在训练模型之前,让我们生成一个示例响应,看看模型在默认设置下的行为。为此,我们需要对一个示例提示进行标记,然后使用模型生成响应:

messages = [
    {"role": "user", "content": "¿Cuál es el capital de Australia?"},
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)

output_ids = model.generate(input_ids, max_new_tokens=512)
response = tokenizer.batch_decode(output_ids)[0]
print(response)

在此示例中,我们可以看到模型首先用英语推理问题,然后提供西班牙语的最终响应。这是模型的默认行为,但让我们看看是否可以通过一些微调来更改它。

为此,我们将使用一种称为 LoRA(低秩自适应)的技术来微调模型。这项技术允许我们调整模型的少数特定层,这对于像 openai/gpt-oss-20b 这样的大型模型特别有用。

首先,我们需要将模型包装为 PeftModel 并定义 LoRA 配置。我们将使用 PEFT 库中的 LoraConfig 类来完成此操作:

from peft import LoraConfig, get_peft_model

peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules="all-linear",
    target_parameters=[
        "7.mlp.experts.gate_up_proj",
        "7.mlp.experts.down_proj",
        "15.mlp.experts.gate_up_proj",
        "15.mlp.experts.down_proj",
        "23.mlp.experts.gate_up_proj",
        "23.mlp.experts.down_proj",
    ],
)
peft_model = get_peft_model(model, peft_config)
peft_model.print_trainable_parameters()

这里我们使用了一些基本的 LoRA 超参数,但您可以尝试不同的值来查看它们如何影响模型的性能。例如,如果您增加 r,您将启用更多的可训练参数,这可能会产生更好的模型,但需要更多的 VRAM 和训练时间。

注意: openai/gpt-oss-20b 模型是 混合专家 (MoE) 架构。除了目标注意力层 (target_modules="all-linear") 之外,还包括专家模块内的投影层很重要。PEFT 通过 target_parameters 参数来实现这一点,该参数允许您指定专家特定的层,例如 mlp.experts.down_projmlp.experts.gate_up_proj。在此示例中,我们针对这些投影层的一个子集,但鼓励您尝试不同的配置。

现在我们已经准备好了模型和数据集,我们可以定义训练的超参数。

微调

TRL 提供了一种使用 SFTConfig 类定义训练超参数的便捷方法。我们将如下设置学习率、批处理大小、训练轮数和其他参数:

from trl import SFTConfig

training_args = SFTConfig(
    learning_rate=2e-4,
    gradient_checkpointing=True,
    num_train_epochs=1,
    logging_steps=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    max_length=2048,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine_with_min_lr",
    lr_scheduler_kwargs={"min_lr_rate": 0.1},
    output_dir="gpt-oss-20b-multilingual-reasoner",
    report_to="trackio",
    push_to_hub=True,
)

请注意,per_device_train_batch_size 设置为 4,gradient_accumulation_steps 设置为 4。这意味着我们将在 1 个 GPU 上有效地拥有一个 4 x 4 = 16 的批处理大小。您可能需要根据您的硬件设置调整这些值。我们还使用 Trackio 来记录训练进度和指标,但您可以使用任何其他您选择的日志记录库。

现在我们拥有了训练模型所需的所有组件。我们将使用 TRL 中的 SFTTrainer 类来处理训练过程。该训练器将负责格式化数据集、应用聊天模板和训练模型:

from trl import SFTTrainer

trainer = SFTTrainer(
    model=peft_model,
    args=training_args,
    train_dataset=dataset,
    processing_class=tokenizer,
)
trainer.train()

在 H100 GPU 上,这大约需要 18 分钟才能完成训练,但根据您的硬件可能会花费更长的时间。

保存模型并推送到 Hugging Face Hub

最后,您可以将微调后的模型推送到您的 Hub 存储库,与社区共享:

trainer.save_model(training_args.output_dir)
trainer.push_to_hub(dataset_name="HuggingFaceH4/Multilingual-Thinking")

注意:为避免内存不足 (OOM) 错误,我们建议在此处重新启动内核。训练后的模型仍然占用 GPU 内存,但不再需要它。

推理

模型上传到 Hub 后,我们就可以使用它进行推理了。为此,我们首先初始化原始基础模型及其分词器。接下来,我们需要将微调后的权重与基础模型合并以实现快速推理:

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

# 首先加载原始模型
model_kwargs = dict(attn_implementation="eager", torch_dtype="auto", use_cache=True, device_map="auto")
base_model = AutoModelForCausalLM.from_pretrained("openai/gpt-oss-20b", **model_kwargs).cuda()

# 合并微调权重与基础模型
peft_model_id = "gpt-oss-20b-multilingual-reasoner"
model = PeftModel.from_pretrained(base_model, peft_model_id)
model = model.merge_and_unload()

现在模型已加载,最后一步是从它生成一些标记!这里我们使用模型的 generate 方法根据输入提示生成输出。让我们首先定义提示:

现在我们可以对提示进行标记并生成输出了。最后,我们可以解码输出标记以获得最终响应:

REASONING_LANGUAGE = "German"
SYSTEM_PROMPT = f"reasoning language: {REASONING_LANGUAGE}"
USER_PROMPT = "¿Cuál es el capital de Australia?"  # 西班牙语,意为“澳大利亚的首都是什么?”

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": USER_PROMPT},
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)

gen_kwargs = {"max_new_tokens": 512, "do_sample": True, "temperature": 0.6, "top_p": None, "top_k": None}

output_ids = model.generate(input_ids, **gen_kwargs)
response = tokenizer.batch_decode(output_ids)[0]
print(response)

让我们也尝试使用模型没有明确微调过的语言,例如中文和印地语:

REASONING_LANGUAGE = "Chinese"  # 或 Hindi,或其他任何语言...
SYSTEM_PROMPT = f"reasoning language: {REASONING_LANGUAGE}"
USER_PROMPT = "What is the national symbol of Canada?"

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": USER_PROMPT},
]

input_ids = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)

output_ids = model.generate(input_ids, **gen_kwargs)
response = tokenizer.batch_decode(output_ids)[0]
print(response)

太好了,它奏效了——我们现在已经微调了 openai/gpt-oss-20b 以使其能够进行多种语言的推理!

结论

恭喜!您已成功使用 TRL 库和 LoRA 微调了一个多语言推理模型。此笔记本中的步骤可以改编为在 Hugging Face Hub 上的许多其他 数据集 上微调 openai/gpt-oss-20b — 我们期待看到您将构建什么!