code/config/train_gpt253m.py

# config for training GPT-2 (253M)
# launch as the following (e.g. in a screen session) and wait ~5 days:
# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt253m.py

wandb_log = False
wandb_project = 'complex-rope'
wandb_run_name='gpt2-253M'
base_dir = ''

# these make tokens per batch to be ~0.5M
# 8 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 327,680
dataset = 'openwebtext'
batch_size = 8
block_size = 1024
gradient_accumulation_steps = 1 * 8

# model
n_layer = 16
n_head = 16
n_embd = 1024
pos_type = 'rope'
base_freq = 10000
rotate_fraction = 1.0
thetab_init = 'zero'
dropout = 0.0

# this makes total number of tokens be 9B (=owt train split size)
max_iters = 100000
lr_decay_iters = 100000
warmup_iters = 1000

# eval stuff
out_dir = 'out-openwebtext'
eval_interval = 1000
eval_iters = 200
log_interval = 10
always_save_checkpoint = False

# weight decay
weight_decay = 1e-2