code/config/train_gpt124m.py

# config for training GPT-2 (124M) down to very nice loss of ~2.85 on 1 node of 8X A100 40GB
# launch as the following (e.g. in a screen session) and wait ~5 days:
# $ torchrun --standalone --nproc_per_node=8 train.py config/train_gpt124m.py

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

# these make tokens per batch to be ~0.065M
# 8 batch size * 1024 block size * 1 gradaccum * 8 GPUs = 65,536
dataset = 'openwebtext'
batch_size = 8
block_size = 1024
gradient_accumulation_steps = 1 * 8

# model
n_layer = 12
n_head = 12
n_embd = 768
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