gpjt/jax-with-mha-bias-larger-chinchilla-1

Model Card for gpjt/jax-with-mha-bias-larger-chinchilla-1

This model is gpjt/jax-with-mha-bias-larger-chinchilla-1, a trained-from-scratch base model using the GPT-2-style architecture from Sebastian Raschka's book "Build a Large Language Model (from Scratch)".

The model was trained in JAX using a black-box reimplementation of Raschka's original PyTorch code, but the safetensors have been converted back to a PyTorch-compatible format for convenience -- if you run this model, it will be in PyTorch.

This model had its number of parameters scaled up so that if it was trained with the Chinchilla-optimal number of tokens, it would approximately the same compute cost to train as a 2x overtrained model like `gpjt/jax-with-mha-bias-no-dropout-extended` or `gpjt/jax-with-mha-bias-no-dropout-2-epoch`. So it's somewhere between GPT-2 small and GPT-2 medium in its number of parameters.

Model Details

Model Description

• —Developed by: Giles Thomas, based on code by Sebastian Raschka
• —Model type: GPT-2 style transformers-based causal LLM.
• —License: Apache 2
• —Parameters: 235,621,120
• —Context length: 1,024
• —Embedding dimensions: 896
• —MHA heads: 14
• —Layers: 15
• —QKV bias: False
• —Weight tying: False

Don't have high expectations for the model! It has only 163M parameters (the GPT-2 "small" size) and was trained on roughly the Chinchilla-optimal number of tokens (~20x the number of parameters), which means that it doesn't know many facts and is not terribly smart. If you want to do serious work, use a serious model (I like Qwen's). But if you want to build on this and see what you can do with a 2020-vintage LLM, please do feel free to play with it!

Model Sources

• —Repository: https://github.com/gpjt/jax-gpt2-from-scratch for the training code, gpjt/ddp-base-model-from-scratch for the code used to run it from here.
• —Blog post: A quick(ish) Chinchilla check This is the first model trained in the post, referrred to as slightly-larger in the text.

How to Get Started with the Model

You can download and run the model for inference directly:

from transformers import pipeline
pipe = pipeline("text-generation", model="gpjt/jax-with-mha-bias-larger-chinchilla-1", trust_remote_code=True)
out = pipe(
    "Every effort moves you",
    max_new_tokens=20,
    do_sample=True,
    temperature=1.4,
    top_k=25,
)
print(out[0]["generated_text"])

Note that because it uses custom code, you'll need to set trust_remote_code to True.

It supports AutoTokenizer, AutoModel and AutoModelForCausalLM:

>>> from transformers import AutoTokenizer, AutoModel, AutoModelForCausalLM
>>> tokenizer = AutoTokenizer.from_pretrained("gpjt/jax-with-mha-bias-larger-chinchilla-1")
>>> model = AutoModel.from_pretrained("gpjt/jax-with-mha-bias-larger-chinchilla-1", trust_remote_code=True)
>>> llm_model = AutoModelForCausalLM.from_pretrained("gpjt/jax-with-mha-bias-larger-chinchilla-1", trust_remote_code=True)

You can also fine-tune it; this notebook has an example.

Again, don't expect too much from this model! It's a 236M-parameter GPT-2 one, trained on a limited number of tokens. It's both dumb and ignorant ;-)

Training Details

• —Machine type: Local machine with an RTX 3090
• —Tokens: 4,712,422,400 (Chinchilla-optimal of 20x parameters) rounded up to the nearest batch.
• —Dataset: gpjt/fineweb-gpt2-tokens
• —Micro-batch size: 4
• —Global batch size: 96
• —Dropout: 0.0
• —Gradient clipping: 3.5
• —Learning rate: 0.0014
• —Schedule learning rate: True
• —Weight decay: 0.01