Upload 12 files

README.md

@@ -1,67 +1,158 @@
 ---
 library_name: transformers
-model_name: trlm-135m
+license: apache-2.0
+base_model: Shekswess/trlm-stage-2-sft-final-2
 tags:
-- generated_from_trainer
-- grpo
 - trl
-licence: license
+- dpo
+- preference-alignment
+- reasoning
+- generated_from_trainer
+model-index:
+- name: trlm-stage-3-dpo-final-2
+  results: []
 ---
 
-# Model Card for trlm-135m
+# Tiny Reasoning Language Model (trlm-135)
 
-This model is a fine-tuned version of [None](https://huggingface.co/None).
-It has been trained using [TRL](https://github.com/huggingface/trl).
+![image/png](https://github.com/user-attachments/assets/5f453496-8180-4cf4-94da-26ebbe1159d4)
 
-## Quick start
+## Table of Contents
 
-```python
-from transformers import pipeline
+1. [Model Summary](#model-summary)
+2. [Post-Training Pipeline](#post-training-pipeline)
+3. [How to use](#how-to-use)
+4. [Training](#training)
+5. [Evaluation](#evaluation)
+6. [Limitations](#limitations)
+7. [Acknowledgements](#acknowledgements)
+8. [License](#license)
+---
+
+## Model Summary
+
+The **Tiny Reasoning Language Model (trlm-135)** is a **135M parameter** research prototype designed to study how small models can learn step-by-step reasoning.
+It was built on top of [SmolLM2-135M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct) and fine-tuned through a **3-stage pipeline**:
+
+* **[Stage 1 SFT](https://huggingface.co/Shekswess/trlm-stage-1-sft-final-2)**: general instruction tuning (non-reasoning).
+* **[Stage 2 SFT](https://huggingface.co/Shekswess/trlm-stage-2-sft-final-2)**: reasoning traces with `<think>` tags.
+* **[Stage 3 DPO](https://huggingface.co/Shekswess/trlm-stage-3-dpo-final-2)**: preference alignment for reasoning style.
+
+The **code** for everything can be found **[here](https://github.com/Shekswess/tiny-reasoning-language-model/blob/main/README.md)**
+
+---
+
+## Post-Training Pipeline
+<img width="1014" height="563" alt="image" src="https://github.com/user-attachments/assets/195ef389-6aa9-4527-b4f0-bea68c0841ae" />
+
+---
+
+## How to use
 
-question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?"
-generator = pipeline("text-generation", model="Shekswess/trlm-135m", device="cuda")
-output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0]
-print(output["generated_text"])
+```bash
+pip install -U transformers accelerate
 ```
 
-## Training procedure
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "Shekswess/trlm-135m"
+device = "cuda"  # or "cpu"
+
+# Load tokenizer & model
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+).to(device)
+
+# Example prompt
+prompt = "Give me a brief explanation of gravity in simple terms."
+messages = [
+    {"role": "user", "content": prompt}
+]
+
+# Apply chat template
+text = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+)
+
+inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+# Generate
+outputs = model.generate(**inputs, max_new_tokens=256)
+print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+```
 
- 
+> [!TIP]
+> For reasoning-heavy tasks, set `temperature=0.6` and `top_p=0.95`.
 
+---
 
-This model was trained with GRPO, a method introduced in [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://huggingface.co/papers/2402.03300).
+## Training
 
-### Framework versions
+### Model
 
-- TRL: 0.23.0
-- Transformers: 4.56.1
-- Pytorch: 2.8.0+cu126
-- Datasets: 4.0.0
-- Tokenizers: 0.22.0
+* **Architecture**: Decoder-only transformer (SmolLM2 backbone which infact is Llama 3 based model).
+* **Parameters**: ~135M.
+* **Precision**: mix-precision (bfloat16) during training.
 
-## Citations
+### Software & Hardware
 
-Cite GRPO as:
+* **Training Frameworks**: PyTorch (ROCm), Hugging Face Transformers & TRL.
+* **Hardware**: AMD MI300X (192GB VRAM, 224GB RAM).
 
-```bibtex
-@article{shao2024deepseekmath,
-    title        = {{DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models}},
-    author       = {Zhihong Shao and Peiyi Wang and Qihao Zhu and Runxin Xu and Junxiao Song and Mingchuan Zhang and Y. K. Li and Y. Wu and Daya Guo},
-    year         = 2024,
-    eprint       = {arXiv:2402.03300},
-}
+**Special thanks to [@HotAisle](https://x.com/HotAisle)**
 
-```
+### Training Stages
+
+1. **Stage 1 – SFT (non-reasoning)**
+   * ~58k samples, everyday conversations & instruction following.
+2. **Stage 2 – SFT (reasoning)**
+   * ~78k samples with `<think>` segments.
+3. **Stage 3 – DPO (alignment)**
+   * ~50k preference pairs (chosen vs. rejected reasoning traces).
+---
+
+## Evaluation
+
+Evaluation was done with `lm-eval-harness`:
+
+| **Benchmark**        | **Tiny Reasoning Language Model (trlm-135M)**  | **SmolLM2-135M-Instruct** | **Improvements** |
+| -------------------- | ---------------------------- | ------------------------- | ---------------------------- |
+| **ARC Challenge**    | **40.61** (avg)              | 37.3 (avg)                | **+3.31**                    |
+| **BBH**              | **36.80** (3-shot)           | 28.2 (3-shot)             | **+8.6**                     |
+| **BoolQ**            | **62.17**                    | –                         | N/A                          |
+| **GSM8K**            | **2.59** (5-shot)            | 1.4 (5-shot)              | **+1.19**                    |
+| **IFEval**           | **35.49** (avg)              | 29.9 (avg)                | **+5.59**                    |
+| **MMLU**             | **34.95**                    | 29.3                      | **+5.65**                    |
+| **PIQA**             | **64.91**                    | 66.3                      | **–1.39**                    |
+| **HellaSwag**        | –                            | 40.9                      | N/A                          |
+| **MT-Bench**         | –                            | 19.8                      | N/A                          |
+
+---
+
+## Limitations
+
+* **Not production-ready**: hallucinations and logical errors are frequent.
+* **Small size**: limited general knowledge and reasoning depth.
+* **English-only**: multilingual capabilities not explored.
+
+---
+
+## Acknowledgements
+
+- [@HotAisle](https://x.com/HotAisle) for providing the compute resources to train all three stages on a awesome AMD MI300x setup.
+- [@mkurman88](https://x.com/mkurman88) for ideas, feedback and code samples.
+- [HuggingFaceTB team](https://huggingface.co/HuggingFaceTB) for SmolLM2-135M-Instruct model and the Smoltalk2 dataset collection.
+- [@scottgeng00](https://huggingface.co/scottgeng00) for the OLmO-3-Preference-Mix-Deltas dataset.
+- [@eliebakouchi](https://x.com/eliebakouch) for help with the tokenization.
+
+---
+
+## License
+
+[Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0)
 
-Cite TRL as:
-    
-```bibtex
-@misc{vonwerra2022trl,
-	title        = {{TRL: Transformer Reinforcement Learning}},
-	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
-	year         = 2020,
-	journal      = {GitHub repository},
-	publisher    = {GitHub},
-	howpublished = {\url{https://github.com/huggingface/trl}}
-}
-```
+---

config.json

@@ -32,7 +32,7 @@
       "q4f16": "float16"
     }
   },
-  "transformers_version": "4.56.1",
+  "transformers_version": "4.56.2",
   "use_cache": true,
   "vocab_size": 49154
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6df764ae38de74036a73950a26c4ce2eabe6461855cb96ee3668e6c6c0d1b15a
+oid sha256:f72be34d84f9d95d2f1ba9e6aa5d352ec841ce1e7aed81998135ce5bf96e9a09
 size 269062856

tokenizer.json

@@ -1,14 +1,7 @@
 {
   "version": "1.0",
   "truncation": null,
-  "padding": {
-    "strategy": "BatchLongest",
-    "direction": "Left",
-    "pad_to_multiple_of": null,
-    "pad_id": 2,
-    "pad_type_id": 0,
-    "pad_token": "<|im_end|>"
-  },
+  "padding": null,
   "added_tokens": [
     {
       "id": 0,

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7a4eaa81d993d68c33c29824d61c91baead1e98508a5a625eb153ada959ab125
-size 7185
+oid sha256:247b792c868f8245ceddd15c2b2486a99317401202045e562ed34e945a36ed82
+size 6865