xxrickyxx/Ailo340m-v4

AILO-340M-v4

A 341M-parameter English chat model (32 layers, GPT-2 tokenizer) trained end-to-end on a single consumer AMD GPU (RX 6750 XT, 12 GB) with a hand-written backward engine — PyTorch autograd is unstable on this hardware for deep transformers.

Small enough to run on a Raspberry Pi Zero 2 W (Q4, ~220 MB) at a few tokens per second.

How it was built

A four-stage pipeline, each stage adding one capability:

1. Logit distillation — a 16-layer student distilled from GPT-J 6B at the logit level (shared GPT-2 vocabulary; teacher top-64 logits precomputed offline on CPU because the teacher does not fit in 12 GB of VRAM). Held-out perplexity 39.9 → 32.1.
2. Discursive chat SFT — fine-tuned on ~2,500 conversational Q&A generated by Gemma 3 4B, loss masked to answer tokens only.
3. Depth expansion — 16 → 32 layers via LLaMA-Pro identity blocks (the expanded model starts mathematically identical), then re-trained on a mixed corpus: OpenAssistant, No-Robots, Dolly, Python code, MBPP, and SQuAD (which teaches exact extraction of facts and figures from a provided passage).
4. Math & instruction-following — arithmetic drills, orca-math word problems, GSM8K, and constraint-following examples.

Usage

Ollama (GGUF included)

ollama run Alieno/ailo-340m-v4

Transformers

from transformers import AutoModelForCausalLM, GPT2TokenizerFast
import torch

m = AutoModelForCausalLM.from_pretrained("xxrickyxx/Ailo340m-v4", trust_remote_code=True).eval()
tok = GPT2TokenizerFast(vocab_file="vocab.json", merges_file="merges.txt")

prompt = "<|user|>\nWhy is the sky blue?\n<|assistant|>\n"
ids = tok(prompt, return_tensors="pt").input_ids
print(tok.decode(m.generate(ids, max_new_tokens=120)[0][ids.shape[1]:], skip_special_tokens=True))

Chat format: <|user|>\n{question}\n<|assistant|>\n, ending with <|end|>. Recommended sampling: temperature 0.3, top_k 20, top_p 0.9, repeat_penalty 1.25.

What it is good at

• Conversational answers — fluent, on-topic, natural register.
• Extracting facts from a provided passage, including numbers and dates. Given a Wikipedia paragraph it will correctly return "330 metres tall" or "7 November 1867". This makes it a reasonable reader/summariser in a retrieval pipeline.
• Following output constraints — e.g. "summarise in exactly five words" (where some larger models fail).
• Basic Python structure — writes syntactically plausible functions, though often with bugs.

What it is NOT good at — please read

• Arithmetic is unreliable. Measured exact-match accuracy on unseen operands: addition 0%, subtraction 0%, percentages 0%. It will confidently produce wrong numbers (25 + 17 = 51).
• Multi-step reasoning fails on word problems and sequences.
• Factual recall from memory is unreliable — it invents dates, nationalities and details. It is markedly better when facts are supplied in the prompt.

The gap against Qwen2.5-0.5B is not mainly architectural: that model saw on the order of 18 trillion pre-training tokens, while AILO's base comes from distillation over a corpus many orders of magnitude smaller. Supervised fine-tuning shapes behaviour; it does not substitute for pre-training scale.

Use it as a small, local conversational reader — with retrieval for facts and a calculator for arithmetic. Do not use it as a source of truth.

Specs

Parameters	341M (32 layers, 768 hidden, 12 heads)
Tokenizer	GPT-2 BPE (50257)
Context	1024
Teachers	GPT-J 6B (logit-KD) · Gemma 3 4B (chat SFT)
Trained on	1× AMD RX 6750 XT (12 GB), DirectML, manual backward

License

CC BY-NC-SA 4.0 (non-commercial). Distilled in part from GPT-J (Apache-2.0) and Gemma 3 (Gemma Terms of Use); downstream use must respect the teachers' terms. Training data includes OpenAssistant (Apache-2.0), Dolly (CC BY-SA 3.0), No-Robots (CC BY-NC 4.0), SQuAD (CC BY-SA 4.0), GSM8K (MIT) and orca-math (MIT).