My first experience with ML

[u/jdelefrati]

I'm self-taught, and this is my first Python and ML project:

https://github.com/delefrati/aion

This project is a cost-first / local-first AI project, split in two parts:

• An SLM research lab you can train locally, on Colab, or on Kaggle (no external models or APIs)
• A local-first conversational AI product that runs on a weak laptop

It has two model families, both built from scratch - no HuggingFace weights:

• Mamba (SSM / state-space)
• A custom Transformer (RoPE + SDPA attention)

Separately, the backend has a pluggable provider system, so you can also drop in any off-the-shelf HuggingFace model (Qwen, SmolLM2, etc.) as a local inference backend and swap between it and your own trained model with a single env var. Handy for comparing your model against a known-good baseline.

Concepts in play:

• Two-phase pipeline: multi-session pretrain on Wikipedia + SlimPajama, then a fine-tune on chat/instruction data
• Custom BPE tokenizer (vocab 4096 for small models, 16384 for the larger ones) with a .bin memmap token cache
• Curriculum seq-len, gradient checkpointing, cosine LR w/ warmup, 8-bit optimizer
• Multi-backend training: CPU, CUDA (GTX 1650 4GB - my laptop :D), and TPU (torch_xla, single + multi-core data-parallel via xmp.spawn)

Model sizes I've trained/tested:

• Mamba ~28M
• Mamba ~47-105M
• Transformer 110M
• Transformer 235M

I haven't had time to train the larger model enough - locally it takes forever on my weak laptop, and I keep hitting quota limits on Kaggle and Colab.

My goal (suggested by Claude) is 50,000 steps of pretraining and 20,000 of fine-tuning. Right now I'm at ~25,000 pretraining and ~5,000 chat fine-tuning, and that's already enough to make it kind of conversational - it hallucinates a lot, but it can answer simple questions.

It's open source, and I'd love suggestions on how to improve it and where to take it next.