Atom-1.0-Instruct

Atom-1.0-Instruct is a 29.1M-parameter decoder-only language model, instruction-tuned to write short children's stories from a list of words and story features. It is the supervised fine-tuned (SFT) version of viknesh100/Atom-1.0, trained on TinyStoriesInstruct.

It is small enough to run on a laptop CPU, and it is a hobby-scale model: good for experiments, demos and learning, not for production use.

Model details

Developer	Viknesh (viknesh100)
Model type	Decoder-only transformer (custom atom-gpt architecture)
Parameters	29.1M
Layers / heads / width	8 layers, 8 heads, d_model = 512 (head dim 64)
Feed-forward	SwiGLU, hidden size 1344
Normalisation	RMSNorm (eps 1e-6), pre-norm
Positional encoding	Rotary (RoPE, base 10000, rotate-half)
Embeddings	Tied input/output embeddings, no biases
Context length	512 tokens
Vocabulary	8,192 tokens (tokenizer.json, tokenizers library)
End-of-text token	<|endoftext|> (id 0)
Language	English
Base model	viknesh100/Atom-1.0 (pretrained, step 1525)
Format	safetensors

Note: atom-gpt is a custom architecture. It is not supported by transformers' AutoModel classes. Use the standalone code below.

Intended use

• Generating short, simple, child-level English stories from a set of required words and story features.
• Studying, probing or fine-tuning a small language model.
• Teaching and demos of a from-scratch GPT-style (RMSNorm + RoPE + SwiGLU) model.

Out of scope: factual question answering, chat, code, other languages, or anything where correctness or safety matters. The model only knows the narrow style of TinyStories-like text.

Prompt format

The model was trained on one fixed template. Use it exactly, including the trailing newline after ### Story:.

### Instruction:
Write a short story.
Words: run, shiny, happy
Features: Dialogue
### Story:

• Words: is a comma-separated list of words that should appear in the story.
• Features: is a comma-separated list of story features, written as they appear in the dataset, e.g. Dialogue, BadEnding, MoralValue, Foreshadowing.
• Two optional lines can be added, in this order, after Features:: Summary: ... (a one- or two-sentence plot summary) and Random sentence: ... (a sentence the story should include).
• Generation ends when the model emits <|endoftext|>.

How to use

Install the dependencies:

pip install torch safetensors tokenizers huggingface_hub

Then run the following. The model classes are the same ones used in training.

import json
import torch, torch.nn as nn, torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_model
from tokenizers import Tokenizer

REPO = "viknesh100/Atom-1.0-Instruct"
cfg = json.load(open(hf_hub_download(REPO, "config.json")))
D, L, H, T, VOCAB = cfg["D"], cfg["L"], cfg["H"], cfg["T"], cfg["VOCAB"]
EOS_ID = cfg["eos_token_id"]


class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__(); self.w = nn.Parameter(torch.ones(d)); self.eps = eps
    def forward(self, x):
        xf = x.float()
        return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)).type_as(x) * self.w

def rope_cache(T, hd, base=10000.0):
    inv = 1.0 / (base ** (torch.arange(0, hd, 2).float() / hd))
    f = torch.outer(torch.arange(T).float(), inv)
    return f.cos(), f.sin()

def apply_rope(x, cos, sin):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)

class Attn(nn.Module):
    def __init__(self, d, h):
        super().__init__(); self.h = h
        self.qkv = nn.Linear(d, 3 * d, bias=False); self.o = nn.Linear(d, d, bias=False)
    def forward(self, x, cos, sin):
        B, T_, C = x.shape
        q, k, v = self.qkv(x).view(B, T_, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
        q, k = apply_rope(q, cos[:T_], sin[:T_]), apply_rope(k, cos[:T_], sin[:T_])
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o(y.transpose(1, 2).reshape(B, T_, C))

class SwiGLU(nn.Module):
    def __init__(self, d):
        super().__init__(); hid = int(8 * d / 3 / 64) * 64
        self.w1 = nn.Linear(d, hid, bias=False); self.w3 = nn.Linear(d, hid, bias=False)
        self.w2 = nn.Linear(hid, d, bias=False)
    def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x))

class Block(nn.Module):
    def __init__(self, d, h):
        super().__init__()
        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)
        self.attn, self.mlp = Attn(d, h), SwiGLU(d)
    def forward(self, x, cos, sin):
        x = x + self.attn(self.n1(x), cos, sin)
        return x + self.mlp(self.n2(x))

class GPT(nn.Module):
    def __init__(self, vocab, d, L, h, T):
        super().__init__()
        self.emb = nn.Embedding(vocab, d)
        self.blocks = nn.ModuleList([Block(d, h) for _ in range(L)])
        self.norm = RMSNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
        self.head.weight = self.emb.weight  # tied embeddings
        cos, sin = rope_cache(T, d // h)
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
    def forward(self, idx):
        x = self.emb(idx)
        for b in self.blocks: x = b(x, self.cos, self.sin)
        return self.head(self.norm(x))


device = "cuda" if torch.cuda.is_available() else "cpu"
model = GPT(VOCAB, D, L, H, T)
load_model(model, hf_hub_download(REPO, "model.safetensors"))
model.to(device).eval()
tok = Tokenizer.from_file(hf_hub_download(REPO, "tokenizer.json"))


@torch.no_grad()
def generate(prompt, max_new_tokens=400, temperature=0.7, top_k=50):
    x = torch.tensor([tok.encode(prompt).ids], device=device)
    out = []
    for _ in range(max_new_tokens):
        logits = model(x[:, -T:])[:, -1, :] / temperature
        v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
        logits[logits < v[:, [-1]]] = float("-inf")
        nxt = torch.multinomial(F.softmax(logits, dim=-1), 1)
        if nxt.item() == EOS_ID:
            break
        out.append(nxt.item())
        x = torch.cat([x, nxt], dim=1)
    return tok.decode(out)


prompt = """### Instruction:
Write a short story.
Words: run, shiny, happy
Features: Dialogue
### Story:
"""
print(generate(prompt))

Example output

Prompt: Words: run, shiny, happy / Features: Dialogue

Once upon a time, there was a little girl named Lily. She had a shiny toy car that she loved to play with. One day, Lily went to the park to play with her toys. She saw a boy with a shiny red car.

"Hi, what's your name?" asked Lily. "My name is Timmy," said the boy. "That's cool! Can I play with you?" asked Lily. "Sure! Let's run around the park and have fun," said Timmy.

Lily and Timmy ran around the park and played on the slide. They had so much fun together. When it was time to go home, Lily said goodbye to Timmy and promised to come back and play again.

The model also appended a Summary: ... line after this story before ending. See Limitations.

Training

Pipeline

1. Pretraining: viknesh100/Atom-1.0, a base model trained from scratch (checkpoint at step 1525, validation loss 1.587 on its own pretraining validation set).
2. Supervised fine-tuning (this model): full-parameter fine-tuning of all 29.1M weights on instruction-formatted stories.

Data

• Source: roneneldan/TinyStoriesInstruct (TinyStories-Instruct-train.txt and -valid.txt), a synthetic dataset of short stories written for young children, each annotated with required words, features, a summary and sometimes a random sentence.
• Examples: the first 200,000 stories of the train file were parsed. Stories whose prompt plus story exceeded 512 tokens were dropped, leaving 191,031 training examples (51.9M tokens, average 272 tokens per example). 2,000 examples from the validation file were used for validation.
• Format: each story was rewritten into the prompt template above, followed by the story text and <|endoftext|>.
• Loss masking: the loss was computed on the response tokens only (the story and the end token, about 74% of all tokens). Prompt tokens were ignored.

Hyperparameters

Epochs	2 (5,968 optimizer steps, 2,984 per epoch)
Batch size	64 sequences per step (2 micro-batches of 32)
Optimizer	AdamW, betas (0.9, 0.95), weight decay 0.01 on matrices only
Learning rate	1e-4 peak, cosine decay to 1e-5, 50 warm-up steps
Gradient clipping	1.0
Precision	fp16 mixed precision with GradScaler
Sequence length	512 (padded, padding masked from the loss)
Shuffling	Seeded per epoch (resumable)
Hardware	1x NVIDIA Tesla T4 (Google Colab), about 0.8 s per optimizer step

Evaluation

Metric	Value
Validation loss (response tokens, 2,000 held-out examples)	1.2423
Validation perplexity	about 3.46
Best checkpoint	step 5968 (the last step)

The best checkpoint was selected on the same validation set that is reported here, so the number is slightly optimistic. No benchmark beyond this held-out loss was run, and word and feature adherence has not been measured systematically.

Limitations and biases

• Tiny model, narrow domain. At 29M parameters it writes simple, repetitive, TinyStories-style prose. It does not know facts and will not follow general instructions.
• Weak instruction following. Required words and features are often used, but not guaranteed.
• Trailing Summary: line. Because the training data pairs stories with summaries, the model sometimes writes a Summary: ... line after the story and before <|endoftext|>. Cut the output at the first \nSummary: if you only want the story.
• Context limit. 512 tokens in total, prompt included. Longer stories will be cut off or degrade.
• Synthetic-data bias. The training stories were machine-generated, so the model inherits their style, vocabulary, names and typical plots. Features like BadEnding can produce bleak or moralising endings.
• Not safety-tuned. No filtering, red-teaming or alignment was done beyond what the source dataset provides.

Files in this repository

File	Purpose
model.safetensors	Final weights (best checkpoint, step 5968). Needed for inference.
config.json	Architecture settings (D, L, H, T, VOCAB), EOS id, prompt template, training step and validation loss. Needed.
tokenizer.json	Tokenizer. Needed.
last.safetensors, last_config.json	Last resumable checkpoint (same step as the final weights).
train_state.pt	Optimizer and scaler state, only for resuming training.
ckpt_best.pt, ckpt_last.pt	Original pickled training checkpoints.
sft_train_*, sft_val_*	Pre-tokenized fine-tuning data and loss masks.

For inference you only need model.safetensors, config.json and tokenizer.json (about 120 MB).

Acknowledgements and citation

The training data comes from the TinyStories work:

@article{eldan2023tinystories,
  title   = {TinyStories: How Small Can Language Models Be and Still Speak Coherent English?},
  author  = {Eldan, Ronen and Li, Yuanzhi},
  journal = {arXiv preprint arXiv:2305.07759},
  year    = {2023}
}

If you use Atom-1.0-Instruct, please link back to this repository.