针对对象存储的 Deepseek 式强化学习

tl;dr：我们训练一个小LLM玩家，让他们擅长使用强化学习（类似于导致 Deepseek R1 的过程）进行推理，所有这些都针对本地模型存储库 AIStor AIHub。基于 Will brown 的出色 GRPO 演示。

赋予动机：

对团队日益增长的要求是需要一个有序、安全的“单一事实来源”来存储所有模型和数据集。这与公共模型和数据集存储库和 API（最值得注意的是 HuggingFace）的广泛采用形成鲜明对比。AIStor AIHub 就是考虑到这个难题而创建的。AIHub 允许企业在私有云（或气隙环境中）创建和使用自己的数据集和模型存储库，而无需更改任何代码。换句话说，通过为自己提供一个通用的模型和数据集“动物园”，组织可以大幅减少实验开销、AI 价值实现时间以及关键训练和推理工作负载的外部依赖关系。

但是，为什么团队真的希望他们的所有模型和数据集都有一个“有序、安全、单一的事实来源”呢？

1 . 数据集通常是专有的（尤其是对于调整模型）。

2 . 快速实验取决于对整个组织中不同模型和数据集的可靠访问。

3 . 用于这些实验和从这些实验中产生的模型通常是专有的。

4 . 希望减少对推理服务基础设施的第三方依赖。

在这篇文章中，我们将介绍一个非常常见的工作流程，适用于处理推理模型的团队。我们调整了一个小型开源LLM，使其能够更好地使用 RL 针对本地存储在 AIHub 上的数据集和模型进行推理。我们使用 trl 和 transformers 等常用库来完成这一切（参见：无需更改代码）。

AIHub 的工作原理

要使 RL 实验代码指向 AIHub（本地模型和数据集存储库），唯一需要的更改是设置环境变量。具体说来：

HF_ENDPOINT=<your-aistor-aihub-endpoint>

AIHub 也支持身份验证，在这种情况下，您还必须设置以下环境变量：

HF_TOKEN=<your-aihub-token>

注意：您可以使用您的 AIStor 凭证向 AIHub 发出经过身份验证的请求以获取令牌。

设置这些环境变量后，使用 transformers、datasets、trl、vllm 等（使用 HuggingFace 的库）的代码将指向您在 AIStor 上构建的本地模型和数据集存储库。开发人员友好。企业友好。

AIHub 驻留在您选择的存储桶中。在我的 AIStor 部署中，AIHub 是在我命名为“ai-central”的存储桶上初始化的。

AIHub 驻留在您选择的存储桶中。在我的 AIStor 部署中，AIHub 是在我命名为“ai-central”的存储桶上初始化的。

当您的客户端代码（即 transformers、trl、datasets、vllm 等）调用 AIHub 上尚不存在的模型或数据集时，将从公共模型存储库中一次性获取该模型或数据集。但是，加载该模型或数据集的后续调用将仅在您的客户端代码和本地 AIHub 之间进行。这最终能够：（1） 使用所有最新的模型和数据集进行快速实验，（2） 重用和使用专有模型/数据集的开发人员体验与使用公共模型/数据集相同，以及 （3） 在加载/提供模型进行推理时没有第三方依赖。

方法

现在，以下是我们将采用的基本方法：

1 . 将 AIHub 的基本模型加载到 GPU 上

2 . 预处理数据集

3 . 定义强化学习的奖励

4 . 配置训练参数

5 . Train

6 . 将新检查点保存到 AIHub

训练推理器

注意：此代码基于 Will brown 的 GRPO 演示。因此，这也适用于 HuggingFace（只需取消设置 HF_ENDPOINT）。

在这里，我们在 GSM8K 数据集上训练 Qwen2.5-0.5B-Instruct 模型（有关此数据集的更多详细信息见下文），以我们想要的格式开发推理行为。通常，以全精度训练的 LLM GPU 内存要求的启发式方法是每 10 亿个参数 ~20 GB GPU 内存。因此，这是在配备 NVIDIA A100 GPU 的机器上运行的。本文的代码在此处的存储库中提供。这是整个训练脚本，然后是有关此处发生情况的一些说明。

使用 HF_ENDPOINT= 指向 AIHub 终端节点的环境变量运行此脚本。

import re
import torch
from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import GRPOConfig, GRPOTrainer

SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""

XML_COT_FORMAT = """\
<reasoning>
{reasoning}
</reasoning>
<answer>
{answer}
</answer>
"""

def extract_xml_answer(text: str) -> str:
    answer = text.split("<answer>")[-1]
    answer = answer.split("</answer>")[0]
    return answer.strip()

def extract_hash_answer(text: str) -> str | None:
    if "####" not in text:
        return None
    return text.split("####")[1].strip()

# uncomment middle messages for 1-shot prompting
def get_gsm8k_questions(split = "train") -> Dataset:
    data = load_dataset('openai/gsm8k', 'main')[split] # type: ignore
    data = data.map(lambda x: { # type: ignore
        'prompt': [
            {'role': 'system', 'content': SYSTEM_PROMPT},
            {'role': 'user', 'content': x['question']}
        ],
        'answer': extract_hash_answer(x['answer'])
    }) # type: ignore
    return data # type: ignore

dataset = get_gsm8k_questions()[:3]

# Reward functions
def correctness_reward_func(prompts, completions, answer, **kwargs) -> list[float]:
    responses = [completion[0]['content'] for completion in completions]
    q = prompts[0][-1]['content']
    extracted_responses = [extract_xml_answer(r) for r in responses]
    print('-'*20, f"Question:\n{q}", f"\nAnswer:\n{answer[0]}", f"\nResponse:\n{responses[0]}", f"\nExtracted:\n{extracted_responses[0]}")
    return [2.0 if r == a else 0.0 for r, a in zip(extracted_responses, answer)]

def int_reward_func(completions, **kwargs) -> list[float]:
    responses = [completion[0]['content'] for completion in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    return [0.5 if r.isdigit() else 0.0 for r in extracted_responses]

def strict_format_reward_func(completions, **kwargs) -> list[float]:
    """Reward function that checks if the completion has a specific format."""
    pattern = r"^<reasoning>\n.*?\n</reasoning>\n<answer>\n.*?\n</answer>\n$"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]

def soft_format_reward_func(completions, **kwargs) -> list[float]:
    """Reward function that checks if the completion has a specific format."""
    pattern = r"<reasoning>.*?</reasoning>\s*<answer>.*?</answer>"
    responses = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, r) for r in responses]
    return [0.5 if match else 0.0 for match in matches]

def count_xml(text) -> float:
    count = 0.0
    if text.count("<reasoning>\n") == 1:
        count += 0.125
    if text.count("\n</reasoning>\n") == 1:
        count += 0.125
    if text.count("\n<answer>\n") == 1:
        count += 0.125
        count -= len(text.split("\n</answer>\n")[-1])*0.001
    if text.count("\n</answer>") == 1:
        count += 0.125
        count -= (len(text.split("\n</answer>")[-1]) - 1)*0.001
    return count

def xmlcount_reward_func(completions, **kwargs) -> list[float]:
    contents = [completion[0]["content"] for completion in completions]
    return [count_xml(c) for c in contents]

model_name = "Qwen/Qwen2.5-0.5B-Instruct"

output_dir="outputs/Qwen-0.5B-GRPO"
run_name="reasoner/Qwen-0.5B-GRPO-gsm8k"

training_args = GRPOConfig(
    output_dir=output_dir,
    run_name=run_name,
    learning_rate=5e-6,
    adam_beta1 = 0.9,
    adam_beta2 = 0.99,
    weight_decay = 0.1,
    warmup_ratio = 0.1,
    lr_scheduler_type='cosine',
    logging_steps=1,
    bf16=True,
    per_device_train_batch_size=16,
    gradient_accumulation_steps=4,
    num_generations=16,
    max_prompt_length=256,
    max_completion_length=200,
    num_train_epochs=1,
    save_steps=100,
    max_grad_norm=0.1,
    log_on_each_node=False,
    use_vllm=True,
    vllm_gpu_memory_utilization=.3,
    vllm_device="cuda:0",
    report_to="none",
    push_to_hub=True,
    hub_model_id=f"{run_name}",
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
).to("cuda")

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[
        xmlcount_reward_func,
        soft_format_reward_func,
        strict_format_reward_func,
        int_reward_func,
        correctness_reward_func],
    args=training_args,
    train_dataset=dataset,
)
print(trainer.accelerator.num_processes)
trainer.train()

培训目标

我们的目标是训练我们LLM更好地阐明提出特定答案的原因。换句话说，我们希望 to LLM 回答一个问题，后面有推理，然后是答案。OpenAI o1 和 Deepseek-R1（及其前身 R1-Zero）等模型不仅是一种非常酷的行为，还表明推理导致了强大的涌现行为，如 “自我验证、反思和生成长 CoT” （Deepseek）。这是一个活跃的研究领域，但在这里我们专门使用 Deepseek 风格的组相对策略优化 （GRPO）。

数据预处理

我们用来训练推理器的数据集是来自 OpenAI 的 GSM8K 数据集。在这里查看。该数据集包含成对的高质量小学数学单词问题及其相应的详细答案。GSM8K 数据集中的答案格式如下：

<steps> #### <final answer>

我们在 get_gsm8k_questions 中进行一些温和的预处理，以提取最终答案以用于 RL 例程

奖励定义（关于 GRPO 的旁白）

接下来，我们定义将用于指导 RL 例程的奖励函数。你可以将其视为定义我们希望在训练期间奖励的行为，在期望的行为发生时给予 “分数”，否则不给予分数。后一点是 GRPO 方法与以前方法相比的一大差异。在 GRPO 中，奖励是严格的（严格满足期望的行为时得分，否则不得分）。例如，我们想要引导模型朝着一个行为方向发展，生成格式如下的响应：

<reasoning>
...
</reasoning>
<answer>
...
</answer>

在训练期间，如果模型的输出与此格式匹配，我们希望用 strict_format_reward_func（） 中定义的 0.5 “积分”作为奖励。但是，我们还希望确保模型继续保持正确性，我们按照 correctness_reward_func（） 中的定义慷慨地奖励了这一点。

初始模型

我们将训练的初始模型是 Qwen2.5-0.5B-Instruct。我在我组织的私有云中的先前实验中使用过这个模型，因此它自动驻留在 AIHub 上。因此，现在将从您的私有云提供进一步的模型加载请求，无论是用于推理训练。

通过指向 AIHub 的 HF_ENDPOINT，我们将模型从 AIHub 获取到我们的进程中，并使用标准转换器代码将其加载到 GPU 上，如下所示：

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
).to("cuda")

创建并运行它！

trl 库为配置和运行 GRPO RL 例程（GRPOConfig 和 GRPOTrainer）提供了有用的抽象。为了确保将生成的推理器模型保存到 AIHub 中以供进一步实验或部署以进行推理，我设置了 push_to_hub=True。这就是它的全部内容。如上所述，HF_ENDPOINT 环境变量必须指向 AIHub，否则此例程将尝试将模型推送到 HuggingFace 上的公有云存储库。完成所有这些后，我们可以开始 RL 训练程序：

~/sid/dev$ HF_ENDPOINT=<my-aihub-endpoint> HF_TOKEN=<my-token> uv run grpo_script.py
INFO 02-28 01:46:17 init.py:207] Automatically detected platform cuda.
INFO 02-28 01:46:25 config.py:549] This model supports multiple tasks: {'embed', 'classify', 'score', 'reward', 'generate'}. Defaulting to 'generate'.
INFO 02-28 01:46:25 llm_engine.py:234] Initializing a V0 LLM engine (v0.7.3) with config: model='Qwen/Qwen2.5-0.5B-Instruct'...
...

为了演示我们新创建的模型最终学习的行为，这里有一个例子。

在训练例程开始时，模型并不热衷于遵循所需的推理格式：

Question:
Every Monday to Friday morning, Alice walks 10 miles through a large grass field to get to school. To avoid the hot afternoon sun on her way home, she walks 12 miles through a forest. When she gets home on Friday, how many miles will she have walked that week?

Answer:
110

Response:
Alice walks 10 miles each morning from Monday to Friday, so she walks a total of 5 * 10 = 50 miles by Friday morning.
The forest walks are different, but 12 miles through the forest are covered on Friday evening and the following Friday morning, converting this to miles, she walks 12 * 2 = 24 miles by Friday evening after the forest walks.
By the following Friday, Alice walks the 10 miles from Monday to Friday over a total of 5 + 2 = 7 days.
So over the entire week, Alice will have walked a total of 24 miles from Friday evening to Friday morning.

请注意，模型的响应不遵循我们希望它开发的 reasoning-s-by-answer 格式。此外，我们注意到答案离得很远。

相比之下，这是训练例程结束时的示例：

Question:
Steve finds 100 gold bars while visiting Oregon. He wants to distribute his gold bars evenly to his 4 friends. If 20 gold bars were lost on the way back to San Diego, how many gold bars will each of his 4 friends get when he returns?

Answer:
20

Response:
<reasoning>
Steve initially finds 100 gold bars and wants to distribute them evenly to his 4 friends. After 20 gold bars are lost, he has \(100 - 20 = 80\) gold bars left to distribute.
Therefore, each of Steve's 4 friends will get \(80 \div 4 = 20\) gold bars.
So, each of his 4 friends will get \(20\) gold bars when he returns.
</reasoning>
<answer>
20
</answer>

请注意，所需的格式已完全满足，答案是正确的。大不相同！

新模式

.
..
...
....
.....
{'loss': 0.1154, 'grad_norm': 18.625, 'learning_rate': 0.0, 'rewards/xmlcount_reward_func': 0.19140625, 'rewards/soft_format_reward_func': 0.0, 'rewards/strict_format_reward_func': 0.0, 'rewards/int_reward_func': 0.1953125, 'rewards/correctness_reward_func': 0.21875, 'reward': 0.60546875, 'reward_std': 0.7295257151126862, 'completion_length': 118.921875, 'kl': 2.883904218673706, 'epoch': 1.0}
{'train_runtime': 7453.3799, 'train_samples_per_second': 1.003, 'train_steps_per_second': 0.251, 'train_loss': 0.07686217568939345, 'epoch': 1.0}
100%|███████████████████████████████████████████████████████████| 1868/1868 [2:04:13<00:00,  3.99s/it]
~/sid/dev$

RL 例程已结束。由于 HF_ENDPOINT 指向您的 AIHub，这确保了在 AIStor 上创建新模型存储库，并将新检查点的文件推送到那里。让我们通过检查 AIHub 所在的存储桶 （“ai-central”） 来验证它，该存储桶位于与此新模型 （ ） 对应的前缀下 "models/sid/Qwen-0.5B-GRPO-gsm8k/main" ：

现在，整个组织都可以使用这个推理器模型进行实验或加载以进行推理，所有这些都不会打破您的私有云或丧失模型服务过程的所有权。

训练后扩展的好处正在推动组织根据专有数据创建自己的推理模型。想象一下，一个组织在其客户服务通话记录上调整开源模型，以便为客户服务团队打造一个能力强大、能够解决问题的副驾驶。那么想象一下，他们可靠地加载这些新保存的模型进行推理，而无需担心公有云模型存储库中断。与此同时，想象一下，这个组织可以使用开发人员喜欢的库（如 transformers、trl 等）来使他们的模型易于试验和构建应用程序。有了 AIHub，我们可以做的不仅仅是想象。AIHub 使组织能够将其 AI 实验和基础设施保留在本地，同时仍为开发人员提供已建立框架的自由和功能。