File size: 4,406 Bytes
0bbaa80
 
 
 
 
 
 
 
 
 
 
 
eed943d
bda6da5
 
0bbaa80
b06def8
 
0bbaa80
 
 
 
 
 
eed943d
bda6da5
0bbaa80
bda6da5
 
 
0bbaa80
 
 
b06def8
0bbaa80
bda6da5
b06def8
0bbaa80
 
 
 
 
b06def8
 
 
0bbaa80
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bda6da5
0bbaa80
 
 
 
 
 
 
 
 
 
 
bda6da5
0bbaa80
 
 
eed943d
 
 
 
0bbaa80
 
 
 
bda6da5
0bbaa80
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
---
language: en
library_name: transformers
pipeline_tag: text-generation
tags:
  - mini-llm
  - from-scratch
  - small-language-model
---

# Mini-LLM Assistant 40M

A ~40M parameter conversational assistant with a full training pipeline implemented
end-to-end in PyTorch: BPE tokenizer, pretraining and supervised fine-tuning, no
pretrained backbone.

This repository hosts the SFT assistant. The base model it was fine-tuned from is
published separately: [GautierP/MiniLLM-40M-base](https://huggingface.co/GautierP/MiniLLM-40M-base).

## Model details

|                |                                                                                                                                          |
| -------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
| Parameters     | 42,476,032 (input embedding shared with the output head)                                                                                 |
| Architecture   | Transformer decoder: RMSNorm, RoPE, causal attention, SwiGLU                                                                              |
| Context length | 1,024 tokens, twice the 512 of the earlier 1.0-generation runs                                                                        |
| Vocabulary     | 32,768 (BPE)                                                                                                                              |
| Pretraining    | ~4.60B tokens (FineWeb-Edu + Wikipedia), 35,096 steps (~115 tokens/param), base val loss 3.202                                           |
| Fine-tuning    | 2 epochs of SFT on Dolly-15k, No Robots, OpenAssistant oasst2, alpaca-cleaned and identity/small-talk examples, val loss 3.058             |
| Efficiency     | ~115 tokens/param, zero-shot scores at or above the family 70M (1.44B tokens) with 40% fewer parameters                                   |

## Base model

The base repository contains the pretrained checkpoint
(`mini-llm-base-40M-35096steps.pt`) before SFT: 42,476,032 parameters, 32,768-token
BPE vocabulary, ~4.60B tokens (FineWeb-Edu + Wikipedia), 35,096 steps, base val loss
3.202. It is a raw completion model, not an assistant: the base ships without a chat
template and without the assistant's sampling defaults.

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "GautierP/MiniLLM-40M-base"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)

inputs = tok("The capital of France is", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=64, do_sample=True, temperature=0.8, top_k=50)
print(tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))
```

## Usage

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "GautierP/MiniLLM-40M"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)

messages = [{"role": "user", "content": "What is the capital of France?"}]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
out = model.generate(inputs, max_new_tokens=128)
print(tok.decode(out[0][inputs.shape[1]:], skip_special_tokens=True))
```

`trust_remote_code=True` is required, the architecture ships with the repository
(`modeling_mini_llm.py`, `configuration_mini_llm.py`). Sampling defaults
(temperature 0.5, top_k 40, repetition penalty 1.15) live in `generation_config.json`.

The training format is single-turn and fixed; `chat_template.jinja` implements it
exactly:

```
User: <question>
Assistant: <answer>
```

Each question is processed independently, no conversation history is attached.

## Limitations

**Factual accuracy** is limited by model scale: fluent, well-formatted answers do not
imply verified facts, and any factual statement should be treated as unconfirmed.
Generation quality is tracked by the
[Mini-Eval](https://github.com/Mini-llm-lab/Mini-Eval) suite. English only.

## Links

- Base model weights: https://huggingface.co/GautierP/MiniLLM-40M/tree/main/base
- Code, training pipeline and full documentation: https://github.com/Mini-llm-lab/
- Evaluation protocol and scores: https://github.com/Mini-llm-lab/Mini-Eval