axiomic banner

GPT-S-1.4M

GPT-S-1.4M a first-generation model in the GPT-S small-model family: 1.4M parameters, 6B training tokens, a custom 4K tokenizer, 5 layers, and all new Exclusive Grouped-query Attention (XGQA), trained from scratch on a 5-source corpus.

See how it compares to similar models here: Open SLM Leaderboard

Benchmarks

All evaluations use zero-shot multiple-choice scoring. Normalized accuracy is reported where available.

Benchmark Score
HellaSwag 26.89%
ARC-Easy 31.57%
ARC-Challenge 21.93%
ARC Average 26.75%
PIQA 55.17%
ArithMark 2 25.16%

Architecture

Component Details
Position encoding RoPE, theta=2,500
Normalization RMSNorm
Feed-forward SwiGLU
Attention Exclusive Grouped-query attention, 4 query heads / 2 KV heads
Embeddings Weight tied
Context length 384 tokens

Config

vocab_size       = 4,096
hidden_size      = 128
num_layers       = 5
num_heads        = 4
num_kv_heads     = 2
head_dim         = 32
intermediate     = 341
block_size       = 384
rope_theta       = 2,500

Training

GPT-S-1.4M was trained from scratch for 6B tokens on a mixed English corpus built around educational web text, synthetic textbook-style material, and higher-quality web text.

Source Dataset Mix Purpose
FineWeb-Edu HuggingFaceFW/fineweb-edu 15% Primary educational web text
Cosmopedia v2 HuggingFaceTB/smollm-corpus 30% Synthetic textbook-style coverage
FineWeb-HQ epfml/FineWeb-HQ 20% Higher-quality general web text
Ultra-FineWeb QA openbmb/Ultra-FineWeb L3 English QA slice 20% Question-answer style web text
Ultra-FineWeb Multi-style openbmb/Ultra-FineWeb L3 English multi-style slice 15% Broader writing-style coverage

Hyperparameters

Hyperparameter Value
Optimizer AdamW
Adam betas 0.9 / 0.95
Weight decay 0.01
Peak learning rate 3.5e-3
Minimum learning rate 0
LR schedule Warmup-stable-decay
Warmup steps 2,000
Decay start 80% of configured training run
Training tokens 6B
Total batch size 294,912 tokens
Microbatch 256 x 384 tokens
Gradient accumulation steps 3
Gradient clipping 1.0
Precision bfloat16 autocast

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "AxiomicLabs/GPT-S-1.4M"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

prompt = "The future of AI is"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    output = model.generate(
        **inputs,
        max_new_tokens=80,
        do_sample=True,
        temperature=0.8,
        top_p=0.95,
        repetition_penalty=1.1,
        no_repeat_ngram_size=4,
    )

print(tokenizer.decode(output[0], skip_special_tokens=True))

Limitations

This is a very small base language model. It is not instruction tuned, has limited factual capacity, and uses a 384-token context window.

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

Datasets used to train AxiomicLabs/GPT-S-1.4M

Space using AxiomicLabs/GPT-S-1.4M 1