在 48GB 显卡上仅用 12 小时重现DeepSeek“Wait!/Aha”时刻-电子发烧友网

作者：算力魔方创始人/英特尔创新大使刘力

近日，荷兰科学家Raz发布了Reinforce-Lite算法，实现了在 48GB显存的显卡上仅用 12 小时在3B模型上重现DeepSeek“Wait!/Aha”时刻。

原文链接：https://medium.com/@rjusnba/overnight-end-to-end-rl-training-a-3b-model-on-a-grade-school-math-dataset-leads-to-reasoning-df61410c04c6
滑动查看更多

一，Reinforce-Lite算法的显存要求

Raz通过移除KL，移除替代比率，去掉评论模型，使用组相对奖励（DeepSeek的GRPO风格）进行优势计算，提出了一种更简单、更稳定、更高效的轻量级强化学习方法：Reinforce-Lite，使得显存需求，从72GB下降到36GB！下表是：端到端的用强化学习训练 3B 模型的显存需求。

二，Reinforce-Lite算法的PyTorch实现

Reinforce-lite算法的PyTorch实现如下所示：

第一步，初始化一个指令微调的LLM，并适当提示以将其推理步骤包含在标签中。

第二步，定义一个奖励函数用于模型输出（例如，GSM8K数学推理任务中的正确性）。通过正则表达式提取标签中的数值，并与数据集中的实际答案进行比较。

第三步，通过直接计算相对于奖励的梯度来优化策略，而不需要替代损失。

第四步，使用组相对归一化进行优势计算，消除了对评论模型的需求。我们使用组大小为10。

第五步，使用标准对数概率梯度更新模型。

def reinforce_lite(batch, policy_model, tokenizer, device, step, save_dir):
"""
使用强化学习方法训练策略模型。

Args:
batch (list of tuples): 包含提示和目标句子的列表。
policy_model (torch.nn.Module): 策略模型，用于生成响应。
tokenizer (transformers.PreTrainedTokenizer): 用于处理文本的标记器。
device (torch.device): 指定模型运行的设备。
step (int): 当前训练步数。
save_dir (str): 保存模型的目录。

Returns:
tuple: 包含策略损失、平均奖励、策略损失项、0.0、第一个响应和所有响应的长度。

"""
# 设置模型为训练模式
policy_model.train()
# 解包输入数据
prompts, targets = zip(*batch)
# 获取批量大小
batch_size = len(prompts)
# 初始化评估组索引
evaluated_group = 0

# 初始化存储列表
all_logprobs = []
all_rewards = []
all_responses = []
all_lengths = []

for group_idx in range(config.GROUP_SIZE):
# 格式化提示
formatted_prompts = [format_prompt(p, tokenizer) for p in prompts]
# 将提示转换为模型输入
inputs = tokenizer(
formatted_prompts,
return_tensors="pt",
padding=True,
truncation=True,
max_length=config.MAX_SEQ_LENGTH
).to(device)

# 生成参数
generate_kwargs = {
**inputs,
"max_new_tokens": config.MAX_NEW_TOKENS,
"do_sample": True,
"temperature": 0.7,
"top_p": 0.9,
"pad_token_id": tokenizer.pad_token_id,
"return_dict_in_generate": True,
}

# 判断当前组是否为评估组
if group_idx == evaluated_group:
# 生成响应
generated = policy_model.generate(**generate_kwargs)
# 获取生成的响应ID
generated_ids = generated.sequences
# 获取模型输出
outputs = policy_model(
generated_ids,
attention_mask=(generated_ids != tokenizer.pad_token_id).long()
)
# 获取提示长度和响应长度
prompt_length = inputs.input_ids.shape[1]
response_length = generated_ids.shape[1] - prompt_length
# 计算对数概率
if response_length > 0:
logits = outputs.logits[:, prompt_length-1:-1, :]
response_tokens = generated_ids[:, prompt_length:]
log_probs = torch.log_softmax(logits, dim=-1)
token_log_probs = torch.gather(log_probs, -1, response_tokens.unsqueeze(-1)).squeeze(-1)
sequence_log_probs = token_log_probs.sum(dim=1)
else:
sequence_log_probs = torch.zeros(batch_size, device=device)
else:
# 在不计算梯度的情况下生成响应
with torch.no_grad():
generated = policy_model.generate(**generate_kwargs)
sequence_log_probs = torch.zeros(batch_size, device=device)

# 解码生成的响应
responses = tokenizer.batch_decode(
generated.sequences[:, inputs.input_ids.shape[1]:],
skip_special_tokens=True
)
# 计算奖励
rewards = torch.tensor([get_reward(resp, tgt) for resp, tgt in zip(responses, targets)], device=device)

# 存储结果
all_responses.extend(responses)
all_rewards.append(rewards)
all_logprobs.append(sequence_log_probs)
all_lengths.extend([len(r.split()) for r in responses])

# 堆叠奖励和对数概率
rewards_tensor = torch.stack(all_rewards)
logprobs_tensor = torch.stack(all_logprobs)

# 分离评估组的奖励和其他组的奖励
evaluated_rewards = rewards_tensor[evaluated_group]
others_rewards = torch.cat([
rewards_tensor[:evaluated_group],
rewards_tensor[evaluated_group+1:]
], dim=0)

# 计算基线值
baseline = others_rewards.mean(dim=0)
# 计算优势
advantages = (evaluated_rewards - baseline) / (others_rewards.std(dim=0) + 1e-8)
advantages = torch.clamp(advantages, -2.0, 2.0)

# 计算策略损失
policy_loss = -(logprobs_tensor[evaluated_group] * advantages.detach()).mean()

return policy_loss, rewards_tensor.mean().item(), policy_loss.item(), 0.0, all_responses[0], all_lengths
滑动查看更多

三，Reinforce-Lite算法的数据集：GSM 8K

Reinforce-Lite使用GSM8K数据集：这是一个小学数学数据集，包含数学问题及其答案，格式如下：