code/config/train_wikitext103.py

# train a 124M GPT2 sized model
# needs atleast 5 GPUs to run

out_dir = 'out-wikitext103'
# eval stuff
eval_interval = 250 # keep frequent because we'll overfit
eval_iters = 200
log_interval = 10 # don't print too too often

# we expect to overfit on this small dataset, so only save when val improves
always_save_checkpoint = False

wandb_log = False # override via command line if you like
wandb_project = 'complex-rope'
wandb_run_name = 'gpt2-124M'

dataset = 'wikitext103'
# 16 batch_size * 1024 block_size * 1 gradaccum * 1 GPUs = 16,384
gradient_accumulation_steps = 1
batch_size = 16
block_size = 1024 # context length

# GPT-2 sized model :)
n_layer = 12
n_head = 12
n_embd = 768
bias = True
pos_type = 'rope'

max_iters = 10000
lr_decay_iters = 10000 # make equal to max_iters usually
warmup_iters = 100

# on macbook also add
# device = 'cpu'  # run on cpu only
# compile = False # do not torch compile the model