AtomLM

A decoder-only Transformer language model built completely from scratch in PyTorch for learning, research, and experimentation.

AtomLM is a research project exploring how modern Large Language Models are built from tokenization and attention mechanisms to distributed training, instruction tuning, reasoning, and tool use.

Unlike projects that simply fine-tune existing LLMs, AtomLM is trained from scratch with a custom architecture and training pipeline.

Features

• Decoder-only Transformer
• RoPE (Rotary Positional Embeddings)
• FlashAttention
• Grouped Query Attention (GQA)
• RMSNorm
• SwiGLU Feed Forward Network
• Mixed Precision (FP16)
• Distributed Data Parallel (DDP)
• Byte-Level BPE tokenizer
• Cosine Learning Rate Scheduler
• Gradient Accumulation
• Checkpointing & Resume Training

Documentation

New here? Start with docs/GETTING_STARTED.md — it walks you from clone to a trained model in a few commands.

Want to change the model size, datasets, or mixing ratios? See docs/CONFIGURATION.md.

The whole pipeline is driven from one CLI entry point:

python main.py download     # download datasets
python main.py tokenize     # train the BPE tokenizer
python main.py build-data   # build the token stream
python main.py train        # pretrain the model (--smoke, --nproc N)
python main.py eval         # evaluation suite
python main.py chat         # interactive chat
python main.py all          # run everything in order

Current Base Model

Property	Value
Parameters	52.3M
Layers	16
Hidden Size	512
Attention Heads	8
KV Heads	2 (GQA)
Context Length	1024
Vocabulary	8K Byte-Level BPE
Framework	PyTorch

Current Status

The current checkpoint is a base pretrained model.

It has learned:

• English grammar
• Story generation
• Basic mathematical reasoning patterns
• Long-form text generation
• Structured reasoning format

It has not yet been instruction-tuned or aligned, so factual accuracy, reasoning quality, and instruction following remain limited.

Roadmap

Foundation

• Prototype Language Model (2.2M parameters)
• Transformer Baseline (8.8M parameters, TinyStories pretraining)
• AtomLM Base (52M parameters)
  • Decoder-only Transformer
  • RoPE
  • FlashAttention
  • Grouped Query Attention (GQA)
  • RMSNorm
  • SwiGLU
  • Pretrained on 343M tokens

Model Training

• Continued Pretraining
• Supervised Instruction Fine-Tuning (SFT)
• Reasoning Fine-Tuning
• Tool-Use Training (Search, Python, Calculator)
• Preference Alignment (RLHF / DPO / GRPO)

Optimization

• Quantization & Efficient Inference
• Mobile Deployment
• Long-Context Support
• Retrieval-Augmented Generation (RAG)

Scaling

• AtomLM-100M
• AtomLM-300M
• AtomLM-1B

Training logs, generations, loss curves, and experiment history are available in:

results/LOGS.md

Project Structure

AtomLM/
├── main.py                    # CLI entry point — run the whole pipeline
├── requirements.txt
├── docs/
│   ├── GETTING_STARTED.md     # start here: install & run from scratch
│   └── CONFIGURATION.md       # edit config, datasets, mixing ratios
├── src/                       # core pipeline (a Python package)
│   ├── config.py              # model + training + dataset settings
│   ├── download_data.py       # stage 1: download datasets
│   ├── tokenizer.py           # stage 2: train the BPE tokenizer
│   ├── build_data.py          # stage 3: build the token stream
│   ├── dataset.py             # PyTorch dataset
│   ├── model.py               # the Transformer (RoPE, GQA, SwiGLU)
│   ├── train.py               # stage 4: pretrain (DDP, resume)
│   ├── evaluate.py            # stage 5: evaluation suite
│   ├── inference.py           # chat
│   └── utils.py               # shared helpers (seed, device, DDP)
├── cpt/                       # continued pretraining scripts
├── checkpoints/               # model checkpoints (gitignored)
├── tokenizer/                 # trained vocab.json / merges.txt
├── data/                      # raw + processed data (gitignored)
├── results/                   # training logs (LOGS.md)
└── README.md

Author

Ganesh Kumar
MIT