liangzhidanta/Qwen3-4B-GRPO-119

<div align="center">

Qwen3-4B-GRPO-119

[English](#-english-version) ｜ [中文](#-中文版本)

</div>

<div align="center">

<a name="-english-version"></a>

English Version

GRPO Reinforcement Learning on Math Reasoning — trained end-to-end on 8×RTX 4090 with [slime](https://github.com/THUDM/slime)

 ![Algorithm]()  ![License]() ![Hardware]()

AIME 2024: 57.5% → 65.0% (+7.5pp) · response length −49% · truncation −54%

</div>

📈 Training Curves

<div align="center">

<img src="curves.png" width="100%" alt="Training curves: AIME eval, reward/length, truncation, timing">

Top-left: AIME 2024 pass@1 every 20 iterations (gray band = ±1pp test-retest noise, calibrated by 3 baseline runs). Top-right: train-batch reward rises while mean response length falls — accuracy and efficiency improve together. Bottom-left: truncation ratio halves on both train (8k cap) and eval (16k cap). Bottom-right: rollout dominates ~85–96% of each ~23-min iteration.

</div>

🏆 Key Results (120 GRPO iterations)

*The headline finding: the model became more accurate and more concise simultaneously* — no length-based reward hacking, entropy stays healthy (0.34 → 0.29).

<details> <summary><b>📊 Full AIME eval history (7 points)</b></summary>

• —Test-retest noise: ±1pp (baseline measured 3×: 57.5 / 57.9 / 58.5%)
• —Fast gains in the first 80 iterations (+8.5pp), then a 65%±1pp plateau — training stopped at 120 based on this

</details>

⚙️ Training Setup

<details> <summary><b>🛠️ Engineering notes — making an H100 recipe fit 24GB consumer GPUs</b></summary>

The official recipe targets 8×H100 80GB. Adapting it to 8×RTX 4090 24GB required (full story in EXPERIMENT_REPORT.md):

1. 1.TP 2→4: a single truncated sample (~8.4k tokens) materializes a ~5GiB fp32 full-vocab logits block in the log-prob path — halving the vocab shard per GPU was the decisive fix;
2. 2.Token budget 9216→4096 + log-probs-chunk 1024: logits buffers scale linearly with microbatch tokens; chunking bounds softmax temporaries;
3. 3.Disabled the (coefficient-zero) KL reference forward — pure waste of memory and time;
4. 4.expandable_segments:True against fragmentation (requires the +slime fork of torchmemorysaver);
5. 5.Session/cluster survival (loginctl enable-linger + tmux) and checkpoint retention (53GB each, keep-latest) for a multi-day run.

</details>

🚀 Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "liangzhidanta/Qwen3-4B-GRPO-119",
    torch_dtype=torch.bfloat16, device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("liangzhidanta/Qwen3-4B-GRPO-119")

messages = [{"role": "user",
             "content": "Find the number of ordered pairs of prime numbers $(p, q)$ "
                        "such that $3p + 5q = 102$. Put the answer in \\boxed{}."}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
).to(model.device)

out = model.generate(**inputs, max_new_tokens=8192,
                     do_sample=True, temperature=1.0)
print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Best served with thinking mode enabled (default chat template) and a generous generation budget; the model was trained at temperature 1.0.

⚠️ Limitations

• —Trained and evaluated on English math reasoning only (DAPO-Math-17K / AIME 2024);
• —21% of training rollouts still hit the 8192-token cap;
• —Checkpoint at iteration 119 (plateau representative); the iter-79 peak (66.0%) is within the ±1pp noise band and was not retained;
• —~2/3 of prompt groups produced zero-variance rewards each iteration (a DAPO dynamic-sampling opportunity — see report §8).

<div align="center">

<a name="-中文版本"></a>

中文版本

GRPO 数学推理强化学习 —— 基于 [slime](https://github.com/THUDM/slime) 框架、在 8×RTX 4090 上端到端训练

 ![Algorithm]()  ![License]() ![Hardware]()

AIME 2024：57.5% → 65.0%（+7.5 个百分点）· 响应长度 −49% · 截断率 −54%

</div>

📈 训练曲线

<div align="center">

<img src="curves.png" width="100%" alt="训练曲线：AIME 评估、奖励/长度、截断率、耗时">

左上：每 20 轮一次的 AIME 2024 pass@1（灰色带 = ±1pp 复测噪声，由 3 次基线重复测量标定）。右上：训练批奖励上升的同时平均响应长度下降——正确率与效率同步改善。左下：训练（8k 上限）与评估（16k 上限）的截断率均减半。右下：每轮约 23 分钟中 rollout 占 85–96%。

</div>

🏆 核心结果（120 轮 GRPO 训练）

核心发现：模型变得更准的同时变得更简洁——没有出现"靠写长碰答案"的 reward hacking，策略熵保持健康（0.34 → 0.29）。

<details> <summary><b>📊 完整 AIME 评估历史（7 个评估点）</b></summary>

• —复测噪声：±1pp（基线测了 3 遍：57.5 / 57.9 / 58.5%）
• —前 80 轮贡献几乎全部增益（+8.5pp），之后进入 65%±1pp 平台期——据此在第 120 轮停止训练

</details>

⚙️ 训练配置

<details> <summary><b>🛠️ 工程笔记——让 H100 配方塞进 24G 消费级显卡</b></summary>

官方配方按 8×H100 80G 设计，适配 8×RTX 4090 24G 做了以下改造（完整记录见 EXPERIMENT_REPORT.md）：

1. 1.TP 2→4：单条截断样本（约 8.4k token）在 log-prob 路径会物化约 5GiB 的 fp32 全词表 logits——把词表分片再对半是决定性修复；
2. 2.token 预算 9216→4096 + log-probs-chunk 1024：logits 缓冲与 microbatch token 数成正比，分块限制 softmax 临时缓冲；
3. 3.关闭（系数为零的）KL 参考前向——纯粹的显存与时间浪费；
4. 4.expandable_segments:True 对抗碎片（需要 torchmemorysaver 的 +slime fork）；
5. 5.多日训练的会话/集群存活（loginctl enable-linger + tmux）与 checkpoint 保留策略（每份 53GB，只留最新）。

</details>

🚀 使用方法

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "liangzhidanta/Qwen3-4B-GRPO-119",
    torch_dtype=torch.bfloat16, device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("liangzhidanta/Qwen3-4B-GRPO-119")

messages = [{"role": "user",
             "content": "求满足 3p + 5q = 102 的质数有序对 (p, q) 的个数。答案写在 \\boxed{} 中。"}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
).to(model.device)

out = model.generate(**inputs, max_new_tokens=8192,
                     do_sample=True, temperature=1.0)
print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

建议开启思考模式（默认 chat template）并给足生成长度；模型训练温度为 1.0。

⚠️ 局限性

• —仅在英文数学推理数据上训练和评估（DAPO-Math-17K / AIME 2024）；
• —21% 的训练 rollout 仍会触及 8192 token 上限；
• —采用第 119 轮 checkpoint（平台期代表）；第 79 轮峰值（66.0%）在 ±1pp 噪声带内、未保留；
• —每轮约 2/3 的 prompt 组奖励零方差（DAPO dynamic sampling 的改进方向——见报告 §8）。

🙏 致谢

• —slime——RL 训练框架；
• —Qwen 团队的优秀基座模型；
• —DAPO 与 DeepScaler 论文提供的数据、奖励设计与非对称裁剪配方。

📁 仓库文件