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.

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.

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

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 suite. English only.

Links

Downloads last month
874
Safetensors
Model size
42.5M params
Tensor type
F32
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Collection including GautierP/MiniLLM-40M