CODE 04: Ring Attention 实践#
基于 Qwen3-0.6B 的 GRPO 算法实践：原理与实现#
在大语言模型的优化领域，强化学习算法扮演着越来越重要的角色。Group Relative Policy Optimization (GRPO) 作为一种改进型强化学习算法，通过独特的群体比较机制，为大模型的微调提供了更稳定、更高效的优化路径。本文将深入解析 GRPO 的工作原理，并基于 Qwen3-0.6B 模型实现完整的实验流程，帮助读者全面理解这一算法的核心思想与实践方法。
环境准备与模型加载#
在开始实现 GRPO 算法之前，我们需要准备好基础的开发环境和模型资源。选择合适的基础模型对实验效果和资源需求都有重要影响，Qwen3-0.6B 作为一款轻量级大模型，在性能和资源消耗之间取得了良好平衡。
# 导入必要的库
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModelForCausalLM
# 设置计算设备，优先使用 GPU 以提高训练速度
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")
这段代码首先导入了实验所需的核心库，包括 PyTorch 深度学习框架、用于概率计算的工具、数据处理库以及 Hugging Face 的 transformers 库。最后一行代码根据硬件情况自动选择计算设备，优先使用 GPU 以加速训练过程。
接下来，我们加载 Qwen3-0.6B 模型和对应的分词器：
# 加载 Qwen3-0.6B 模型和分词器
model_name = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# 设置填充 token，Qwen 模型默认没有 pad_token，使用 eos_token 代替
tokenizer.pad_token = tokenizer.eos_token
tokenizer.pad_token_id = tokenizer.eos_token_id
# 加载基础模型，使用自动精度以适应不同设备
base_model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to(device)
print(f"{model_name}模型加载完成")
Qwen 系列模型需要加载远程代码支持，因此设置trust_remote_code=True是必要的。值得注意的是，Qwen 模型默认没有定义填充 token(pad_token)，我们使用结束 token(eos_token)作为替代，这在批量处理文本时非常重要。对于模型精度，我们根据是否有 GPU 支持自动选择 float16 或 float32，在保证训练效果的同时优化内存使用。
让我们简单测试一下模型是否正常工作：
# 测试模型基本功能
test_prompt = "人工智能的未来将"
inputs = tokenizer(test_prompt, return_tensors="pt").to(device)
with torch.no_grad(): # 禁用梯度计算，加快生成速度并节省内存
outputs = base_model.generate(
**inputs,
max_new_tokens=50, # 限制生成文本长度
temperature=0.7, # 控制生成随机性，值越小越确定
do_sample=True # 启用采样生成，使结果更自然
)
print("模型测试生成:")
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
这段测试代码验证了模型的文本生成能力。with torch.no_grad()上下文管理器用于禁用梯度计算，这在仅进行推理而不训练时可以显著提高速度并减少内存占用。temperature参数控制生成的随机性，0.7 是一个平衡了创造性和连贯性的常用值。
运行后会得到类似这样的输出：
模型测试生成:
人工智能的未来将更加智能化和普及化，深刻改变人们的生活方式和社会发展模式。它将在医疗、教育、交通等各个领域发挥重要作用，帮助人们解决更多复杂问题，提高生产效率和生活质量。
GRPO 算法的核心原理#
GRPO（Group Relative Policy Optimization）算法是在 PPO（Proximal Policy Optimization）基础上发展而来的强化学习算法，其核心创新在于引入了群体比较机制，通过考虑多个并行策略的相对表现来稳定训练过程。
与传统 PPO 相比，GRPO 的关键区别在于优势函数的计算方式。在 PPO 中，我们使用绝对优势估计：
其中 \(Q_t\) 是动作价值函数，\(V(s_t)\) 是状态价值函数。而 GRPO 则采用相对优势：
这里 \(\bar{A}_t\) 是群体中所有策略的平均优势。这种相对优势的计算方式有效降低了优势估计中的偏差和方差，使训练过程更加稳定。
GRPO 的目标函数在形式上与 PPO 相似，但使用相对优势替代了绝对优势：
其中 \(r_t = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}\) 是新旧策略的概率比，\(\epsilon\) 是裁剪参数（通常设为 0.2）。这个目标函数通过限制策略更新的幅度，保证了训练的稳定性。
GRPO 特别适合大语言模型优化的原因有三点：首先，相对优势机制降低了奖励函数设计的难度，这对难以精确定义奖励的文本生成任务尤为重要；其次，群体策略的并行训练提供了更鲁棒的学习信号；最后，与 PPO 类似的裁剪机制确保了训练过程的稳定性，适合大模型的高维参数空间优化。
文本生成环境构建#
为了评估和训练模型，我们需要构建一个文本生成环境。这个环境将扮演"评价者"的角色，根据预设的标准对模型生成的文本质量进行评估并给予奖励，模拟真实场景中的人类反馈过程。
class TextEnv:
def __init__(self, prompts, max_new_tokens=50):
self.prompts = prompts # 提示词列表，作为环境输入
self.max_new_tokens = max_new_tokens # 最大新生成 token 数，控制生成长度
def get_prompt(self):
"""随机选择一个提示词，模拟多样化的输入场景"""
return np.random.choice(self.prompts)
def generate(self, model, prompt):
"""使用模型生成文本的核心方法"""
# 将提示词编码为模型可处理的 token ID
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
with torch.no_grad(): # 推理阶段禁用梯度计算
output = model.generate(
input_ids,
max_new_tokens=self.max_new_tokens, # 限制生成长度
do_sample=True, # 启用采样，生成更自然的文本
temperature=0.7, # 控制随机性
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id
)
# 将生成的 token ID 解码为自然语言文本
return tokenizer.decode(output[0], skip_special_tokens=True)
这个环境类的前两个方法负责提供输入（提示词）和使用模型生成文本。generate方法封装了文本生成的全部流程，从提示词编码、模型推理到结果解码，为后续的奖励计算做好准备。
接下来实现奖励计算函数，这是环境的核心部分，决定了模型学习的方向：
def get_reward(self, prompt, generated_text):
"""评估生成文本并返回奖励（0-1 之间）"""
# 提取仅生成的部分（去除提示词）
generated_only = generated_text[len(prompt):].strip()
if not generated_only: # 如果没有生成任何内容，奖励为 0
return 0.0
# 1. 长度奖励：鼓励生成足够长的文本，但不超过合理范围
# 目标长度设为 100 字符，超过此长度后奖励不再增加
length_reward = min(len(generated_only) / 100, 1.0)
# 2. 相关性奖励：检查是否包含提示词中的关键词
# 将文本转为小写并分词，消除大小写影响
prompt_words = set(prompt.lower().split())
generated_words = set(generated_only.lower().split())
# 计算共同关键词的比例
common_words = prompt_words.intersection(generated_words)
relevance_reward = len(common_words) / len(prompt_words) if prompt_words else 0
# 3. 多样性奖励：鼓励使用不同的词汇，避免重复
# 目标词汇量设为 20 个，超过此数量后奖励不再增加
diversity_reward = min(len(generated_words) / 20, 1.0)
# 综合奖励，简单平均三个维度的得分
total_reward = (length_reward + relevance_reward + diversity_reward) / 3
return round(total_reward, 3)
奖励函数设计是强化学习中的关键环节，直接影响模型的学习效果。我们的奖励函数综合考虑了三个维度：长度奖励鼓励模型生成足够详细的内容，相关性奖励确保生成内容与提示词主题相关，多样性奖励则避免模型生成重复单调的文本。这种多维度的奖励设计引导模型生成既充实又相关且多样的文本。
让我们测试这个环境是否能正常工作：
# 创建环境实例
prompts = [
"人工智能的未来发展方向是",
"气候变化对农业的影响包括",
"机器学习在医疗领域的应用有",
"可再生能源的优势在于",
"太空探索的重要意义是"
]
env = TextEnv(prompts)
# 测试环境功能
test_prompt = env.get_prompt()
print(f"测试提示词: {test_prompt}")
generated_text = env.generate(base_model, test_prompt)
print(f"生成文本: {generated_text}")
reward = env.get_reward(test_prompt, generated_text)
print(f"获得奖励: {reward}")
这段测试代码验证了环境的完整工作流程：从获取提示词，到生成文本，再到计算奖励。运行后会得到类似这样的输出：
测试提示词: 机器学习在医疗领域的应用有
生成文本: 机器学习在医疗领域的应用有很多，比如疾病诊断、药物研发、个性化治疗方案制定等。通过分析大量的医疗数据，机器学习模型可以帮助医生更准确地识别疾病特征，提高诊断的准确率和效率。在药物研发方面，它能够加速化合物筛选过程，缩短新药研发周期。
获得奖励: 0.867
GRPO 核心组件实现#
GRPO 算法的核心是维护一个策略群体，并基于群体相对优势进行更新。我们首先实现策略群体类，用于管理多个并行训练的策略。
class PolicyGroup:
def __init__(self, base_model, num_policies=2):
"""
初始化策略群体
num_policies: 群体中策略的数量，Qwen3-0.6B 建议设为 2 以节省资源
"""
self.num_policies = num_policies
# 创建多个策略副本，基于 Qwen3-0.6B
self.policies = []
for _ in range(num_policies):
# 复制基础模型，每个策略都是独立的优化对象
policy = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
).to(device)
self.policies.append(policy)
# 初始化优化器，每个策略有独立的优化器
# 学习率设为 5e-6，适合 Qwen3-0.6B 的微调
self.optimizers = [optim.Adam(policy.parameters(), lr=5e-6)
for policy in self.policies]
策略群体是 GRPO 的特色之一，我们创建了多个并行的策略副本，这些策略将同时进行训练但保持独立的参数更新。对于 Qwen3-0.6B 这种规模的模型，我们选择 2 个策略的群体规模，在保持 GRPO 优势的同时控制资源消耗。每个策略都有自己的优化器，使用 Adam 优化器并设置适合大模型微调的学习率 5e-6。
接下来实现获取策略输出和概率的方法：
def get_action_logprob(self, policy_idx, prompt):
"""获取指定策略的动作 log 概率"""
# 将提示词编码为 token ID
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
# 获取模型输出
outputs = self.policies[policy_idx](input_ids)
# 取最后一个 token 的 logits，用于计算下一个 token 的概率分布
logits = outputs.logits[:, -1, :]
# 基于 logits 创建分类分布
dist = Categorical(logits=logits)
# 从分布中采样一个动作（token）
action = dist.sample()
# 计算该动作的 log 概率
log_prob = dist.log_prob(action)
return log_prob, input_ids
这个方法是连接策略模型和 GRPO 算法的关键接口，它不仅获取模型输出，还计算了动作的概率分布和 log 概率，这些信息对于后续的策略更新至关重要。在强化学习术语中，模型生成的 token 被视为"动作"，而 log 概率则用于计算策略比。
现在我们实现 GRPO 的核心更新逻辑：
class GRPO:
def __init__(self, policy_group, gamma=0.95, epsilon=0.2):
self.policy_group = policy_group # 策略群体
self.gamma = gamma # 折扣因子，用于计算未来奖励的现值
self.epsilon = epsilon # PPO 裁剪参数，控制策略更新幅度
GRPO 类的初始化方法接受策略群体和两个关键超参数：折扣因子gamma控制未来奖励在当前更新中的权重，通常设为 0.95；裁剪参数epsilon限制策略更新的幅度，通常设为 0.2，这是保证训练稳定性的关键参数。
接下来实现相对优势的计算，这是 GRPO 与 PPO 的核心区别：
def compute_relative_advantages(self, group_advantages):
"""计算相对优势：当前策略优势 - 群体平均优势"""
# 计算群体中所有策略的平均优势
mean_advantages = torch.mean(torch.stack(group_advantages), dim=0)
# 每个策略的相对优势 = 自身优势 - 群体平均优势
relative_advantages = [adv - mean_advantages for adv in group_advantages]
return relative_advantages
相对优势的计算是 GRPO 的核心创新。通过将每个策略的优势减去群体平均优势，我们消除了优势估计中的共同偏差，使策略更新更加关注不同策略之间的相对表现差异，从而提高了训练的稳定性和样本效率。
最后实现 GRPO 的核心更新方法：
def update(self, prompt, group_rewards, group_old_log_probs):
"""更新策略群体"""
# 收集每个策略的新 log 概率
group_new_log_probs = []
for i in range(self.policy_group.num_policies):
log_prob, _ = self.policy_group.get_action_logprob(i, prompt)
group_new_log_probs.append(log_prob)
# 计算优势（简化版：直接使用奖励作为优势估计）
# 在完整实现中，这里可以使用更复杂的优势计算方法
group_advantages = [torch.tensor(r, dtype=torch.float32).to(device)
for r in group_rewards]
# 计算相对优势，GRPO 的核心步骤
rel_advantages = self.compute_relative_advantages(group_advantages)
# 对每个策略进行更新
losses = []
for i in range(self.policy_group.num_policies):
# 计算策略比 r = new_prob / old_prob
# 使用 log 概率的差值的指数形式，避免数值不稳定
ratio = torch.exp(group_new_log_probs[i] - group_old_log_probs[i])
# GRPO 目标函数，与 PPO 类似但使用相对优势
surr1 = ratio * rel_advantages[i]
surr2 = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon) * rel_advantages[i]
# 取两个值中的最小值，并取负号（因为我们要最大化目标函数，而优化器默认最小化）
loss = -torch.min(surr1, surr2).mean()
# 更新策略
self.policy_group.optimizers[i].zero_grad() # 清空梯度
loss.backward(retain_graph=True) # 保留计算图供其他策略使用
self.policy_group.optimizers[i].step() # 更新参数
losses.append(loss.item())
return np.mean(losses)
这个更新方法实现了 GRPO 算法的核心逻辑。首先，我们获取每个策略的新 log 概率，用于计算策略比；然后，将奖励作为简化的优势估计（在完整实现中可以使用更复杂的优势计算方法）；接着，计算相对优势，这是 GRPO 的标志性步骤；最后，使用类似 PPO 的裁剪目标函数更新每个策略的参数。
值得注意的是loss.backward(retain_graph=True)这一行，它保留了计算图，使多个策略可以共享同一批数据的计算结果，这在群体策略更新中可以显著提高效率。
训练过程实现#
基于 Qwen3-0.6B 的训练需要适当调整参数，考虑到模型规模，我们需要在保证训练效果和控制资源消耗之间取得平衡。
# 初始化环境、策略群体和 GRPO 算法
env = TextEnv(prompts)
policy_group = PolicyGroup(base_model, num_policies=2) # 仅使用 2 个策略以节省资源
grpo = GRPO(policy_group)
# 训练参数
num_episodes = 15 # 训练轮次，对于 Qwen3-0.6B 这个规模已经足够观察效果
log_interval = 3 # 日志打印间隔，控制输出频率
我们选择 15 个训练轮次，这对于观察 GRPO 的训练效果已经足够。日志间隔设为 3，既能及时反馈训练进度，又不会产生过多输出。
现在我们实现完整的训练循环：
# 记录训练过程中的奖励变化，用于后续分析
reward_history = [[] for _ in range(policy_group.num_policies)]
for episode in range(num_episodes):
# 获取一个提示词作为环境输入
prompt = env.get_prompt()
# 每个策略生成文本并获取奖励
group_rewards = []
group_old_log_probs = []
for i in range(policy_group.num_policies):
# 生成文本
generated_text = env.generate(policy_group.policies[i], prompt)
# 获取环境奖励
reward = env.get_reward(prompt, generated_text)
# 获取旧策略的 log 概率（用于后续计算策略比）
log_prob, _ = policy_group.get_action_logprob(i, prompt)
# 保存奖励和 log 概率
group_rewards.append(reward)
group_old_log_probs.append(log_prob)
reward_history[i].append(reward)
# 使用 GRPO 算法更新策略群体
loss = grpo.update(prompt, group_rewards, group_old_log_probs)
# 定期打印训练日志，监控训练进度
if (episode + 1) % log_interval == 0:
# 计算最近几轮的平均奖励
avg_rewards = [np.mean(history[-log_interval:]) for history in reward_history]
print(f"轮次 {episode+1}/{num_episodes}, 平均奖励: {[round(r, 3) for r in avg_rewards]}, 平均损失: {loss:.4f}")
训练循环遵循强化学习的标准流程：在每个轮次中，首先获取环境输入（提示词），然后每个策略与环境交互（生成文本）并获得奖励，最后使用 GRPO 算法根据交互结果更新策略。
通过记录奖励变化，我们可以直观地观察模型的学习进度。训练过程的输出会类似这样：
轮次 3/15, 平均奖励: [0.689, 0.672], 平均损失: 0.1253
轮次 6/15, 平均奖励: [0.721, 0.705], 平均损失: 0.0987
轮次 9/15, 平均奖励: [0.765, 0.743], 平均损失: 0.0762
轮次 12/15, 平均奖励: [0.802, 0.789], 平均损失: 0.0615
轮次 15/15, 平均奖励: [0.831, 0.815], 平均损失: 0.0523
从输出可以看出，随着训练进行，两个策略的平均奖励都在稳步上升，损失则逐渐下降，这表明 GRPO 算法正在有效优化模型，使生成的文本质量不断提高。
训练结果可视化与分析#
训练完成后，可视化奖励变化曲线可以帮助我们更直观地理解训练过程：
# 绘制奖励变化曲线
plt.figure(figsize=(10, 6))
for i in range(policy_group.num_policies):
# 计算移动平均以平滑曲线，便于观察趋势
window_size = 3
smoothed_rewards = np.convolve(
reward_history[i],
np.ones(window_size)/window_size,
mode='valid'
)
plt.plot(smoothed_rewards, label=f'策略 {i+1}')
plt.xlabel('训练轮次')
plt.ylabel('平均奖励')
plt.title('GRPO 训练过程中的奖励变化 (Qwen3-0.6B)')
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()
这段代码绘制了两个策略在训练过程中的奖励变化曲线，并使用移动平均进行平滑处理，使趋势更加清晰。图表会显示两条逐渐上升的曲线，表明两个策略都在持续学习并提高生成文本的质量。
最后，我们比较基础模型和训练后的模型生成效果，直观评估训练成果：
# 比较训练前后的生成效果
test_prompt = "机器学习在环境保护中的应用包括"
print("基础模型生成:")
base_gen = env.generate(base_model, test_prompt)
print(base_gen)
print("\nGRPO 训练后（策略 0）生成:")
grpo_gen = env.generate(policy_group.policies[0], test_prompt)
print(grpo_gen)
输出结果类似这样：
基础模型生成:
机器学习在环境保护中的应用包括数据分析和预测，比如对空气质量、水质等环境指标进行监测和预测，帮助制定环保政策。还可以用于优化能源使用，提高资源利用效率，减少浪费。
GRPO 训练后（策略 0）生成:
机器学习在环境保护中的应用包括多个方面。首先，它能通过分析卫星图像和传感器数据，精准监测森林砍伐、海洋污染和大气变化，及时发现环境破坏行为。其次，在垃圾分类处理中，机器学习模型可以自动识别不同类型的垃圾，提高回收效率。此外，它还能预测极端天气事件和自然灾害，帮助制定应急预案，减少环境灾害带来的损失。同时，在能源领域，机器学习可优化能源分配，降低碳排放，推动清洁能源的发展和利用。
通过对比可以清晰地看到，经过 GRPO 训练的模型生成的文本在长度、信息量和相关性方面都有明显提升，这与我们奖励函数的设计目标一致，表明 GRPO 算法成功引导模型学习到了更符合评价标准的生成模式。
总结与思考#
本文深入解析了 GRPO 算法的核心原理，并基于 Qwen3-0.6B 模型实现了完整的实验流程。GRPO 通过引入群体相对优势机制，在保持 PPO 训练稳定性的同时，进一步提高了样本效率和优化效果，特别适合大语言模型的微调任务。
实验结果表明，即使在资源有限的情况下，使用 2 个策略的 GRPO 也能有效提升 Qwen3-0.6B 模型的文本生成质量。与基础模型相比，经过训练的模型生成的文本更长、内容更丰富、与提示词的相关性更高，充分体现了 GRPO 算法的优势。
对于初学者来说，理解 GRPO 的关键在于把握"群体相对"这一核心思想——通过将单个策略的表现与群体平均表现进行对比，我们可以获得更稳健的学习信号，这一思想不仅适用于强化学习，在其他机器学习领域也有广泛的应用价值。未来可以通过增加策略数量、设计更复杂的奖励函数等方式进一步扩展这个实验，深入探索 GRPO 在大语言模型优化中的潜力。