LLM training in simple, pure C/CUDA. There is no need for 245MB of PyTorch or 107MB of cPython. For example, training GPT-2 (CPU, fp32) is ~1,000 lines of clean code in a single file. It compiles and runs instantly, and exactly matches the PyTorch reference implementation. I chose GPT-2 as the first working example because it is the grand-daddy of LLMs, the first time the modern stack was put together.

Currently, I am working on:

- direct CUDA implementation, which will be significantly faster and probably come close to PyTorch.

- speed up the CPU version with SIMD instructions, AVX2 on x86 / NEON on ARM (e.g. Apple Silicon).

- more modern architectures, e.g. Llama2, Gemma, etc.

For the repo, I'd like to maintain both clean, simple reference implementations alongside a also lot more optimized versions that can come close to PyTorch, but in a tiny fraction of the code and dependencies.

Download and tokenize a dataset. The [tinyshakespeare](https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt) dataset is the fastest to download and tokenize:

python prepro_tinyshakespeare.py

This prints:

The .bin files are raw byte streams of int32 numbers indicating the token ids with the GPT-2 tokenizer. Alternatively you could also tokenize the [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) dataset with `prepro_tinystories.py`.

In principle we'd be ready to the train the model right here. However the baseline CPU/fp32 reference code is so inefficient that it's not practical to train these models from scratch yet. Instead, we initialize with the GPT-2 weights released by OpenAI and just do finetuning. For that, we have to download the GPT-2 weights and save them as a checkpoint we can load in C:

python train_gpt2.py

You'll recognize this code from nanoGPT as a simple GPT-2 reference implementation in PyTorch. This script will download the GPT-2 (124M) model, overfit a single batch of data for 10 iterations, run a few steps of generation, and most importantly it will save two files: 1) the `gpt2_124M.bin` file that contains the raw model weights for loading in C, and `gpt2_124M_debug_state.bin`, which also contains more debug state: the inputs, targets, logits and loss. This is very useful for debugging C code, for unit testing, and making sure we're exactly matching the PyTorch reference implementation. For now all we care about are the model weights in `gpt2_124M.bin`. We can now initialize with them and train in raw C. First compile the code:

make train_gpt2

You can have a look inside the `Makefile` and its comments. It will try to autodetect if OpenMP is available on your system, which is very helpful for speeding up the code at very low cost of code complexity. Once `train_gpt2` is compiled, you can run it:

OMP_NUM_THREADS=8 ./train_gpt2

You should tune the number of threads depending on how many cores your CPU has. The program will load the model weights, the tokens, it will run a finetuning loop for a few iterations with Adam lr 1e-4, and then generate a sample from the model. The file is (I think) very readable and you should have a look. Simply, there are implementations for the forward and backward pass of all the layers, and they get strung together into a large, manual, forward/backward/update loop. The output looks like this on my MacBook Pro (Apple Silicon M3 Max):

The generation just gives you the token ids for now, which we have to decode back to text. We can implement this in C quite easily also, because decoding is very straight-forward, it's just string chunk lookups and prints. For now we can use tiktoken:

import tiktoken

enc = tiktoken.get_encoding("gpt2")

print(enc.decode(list(map(int, "50256 16773 18162 21986 11 198 13681 263 23875 198 3152 262 11773 2910 198 1169 6002 6386 2583 286 262 11858 198 20424 428 3135 7596 995 3675 13 198 40 481 407 736 17903 11 329 703 6029 706 4082 198 42826 1028 1128 633 263 11 198 10594 407 198 2704 454 680 1028 262 1027 28860 286 198 3237 323".split()))))

which prints:

I like how Netflix comes up, it's clear that the shadow of the training past is still lurking in the model. I did not attempt to tune the finetuning hyperparameters so it's quite likely this can be improved quite a bit, most likely especially if one was to train a bit longer.

I am also attaching a simple unit test for making sure our C code agrees with the PyTorch code. Compile and run with:

This now loads the `gpt2_124M_debug_state.bin` file, runs a forward pass, compares the logits and loss with the PyTorch reference implementation, then it does 10 iterations of training with Adam and makes sure the losses match PyTorch.

import os

import requests

from tqdm import tqdm

import tiktoken

import numpy as np

DATA_CACHE_DIR = "data"

enc = tiktoken.get_encoding("gpt2")

encode = lambda s: enc.encode(s, allowed_special={'<|endoftext|>'})

def download_file(url: str, fname: str, chunk_size=1024):

"""Helper function to download a file from a given url"""

resp = requests.get(url, stream=True)

total = int(resp.headers.get("content-length", 0))

with open(fname, "wb") as file, tqdm(

desc=fname,

total=total,

unit="iB",

unit_scale=True,

unit_divisor=1024,

) as bar:

for data in resp.iter_content(chunk_size=chunk_size):

size = file.write(data)

bar.update(size)

def download():

"""Downloads the TinyShakespeare dataset to DATA_CACHE_DIR"""

os.makedirs(DATA_CACHE_DIR, exist_ok=True)

# download the TinyStories dataset, unless it's already downloaded

data_url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"

data_filename = os.path.join(DATA_CACHE_DIR, "tiny_shakespeare.txt")

if not os.path.exists(data_filename):

print(f"Downloading {data_url} to {data_filename}...")

download_file(data_url, data_filename)

else:

print(f"{data_filename} already exists, skipping download...")

def tokenize():

eot = enc._special_tokens['<|endoftext|>'] # end of text token

data_filename = os.path.join(DATA_CACHE_DIR, "tiny_shakespeare.txt")

text = open(data_filename, 'r').read()

# let's treat every person's statement in the dialog as a separate document

text = "<|endoftext|>" + text

text = text.replace('\n\n', '\n\n<|endoftext|>')

# encode the text

tokens = encode(text)

tokens_np = np.array(tokens, dtype=np.int32)

# let's take the first 32,768 tokens as the validation split (~10%)

val_tokens_np = tokens_np[:32768]

train_tokens_np = tokens_np[32768:]

# save to file

val_filename = os.path.join(DATA_CACHE_DIR, "tiny_shakespeare_val.bin")

train_filename = os.path.join(DATA_CACHE_DIR, "tiny_shakespeare_train.bin")

with open(val_filename, "wb") as f:

f.write(val_tokens_np.tobytes())

with open(train_filename, "wb") as f:

f.write(train_tokens_np.tobytes())

# prints

print(f"Saved {len(val_tokens_np)} tokens to {val_filename}")

print(f"Saved {len(train_tokens_np)} tokens to {train_filename}")

if __name__ == "__main__":

download()

tokenize()

import os

import glob

import json

import random

import requests

from tqdm import tqdm

from concurrent.futures import ProcessPoolExecutor, as_completed

import tiktoken

import numpy as np

DATA_CACHE_DIR = "data"

enc = tiktoken.get_encoding("gpt2")

encode = lambda s: enc.encode_ordinary(s)

def download_file(url: str, fname: str, chunk_size=1024):

"""Helper function to download a file from a given url"""

resp = requests.get(url, stream=True)

total = int(resp.headers.get("content-length", 0))

with open(fname, "wb") as file, tqdm(

desc=fname,

total=total,

unit="iB",

unit_scale=True,

unit_divisor=1024,

) as bar:

for data in resp.iter_content(chunk_size=chunk_size):

size = file.write(data)

bar.update(size)

def download():

"""Downloads the TinyStories dataset to DATA_CACHE_DIR"""

os.makedirs(DATA_CACHE_DIR, exist_ok=True)

# download the TinyStories dataset, unless it's already downloaded

data_url = "https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStories_all_data.tar.gz"

data_filename = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data.tar.gz")

if not os.path.exists(data_filename):

print(f"Downloading {data_url} to {data_filename}...")

download_file(data_url, data_filename)

else:

print(f"{data_filename} already exists, skipping download...")

# unpack the tar.gz file into all the data shards (json files)

data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")

if not os.path.exists(data_dir):

os.makedirs(data_dir, exist_ok=True)

print(f"Unpacking {data_filename}...")

os.system(f"tar -xzf {data_filename} -C {data_dir}")

else:

print(f"{data_dir} already exists, skipping unpacking...")

# print a single example just for debugging and such

shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))

with open(shard_filenames[0], "r") as f:

data = json.load(f)

print("Download done.")

print(f"Number of shards: {len(shard_filenames)}")

#print(f"Example story:\n{data[0]}")

def process_shard(shard_index, shard_filename):

with open(shard_filename, "r") as f:

data = json.load(f)

eot = enc._special_tokens['<|endoftext|>'] # end of text token

rng = random.Random(1337 + shard_index)

rng.shuffle(data)

all_tokens = []

for example in data:

text = example["story"]

text = text.strip() # get rid of leading/trailing whitespace

tokens = encode(text)

all_tokens.append(eot)

all_tokens.extend(tokens)

return all_tokens

def tokenize():

# shard 0 will be the val split, rest is train

data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")

shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))

val_shards = [shard_filenames[0]]

train_shards = shard_filenames[1:]

for split_name, split_shards in [("val", val_shards), ("train", train_shards)]:

print(f"Tokenizing {split_name} split...")

all_tokens = []

with ProcessPoolExecutor() as executor:

futures = [executor.submit(process_shard, shard_index, shard_filename)

for shard_index, shard_filename in enumerate(split_shards)]

for future in as_completed(futures):

all_tokens.extend(future.result())

all_tokens_np = np.array(all_tokens, dtype=np.int32)

split_filename = os.path.join(DATA_CACHE_DIR, f"TinyStories_{split_name}.bin")

with open(split_filename, "wb") as f:

f.write(all_tokens_np.tobytes())

print(f"Saved {len(all_tokens_np)} tokens to {split_filename}")

if __name__ == "__main__":

download()

tokenize()

# Prints:

# Tokenizing val split...

# Saved 19043638 tokens to data/TinyStories_val.bin

# Tokenizing train split...

# Saved 925653391 tokens to data/TinyStories_train.bin