Training in progress, step 10

README.md

@@ -1,158 +1,67 @@
 ---
 library_name: transformers
-license: apache-2.0
-base_model: Shekswess/trlm-stage-2-sft-final-2
+model_name: trlm-135m
 tags:
-- trl
-- dpo
-- preference-alignment
-- reasoning
 - generated_from_trainer
-model-index:
-- name: trlm-stage-3-dpo-final-2
-  results: []
----
-
-# Tiny Reasoning Language Model (trlm-135)
-
-![image/png](https://github.com/user-attachments/assets/5f453496-8180-4cf4-94da-26ebbe1159d4)
-
-## Table of Contents
-
-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)**
-
+- grpo
+- trl
+licence: license
 ---
 
-## Post-Training Pipeline
-<img width="1014" height="563" alt="image" src="https://github.com/user-attachments/assets/195ef389-6aa9-4527-b4f0-bea68c0841ae" />
+# Model Card for trlm-135m
 
----
+This model is a fine-tuned version of [None](https://huggingface.co/None).
+It has been trained using [TRL](https://github.com/huggingface/trl).
 
-## How to use
-
-```bash
-pip install -U transformers accelerate
-```
+## Quick start
 
 ```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`.
+from transformers import pipeline
 
----
-
-## Training
-
-### Model
-
-* **Architecture**: Decoder-only transformer (SmolLM2 backbone which infact is Llama 3 based model).
-* **Parameters**: ~135M.
-* **Precision**: mix-precision (bfloat16) during training.
-
-### Software & Hardware
-
-* **Training Frameworks**: PyTorch (ROCm), Hugging Face Transformers & TRL.
-* **Hardware**: AMD MI300X (192GB VRAM, 224GB RAM).
-
-**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`:
+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"])
+```
 
-| **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                          |
+## Training procedure
 
----
+ 
 
-## Limitations
 
-* **Not production-ready**: hallucinations and logical errors are frequent.
-* **Small size**: limited general knowledge and reasoning depth.
-* **English-only**: multilingual capabilities not explored.
+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).
 
----
+### Framework versions
 
-## Acknowledgements
+- TRL: 0.23.0
+- Transformers: 4.56.1
+- Pytorch: 2.8.0+cu126
+- Datasets: 4.0.0
+- Tokenizers: 0.22.0
 
-- [@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.
+## Citations
 
----
+Cite GRPO as:
 
-## License
+```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},
+}
 
-[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.2",
+  "transformers_version": "4.56.1",
   "use_cache": true,
   "vocab_size": 49154
 }

model.safetensors

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

runs/Sep28_20-28-56_8cb6b1a61c6a/events.out.tfevents.1759091343.8cb6b1a61c6a.390.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a4b028d223ad831a7a83a1f6912243cf65cad3d4a573fc289cca3d3f8c4016c
+size 8198

tokenizer.json

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

training_args.bin

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