NanoRL极简实现:用1800行代码读懂LLM强化学习训练核心原理 我本来以为强化学习训练 LLM 是一件“只能看论文 用成熟框架”的事情。直到我看到 NanoRL 这个项目用大约 1,800 行 Python 代码实现了 RL 训练 LLM 的完整链路。这个体量放在当前 LLM 框架动辄几万行代码的环境里确实非常吸引人。它不一定能直接替代生产级框架但非常适合用来理解 RL 训练的核心原理也适合做二次开发和教学实验。如果你和我一样之前对 RLHF、PPO、GRPO 这些术语停留在“听过但没跑通”的阶段那这篇文章就是写给你的。我会从 RL 训练 LLM 的基本概念讲起剖析 NanoRL 这类极简实现的关键模块再结合一个可操作的实战思路把完整链路拆开来看。我们不去追逐复杂的大规模分布式训练而是先用手能伸进去改代码的最小实现把原理跑明白。1. 背景与核心概念1.1 为什么需要“极简 RL 训练框架”大语言模型LLM的训练大致分为三个阶段预训练Pre-training、监督微调SFTSupervised Fine-Tuning和偏好对齐。前两个阶段解决的是“模型能生成内容、能听懂指令”的问题而第三个阶段解决的是“模型生成的内容符合人类偏好、符合事实和安全要求”的问题。第三个阶段的常用手段之一就是强化学习RLReinforcement Learning。我们常听到的 RLHFReinforcement Learning from Human Feedback就是其中的代表流程先让模型生成多个候选回答再由人类或奖励模型打分最后用强化学习算法更新模型参数使高分回答的出现概率越来越高。但 RL 训练 LLM 涉及多个组件策略模型Actor、参考模型Reference Model、奖励模型Reward Model、价值模型CriticPPO 需要还要处理生成、采样、KL 散度控制、优势估计等细节。大型框架把这些逻辑封装得很深新手打开源码往往不知道从哪里看起。NanoRL 的价值就在于它把这条链路精简到 1,800 行左右每个模块都短小直接适合按行阅读。你可以看清一次完整的 RL 迭代里数据是怎么流动的、梯度是怎么计算的、KL 惩罚是怎么加进去的。这类项目很像 Karpathy 的 nanoGPT 带来的启示先把规模缩小把原理讲透再走向规模化方案。1.2 NanoRL 是什么NanoRL 是一个用于 LLM 强化学习训练的极简实现项目核心代码量约 1,800 行。它并不追求功能大而全而是把 RL 训练的关键路径清晰呈现出来。你可以把它看成一份“可运行的 RL 训练源码级教程”。它的核心定位可以理解为教学友好代码结构清晰注释和模块划分适合学习。易于修改没有大规模分布式抽象改策略、改奖励函数都相对直接。快速验证适合在小规模模型、小数据集上验证 RL 算法实验。同时也要说明NanoRL 不是一个面向生产环境的框架。它更适合研究算法、验证想法、教学演示或者作为二次开发的骨架。如果要做大规模 RLHF仍然建议基于 TRL、OpenRLHF、DeepSpeed-Chat 等更成熟的框架。1.3 RL 训练 LLM 的整体链路先看一张简化的流程图它描述了 RL 训练 LLM 时一个 Step 的完整过程提示词输入 ↓ Actor 模型生成回答多个采样 ↓ Reward Model 或规则奖励计算得分 ↓ 结合 KL 散度控制模型不要偏离参考模型太远 ↓ PPO/GRPO 计算优势函数与损失 ↓ 更新 Actor 模型参数这里要理解几个关键点Actor策略模型是需要训练的模型一边生成回答一边更新参数。Reward Model奖励模型给生成结果打分可以理解成一个“评分老师”。KL 散度惩罚防止模型为了刷分而生成完全偏离原始分布的内容保持生成质量稳定。PPOProximal Policy Optimization是 RL 训练中最常用的算法之一它通过裁剪clip限制每次参数更新的幅度提升训练稳定性。2. 环境准备与版本说明在动手跑 NanoRL 或类似极简 RL 训练项目之前先把环境准备步骤理顺。很多 RL 训练跑不起来不是因为代码问题而是 PyTorch、CUDA、transformers 版本不一致导致的。2.1 运行环境建议由于 RL 训练需要反复进行前向和反向计算对显存和 CUDA 环境要求偏高。建议按如下环境准备环境项建议配置操作系统Ubuntu 20.04 / 22.04Windows 也可以但推荐 WSL2Python3.10 或 3.11CUDA11.8 或 12.1按显卡驱动选择PyTorch2.x 版本需与 CUDA 版本匹配GPU 显存起步建议 16GB越大越方便以下命令中的具体 PyTorch 版本需要依据你的 CUDA 环境调整安装方式建议参考 PyTorch 官网的生成命令。# 创建虚拟环境 python -m venv nano-rl-env source nano-rl-env/bin/activate # 安装 PyTorch示例为 CUDA 12.1请按实际环境选择 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 # 安装常用依赖 pip install transformers datasets accelerate2.2 项目结构当拿到一个 NanoRL 风格的项目时它的目录结构通常可以归纳为以下几块nanorl/ ├── configs/ # 训练配置文件 ├── data/ # 训练数据 ├── nanorl/ │ ├── models/ # Actor、Reward 等模型定义 │ ├── rl/ # RL 算法核心如 PPO、GRPO │ ├── utils/ # KL 计算、日志、采样等工具 │ └── trainer.py # 训练主流程 ├── scripts/ # 启动脚本 └── requirements.txt实际的代码组织会根据作者习惯有所不同但核心模块大致是这些。在阅读代码时建议优先看trainer.py和rl/目录下的文件这两个地方定义了算法的主循环和策略更新方式。3. 核心原理拆解从 RLHF 到 PPO/GRPO在阅读 NanoRL 源码之前先要把几个核心算法概念弄明白。很多教程一上来就贴 PPO 公式容易劝退新手。下面我用“要解决什么问题 怎么解决”的思路拆开讲。3.1 为什么标准 SFT 不够SFT 阶段虽然用人工标注数据训练模型但它本质上是“模仿学习”。模型学习的是“给定指令输出标准答案”的条件概率分布。这种训练方式有两个问题模型的回答缺乏多样性遇到训练集之外的输入时容易表现不稳定。SFT 很难直接优化“人类偏好”这种非结构化目标因为标注数据无法覆盖所有生成结果。RL 阶段解决的是第二个问题我们不直接告诉模型“哪句话是对的”而是给模型一个“评分机制”让模型在探索中持续调整让得分更高的回答更可能出现。3.2 PPO 的核心思想PPO 是 OpenAI 在 2017 年提出的策略梯度算法后来成为 RLHF 的主流算法。对 LLM 来说PPO 中最重要的几个概念是策略策略指的是 Actor 模型在给定提示词下生成各个 token 的概率分布。强化学习的“策略网络”就是我们的 LLM。优势函数Advantage优势函数衡量的是“当前动作相对于平均水平好多少”。在 LLM 场景下可以理解为这次生成的结果比平均结果好多少。如果优势为正就增加这个回答的生成概率如果为负就降低概率。截断目标函数Clipped ObjectivePPO 通过限制新旧策略的比值避免单次更新过大。这一点非常重要因为如果一次更新权重太大模型可能直接产生灾难性退化生成质量迅速崩溃。KL 惩罚KL 散度衡量两个分布之间的差异。在 RLHF 中Actor 模型不能离 Reference 模型太远否则模型会为了奖励“钻空子”生成语义不稳定甚至退化的文本。通常会在奖励上减去一个 KL 惩罚项。用一句话概括 PPO 训练 LLM根据模型生成的回答得分决定下一步是鼓励还是抑制这类回答并且在更新时保持一个“安全距离”。3.3 GRPO 带来的简化思路GRPOGroup Relative Policy Optimization是 DeepSeek 等工作中采用的一种策略优化方法它的核心思想是做组内相对比较。具体来说对同一个提示词采样多个回答比如 8 个或 16 个。用奖励模型的得分计算这组回答的均值和标准差。每个回答的优势函数用它与组内均值的差值来刻画。不需要单独的 Critic价值模型从而省掉一个大规模模型及其训练开销。GRPO 与 PPO 相比最大的优势是简洁。不需要训练价值模型计算量更小在开源社区的应用越来越广泛。NanoRL 这类极简框架很可能倾向于实现 GRPO因为它更简单、更容易用 1,800 行代码表达核心思想。3.4 奖励模型与策略模型的交互在完整 RLHF 流程中奖励模型通常先用人类偏好数据训练出来。而在 GRPO 简化流程中有时可以用规则奖励来代替例如数学题奖励与标准答案的一致性。代码题奖励代码能否通过测试用例。安全对齐奖励模型给有害内容打低分。NanoRL 类项目为了方便实验通常会支持多种奖励来源。理解这一点很重要RL 训练时模型优化的目标完全由奖励函数决定所以奖励函数必须谨慎设计。4. 完整实战案例用 NanoRL 训练一个小型 LLM下面进入实战环节。由于 NanoRL 属于教学型极简实现不同仓库的接口可能不同下面给出一个通用可操作的落地思路和核心代码模式。你需要根据自己的实际代码做调整但整体流程是通用的。4.1 创建项目结构我们模仿 NanoRL 的设计思路搭建一个最小项目my-nanorl/ ├── configs/ │ └── train.yaml ├── data/ │ └── prompts.jsonl ├── my_nanorl/ │ ├── __init__.py │ ├── model.py │ ├── rl_trainer.py │ └── utils.py ├── train.py └── requirements.txt这里的原则是每个文件只做一件事方便阅读和修改。4.2 准备提示词数据集RL 训练的第一步是准备一组提示词prompts让模型针对这些提示词生成回答。数据格式可以采用 JSON Lines每行一个 JSON 对象。{prompt: 一个三角形的三条边分别是 3、4、5这个三角形的面积是多少} {prompt: 简述强化学习中的探索与利用是什么意思。} {prompt: 写一段 Python 代码判断一个字符串是否为回文。}数据不需要像 SFT 那样包含标准答案因为 RL 阶段会在训练过程中通过奖励函数计算得分。建议提示词数据集覆盖模型可能面对的任务类型数量不用太多先跑通流程再说。4.3 编写策略模型与参考模型核心模型部分我们可以使用 transformers 库加载一个小型模型比如 Qwen2-0.5B、Phi-3-mini 或你手头已有的模型。关键代码如下# 文件路径my_nanorl/model.py from transformers import AutoModelForCausalLM, AutoTokenizer def load_actor_and_ref(model_name: str, device: str cuda): 加载 Actor 模型和 Reference 模型。 Actor 模型需要计算梯度并更新参数。 Reference 模型冻结参数只用于计算 KL 散度。 actor AutoModelForCausalLM.from_pretrained( model_name, torch_dtypeauto, ).to(device) ref_model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypeauto, ).to(device) # 冻结参考模型参数 for param in ref_model.parameters(): param.requires_grad False tokenizer AutoTokenizer.from_pretrained(model_name) if tokenizer.pad_token is None: tokenizer.pad_token tokenizer.eos_token return actor, ref_model, tokenizer需要注意同一个模型需要加载两份一份参与训练Actor一份冻结Reference。这样做是为了计算 KL 散度约束 Actor 模型的更新幅度。4.4 定义规则奖励函数在 RL 训练中奖励函数决定模型学习方向。下面演示一个简单的规则奖励函数检查模型生成结果中是否包含正确答案。# 文件路径my_nanorl/utils.py import re def simple_rule_reward(prompt: str, response: str) - float: 一个非常简单的规则奖励示例 如果回答中包含数字 6认为回答可能正确给 1 分 否则给 0 分。 注意实际项目必须按任务设计可靠的奖励函数 if 6 in response: # 这里只做示例真实场景需要更严谨的规则。 return 1.0 return 0.0 def extract_assistant_text(seq: str, eos_token: str) - str: 截取生成部分去掉 prompt 前缀。 # 思路找到与 prompt 等长的位置之后的内容视为模型输出 # 实际实现需要根据 tokenizer 和模板调整。 return seq.split(eos_token)[0]这里的示例奖励函数非常简单真实场景下不建议直接使用。更好的方式是对于数学题用规则解析最终答案并对比对于代码题运行测试用例对于通用场景使用训练好的奖励模型。4.5 RL 训练主循环GRPO 思路下面给出 GRPO 风格训练主循环的伪代码结构它展示了 RL 训练一次迭代的核心步骤。# 文件路径my_nanorl/rl_trainer.py import torch from torch.optim import AdamW def grpo_train_step(actor, ref_model, tokenizer, prompts, reward_fn, optimizer, max_new_tokens128, group_size8, kl_coef0.1): 一个简化的 GRPO 训练步 1. 对每个 prompt 采样 group_size 个回答。 2. 计算每个回答的奖励。 3. 计算组内优势。 4. 用策略梯度损失更新 actor。 actor.train() ref_model.eval() all_prompt_ids [] all_response_ids [] all_rewards [] # Step 1: 采样 for prompt in prompts: inputs tokenizer(prompt, return_tensorspt).to(actor.device) with torch.no_grad(): outputs actor.generate( **inputs, max_new_tokensmax_new_tokens, num_return_sequencesgroup_size, do_sampleTrue, top_p0.9, temperature0.8, pad_token_idtokenizer.pad_token_id, ) # 收集生成结果和奖励 for seq in outputs: response_text tokenizer.decode(seq, skip_special_tokensTrue) reward reward_fn(prompt, response_text) all_rewards.append(reward) all_prompt_ids.append(inputs[input_ids].repeat(group_size, 1)) all_response_ids.append(outputs) # 组合成一个批次 prompt_ids torch.cat(all_prompt_ids, dim0) response_ids torch.cat(all_response_ids, dim0) rewards torch.tensor(all_rewards, dtypetorch.float32, deviceactor.device) # Step 2: 组内优势 # 对每个 prompt 独立计算组内均值和标准差 rewards rewards.view(len(prompts), group_size) mean_rewards rewards.mean(dim-1, keepdimTrue) std_rewards rewards.std(dim-1, keepdimTrue) 1e-6 advantages (rewards - mean_rewards) / std_rewards # Step 3: 策略梯度损失 # 将 prompt 和 response 拼接计算整个序列的 logits input_ids torch.cat([prompt_ids, response_ids[:, prompt_ids.shape[1]:]], dim-1) logits actor(input_idsinput_ids).logits # 计算新策略下 response token 的 log 概率 response_logps gather_log_probs(logits, response_ids) # 计算旧策略下的 log 概率简化起见用 ref_model 作为旧策略示例 with torch.no_grad(): ref_logits ref_model(input_idsinput_ids).logits ref_logps gather_log_probs(ref_logits, response_ids) # KL 惩罚项分布偏离参考模型越多惩罚越大 kl_penalty kl_coef * torch.exp(ref_logps - response_logps) # GRPO 损失让优势为正的响应概率上升 loss -torch.mean(advantages.flatten() * response_logps - kl_penalty) # Step 4: 反向传播 optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(actor.parameters(), max_norm1.0) optimizer.step() return loss.item() def gather_log_probs(logits, response_ids): 根据生成的 token id 从 logits 中取出对应的 log 概率。 真实实现需要处理 prompt 部分的 logits只保留 response 部分。 # 示例逻辑真实代码需要结合 padding 和 mask 处理。 log_probs torch.log_softmax(logits, dim-1) return log_probs.gather(-1, response_ids.unsqueeze(-1)).squeeze(-1)注意上面的代码是教学示例省略了 padding mask、loss 掩码、旧策略保存等细节。在实际 NanoRL 源码中这些细节都会以更严谨的方式实现。但它能帮助我们理解 GRPO 的核心逻辑采样 → 打分 → 组内优势 → 更新策略。4.6 编写训练启动脚本最后写一个训练入口脚本把整个流程串起来。# 文件路径train.py import argparse import json import torch from torch.optim import AdamW from my_nanorl.model import load_actor_and_ref from my_nanorl.utils import simple_rule_reward from my_nanorl.rl_trainer import grpo_train_step def main(): parser argparse.ArgumentParser() parser.add_argument(--model_name, typestr, defaultQwen/Qwen2-0.5B) parser.add_argument(--data_path, typestr, defaultdata/prompts.jsonl) parser.add_argument(--epochs, typeint, default10) parser.add_argument(--lr, typefloat, default1e-5) parser.add_argument(--device, typestr, defaultcuda) args parser.parse_args() # 加载模型 actor, ref_model, tokenizer load_actor_and_ref(args.model_name, args.device) # 读取提示词 prompts [] with open(args.data_path, r, encodingutf-8) as f: for line in f: prompts.append(json.loads(line.strip())[prompt]) # 优化器只更新 actor 参数 optimizer AdamW(actor.parameters(), lrargs.lr) # 训练循环 for epoch in range(args.epochs): loss grpo_train_step( actoractor, ref_modelref_model, tokenizertokenizer, promptsprompts, reward_fnsimple_rule_reward, optimizeroptimizer, ) print(fEpoch {epoch 1}/{args.epochs}, loss: {loss:.4f}) # 保存模型 actor.save_pretrained(output/actor_model) tokenizer.save_pretrained(output/actor_model) if __name__ __main__: main()4.7 运行与验证启动训练的命令python train.py --model_name Qwen/Qwen2-0.5B --epochs 10 --lr 1e-5预期输出大致如下数值不固定Epoch 1/10, loss: -0.5321 Epoch 2/10, loss: -0.7856 Epoch 3/10, loss: -0.9102 ...训练结束后可以用一个简单的脚本验证模型生成质量是否发生变化from transformers import AutoModelForCausalLM, AutoTokenizer model AutoModelForCausalLM.from_pretrained(output/actor_model) tokenizer AutoTokenizer.from_pretrained(output/actor_model) prompt 一个三角形的三条边分别是 3、4、5这个三角形的面积是多少 inputs tokenizer(prompt, return_tensorspt) outputs model.generate(**inputs, max_new_tokens64, do_sampleFalse) print(tokenizer.decode(outputs[0], skip_special_tokensTrue))对比训练前后的生成结果就能直观看到 RL 训练对模型输出风格和准确率的影响。5. 常见问题与排查思路RL 训练 LLM 比普通微调更考验细节。以下是我认为最容易踩坑的几个问题整理成表格供排查参考。问题现象常见原因解决思路训练初期 loss 持续为负数且异常大日志计算方式不正确或优势函数没有归一化检查 loss 符号和 KL 项系数建议先固定随机种子复现生成结果全是 EOS token 或空白生成长度设置过短、EOS 采样概率过高调整 max_new_tokens、temperature、top_preward 一直为 0奖励函数设计不合理或生成格式与奖励解析规则不匹配先单独测试奖励函数输入固定文本验证得分模型很快退化输出重复或无意义内容KL 系数设置太小模型为了刷奖励钻空子调大 kl_coef或对参考模型增加约束显存不足OOMbatch size 太大、group_size 太大、序列太长减小 group_size 或 max_new_tokens使用梯度累积PPO/GRPO 中 loss 为 NaN精度溢出、学习率过大、logits 中出现 inf检查是否用了 fp16/bf16是否关闭了 grad scaling降低学习率在排查问题时我建议遵循一个原则先分别验证每个环节再集成测试。比如先用固定输入测试奖励函数再单独测试采样生成最后再跑训练循环。这样能把问题定位在单一模块避免“全都出问题但不知道从哪查起”。6. 最佳实践与工程建议通过阅读和修改 NanoRL 这类项目我沉淀了一些适用于实际项目的工程建议分享给你。6.1 奖励函数设计要可验证很多人第一次跑 RL 训练时喜欢写一个“看起来很合理”的奖励函数然后发现训练不收敛。问题往往不在算法而在奖励函数本身。设计奖励函数时要注意奖励函数必须是可计算的不能有模糊地带。生成的文本要先解析成结构化结果再进行打分。要准备一批验证集定期检查奖励是否符合预期。奖励值范围不要过大建议归一化到 [-1, 1] 或 [0, 1]。6.2 混合精度训练要谨慎训练 LLM 时常用 fp16 或 bf16 来节省显存。但在 RL 训练中策略损失对数值稳定性更敏感。建议初期先用 FP32 跑通流程再切换到混合精度。如果使用 fp16注意梯度溢出问题必要时保留GradScaler。bf16 通常更稳定但需要较新的 GPU 支持。以常见精度选择为例做一个小表格精度优点缺点适用场景FP32数值稳定显存占用高、训练慢调试、小模型FP16显存占用小、速度快容易出现精度溢出显存紧张时使用BF16数值范围大、较稳定低精度下精度略低现代 GPU 上优先推荐6.3 日志与监控要完整RL 训练不只看 loss还要看很多指标。建议至少记录以下几类平均奖励、奖励方差。KL 散度均值。生成回答的文本样例定期打印几条。Token 级别的概率分布变化。这些日志能帮助你在模型崩溃之前发现问题而不是等到生成结果全部变乱码才回头查。6.4 超参数调优建议RL 训练的超参数比 SFT 更敏感。主要的调节方向如下学习率一般比 SFT 小 1 到 2 个数量级1e-6 到 5e-6 是常见区间。KL 系数KL 系数控制着对模型漂移的惩罚力度。过大则模型学不到新东西过小则模型可能崩坏。group_sizeGRPO 中每个 prompt 的采样数量。越大优势估计越稳定但计算成本也越高。temperature采样温度。RL 训练阶段通常比推理时高一些保持探索性。建议在修改超参数时每次只改一个变量并记录实验效果形成自己的实验笔记。6.5 安全与合规提醒RL 训练中奖励函数事实上在“塑造”模型价值观和行为边界。如果你是做实际业务请务必注意不要在未授权环境中收集用户数据。对模型生成内容要有审核和过滤机制。对于医疗、金融等高风险领域不要用简单规则奖励代替专业审核。训练完成后应在测试集上做安全评估不要只看奖励分数。7. 总结与学习路线用极简代码理解 RL 训练 LLM是非常高效的学习路径。NanoRL 这类项目用 1,800 行代码把所有核心模块暴露在你的面前你能直接看到策略模型如何生成、奖励如何计算、优势如何估计、参数如何更新。读完它你再看 TRL、OpenRLHF 这些工程框架时就不容易迷失在抽象层里。如果你准备继续深入我建议按这个顺序走先跑通 NanoRL 或你自己搭的最小实现不要追求大模型先让代码循环跑起来。修改奖励函数换一个自己熟悉的任务观察模型行为变化。对比 PPO 和 GRPO感受一个价值模型对训练稳定性的影响。再读 TRL 源码理解工程框架中如何处理 batch、padding、多卡并行等问题。最后扩展到更大模型和分布式环境。RL 训练 LLM 的坑确实比普通微调多但只要你手里有一个能看懂、能改、能跑通的最小实现剩下的事情就有迹可循。希望这篇文章能把你的第一步垫稳剩下的路就差你亲手跑一条训练曲线了。