code/config/train_hrg205m.py

# config for training a 205M GPT on Human Reference Genome
# launch as the following (e.g. in a screen session) and wait ~5 days:
# $ torchrun --standalone --nproc_per_node=8 train.py config/train_hrg205m.py

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

# these make tokens per batch to be ~0.064M
# 8 batch size * 1000 block size * 1 gradaccum * 8 GPUs = 64,000
dataset = 'hrg'
n_workers = 0
persistent = False
batch_size = 8
block_size = 1000
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 ~6.4B
max_iters = 100000
lr_decay_iters = 100000
warmup_iters = 4000
learning_rate = 2.5e-4
min_lr = 2.5e-5
beta2 = 0.999

# eval stuff
out_dir = 'out-hrg'
eval_interval = 1000
eval_iters = 500 # val: 7784 test: 8469
log_interval = 10
always_save_checkpoint = False

# weight decay
weight_decay = 1e-2