GPT 2B - Custom Trained Model

This is a custom GPT architecture trained from scratch for ~2 Billion tokens on English text. It was trained in two phases, starting with a 256 context window for the first billion tokens, and seamlessly expanding to a 512 context window for the second billion tokens.

Model Details

• Architecture: Custom GPT (Transformer Decoder)
• Parameters: 124M
• Context Length: 512 tokens
• Training Tokens: ~2 Billion
• Embedding Dimension: 768
• Heads: 12
• Layers: 12

Training Loss

The model was trained on RunPod. Below is the loss curve for the second phase of training (1B to 2B tokens).

How to use

You can load and use this model locally by cloning this repository. You will need the model.py and test.py scripts included in the repo.

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

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

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

# Clean up torch.compile prefixes if present
unwanted_prefix = '_orig_mod.'
for k,v in list(state_dict.items()):
    if k.startswith(unwanted_prefix):
        state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)

model.load_state_dict(state_dict, strict=False)
model.eval()
model.to("cuda")

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

# Generate
def generate_and_print_sample(model, tokenizer, device, start_context):
    context_size = model.pos_emb.weight.shape[0]
    encoded = tokenizer(start_context, return_tensors="pt")["input_ids"].to(device)
    
    token_ids = generate_text(
        model=model,
        idx=encoded,
        max_new_tokens=512,
        context_size=context_size,
        temperature=0.65,
        top_k=30,
        top_p=0.9,
        repetition_penalty=1.15
    )
    
    decoded_text = tokenizer.decode(token_ids.squeeze(0).tolist())
    return decoded_text

print(generate_and_print_sample(model, tokenizer, "cuda", "once upon a time"))

Sample Output

Prompt: once upon a time

Generation:

once upon a time, and he had been a very good man. He was a great man; but his life seemed to be over; and when he died he would die at the hands of those who knew him better than any other man. But if he had lived so long as he could, we would have nothing to say about him. He is always in danger; but that which he will not believe he will not be able to live more than once; for he has no power to do what he wants. For it is true that he is in danger. If he does not know what he will do or why he will not do it, he will never make himself known. And if he cannot understand it he must be afraid; and then he may even have thought he must know what he will do. In short, the one thing he must do is to be afraid: and he should fear that this thing might happen again. So, the first step towards understanding the nature of God's existence was the realization that all things are made of matter, and that all things are made of them. The word "matter" comes from the Greek words meaning "life", and the Latin word means "earth". It is only in the sense of being like air, water, air, water, fire, fire, fire, etc., that the soul can perceive itself, and the mind can comprehend its own existence. Nowhere else can there be an un-existent world, or, in this case, heaven, earth. The first stage is that of the universe, and there is no beginning; it is only by the act of creating that part of the world that the soul can experience the essence of its being. This is the point of the second stage. When the soul is fully formed, it becomes aware of its own being, and then perceives itself as the whole of everything that exists in its place. But when the soul is fully formed, it becomes aware of its own being, and knows its own being, and is conscious of its own being. That is why the first stage is that of the first stage, that of the last stage, of the second stage. The third stage of the third stage is that of the third stage, which is the one which is the most important. All things are made of atoms, and all things are made of atoms. There are two kinds of atoms: those which are not atoms, and those which are not atoms. Some are not atoms, but some molecules. They are particles which are