Sherwinroger002/gpt-124m-param-1.6b-tokens

GPT 124M — Pretrained on 1.6 Billion Tokens

A 124 million parameter GPT-2 style language model pretrained from scratch on 1.6 billion tokens from the FineWeb-Edu dataset.

Model Details

Training Loss

[image]

Training

This model was pretrained in three phases:

1. 1.Phase 1 (100M tokens): Initial pretraining from scratch on FineWeb-Edu
2. 2.Phase 2 (500M tokens): Continued pretraining with cosine LR schedule, torch.compile, and FlashAttention
3. 3.Phase 3 (1B tokens): Final phase with cosine LR warmup and decay over 1 billion fresh tokens

All phases used:

• —Optimizer: AdamW (lr=3e-4, weight_decay=0.1)
• —Precision: Mixed precision (FP16) with GradScaler
• —LR Schedule: Cosine with 5% warmup
• —Gradient Clipping: Max norm 1.0

Usage

⚠️ This model uses a custom architecture. You must download model.py from this repo to use it.

import json
import torch
from safetensors.torch import load_file
from transformers import AutoTokenizer
from model import GPTModel, generate_and_print_sample

# Load config
with open("config.json") as f:
    cfg = json.load(f)

# Create and load model
model = GPTModel(cfg)
state_dict = load_file("model.safetensors")

# Remove torch.compile prefix if present
unwanted_prefix = '_orig_mod.'
for k in list(state_dict.keys()):
    if k.startswith(unwanted_prefix):
        state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)

model.load_state_dict(state_dict)
model.to("cuda")
model.eval()

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("./")

# Generate text
generate_and_print_sample(model, tokenizer, "cuda", "The future of artificial intelligence")

Architecture

The model follows the GPT-2 architecture with the following components:

• —Multi-Head Attention with Flash Attention (scaled_dot_product_attention)
• —Pre-norm Transformer blocks (LayerNorm before attention and FFN)
• —GELU activation in the feed-forward network
• —Learned positional embeddings (max 256 positions)

Limitations

• —Short context window: The model can only process 256 tokens at a time due to fixed positional embeddings
• —Small model: At 124M parameters, this model is primarily for educational/research purposes
• —Training data: Trained only on English educational web content (FineWeb-Edu)
• —No instruction tuning: This is a base model — it completes text but does not follow instructions

Files

License

MIT