GPT-X2-125M

GPT-X2-125M is the second generation GPT-X model built on top of the T-X2 architecture. 125M parameters, 75B tokens, custom 32K tokenizer, 30 layers. Trained from scratch using a 4-source progressive curriculum with AST-normalized code, achieving near state-of-the-art performance on both natural language and structured reasoning benchmarks at the 125M scale on a fraction of the training data.

Results

GPT-X2-125M achieves competitive performance with leading models despite using significantly less data. Most notably, it matches SmolLM2-135M within ~1 point on aggregate while using ~27x fewer tokens and having ~8% fewer parameters.

Evaluated with an internal harness modeled on EleutherAI/lm-eval-harness; all benchmarks are zero-shot.

UPDATE For future models we will be dropping Winogrande from our benchmark suite as it fails to meet our standards for effective evaluation at this scale

LogicMark and ArithMark are procedural benchmarks designed to evaluate structured reasoning and arithmetic generalization across increasing difficulty levels.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_name = "AxiomicLabs/GPT-X2-125M"

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    trust_remote_code=True,
    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=120,
        do_sample=True,
        temperature=0.8,
        top_p=0.95,
        repetition_penalty=1.1,
        no_repeat_ngram_size=4,
    )

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

What's New in v2

Architecture

Config

vocab_size     = 32,768    (custom BPE trained on FineWeb-Edu)
n_layer        = 30
n_head         = 9         (query heads)
n_kv_heads     = 3         (key-value heads, 3:1 GQA)
n_embd         = 576
head_dim       = 64
intermediate   = 1,536     (SwiGLU, 2.67x ratio)
block_size     = 1,024
rope_theta     = 100,000
total params   = 125,081,664

Parameter Breakdown

Training

Data

The curriculum gradually introduces specialized data (math/code) after the model learns core language, improving reasoning without harming fluency. Four data sources mixed via a progressive curriculum:

• Tokens: 75B (143,051 steps x 524,288 tokens/step)
• Tokenizer: Custom 32K BPE trained on 50GB of FineWeb-Edu
• Final val loss: 2.7525 (FineWeb-Edu held-out)

Progressive Data Curriculum

Rather than a fixed mixture, data sources are introduced progressively to let the model build foundational language capabilities before adding specialized data:

FineMath and code are drip-fed at 1% from the start so the model sees early examples, then ramped to peak during the stable LR phase for maximum learning, and gradually tapered back toward primary sources during the LR decay phase.

AST-Based Code Normalization

Raw Python code is converted to a compact pseudocode representation (TinyDSL) before tokenization using AST parsing. This achieves ~1.25x token compression compared to raw code, letting the model learn programming reasoning without the need for code-specific tokens and reducing understanding overhead.

The normalizer:

• Strips comments, docstrings, whitespace, and syntactic noise
• Replaces Python builtins with full English words (len -> length, str -> string, etc.)
• Uses natural language keywords with spaces (end function, for else, list comprehension, etc.)
• Expands comprehensions, decorators, exception handling, and all Python AST node types

Example:

# Raw Python
def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n - 1) + fibonacci(n - 2)

# Normalized TinyDSL
# function fibonacci n 
# begin
# if n <= 1
# begin
# return n
# end
# return call fibonacci n - 1 + call fibonacci n - 2 
# end function

Optimization

• Optimizer: AdamW (betas=0.9/0.95, weight_decay=0.1 < 31Bt then 0.01)
• Learning rate: 1.5e-3 max, decays to 0
• Schedule: WSD -- 2,000 step warmup, stable phase (80%), linear decay to 0 over final 20%
• Batch size: 524,288 tokens (micro_batch=8, seq_len=1024, 64 grad accumulation)
• Precision: bfloat16 mixed precision
• Gradient clipping: 1.0

Hardware

• 1x RTX 3080 Ti
• Training time: ~500 hours

Design Decisions

• 30 layers x 576 hidden -- 3 more layers than v1, made possible by smaller 32K vocab. Depth is the primary driver of quality at 125M scale (SmolLM, MobileLLM).
• Custom 32K tokenizer -- Trained on FineWeb-Edu for ~9% better compression than GPT-2 BPE. Fewer vocab entries = smaller embedding table = more params for transformer layers.
• rope_theta=100K -- Matches SmolLM2-135M. Better extrapolation and long-range dependency modeling.
• 1.5e-3 learning rate -- GPT-X v1 used 6e-4 which was too conservative. SmolLM proved 3e-3 @ 1M batch works at 135M scale; thus used 1.5e-3 at 524288 batch.
• GQA 3:1 -- 9 query heads, 3 KV heads. Saves attention parameters reinvested into SwiGLU capacity.
• QK-Norm -- RMSNorm on Q and K before RoPE.
• z-loss -- Prevents logit magnitude drift (PaLM, T5) used early.
• Progressive curriculum -- FineMath and NPset-Python introduced at 1% early, ramped to peak during stable LR, tapered during decay. Lets the model build language foundations first, then learn specialized reasoning and logic patterns.
• AST code normalization (NPset-Python) -- 1.25x token compression via Python AST to TinyDSL conversion. Strips syntactic noise and normalizes identifiers so the model learns programming structure rather than memorizing variable names.

Limitations

• Hardware: Total training took roughly 500 hours for 75B tokens; as a result, ablations were not feasible, and VRAM capped the context window at 1024
• Small model: 125M parameters limits reasoning and factual recall
• Educational data only: Primarily trained on educational datasets; not representative of general web text
• Not instruction-tuned: Base model only, not aligned for chat
• English only
• 1024 context window

Citation

@misc{gptx2_2025,
  title={GPT-X2: Data-Efficient Language Modeling at 125M Scale},
  author={Axiomic Labs},
  year={2025},
  howpublished={\url{https://huggingface.co/AxiomicLabs/GPT-X2-125M}},
  note={Trained on 75B tokens with a progressive curriculum and custom tokenizer}
}