Files
fluxion-370m-instruct/README.md

364 lines
13 KiB
Markdown
Raw Permalink Normal View History

---
library_name: transformers
tags:
- code
- math
- text-generation-inference
- slm
- instructmodel
license: apache-2.0
language:
- en
---
# ⚡ Fluxion-370M-Instruct
### *Trained on 0.227T tokens. Competing with models trained on 28T.*
<br/>
[![arXiv](https://img.shields.io/badge/arXiv-coming%20soon-b31b1b?style=flat-square&logo=arxiv)](<!-- ADD -->)
[![HuggingFace](https://img.shields.io/badge/🤗%20Model-nareshmeena12%2Ffluxion--370m--instruct-yellow?style=flat-square)](https://huggingface.co/nareshmeena12/fluxion-370m-instruct)
[![GitHub](https://img.shields.io/badge/GitHub-Code-black?style=flat-square&logo=github)](<!-- ADD -->)
[![License](https://img.shields.io/badge/License-Apache%202.0-blue?style=flat-square)](https://www.apache.org/licenses/LICENSE-2.0)
<div align="center">
![Math and Code Benchmarks](https://github.com/nareshmeena12/fluxion-370m-instruct/blob/main/assests/math%20and%20code%20benchmarks.png)
</div>
---
## What is Fluxion?
**Fluxion-370M-Instruct** is a 370M-parameter instruction-tuned language model built from scratch, focused on **mathematical reasoning and code generation**. Its core design philosophy is simple: *do more with less*.
Trained on just **227B tokens** — a fraction of what comparable models consume — Fluxion achieves competitive performance on math and coding benchmarks against models that have seen **1828× more data**. On token efficiency (MATH-500 score per trillion training tokens), Fluxion ranks **1st** across all models in its class by a dominant margin.
The model went through three training stages: large-scale pretraining, supervised fine-tuning for instruction following, and GRPO reinforcement learning specifically targeting mathematical reasoning with rule-based rewards.
---
## ⚡ Token Efficiency — The Core Story
| Model | Training Tokens | MATH-500 | **MATH-500 / Trillion Tokens** |
|:---|:---:|:---:|:---:|
| **Fluxion-370M** | **0.227T** | **17.40%** | **🥇 76.7×** |
| MobileLLM-R1-360M | 4.2T | 25.60% | 6.1× |
| Gemma3-270M | 6T | 8.20% | 1.4× |
| Granite-350M | 11T | 19.20% | 1.7× |
| Qwen2.5-0.5B | 18T | 29.80% | 1.7× |
| LFM2.5-350M | 28T | 13.40% | 0.5× |
| SmolLM2-360M | 4T | 3.40% | 0.9× |
**Fluxion delivers 76.7× higher MATH-500 score per trillion training tokens than any other model in this comparison.** This is a structural advantage from architecture choices, custom tokenization, and a high-signal data pipeline — not a measurement artifact.
---
## 🏗️ Architecture
Fluxion is a clean, modern decoder-only transformer. No hybrid layers, no custom operators — just well-chosen standard components assembled deliberately.
```
FluxionForCausalLM
├── Embedding (vocab_size=50,000 · hidden_size=1024)
├── 26 × DecoderBlock
│ ├── RMSNorm (ε = 1e-6)
│ ├── GroupedQueryAttention
│ │ ├── 16 Query heads · 8 KV heads (GQA 2:1)
│ │ └── RoPE (θ = 500,000) + NTK scaling
│ ├── RMSNorm (ε = 1e-6)
│ └── SwiGLU FFN (hidden → 3072 → hidden)
└── RMSNorm → LM Head (weight-tied to embedding)
```
### Hyperparameters
| Parameter | Value |
|:---|:---:|
| **Total Parameters** | ~370M |
| **Layers** | 26 |
| **Hidden Size** | 1024 |
| **FFN Intermediate Size** | 3072 |
| **Attention Heads (Q)** | 16 |
| **KV Heads (GQA)** | 8 |
| **GQA Group Ratio** | 2 : 1 |
| **Head Dimension** | 64 |
| **Context Length** | 4096 |
| **Vocabulary Size** | 50,000 |
| **RoPE Base θ** | 500,000 |
| **RoPE Scaling** | NTK |
| **Normalization** | RMSNorm (ε = 1e-6) |
| **Activation** | SwiGLU |
| **Attention** | Flash Attention |
| **Weight Tying** | Yes (embed ↔ lm\_head) |
| **Attention Bias** | No |
| **FFN Bias** | No |
| **dtype** | bfloat16 |
### Design Choices & Rationale
**Grouped Query Attention (16Q / 8KV):** Halves KV cache memory vs standard MHA with no meaningful quality loss at this scale. Larger effective batch sizes at both training and inference time.
**RoPE with θ = 500,000 + NTK scaling:** High base frequency improves length generalisation beyond the training window. NTK scaling further stabilises attention patterns at longer sequences without fine-tuning.
**SwiGLU activation:** Consistently stronger than GeLU/ReLU on reasoning tasks. The gated structure gives the FFN richer expressiveness per parameter.
**Custom BPE tokenizer (50k vocab):** Trained from scratch on 100GB of domain-focused text — a curated subset of the pretraining corpus weighted toward mathematics, code, and English prose. A smaller, domain-tuned vocabulary means fewer tokens per math expression and code snippet, directly improving effective context utilisation and reducing sequence lengths vs generic 128k tokenizers.
**Weight tying (embed ↔ lm\_head):** Saves ~50M parameters, keeps the model compact, and forces the input and output representation spaces to align.
**Flash Attention:** Full Flash Attention used throughout training for memory efficiency and throughput. TF32 enabled for matrix multiplications.
---
## 📊 Benchmarks
All results are **0-shot** unless noted. Evaluation is strict.
### Math & Code
| Model | Tokens | GSM8K | MATH-500 | HumanEval | MBPP | **Avg** |
|:---|:---:|:---:|:---:|:---:|:---:|:---:|
| Qwen2.5-0.5B | 18T | 43.44% | 29.80% | 30.49% | 38.91% | 35.66% |
| Granite-350M | 11T | 31.69% | 19.20% | 29.88% | 42.02% | 30.70% |
| LFM2.5-350M | 28T | 32.68% | 13.40% | 10.37% | 5.45% | 15.48% |
| MobileLLM-R1-360M | 4.2T | 23.81% | 25.60% | 19.51% | 17.51% | 21.61% |
| **Fluxion-370M** | **0.227T** | **29.57%** | **17.40%** | **16.46%** | **20.23%** | **20.92%** |
| SmolLM2-360M | 4T | 8.87% | 3.40% | 17.68% | 35.02% | 16.24% |
| Gemma3-270M | 6T | 7.58% | 8.20% | 11.59% | 19.07% | 11.61% |
### General Reasoning
| Model | ARC | PIQA | OBQA | CSQA | MMLU | Wino | IFEval | **Avg** |
|:---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| Qwen2.5-0.5B | 50.9% | 70.1% | 36.2% | 42.3% | 45.9% | 51.1% | 74.1% | 52.9% |
| Granite-350M | 51.8% | 68.9% | 37.8% | 40.5% | 33.8% | 53.8% | 83.2% | 52.8% |
| LFM2.5-350M | 49.4% | 67.0% | 31.6% | 33.4% | 42.2% | 52.0% | 88.3% | 52.0% |
| SmolLM2-360M | 45.9% | 71.5% | 35.2% | 29.0% | 24.9% | 54.8% | 77.5% | 48.4% |
| Gemma3-270M | 43.7% | 66.6% | 30.6% | 36.9% | 24.7% | 49.7% | 71.5% | 46.2% |
| MobileLLM-R1-360M | 39.1% | 58.6% | 26.2% | 27.6% | 26.7% | 50.8% | 71.2% | 42.9% |
| **Fluxion-370M** | **36.3%** | **61.6%** | **29.4%** | **27.3%** | **25.3%** | **51.9%** | **66.9%** | **42.7%** |
### Rankings Summary
| Category | Score | Rank |
|:---|:---:|:---:|
| Math Average (GSM8K + MATH-500) | 23.49% | **4th** |
| Code Average (HumanEval + MBPP) | 18.35% | **5th** |
| Math + Code Combined | 20.92% | **4th** |
| General Average | 39.19% | 7th |
| **Token Efficiency (MATH-500 / T tokens)** | **76.7×** | **🥇 1st** |
> Fluxion's design target is math and code under tight data constraints. On general benchmarks like IFEval and MMLU it trails models trained on significantly more diverse data — this is expected and intentional. On token efficiency it has no close competitor.
---
## 💬 Chat Template & Tokenizer
Fluxion-Instruct uses the **ChatML** format:
```
<|im_start|>system
You are a helpful assistant.<|im_end|>
<|im_start|>user
{your prompt here}<|im_end|>
<|im_start|>assistant
```
### Special Tokens
| Token | ID | Role |
|:---|:---:|:---|
| `<\|im_start\|>` | 100278 | Turn start marker |
| `<\|im_end\|>` | 100279 | Turn end marker |
| `<\|pad\|>` | — | Padding token |
### Stopping Token IDs
Always pass `eos_token_id=[1, 4]` during generation. Without these, the model may continue generating beyond the intended response boundary:
```python
output = model.generate(
**inputs,
eos_token_id=[1, 4],
...
)
```
Token ID `1` is the primary EOS token and token ID `4` is the `<|im_end|>` marker. Including both ensures clean turn termination in all inference scenarios.
### Tokenizer
Fluxion uses a **custom BPE tokenizer** trained from scratch on ~100GB of domain-focused text — a curated subset of the pretraining corpus weighted toward mathematics, code, and English prose. With a **50,000 token vocabulary**, it is significantly more compact than general-purpose tokenizers (32k128k). This means:
- Fewer tokens per math expression and code snippet
- Better effective context utilisation within the 4096 token window
- Faster embedding lookups and a smaller lm\_head projection
---
## 🚀 Quick Start
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained(
"nareshmeena12/fluxion-370m-instruct",
torch_dtype=torch.bfloat16,
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(
"nareshmeena12/fluxion-370m-instruct"
)
messages = [
{"role": "system", "content": "You are a helpful assistant that solves math problems step by step."},
{"role": "user", "content": "What is the sum of the first 100 natural numbers?"}
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
with torch.no_grad():
output = model.generate(
**inputs,
max_new_tokens=512,
eos_token_id=[1, 4],
do_sample=True,
top_p=0.9,
)
response = tokenizer.decode(
output[0, inputs["input_ids"].shape[1]:],
skip_special_tokens=True
)
print(response)
```
> **Note:** Always include `eos_token_id=[1, 4]` in all `generate()` calls regardless of quantization mode.
---
## 🗂️ Training
Fluxion was trained in three stages.
### Stage 1 — Pretraining
Large-scale unsupervised pretraining on a curated mix of math, code, and English text.
| Parameter | Value |
|:---|:---:|
| **Total Tokens** | ~227B (0.227T) |
| **Steps** | 264,000 |
| **Batch Size** | 10 sequences |
| **Gradient Accumulation** | 21 steps |
| **Effective Batch** | ~860K tokens/step |
| **Learning Rate** | 3e-4 |
| **Min Learning Rate** | 9e-5 |
| **LR Schedule** | Cosine decay |
| **Warmup Steps** | 2,000 |
| **Sequence Length** | 4,096 |
| **Optimizer** | AdamW |
| **Precision** | bfloat16 + TF32 |
| **Attention** | Flash Attention |
| **Positional Encoding** | RoPE (θ=500k) + NTK scaling |
**Pretraining data:** *(to be added)*
### Stage 2 — Supervised Fine-Tuning (SFT)
Instruction tuning using ChatML format to teach the model to follow user instructions across math, code, and general tasks.
| Parameter | Value |
|:---|:---:|
| **Format** | ChatML |
| **Data** | *(to be added)* |
### Stage 3 — GRPO (Math Reasoning)
Group Relative Policy Optimisation applied specifically to strengthen mathematical reasoning.
| Parameter | Value |
|:---|:---:|
| **Algorithm** | GRPO |
| **Reward Signal** | Rule-based answer correctness |
| **Target Domain** | Mathematics |
| **Data** | *(to be added)* |
Rule-based rewards verify final answer correctness against ground truth — no LLM judge, no model-in-the-loop. This keeps the reward signal clean, fast, and free from reward hacking through verbosity or style.
---
## 🎯 Intended Use
**Fluxion-370M-Instruct is designed for:**
- Mathematical problem solving and step-by-step reasoning
- Code generation and explanation (Python, general purpose)
- Instruction-following tasks in math and code domains
- Fine-tuning base for domain-specific math/code applications
- Efficient deployment — capable performance on modest hardware
- Research into token-efficient pretraining and RL for reasoning
**Fluxion-370M-Instruct is NOT designed for:**
- Knowledge-intensive QA requiring broad world knowledge
- Long-document tasks beyond 4096 tokens
- Multilingual use cases
- Safety-critical production deployments without additional alignment
---
## ⚠️ Limitations
- **Narrow training data scope.** 0.227T tokens is intentionally small and domain-focused. The model has limited world knowledge outside math, code, and general English.
- **Context length is 4096 tokens.** Not suitable for long-document tasks without chunking strategies.
- **General benchmarks lag.** IFEval (66.9%) and MMLU (25.3%) reflect that the model was not optimised for instruction following diversity or broad knowledge recall — this is a deliberate trade-off.
- **Custom tokenizer.** The 50k BPE vocabulary is domain-tuned. Rare tokens, non-English text, and highly specialised symbols outside the training distribution may tokenise suboptimally.
- **Small scale.** At 370M parameters, the model is outperformed by larger models on complex multi-step reasoning chains requiring deep world knowledge.
- **Always set `eos_token_id=[1, 4]`.** Without explicit stopping token IDs the model may continue generating past the intended response boundary.
---
## 📰 News
- `<!-- ADD DATE -->` — Fluxion-370M-Instruct released on Hugging Face
- `<!-- ADD DATE -->` — Training code released on GitHub
- `<!-- ADD DATE -->` — Technical report available on arXiv
---
## 📄 Citation
```bibtex
@misc{fluxion2025,
title = {Fluxion-370M: Token-Efficient Pretraining and Reinforcement Learning for Mathematical and Code Reasoning},
author = {Meena, Naresh},
year = {2025},
note = {<!-- ADD: arXiv link or report URL -->}
}
```
---
## 📬 Contact
<!-- ADD: your contact details / social links -->
---
<div align="center">
*Less data. Less compute. More signal.*
**[🤗 Model](https://huggingface.co/nareshmeena12/fluxion-370m-instruct) · [💻 Code](<!-- ADD -->) · [📄 Paper](<!-- ADD -->)**
</div>