从零开始的LLM 4.微调与RL （代码实战篇）

前言

从零开始学习ai文章系列已完成《动手学深度学习》和《磨菇书》两本书的学习，新开的LLM系列课本来自《Happy-LLM》和《Hello-Agents》，但是内容和排版等个人重新进行整理，因此不会按照原来课本中的章节来写。如果哪里有错误的欢迎指正。或者不清晰的可以直接查看原文部分。

本章内容参考自 第十一章 Agentic-RL。不过原作者在讲解时使用的是一套自行封装的框架，对初学者而言抽象程度较高、不易跟随细节，因此我重新整理了一份更加"过程化"的代码实现与配套教程——尽量把每一步的输入输出、中间过程都摊开来讲，方便大家照着跑通、照着改。

在前面的章节中，我们已经系统介绍了大模型的基础知识、预训练方法和微调技术。这里先做一个小小的说明：预训练和微调在代码实现上其实高度相似，核心差异仅在于——预训练用的数据量更大、不需要套用对话模板（chat template）格式；而微调则需要将数据整理成"指令-回复"的对话形式。正因为如此，本章不再重复预训练的内容，而是聚焦于强化学习相关的两个核心方法：监督微调（Supervised Fine-Tuning，SFT） 和 群组相对策略优化（Group Relative Policy Optimization，GRPO）。

文章所有代码都能在：https://github.com/zusda/LLM-learning-from-zero/tree/main 里面找到

---

1. 整体规划

在正式动手之前，先梳理一下整体的实验路线，让大家对"我们要做什么、为什么这么做"有一个清晰的预期。

任务目标：以 Qwen3-0.6B 为基座模型，通过 SFT + GRPO 两阶段训练，使其具备解决 GSM8K 数学推理数据集 中小学应用题的能力，并能按照我们指定的格式输出推理过程与最终答案。

整体流程分为四大块：

1. 数据准备

   1. 下载 GSM8K 原始数据集
   2. 分别转换成 SFT 训练格式（指令-回复对）和 RL 训练格式（仅问题，答案用于奖励计算）

2. 模型准备 下载 Qwen3-0.6B 基座模型，封装好加载、推理、生成等常用接口函数，作为后续 SFT 和 GRPO 训练的统一基础设施。

3. 监督微调（SFT）

   1. 编写训练所需的数据加载、Trainer、损失计算等模块
   2. 启动训练，并通过对比训练前后模型在样例问题上的表现，直观感受 SFT 带来的变化

4. GRPO 强化学习 思路与 SFT 阶段类似：先搭建好奖励函数、采样逻辑、策略更新等核心组件，再启动训练，最后对比 SFT 后模型与 GRPO 后模型的表现差异，看看强化学习阶段到底带来了哪些提升。

2. 数据准备

代码：1. data_preparation.py

2.1 GSM8K 数学推理数据集

数学推理是评估 LLM 推理能力的理想任务。首先，数学问题有明确的标准答案，可以自动评估，无需人工标注或复杂的奖励模型；其次，解决数学问题需要拆解问题、逐步推导，这正是多步推理的典型场景；最后，在数学任务上学到的推理能力具有较强的泛化性，能够迁移到其他领域。相比之下，开放式问答任务（如"如何学习编程？"）的答案质量难以客观评估，往往需要大量人工标注。

GSM8K（Grade School Math 8K）是一个高质量的小学数学应用题数据集。如下图 所示，数据集包含 7,473 个训练样本和 1,319 个测试样本，难度对应小学数学水平（2-8 年级），题型为应用题，每道题需要 2-8 步推理才能得出最终答案。

2.2 下载数据

我们可以通过一行代码，将数据集下载到本地：

from datasets import load_dataset
dataset = load_dataset("openai/gsm8k", "main", cache_dir="data/gsm8k")

这里通过 cache_dir 参数将数据集下载到指定文件夹下（若不指定，默认下载路径为 C:\Users\用户名\.cache\huggingface\datasets）。

\# 查看数据集结构
print(dataset)

\# 查看训练集第一条
print(dataset["train"][0])

{
    'question': 'Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?', 
    'answer': 'Natalia sold 48/2 = <<48/2=24>>24 clips in May.\nNatalia sold 48+24 = <<48+24=72>>72 clips altogether in April and May.\n#### 72'
}

可以看到，原始数据中每条样本包含 question（问题）和 answer（答案，内含完整的解题步骤，以 #### 标记最终数值结果）两个字段。

2.3 转换格式

拿到原始数据后，并不能直接喂给模型训练——SFT 和 RL 这两种训练方式所需要的数据"形态"并不一样，因此需要针对性地做格式转换，如下图所示。

原始格式直接来自数据集，包含问题（question）和答案（answer，含完整解题步骤），更适合人类阅读，但还不能直接用于训练。

2.3.1 SFT格式

SFT 格式用于监督微调，核心思路是：将问题转换为对话格式的 prompt，将完整解答作为 completion。例如：

{
    "prompt": "<|im_start|>user\nNatalia sold clips to 48 of her friends...<|im_end|>\n<|im_start|>assistant\n",
    "completion": "Let me solve this step by step.\n\nStep 1: ...\n\nFinal Answer: 72<|im_end|>"
}

这里的关键在于使用模型自带的对话模板（如 Qwen 的 <|im_start|> / <|im_end|> 标记）：prompt 部分包含用户提出的问题，completion 部分则包含完整的解题过程与最终答案。这样模型在训练中既能学会"按指定格式输出"，也能学会"如何分步推理"。

2.3.2 RL格式

RL 格式用于强化学习，与 SFT 格式不同的是：它只提供问题和正确答案，不提供解题过程。例如：

{
    "prompt": "<|im_start|>user\nNatalia sold clips to 48 of her friends...<|im_end|>\n<|im_start|>assistant\n",
    "ground_truth": "72"
}

关键点在于：prompt 部分与 SFT 格式保持一致，但 ground_truth 只保留最终答案（用于后续计算奖励），不再提供中间推理步骤。这种设计是刻意为之——它迫使模型必须自己生成完整的推理过程，而不是简单地"背答案"。

下面的代码虽然返回了 4 个字段，但实际训练中我们只会用到其中两个：prompt 和 ground_truth，另外两个字段（question、full_answer）主要用于调试和结果展示。

def format_for_rl(self, example: Dict[str, Any]) -> Dict[str, Any]:
    """
    格式化为RL训练格式(Standard Format with Chat Template Applied)

    Returns:
        - prompt: 应用chat template后的文本字符串
        - ground_truth: 正确答案
        - question: 原始问题
        - full_answer: 完整答案
    """
    question = example["question"]
    answer = example["answer"]

    if "####" in answer:
        _, final_answer = answer.split("####")
        final_answer = final_answer.strip()
    else:
        final_answer = answer.strip()

    prompt_content = f"Question: {question}\n\nLet's solve this step by step:"

    # 如果提供了tokenizer,应用chat template
    if self.tokenizer:
        messages = [{"role": "user", "content": prompt_content}]
        prompt_text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
    else:
        prompt_text = prompt_content
    # prompt要进行这个处理，主要是后续需要LLM生成多个答案，用以训练。

    return {
        "prompt": prompt_text,
        "ground_truth": final_answer,
        "question": question,
        "full_answer": answer
    }

2.4 保存与加载本地数据

保存与加载本地数据

将原始数据转换成对应格式后，为了节省时间、避免每次训练前都重新走一遍"下载 + 转换"的流程，我们会把转换好的数据直接缓存到本地磁盘。具体逻辑是：训练开始前先检测本地是否已存在缓存文件，存在则直接加载，不存在则进行转换并保存，供下次复用。

下面以 SFT 对应的数据集类为例，说明这套缓存机制的具体实现。

class GSM8K_SFT_Dataset:
    """GSM8K数学推理数据集 -> SFT格式

    自动处理：下载原始数据 -> 转换为SFT格式 -> 本地缓存为jsonl -> 后续直接加载，
    不用每次都重新走HuggingFace下载+转换的流程。
    """

    def __init__(
        self,
        split: str = "train",
        max_samples: int = None,
        raw_cache_dir: str = "data/gsm8k",
        sft_cache_path: str = None,
    ):
        """
        Args:
            split: 数据集分割 ("train" 或 "test")
            max_samples: 最大样本数
            raw_cache_dir: 原始GSM8K数据集的缓存目录
            sft_cache_path: SFT格式jsonl文件的保存路径，默认根据split自动生成
        """
        self.split = split
        # self.max_samples = max_samples
        # 统一把 None 和 'all' 归一化成 None,表示不限制数量
        if max_samples is None or (isinstance(max_samples, str) and max_samples.lower() == "all"):
            self.max_samples = None
        else:
            self.max_samples = int(max_samples)


        self.raw_cache_dir = raw_cache_dir
        if sft_cache_path:
            self.sft_cache_path=sft_cache_path
        elif max_samples:
            self.sft_cache_path=f"data/gsm8k/gsm8k_sft_{split}_{max_samples}.jsonl"
        else:
            self.sft_cache_path=f"data/gsm8k/gsm8k_sft_{split}_all.jsonl"


        self.data = self.load_or_create()

    # ---------- 单条样本转换逻辑 ----------
    def format_for_sft(self, example):
        question = example["question"]
        answer = example["answer"]

        if "####" in answer:
            reasoning, final_answer = answer.split("####")
            reasoning = reasoning.strip()
            final_answer = final_answer.strip()
        else:
            reasoning = answer
            final_answer = ""

        prompt = f"Question: {question}\n\nLet's solve this step by step:\n"
        completion = f"{reasoning}\n\nFinal Answer: {final_answer}"

        return {
            "prompt": prompt,
            "completion": completion,
            "text": prompt + completion
        }

    # ---------- 下载原始数据并批量转换 ----------
    def create_sft_data(self):
        print(f"📥 加载原始 GSM8K 数据集 (split={self.split})...")
        raw = load_dataset("openai/gsm8k", "main", split=self.split, cache_dir=self.raw_cache_dir)

        if self.max_samples:
            raw = raw.select(range(min(self.max_samples, len(raw))))

        sft_data = [self.format_for_sft(example) for example in raw]    # 对于GSM8K数据集应该不会太慢，一共就9000多条数据
        print(f"⚙️ 已转换 {len(sft_data)} 条样本为SFT格式")
        return sft_data

    # ---------- jsonl 读写 ----------
    def save_to_jsonl(self, data):
        dir_name = os.path.dirname(self.sft_cache_path)
        if dir_name:  # 只有当路径里确实包含目录部分时才创建
            os.makedirs(dir_name, exist_ok=True)
        # os.makedirs(os.path.dirname(self.sft_cache_path), exist_ok=True)
        with open(self.sft_cache_path, "w", encoding="utf-8") as f:
            for item in data:
                f.write(json.dumps(item, ensure_ascii=False) + "\n")
        print(f"💾 已保存至 {self.sft_cache_path}")

    def load_from_jsonl(self):
        data = []
        with open(self.sft_cache_path, "r", encoding="utf-8") as f:
            for line in f:
                data.append(json.loads(line))
        print(f"📂 从本地缓存加载了 {len(data)} 条样本: {self.sft_cache_path}")
        return data

    # ---------- 核心：存在则load，不存在则create+save ----------
    def load_or_create(self):
        if os.path.exists(self.sft_cache_path):
            return self.load_from_jsonl()
        else:
            data = self.create_sft_data()
            self.save_to_jsonl(data)
            return data

    # ---------- 对外接口 ----------
    def to_hf_dataset(self) -> Dataset:
        """转换为HuggingFace Dataset对象，方便配合SFTTrainer等工具使用"""
        return Dataset.from_list(self.data)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

整个类的核心逻辑可以概括为一句话：load_or_create() 负责判断"是加载还是新建"，create_sft_data() 负责"下载 + 转换"，save_to_jsonl() / load_from_jsonl() 负责"读写本地缓存文件"——三者各司其职，组合起来就实现了"首次运行自动生成缓存，之后直接复用"的效果。

使用方法如下：

train_dataset = GSM8K_SFT_Dataset(split="train", max_samples=1000,sft_cache_path='sft_train.jsonl')
# train_dataset = GSM8K_SFT_Dataset(split="test", max_samples=1000,sft_cache_path='sft_train.jsonl')
print(len(train_dataset))
print(train_dataset[2])

几个参数说明一下：

• split：我们下载的原始数据中包含 train 和 test 两部分，需要通过该参数指定读取哪一部分；
• max_samples：指定读取的数据条数；若填写为 None 或不传该参数，则默认读取全部数据；
• sft_cache_path：转换后数据的本地保存路径，若不指定则会根据 split 和 max_samples 自动生成一个文件名（如 gsm8k_sft_train_1000.jsonl）。

数据最终结构如下：

3. 模型准备

代码：2. LLM_preparation.py

首先需要将模型下载到本地。这里以 Qwen3-0.6B 为例，我将模型保存在：.\LLM_models\Qwen_Qwen3-0.6B。

对应的模型 ID 为：

self.model_id="Qwen/Qwen3-0.6B",

3.1 下载模型

当程序首次运行时，本地还没有模型文件，因此需要从 Hugging Face 下载。

def download_model(self):
    print(f"本地未找到模型,从 {os.environ['HF_ENDPOINT']} 下载 {self.model_id} ...")

    self.local_dir.mkdir(parents=True, exist_ok=True)

    model = AutoModelForCausalLM.from_pretrained(
        self.model_id,
        torch_dtype=self.torch_dtype,
        device_map=self.device_map,
    )
    tokenizer = AutoTokenizer.from_pretrained(self.model_id, use_fast=False)

    # 固化保存到本地目录,后续直接从这里加载,不再走网络
    model.save_pretrained(self.local_dir)
    tokenizer.save_pretrained(self.local_dir)

    print(f"下载完成,已保存到 {self.local_dir}")
    return model, tokenizer

这里最核心的代码是：AutoModelForCausalLM.from_pretrained(self.model_id) 。

会根据提供的模型 ID 自动完成以下工作：

1. 从 Hugging Face 下载模型配置文件（config.json）
2. 下载模型权重文件（model.safetensors）
3. 构建对应的模型结构
4. 将权重加载到模型中

随后调用：

model.save_pretrained(self.local_dir)
tokenizer.save_pretrained(self.local_dir)

将模型和 tokenizer 保存到本地目录，避免每次运行时重复下载。

3.2 本地模型目录

模型保存路径由下面的代码生成：

self.local_dir = Path(local_root) / model_id.replace("/", "_")

例如：

model_id = "Qwen/Qwen3-0.6B"

会转换成：LLM_models/Qwen_Qwen3-0.6B/

3.3 加载本地模型

后续程序再次启动时，会先检查模型目录是否存在。

如果目录已经存在，说明模型之前已经下载过，此时无需再访问网络，而是直接从本地加载：

def load_model(self):
    print(f"检测到本地模型,直接从 {self.local_dir} 加载...")

    model = AutoModelForCausalLM.from_pretrained(
        self.local_dir,
        torch_dtype=self.torch_dtype,
        device_map=self.device_map,
    )
    tokenizer = AutoTokenizer.from_pretrained(self.local_dir, use_fast=False)

    return model, tokenizer

这里与下载阶段最大的区别在于：

# 下载时
from_pretrained(self.model_id)

# 加载时
from_pretrained(self.local_dir)

from_pretrained() 既可以接收 Hugging Face 模型 ID，也可以接收本地目录路径。

当传入模型 ID 时，会从 Hugging Face 下载模型；当传入本地路径时，则直接读取本地保存的模型文件。因此模型下载一次后，后续启动都可以直接从本地加载，无需再次联网。

当然，在下载、保存、加载model的时候，tokenizer也是一并一起处理的。

3.4 使用

使用方式非常简单，只需要指定模型 ID 和本地模型存储目录即可：

if __name__ == "__main__":
    # .\LLM_models\Qwen_Qwen3-0.6B
    loader = LLMModelLoader(model_id="Qwen/Qwen3-0.6B", local_root="./LLM_models")
    model_trained, tokenizer = loader.get_model_and_tokenizer()

4. stf训练

代码：4. stf_trainer.py

4.1 确定传入数据

首先确定训练数据格式。 SFT（Supervised Fine-Tuning）的本质是让模型学习：

给定 Prompt，预测 Completion。

因此每条训练样本至少需要能够拆分出：

• Prompt（模型输入）
• Completion（模型目标输出）

下面以 prompt/completion 格式作为示例。

train_data = [
    {
        "prompt": "<|im_start|>user\nNatalia sold clips to 48 of her friends...<|im_end|>\n<|im_start|>assistant\n",

        "completion": "Let me solve this step by step.\n\nStep 1: ...\n\nFinal Answer: 72<|im_end|>"

    },

    {
        "propmpt":"......",
		"completion":"......",
    }

]

4.2 构建 SFTDataset

接下来需要将原始训练数据封装为 Dataset 对象，即重写 Dataset 类。这样才能配合 DataLoader 实现批量加载（Batch）、数据打乱（Shuffle）、多进程读取等功能，并在训练过程中按批次将数据传递给模型。

首先定义数据集类：

class SFTDataset(Dataset):

在初始化函数中传入原始数据、Tokenizer 以及相关配置：

def __init__( self, data, tokenizer, max_length=512, use_chat_template=True ):
    # data：			原始训练数据；
    # tokenizer：	模型对应的分词器；
    # max_length：	最大序列长度；
    # use_chat_template：	是否使用对话模板格式化数据。

---

构建 Prompt ID 与 Full ID

为了后续正确计算 SFT Loss，需要额外实现一个函数，将单条原始数据转换为：

• prompt_ids：仅包含用户输入（Prompt）的 Token ID；
• full_ids：包含 Prompt 与 Completion 的完整 Token ID。

在生成 prompt_ids 时，需要设置：

add_generation_prompt=True

这样会自动在对话末尾补充 Assistant 的起始标记：

<|im_start|>assistant

从而告诉模型：

接下来应该开始生成 Assistant 回复。

而构造完整的 full_ids 时，则需要显式将答案加入对话模板：

messages = [
    {"role": "user", "content": prompt},
    {"role": "assistant", "content": completion},
]

随后再次调用：

full_ids = tokenizer.apply_chat_template( messages, tokenize=True )

此时 completion 会被嵌入到 Assistant 对应的位置，从而得到完整训练样本。

之所以要同时构造 prompt_ids 和 full_ids，是因为在 SFT 训练中：

• Prompt 部分属于输入条件；
• Completion 部分才是模型需要学习预测的内容。

因此后续构造 labels 时，需要将 Prompt 部分屏蔽掉，只让模型在 Completion 部分计算 Loss。

---

重写 getitem

下面实现 Dataset 的核心函数：

def __getitem__(self, idx):
    # 取出第 idx 条原始数据
    # 此时仍然是纯文本格式：
    # {"prompt": "...", "completion": "..."}
    example = self.data[idx]
    prompt_ids, full_ids = self.build_ids(
        example["prompt"], example["completion"]
    )

    # 构造 SFT Labels
    # Prompt 部分不参与 Loss 计算
    labels = (
        [-100] * len(prompt_ids) + full_ids[len(prompt_ids):]
    )

    # 安全截断
    # 防止极长样本导致显存占用过高
    full_ids = full_ids[:self.max_length]
    labels = labels[:self.max_length]

    return {
        "input_ids": torch.tensor(full_ids, dtype=torch.long),
        "labels": torch.tensor(labels, dtype=torch.long),
    }

分成这两部分，主要是我们下面提供input id的时候，我们同时也需要提供labels，告诉后续训练哪些位置才是需要计算loss的，因为prompt部分是不需要计算loss的。我们需要根据这两部分来获取labels。

下面就是我们重写Dataset类的__getitem__函数

其中最重要的是：

labels = [-100] * len(prompt_ids) + full_ids[len(prompt_ids):]

在 Hugging Face 的因果语言模型训练中，label = -100 表示当前位置不参与 Loss 计算。

经过 __getitem__() 处理后，原始数据：

{ 
    "prompt": "...", 
    "completion": "..." }

会被转换为：

{
    "input_ids": [...],
    "labels": [...]
}

其中：

• input_ids 为完整的 prompt + completion 经过 Tokenizer 编码后的结果；
• labels 用于标记哪些 Token 需要参与 Loss 计算。

后续通过 DataLoader 读取数据：

for batch_idx, batch in enumerate(self.dataloader):

每次获得的 batch 都是由多个样本经过 __getitem__() 处理后再拼接得到的。

注意：很多初学者容易误认为 labels 是一个由 0/1 组成的 Mask。实际上并不是。

在 Hugging Face 的 CausalLM 训练中，labels 保存的是真实目标 Token ID，而不是 Mask。

其中：

• -100 表示忽略该位置，不参与 Loss 计算；
• 其它位置保存真实 Token ID，用于监督模型学习下一个 Token 的预测。

因此 SFT 中的 labels 通常类似：

[-100, -100, -100, 151645, 872, 3094, ...]

而不是：

[0, 0, 0, 1, 1, 1, ...]

如果是预训练（Pretraining）场景，由于训练目标是学习语言本身，因此所有 Token 都需要参与预测。

通常直接：

labels = input_ids.copy()

即所有位置都会参与 Loss 计算，不需要对任何部分进行遮蔽。

而 SFT（Supervised Fine-Tuning）则不同。

SFT 数据通常由 Prompt 和 Completion 两部分组成，其中：

• Prompt 用于提供上下文；
• Completion 才是模型需要学习生成的内容。

因此需要将 Prompt 部分设置为 -100。

这也是预训练与 SFT 最本质的区别之一：

4.3 Batch 数据预处理

在上一节中，我们通过 SFTDataset 完成了单条样本的数据预处理，每次调用 __getitem__() 都会得到：

{    
    "input_ids": ...,    
    "labels": ... 
}

但模型训练并不是一次处理一条数据，而是一次处理一个 Batch。因此，在 DataLoader 将多条样本组合成 Batch 时，还需要进行额外处理。

最主要的问题是：

不同样本经过 Tokenizer 编码后长度通常不同，而 Tensor 要求同一维度长度一致，因此需要对 Batch 内的数据进行 Padding（补齐）。

Padding 和 Batch 整理逻辑通常定义在 DataLoader 的 collate_fn 参数中。当 DataLoader 从 Dataset 中取出多个样本准备组成一个 Batch 时，会自动调用我们传入的 collate_fn() 函数，对这些样本进行补齐、拼接等预处理操作，最终生成模型可直接训练的 Batch 数据。

dataset = SFTDataset(train_data, tokenizer, max_length=max_length)
self.dataloader = DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=lambda batch: collate_fn(batch, tokenizer.pad_token_id),
)

处理函数如下，代码注释很清楚因此不过多解释。

def collate_fn(batch, pad_token_id):
    # 把“不同长度的样本”整理成一个标准的训练 batch（对齐 + padding + mask）

    max_len = max(len(item["input_ids"]) for item in batch)

    input_ids = torch.full((len(batch), max_len), pad_token_id, dtype=torch.long)
    labels = torch.full((len(batch), max_len), -100, dtype=torch.long)
    attention_mask = torch.zeros((len(batch), max_len), dtype=torch.long)
    # 初始化空矩阵(以max_len=3，batch=2为例)
    # input_ids:
    # [
    #  [PAD, PAD, PAD],
    #  [PAD, PAD, PAD]
    # ]
    #
    # labels:
    # [
    #  [-100, -100, -100],
    #  [-100, -100, -100]
    # ]
    #
    # attention_mask:
    # [
    #  [0, 0, 0],
    #  [0, 0, 0]
    # ]

    for i, item in enumerate(batch):
        seq_len = len(item["input_ids"])
        input_ids[i, :seq_len] = item["input_ids"]
        labels[i, :seq_len] = item["labels"]
        attention_mask[i, :seq_len] = 1
    # 第二条样本举例
    # item[1]:
    # input_ids = [20, 21]
    # labels    = [20, 21]
    # 变成
    # input_ids
    # [
    #  [20, 21, PAD]
    # ]
    # labels
    # [
    #  [20, 21, -100]
    # ]
    # attention_mask
    # [
    #  [1, 1, 0]
    # ]


    # attention_mask是告诉模型哪些位置是真实token、哪些是padding。在attention计算里,0的位置在计算attention score时会被mask掉。
    # labels是告诉模型哪些位置需要计算loss。

    # 用一个具体的例子,假设两条样本padding后长度都是8:
    # 样本1: [user问话token×3] [assistant回答token×3] [pad×2]
    # 样本2: [user问话token×2] [assistant回答token×4] [pad×2]
    # 两个mask长这样:
    # 样本1
    # attention_mask = [1, 1, 1, 1, 1, 1, 0, 0]  # 后两个是pad,不参与attention
    # labels         = [-100, -100, -100, id, id, id, -100, -100]  # 前3个是prompt不算loss,后两个是pad也不算loss
    #
    # 样本2
    # attention_mask = [1, 1, 1, 1, 1, 1, 0, 0]
    # labels         = [-100, -100, id, id, id, id, -100, -100]

    # Decoder的causal attention本身就自带一个下三角mask,强制每个位置只能看到自己和左边的token。右边的padding天然被这个下三角结构隔断,根本不需要attention_mask额外去屏蔽它。
    # 所以在右padding+decoder模型的训练场景下,attention_mask全为1和正确设置等价,真正起作用的只有labels里的-100。

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

需要注意区分 attention_mask 和 labels 的作用。

attention_mask 作用于 Attention 计算阶段，用于告诉模型哪些位置是真实 Token、哪些位置是 Padding。框架会根据它构造 Attention Mask，从而在构建 Attention 计算图时忽略 Padding 节点。

而 labels 作用于 Loss 计算阶段，用于告诉模型哪些位置需要参与损失计算。对于 SFT 来说，Prompt 部分和 Padding 部分通常会被设置为 -100，从而不参与 Loss 计算。

因此：

• attention_mask 决定哪些 Token 可以参与信息交互；
• labels 决定哪些 Token 参与 Loss 计算。

二者虽然经常同时出现，但作用于模型计算流程中的不同阶段。

不过对于 Qwen、Llama、GPT 等 Decoder-Only 模型，还需要额外说明一点：模型内部的 Causal Attention 本身就自带下三角 Mask，保证每个 Token 只能看到自己及其左侧的 Token，而无法看到未来 Token。

因此在最常见的 右侧 Padding（Right Padding）+ Decoder 模型 的训练场景下，即使将所有位置的 attention_mask 都设置为 1，右侧 Padding Token 通常也无法被前面的真实 Token 看到。在这种情况下，attention_mask 对训练结果的影响往往很小。

真正决定 Prompt 和 Padding 是否参与训练的，实际上是：

labels = -100

因为 Loss 计算阶段会直接忽略这些位置。

当然，在实际工程中仍然建议正确构造 attention_mask，以保证与 Hugging Face、Flash Attention、推理框架以及各种特殊场景保持兼容。

4.4 开始训练

成数据预处理后，就可以正式开始训练了。

整个训练流程如下：

读取Batch
    ↓
Forward
    ↓
计算Loss
    ↓
Backward
    ↓
梯度累积
    ↓
更新参数

首先遍历数据集，并将当前 Batch 中的数据移动到训练设备（GPU 或 CPU）上。

首先遍历数据集，并将当前 Batch 中的数据移动到训练设备（GPU 或 CPU）上。

这里使用的 tqdm 仅用于显示训练进度，对训练逻辑本身没有影响。

for epoch in range(self.num_epochs):
    progress_bar = tqdm(
        enumerate(self.dataloader),
        total=len(self.dataloader),
        desc=f"epoch {epoch}",
    )
    for batch_idx, batch in progress_bar:
    # for batch_idx, batch in enumerate(self.dataloader):
        batch = {k: v.to(self.device) for k, v in batch.items()}
        # {
        #     "input_ids": tensor形状(2, 18),
        #     "labels": tensor形状(2, 18),
        #     "attention_mask": tensor形状(2, 18),
        # }
        # batch.items()遍历这3个键值对。把这3个tensor分别搬到GPU(或者CPU,取决于self.device)

可以看到，最终送入模型的数据仅包含：

input_ids
labels
attention_mask

这三个字段。

随后执行前向传播，获得模型输出并计算 Loss：

for batch_idx, batch in progress_bar:
# for batch_idx, batch in enumerate(self.dataloader):
    batch = {k: v.to(self.device) for k, v in batch.items()}
 
    outputs = self.model(**batch)
    # 训练时仅输出每个位置的向量。并未转换成文本。
    # 推理时你用的model.generate(ids)是HuggingFace在GenerationMixin里实现的一个高层方法,内部封装了自回归循环——每次forward一步,取出logits,采样或者greedy选出下一个token,把这个token拼回input_ids,再forward下一步,直到遇到eos或者达到max_new_tokens。你看到的是"输入一段prompt,输出一段完整的新文本"。
    loss = outputs.loss / self.gradient_accumulation_steps
    loss.backward()
    

由于这里使用了梯度累积（Gradient Accumulation），因此并不会每个 Batch 都更新一次参数，而是累积足够次数后再统一更新。

当达到指定的梯度累积步数后，会执行梯度裁剪、参数更新、学习率更新以及梯度清空等操作。与此同时，还会记录当前 Loss，并按照设定的间隔保存模型检查点（Checkpoint）。

for batch_idx, batch in progress_bar:
# for batch_idx, batch in enumerate(self.dataloader):
    batch = {k: v.to(self.device) for k, v in batch.items()}
    outputs = self.model(**batch)

    loss = outputs.loss / self.gradient_accumulation_steps
    loss.backward()

    if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
        self.optimizer.step()
        self.scheduler.step()
        self.optimizer.zero_grad()
        global_step += 1

        real_loss = loss.item() * self.gradient_accumulation_steps
        progress_bar.set_postfix(step=global_step, loss=f"{real_loss:.4f}")

        if global_step % 10 == 0:
            # real_loss = loss.item() * self.gradient_accumulation_steps
            print(f"epoch {epoch} step {global_step} loss {real_loss:.4f}")

        if global_step % self.save_steps == 0:
            self.save_checkpoint(global_step)

至此，一个完整的 SFT 训练流程就完成了。

4.5 SFT 训练的 Main 函数

前面已经完成了：

• 数据集构建；
• Tokenizer 加载；
• 模型加载；
• Dataset 与 DataLoader 构建；
• SFT Trainer 实现。

因此最终的训练入口函数实际上非常简单。

首先准备训练所需的三个核心对象：

• tokenizer
• model
• train_data

def stf_train_main():
    GSM8K_SFT_Dataset = import_things('1. data_preparation.py','GSM8K_SFT_Dataset')
    LLMModelLoader = import_things('2. LLM_preparation.py','LLMModelLoader')

    loader = LLMModelLoader(model_id="Qwen/Qwen3-0.6B", local_root="./LLM_models")
    model, tokenizer = loader.get_model_and_tokenizer()


    train_data = GSM8K_SFT_Dataset(split="train", max_samples=None)

随后为模型添加 LoRA。

由于笔者使用的是笔记本显卡，无法承担全量参数微调（Full Fine-Tuning）的显存开销，因此这里采用 PEFT（Parameter Efficient Fine-Tuning）中的 LoRA 方案，仅训练少量新增参数。

完成 LoRA 包装后，将模型、Tokenizer 和训练数据传入前面实现的 SimpleSFTTrainer，最后调用：

trainer.train()

即可开始训练。

def stf_train_main():
	......

    # 加载模型之后,传进trainer之前包一层
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        lora_alpha=32,
        lora_dropout=0.1,
        target_modules=["q_proj", "v_proj"],  # Qwen3的attention投影层
        # Transformer的attention层一共有4个投影矩阵:q_proj、k_proj、v_proj、o_proj。
        # 论文做了消融实验,在同样的参数预算下(LoRA的总秩固定),对比"只调q和v"、"调全部4个但每个秩更小"等不同组合的下游任务效果,
        # 发现q_proj+v_proj这个组合性价比最高

    )
    model = get_peft_model(model, lora_config)
    # model.print_trainable_parameters()  # 确认可训练参数量

    trainer = SimpleSFTTrainer(
        model=model,
        tokenizer=tokenizer,
        train_data=train_data,
        output_dir="./checkpoints",
        batch_size=2,
        num_epochs=3,
        save_steps=50,
        gradient_checkpointing=True,  # 8GB显存建议开启
    )
    trainer.train()

在个人笔记本环境下，每个 Epoch 训练时间大约为 30～40 分钟。

为了突出核心原理，文中省略了部分工程化细节，例如混合精度训练、学习率调度器配置、Checkpoint 恢复训练、日志记录等内容。相关实现均已保留在源码及代码注释中，如果需要自行复现，建议结合源码一起阅读。

4.6 训练结果

训练完成后，先观察 Loss 的变化情况。

下图是根据训练日志中的 Loss 数据绘制的曲线图。其中浅色曲线表示每个 Step 的原始 Loss，红色曲线表示经过平滑处理后的平均 Loss。

可以看到，随着训练进行，Loss 快速下降并逐渐趋于稳定。后期平均 Loss 基本维持在 0.4～0.6 之间，说明模型已经较好地拟合了当前训练数据。

当然，仅观察 Loss 并不能完全说明模型是否真正学到了目标能力。对于大语言模型来说，更直观的方法还是实际测试模型输出结果。

因此下面编写一个简单的工具，用于：

• 加载指定 Checkpoint；
• 恢复 LoRA Adapter；
• 封装一个简单的 Chat 接口；
• 对比训练前后的模型输出。

代码如下：

import importlib.util
def import_things(file,class_name):
    spec = importlib.util.spec_from_file_location(
        "data_preparation",
        file
    )
    data_preparation = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(data_preparation)

    # 然后从这个module对象里取出你要的类
    out_class = getattr(data_preparation, class_name)

    return out_class

def load_model_from_checkpoint(checkpoint,merge=False):
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from peft import AutoPeftModelForCausalLM

    adapter_path = f"./checkpoints/checkpoint-{checkpoint}"

    # 之所以能省略base model路径,是因为adapter_config.json里其实已经记录了base_model_name_or_path这个字段
    model = AutoPeftModelForCausalLM.from_pretrained(adapter_path, dtype=torch.bfloat16)
    tokenizer = AutoTokenizer.from_pretrained(adapter_path, use_fast=False)

    if merge:
        merged_model = model.merge_and_unload()
        return merged_model,tokenizer
    return model,tokenizer


import torch

class LLM_wrapper:
    def __init__(self, model, tokenizer):
        model.to("cuda")
        self.model = model.eval()
        self.tokenizer = tokenizer

        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

    def chat(self, query, max_new_tokens=256, do_sample=False, temperature=1.0, top_p=0.95):
        prompt_text = self.tokenizer.apply_chat_template(
            [{"role": "user", "content": query}],
            tokenize=False,
            add_generation_prompt=True,
        )

        inputs = self.tokenizer(
            prompt_text,
            return_tensors="pt",
            add_special_tokens=False,
        ).to(self.model.device)

        gen_kwargs = {
            "max_new_tokens": max_new_tokens,
            "do_sample": do_sample,
            "pad_token_id": self.tokenizer.pad_token_id,
        }
        if do_sample:
            gen_kwargs["temperature"] = temperature
            gen_kwargs["top_p"] = top_p

        with torch.no_grad():
            output_ids = self.model.generate(**inputs, **gen_kwargs)

        # 只解码新生成的部分,去掉重复的prompt
        generated_ids = output_ids[0][inputs["input_ids"].shape[1]:]
        response = self.tokenizer.decode(generated_ids, skip_special_tokens=True)

        return response

随后从测试集中挑选一道题目，同时分别使用：

• SFT 训练后的模型；
• 原始基础模型；

进行推理测试。

if __name__ == '__main__':

    model,tokenizer = load_model_from_checkpoint(2800)
    agent = LLM_wrapper(model,tokenizer)

    query = "Question: Betty is saving money for a new wallet which costs $100. Betty has only half of the money she needs. Her parents decided to give her $15 for that purpose, and her grandparents twice as much as her parents. How much more money does Betty need to buy the wallet?\n\nLet's solve this step by step:\n"
    response = agent.chat(query)
    print(response)

    print('======================')
    LLMModelLoader = import_things('2. LLM_preparation.py','LLMModelLoader')
    loader = LLMModelLoader(model_id="Qwen/Qwen3-0.6B", local_root="./LLM_models")
    model2, tokenizer2 = loader.get_model_and_tokenizer()
    agent2 = LLM_wrapper(model2,tokenizer2)

    response = agent2.chat(query)
    print(response)

测试结果如下：

可以看到，经过 SFT 训练后，模型已经明显学习到了训练数据中的输出模式。

在原始模型中，对于这类 GSM8K 数学推理问题，回答往往不够稳定，或者无法按照数据集要求的格式进行推理。而经过 SFT 后，模型开始主动输出符合训练样本风格的 Step-by-Step 推理过程，并能够按照预期格式组织答案。

当然，这里仅展示了一条测试样例，并不能严格代表模型整体性能。如果需要进行严谨评估，还应该在完整测试集上统计准确率等指标。不过从定性结果来看，本次 SFT 训练已经成功将训练数据中的行为模式迁移到了模型输出中。

5. GRPO训练

代码：5. rl_trainer.py

5.1 GRPO训练概述

在进入代码实现之前，先从整体上理解 GRPO（Group Relative Policy Optimization）的核心思想。

可以将 GRPO 理解为一种基于相对比较的强化学习训练方法：

1. 模型针对同一个问题生成多个回答（Group）。
2. 使用奖励函数（Reward Function）对每个回答进行评分。
3. 在同组回答内部计算相对优劣，即优势（Advantage）。
4. 提高高优势回答的生成概率，降低低优势回答的生成概率。
5. 经过大量训练后，模型会逐渐倾向于生成能够获得更高奖励的回答。

从优化目标来看，GRPO 本质上并不是在向模型灌输新的知识，而是在调整模型的行为策略（Policy）。

换句话说：

SFT 更偏向于学习知识和模仿示例，而 RL 更偏向于学习“什么样的回答更值得生成”。

因此，强化学习阶段的主要作用通常是：

• 提高回答正确率；
• 提高推理质量；
• 约束输出格式；
• 减少幻觉；
• 强化安全策略；
• 对齐人类偏好（Human Preference）。

而不是让模型凭空学会训练数据中从未接触过的新知识。

---

5.1.1 GRPO 中最关键的部分是奖励来源。

通常有两种方式：

方式1：人工设计奖励函数（Rule-based Reward）

例如：

reward = (
    format_reward +
    correctness_reward +
    length_reward
)

奖励函数根据各种规则进行打分：

• 答案是否正确；
• 格式是否符合要求；
• 是否包含指定标签；
• 推理过程是否完整；
• 长度是否合理；

这类奖励函数实现简单、可解释性强，也是多数小规模实验最常见的做法。

方式2：奖励模型（Reward Model）

另一种做法是训练专门的奖励模型（Reward Model）。

训练流程通常为：

人工偏好数据
        ↓
Reward Model
        ↓
GRPO / PPO
        ↓
LLM

奖励模型学习人类对回答质量的偏好：

回答A > 回答B

训练完成后，Reward Model 会输出奖励分数：

reward = reward_model(response)

随后 GRPO 再利用这些奖励计算优势（Advantage），更新策略模型。

因此无论奖励来自：

• 手工规则；
• 奖励模型；
• AI 评审模型（LLM-as-a-Judge）；

最终都会统一转换为一个奖励值，然后进入 Advantage 计算流程。

---

5.1.2 GRPO训练面临的问题

虽然 GRPO 能够有效优化模型行为，但它并不是万能的。

问题1：探索空间不足

假设某个问题生成了 8 个回答：

回答1：错误
回答2：错误
回答3：错误
...
回答8：错误

如果这一组回答全部错误，那么即使计算出相对优势，也只是在：

一堆错误答案中寻找相对较好的错误答案。

此时优化方向仍然可能是错误的。

因此强化学习通常要求：

• 基础模型本身具备一定能力；
• 或者先经过 SFT 训练；
• 或者增加采样数量与探索能力；

否则 RL 很容易陷入局部最优。

问题2：奖励函数决定优化方向

奖励函数本质上只是一个打分器。

它不会告诉模型：

为什么错
应该怎么改
下一步应该学什么

它只会给出：

90分
20分
0分

模型只能通过不断试错来寻找获得高奖励的行为。

因此：

奖励函数的质量，往往直接决定了强化学习训练效果的上限。

如果奖励函数设计存在缺陷，模型甚至可能学会“投机取巧（Reward Hacking）”，通过钻奖励规则的漏洞来获得高分，而不是真正完成任务。

5.2 数据准备

在本项目中，GRPO 阶段使用的数据格式如下：

[
    {
        "prompt": prompt_text,
        "ground_truth": final_answer,
        "question": question,
        "full_answer": answer
    },
    ...
]

与 SFT 不同，GRPO 训练并不需要数据集提供 completion。在 SFT 中，模型学习的是数据集中已经存在的标准答案；而在 GRPO 中，回答（Completion）需要由当前模型实时生成，然后根据奖励函数进行评分。

因此从训练角度来看，真正参与 GRPO 训练的核心字段实际上只有：

• prompt ：于生成回答
• ground_truth ：用于计算正确性奖励（Correctness Reward）

---

与前面的 SFT 数据集不同，GRPO 阶段不需要提前将数据转换成 Token ID。

原因在于后续训练过程中，同一个 Prompt 通常需要生成多个不同回答（Multiple Completions），因此 Tokenize 会统一放到生成阶段进行处理。

Dataset 只负责保留训练所需字段：

class RLPromptDataset(Dataset):
    """
    输入数据格式：
    {
        "prompt": ...,
        "ground_truth": ...,
        "question": ...,
        "full_answer": ...
    }
    """

    def __init__(self, data):
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        example = self.data[idx]

        return {
            "prompt": example["prompt"],
            "ground_truth": example["ground_truth"]
        }

此时 Dataset 中的数据形式如下：

[
    {
        "prompt": p1,
        "ground_truth": t1
    },
    {
        "prompt": p2,
        "ground_truth": t2
    },
    {
        "prompt": p3,
        "ground_truth": t3
    }
]

可以看到，此时保存的仍然是原始字符串，而不是已经编码完成的 Token ID。

接下来定义 collate_fn。

由于后续生成阶段需要对同一个 Prompt 采样多个回答，因此这里同样不进行 Tokenize，而是保持字符串格式。

def rl_collate_fn(batch):
    return {
        "prompts": [item["prompt"] for item in batch],
        "ground_truths": [item["ground_truth"] for item in batch],
    }

然后构建 DataLoader：

self.dataloader = DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    collate_fn=rl_collate_fn
)

假设一个 Batch 中包含三条样本：

[
    {"prompt": p1, "ground_truth": t1},
    {"prompt": p2, "ground_truth": t2},
    {"prompt": p3, "ground_truth": t3}
]

经过 collate_fn 后会被整理为：

{
    "prompts": [p1, p2, p3],
    "ground_truths": [t1, t2, t3]
}

5.3 回答生成

完成数据准备后，GRPO 训练的第一步并不是计算 Loss，而是先让模型生成回答。

这是因为 GRPO 属于在线（Online）强化学习方法，训练所需要的 Completion 并不来自数据集，而是由当前模型实时生成。

对于同一个 Prompt，我们通常会生成多个回答：

Question:
1+1=?

Completion 1: 2
Completion 2: 3
Completion 3: 2
Completion 4: 4

随后对这组回答分别计算 Reward：

Completion 1 → Reward = 1.0
Completion 2 → Reward = 0.0
Completion 3 → Reward = 1.0
Completion 4 → Reward = 0.0

再利用组内奖励计算相对优势（Advantage）：

Advantage = Reward - Group Mean Reward

因此，生成多个回答是整个 GRPO 训练流程的起点。

下面是回答生成的核心代码：

def generate_completions(self, prompts):
    # 每个prompt重复G次,一次性batch生成,效率比for循环高
    repeated_prompts = [p for p in prompts for _ in range(self.num_generations)]

    inputs = self.tokenizer(
        repeated_prompts,
        return_tensors="pt",  # 让tokenizer直接返回PyTorch tensor,而不是普通的python list
        padding=True,   # 因为这16条prompt长度不一样,要padding到同一batch内的最大长度才能拼成矩形tensor喂给模型。结合我们之前设的tokenizer.padding_side = "left",padding会补在左边。只要传了padding=True ,tokenizer就会在padding的同时同步产出一个对应的mask,告诉模型"哪些位置是真实token,哪些是补的pad"

        truncation=True, # 如果某条prompt长度超过tokenizer.model_max_length,会从那头截断,防止个别异常长的样本撑爆显存或者报错。
        add_special_tokens=False,  # chat template已经包含特殊token
    ).to(self.device)

    self.model.eval()
    with torch.no_grad():
        generated = self.model.generate(
            **inputs,
            max_new_tokens=self.max_new_tokens, # 限制最多生成多少个新token(不算prompt长度),防止模型生成停不下来,无限占用显存和时间。
            do_sample=True, # 开启随机采样,而不是贪婪解码(do_sample=False会每次都选概率最高的token,结果完全确定)。这一点对GRPO是必须的——同一个prompt要生成
            temperature=1.0, # 控制采样的"随机程度"。数值越大,生成结果越随机发散;越接近0,越趋向于贪婪解码那种确定性选择。
                             # 1.0是不做额外缩放的原始分布,是比较中性的默认值,
                             # 如果你发现生成内容过于离谱可以调低(比如0.7),
                             # 如果发现生成内容缺乏多样性、组内方差太小可以调高。
            top_p=0.95,
            pad_token_id=self.tokenizer.pad_token_id,
            use_cache=True,  # 必须显式传True,覆盖gradient_checkpointing_enable()强制设置的config.use_cache=False,否则生成会退化成每个新token都重算整个序列,速度暴跌(这才是真正的根因)

        )
    self.model.train()

    prompt_len = inputs["input_ids"].shape[1]
    completion_ids = generated[:, prompt_len:]  # 只取新生成的部分
    completions_text = self.tokenizer.batch_decode(completion_ids, skip_special_tokens=True)

    # inputs["input_ids"] = tensor([
    #     [PAD, PAD, PAD, PAD, tok1, tok2, tok3, tok4, tok5],   # prompt1,前面补4个pad
    #     [tok1, tok2, tok3, tok4, tok5, tok6, tok7, tok8, tok9] # prompt2,刚好9个,不需要pad
    # ])

    # inputs["attention_mask"] = tensor([
    #     [0, 0, 0, 0, 1, 1, 1, 1, 1],   # 前4个pad位置标0,后5个真实token标1
    #     [1, 1, 1, 1, 1, 1, 1, 1, 1]    # 全部都是真实token
    # ])

    # generated = tensor([
    #     [PAD, PAD, PAD, PAD, tok1, tok2, tok3, tok4, tok5, new1, new2, new3, new4, new5, new6],
    #     [tok1, tok2, tok3, tok4, tok5, tok6, tok7, tok8, tok9, new1, new2, new3, new4, new5, new6]
    # ])

    # completions_text = [
    #     "1+1等于2。",
    #     "中国的首都是北京,推理过程是..."
    # ]

    return inputs["input_ids"], inputs["attention_mask"], completion_ids, completions_text

为什么使用左侧 Padding

这里有一个容易被忽略但非常重要的细节：

tokenizer.padding_side = "left"

假设两个 Prompt：

Prompt A:
hello world

Prompt B:
hello world today how are you

经过左侧 Padding 后：

[PAD PAD hello world]
[hello world today how are you]

对应的 Tensor：

[
    [PAD, PAD, t1, t2],
    [t1, t2, t3, t4]
]

生成后会变成：

[
    [PAD, PAD, t1, t2, n1, n2, n3],
    [t1, t2, t3, t4, n1, n2, n3]
]

此时所有样本的 Prompt 长度已经被 Padding 到统一长度：

prompt_len = inputs["input_ids"].shape[1]

因此可以直接：

completion_ids = generated[:, prompt_len:]

一次性截取所有样本的新生成部分。

如果使用右侧 Padding：

[t1 t2 PAD PAD]
[t1 t2 t3 t4]

那么生成结果会混入 Padding 区域，处理起来会复杂得多。

因此在生成任务中通常都会采用左侧 Padding。

经过上述步骤，我们最终获得了针对每条 Query 所生成的多个回答，以及后续 GRPO 训练所需的全部信息：

变量	作用
input_ids	Padding 后的 Prompt Token
attention_mask	Prompt 对应的 Attention Mask
completion_ids	模型生成的 Completion Token
completions_text	解码后的文本回答

5.4 Reward与Advantage计算

生成完多个回答后，接下来需要使用奖励函数（Reward Function）对每个回答进行评分。

例如对于同一个问题：

Question:
1 + 1 = ?

Completion 1: 2
Completion 2: 3
Completion 3: 2
Completion 4: 4

奖励函数可能给出：

Completion 1 → 2.0
Completion 2 → 0.0
Completion 3 → 2.0
Completion 4 → 0.0

但 GRPO 并不会直接使用这些 Reward 进行训练。

因为 Reward 本身只能反映回答的绝对质量，而强化学习真正需要的是：

当前回答相对于同组其他回答来说，到底是更好还是更差。

因此在获得 Reward 后，还需要进一步计算 Advantage（优势）。 \[ 优势 = \frac{ 当前奖励 - 平均奖励 }{ 标准差 + 10^{-4} } \] 对应代码如下：

def compute_advantages(self, completions_text, ground_truths):
    # ground_truths需要按num_generations重复对齐
    repeated_truths = [gt for gt in ground_truths for _ in range(self.num_generations)]
    # 举例：repeated_truths = ["2", "2", "2", "2", "北京", "北京", "北京", "北京", "25", "25", "25", "25"]

    # 根据 生成的结果 与 答案 计算奖励。 这部分的内部操作后续需要修改，因为提取答案的方式似乎跟之前sft训练的不一样
    rewards = self.reward_func(completions_text, ground_truth=repeated_truths)
    rewards = torch.tensor(rewards, dtype=torch.float32, device=self.device)

    # 按组(每num_generations条为一组,对应同一个prompt)做归一化
    # 这是GRPO的核心:不需要value model,直接用组内相对好坏作为advantage
    num_prompts = len(ground_truths)
    rewards = rewards.view(num_prompts, self.num_generations)
    mean = rewards.mean(dim=1, keepdim=True)
    std = rewards.std(dim=1, keepdim=True) + 1e-4  # 防止除0
    advantages = (rewards - mean) / std
    advantages = advantages.view(-1)  # 展平回(num_prompts * num_generations,)

    return advantages, rewards.view(-1)

一开始我也说过了，奖励的好坏是我们训练的关键，由于我这仅仅是做个学习，因此采用原课程中相对简陋的自定义奖励函数：

正如前面提到的，奖励函数的质量往往决定了强化学习训练效果的上限。由于这里的目标主要是学习 GRPO 的整体流程，因此采用课程中的一个相对简单的规则奖励函数（Rule-based Reward）作为示例实现。

奖励函数代码：3. reward_2.py

5.5 计算log概率

前面我们已经获得：

prompt_ids
completion_ids

其中：

prompt_ids
    ↓
模型生成
    ↓
completion_ids

但此时还无法直接计算 GRPO Loss。

因为强化学习训练并不仅仅需要知道模型生成了什么内容，还需要知道：

模型当时到底有多大概率生成这些 Token。

也就是我们常说的：

log_prob

---

为什么要重新 Forward

很多刚接触 RLHF 的读者都会有一个疑问：

既然已经调用 generate() 生成出了回答，

为什么还要再跑一次模型？

原因在于：

generated = model.generate(...)

得到的主要是：

completion_ids

即最终采样出的 Token。

而后续 GRPO Loss 需要的是：

每个生成Token对应的log_prob

因此必须重新将：

prompt_ids
+
completion_ids

拼接回完整序列，再执行一次 Forward。

---

拼接完整序列

首先构造：

full_ids = torch.cat(
    [prompt_ids, completion_ids],
    dim=1
)

例如：

Prompt:
t1 t2 t3

Completion:
t4 t5

拼接后：

full_ids:
t1 t2 t3 t4 t5

然后送入模型：

outputs = self.model(
    input_ids=full_ids,
    attention_mask=full_mask,
    use_cache=False
)

得到：

outputs.logits

形状为：

(batch_size, seq_len, vocab_size)

其中：

vocab_size

表示词表大小。

假设词表大小为 150000，那么每个位置都会产生一个维度为 150000 的向量：

[
    token_0 的logit,
    token_1 的logit,
    token_2 的logit,
    ...
]

表示：

模型认为下一个 Token 是词表中每个 Token 的可能性大小。

---

Logits 与 Token 的对应关系

这里最容易理解错，因此单独举例。

假设：

full_ids: [4, 5, 7]

经过模型计算：

位置0 → logits_0
位置1 → logits_1
位置2 → logits_2

注意：

logits_0

并不是预测当前位置的 Token。

而是在预测：

位置1应该是什么Token

即：

4
 ↓
logits_0
 ↓
预测5

同理：

5
 ↓
logits_1
 ↓
预测7

因此真实对应关系为：

Token:
4   5   7

Logits:
    L1  L2

或者写成：

4 ──► L1 ──► 5
5 ──► L2 ──► 7

其中：

L1[5]

表示：

在位置1的真实答案是 Token 5 时，模型给出的 Logit。

而：

L2[7]

表示：

在位置2的真实答案是 Token 7 时，模型给出的 Logit。

---

为什么要错位

由于语言模型始终在预测：

下一个Token

因此：

logits = outputs.logits[:, :-1, :]
targets = full_ids[:, 1:]

例如：

Token:
t1  t2  t3  t4  t5

Logits:
L2  L3  L4  L5  L6

对应关系为：

t1 → L2 → t2
t2 → L3 → t3
t3 → L4 → t4
t4 → L5 → t5

可以发现：

第一个Token没有对应Logit
最后一个Logit没有对应Token

因此：

logits[:, :-1]
full_ids[:, 1:]

都需要截断。

---

获取真实 Token 对应的 Log Probability

接下来：

log_probs = F.log_softmax(
    logits,
    dim=-1
)

将：

logits

转换为：

log_prob

这里的每个位置都保存了：

词表中所有Token的log概率

例如：

[
    token0的log_prob,
    token1的log_prob,
    token2的log_prob,
    ...
]

但我们只关心：

实际生成出来的那个 Token。

因此：

token_log_probs = (
    log_probs.gather(
        2,
        targets.unsqueeze(-1)
    )
    .squeeze(-1)
)

作用就是：

从整个词表中,只取真实Token对应的那个log_prob

例如：

真实Token = 7

那么：

log_probs[..., 7]

就会被提取出来。

最终得到：

token_log_probs

形状：

(batch_size, seq_len - 1)

其中每一个元素都表示：

模型生成该 Token 时对应的 Log Probability。

---

为什么只保留 Completion

前面计算出的：

token_log_probs

同时包含：

Prompt部分
+
Completion部分

但 GRPO 只关心模型生成的内容。

Prompt 是训练数据提供的：

用户输入

不是模型生成的。

因此：

completion_log_probs =
    token_log_probs[:, prompt_len - 1:]

只保留：

Completion对应的log_prob

例如：

Prompt:
t1 t2 t3

Completion:
t4 t5

最终保留：

t4对应的log_prob
t5对应的log_prob

而t1,t2,t3对应的部分全部丢弃。

---

这里顺便解释一下：

为什么强化学习一定要计算 log_prob？

因为后续 GRPO 的优化目标本质上是在调整：

模型生成当前回答的概率

如果某个回答：

Advantage > 0

说明它比组内平均水平更好。

那么训练会提高：

该回答对应Token的log_prob

反之：

Advantage < 0

则会降低对应 Token 的生成概率。

因此 Log Probability 实际上就是：

强化学习阶段连接「奖励信号」与「模型参数更新」的桥梁。

---

经过上述步骤，我们最终获得：

completion_log_probs
completion_mask_for_loss

其中：

• completion_log_probs：Completion 部分每个 Token 对应的 Log Probability；
• completion_mask_for_loss：对应位置的有效 Mask。

这些数据将在下一节用于计算最终的 GRPO Loss。

以下就是我们的实现代码：

# ---------- 第三步:重新算log prob(generate不保留梯度,必须重新forward) ----------
def compute_log_probs(self, prompt_ids, prompt_mask, completion_ids):
    full_ids = torch.cat([prompt_ids, completion_ids], dim=1)
    completion_mask = (completion_ids != self.tokenizer.pad_token_id).long()
    full_mask = torch.cat([prompt_mask, completion_mask], dim=1)

    outputs = self.model(input_ids=full_ids, attention_mask=full_mask, use_cache=False)
    # outputs.logits的形状是(batch_size, seq_len, vocab_size),最后一维就是词表大小

    # 是把最后一个位置(位置5)的logits砍掉,只保留位置0~4这5个"有真实答案可核对"的预测:
    logits = outputs.logits[:, :-1, :]  # 预测下一个token,所以logits要错位
    # 举例：位置0的token，经过LLM后计算得到logit(0) 这个logint会转换成token，然后放到位置1上面。
    # 那么，full_ids中位置1所在的token就是与位置0的logit匹配的一对。

    # 把full_ids的第一个token(位置0对应的101)砍掉,保留位置1~5这5个token作为"真实答案":
    targets = full_ids[:, 1:]
    # (batch_size, prompt_len + completion_len) id只是个数字，因此只有2维，无需像token一样用第三维表示

    log_probs = F.log_softmax(logits, dim=-1)
    
    # token id 就是token在vocab中的位置
    # 从log_probs这个(batch, seq_len, vocab_size)的大张量里,每个位置只挑出"真实token对应的那一个数值",把维度从词表大小(可能十几万)压缩到1。
    token_log_probs = log_probs.gather(2, targets.unsqueeze(-1)   ).squeeze(-1)
    # token_log_probs的大小是（batch, seq_len-1），里面每个数都是对应答案的log_prob （错位对齐之后的长度-1，因此seq_len-1）

    # 只保留completion部分的log prob(prompt部分不需要算梯度来源)
    prompt_len = prompt_ids.shape[1]

    completion_log_probs = token_log_probs[:, prompt_len - 1 :]
    # 第一个token是没有任何logit与它对应的，组不成token-logit对，因此才会产生之前偏差
    # 因此我们在计算len(prompt)后，需要减去1，这才是 token-logit对 中propmt的长度。

    completion_mask_for_loss = full_mask[:, prompt_len:]
    # 下面是一组token，和一组logits，l2可以通过softmax等方式变成t2，因此l2和t2是一对的
    # 可以发现 仅中间4对构成匹配的数据。
    # 我们以t1，t2,t3为prompt举例，我们要获得completion的mask。则full_mask中将prompt_len长度的前面部分截断即可
    # completion_log_probs中，token_log_probs内容是每个匹配的数据，所以长度是下面重合部分，因此要去除prompt部分的token_log_probs，将prompt_len - 1长度的前面部分截断即可
    # [
    #   t1,t2,t3,  t4,t5
    #      l2,l3,  l4,l5,l6
    # ]

    # 经过以上操作
    # completion_log_probs中仅存在 l4,l5对应的log_prob值
    # completion_mask_for_loss 仅存在 t4,t5 两个位置的mask。（其实和前面completion_mask是等价的）.
    # 由于completion_mask 是从completion ids里面生成的，因此也仅有t4、t5两个位置的mask
    return completion_log_probs, completion_mask_for_loss

5.6 计算 Loss 与参数更新

这部分没什么好说的。对"优势"有什么概念的疑惑，可以在我博客里面搜索“advantage”，找到强化学习 4.策略梯度算法 （应该最开始是这里讲解了，我也有点忘了，也可能是另一篇）

# policy gradient loss: -advantage * log_prob,只在completion有效token上算,按token数做平均
per_token_loss = -advantages.unsqueeze(1) * log_probs
per_token_loss = per_token_loss * completion_mask  # 把pad位置的loss乘以0,直接清零。

然后就是累计指定次数，进行更新

if (batch_idx + 1) % self.gradient_accumulation_steps == 0:
    torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
    self.optimizer.step()
    self.scheduler.step()
    self.optimizer.zero_grad()
    global_step += 1

    mean_reward = raw_rewards.mean().item()
    progress_bar.set_postfix(
        step=global_step,
        loss=f"{loss.item() * self.gradient_accumulation_steps:.4f}",
        reward=f"{mean_reward:.3f}",
    )

    if global_step % 10 == 0:
        print(
            f"epoch {epoch} step {global_step} "
            f"loss {loss.item() * self.gradient_accumulation_steps:.4f} "
            f"mean_reward {mean_reward:.3f}"
        )

    if global_step % self.save_steps == 0:
        self.save_checkpoint(global_step)

5.7 GRPO main函数训练

和之前SFT一样，准备好我们需要的model 、tokenizer、train data 3部分，然后model上面搭上LORA。

from peft import LoraConfig, get_peft_model, TaskType
def rl_train_main():
    model, tokenizer = load_model_from_checkpoint('final',merge=True)
    GSM8K_RL_Dataset = import_things('1. data_preparation.py','GSM8K_RL_Dataset')

    train_data = GSM8K_RL_Dataset(split="train", max_samples=20)

    # 加载模型之后,传进trainer之前包一层
    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        lora_alpha=32,
        lora_dropout=0.1,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ]
    #     GRPO通常参数更新信号很弱。

    )
    model = get_peft_model(model, lora_config)


    trainer = SimpleGRPOTrainer(
        model=model,
        tokenizer=tokenizer,
        train_data=train_data,
        reward_func=reward_func,
        batch_size=1,
        num_generations=4,
        max_new_tokens=612,
        gradient_checkpointing=True,
        save_steps=100,
        gradient_accumulation_steps=1,
        show_time=True
    )
    trainer.train()

5.8 测试与效果

这一部分实际上并没有完成完整测试，主要原因是训练耗时过长，个人用于实验的小笔记本难以支撑长时间的 GRPO 训练。

在前面的实现中，我们设置：

max_new_tokens = 612

这个值看起来比较大，但实际上是必要的。

因为当前训练数据要求模型输出完整的推理过程以及最终答案。如果生成长度限制过小，模型往往还处于推理阶段就已经到达长度上限，被强制截断。

此时模型无法输出最终答案：

Final Answer: ...

奖励函数自然也无法正确判断回答是否正确。

这样就会出现前面提到的问题：

同一组回答全部都是无效回答，Reward 全部接近于 0，模型无法获得有效的优化方向。

因此在实验中只能将 max_new_tokens 设置得相对较大，以保证大部分样本能够完整生成。

但这也带来了另一个问题：

GRPO 的主要时间开销几乎全部集中在回答生成阶段（Generate）。

训练过程中统计得到的耗时如下：

可以看到，相比于 Reward 计算、Loss 计算以及反向传播等步骤，回答生成阶段的耗时远高于其他部分。

其中：

Generate ≈ 76s

单次生成就需要一分钟以上。

而整个实验仅使用约 20 条样本进行训练，总耗时仍达到：

约 36 分钟

这也反映出了 RLHF / GRPO 与传统 SFT 的一个明显区别：

SFT 的主要开销来自反向传播，而 GRPO 的主要开销往往来自回答生成。

因为对于每个 Prompt，模型不仅需要生成回答，而且还需要生成多个回答（Group），随后才能计算 Reward 和 Advantage。

目前工业界解决这一问题最常见的方法，是将训练模型与生成模型分离。

例如：

训练模型
    ↓
同步权重
    ↓
高性能推理引擎
    ↓
批量生成回答

其中较常见的方案就是使用 vLLM 负责生成，而训练框架只负责参数更新。

相比 HuggingFace 原生的 generate()，vLLM 在大批量生成场景下通常能够获得明显更高的吞吐量，因此也是目前许多 RLHF 框架的默认选择。

不过在本系列文章编写时，我的实验环境主要为 Windows，本篇内容更关注 GRPO 流程本身，因此暂未进一步集成 vLLM。

---

由于训练数据规模较小、训练轮数有限，加上缺少系统性的 Benchmark 评测，因此本文并不给出模型效果方面的结论。

本章的重点也并不在于训练出一个高性能模型，而是：

从零实现并跑通一次完整的 GRPO 训练流程，理解 Reward、Advantage、Log Probability 以及策略梯度在 LLM 强化学习中的作用。

至此，一个最小可运行版本的 GRPO 训练框架便已经实现完成。