code/config/train_indirect_idx.py

# dataset
dataset = 'indirect_idx'
max_shift = 15
min_length = 20
max_length = 40

# train a tiny GPT model on the indirect_idx task
out_dir = 'out-indirect-idx'
seed = 0

# 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 = 'indirect_idx'

gradient_accumulation_steps = 1
batch_size = 64
block_size = max_length + 7

# baby GPT model :)
n_layer = 8
n_head = 8
n_embd = 512
pos_type = 'rope'
base_freq = 10000
rotate_fraction = 1.0
thetab_init = 'two_pi'
use_theta_bias = True
dropout = 0.0

learning_rate = 2e-4 # with baby networks can afford to go a bit higher
max_iters = 100000
lr_decay_iters = 100000 # make equal to max_iters usually
min_lr = 2e-5 # learning_rate / 10 usually
beta2 = 0.99 # make a bit bigger because number of tokens per iter is small
grad_clip = 1.0
weight_decay = 1e-2
warmup_iters = 4000 # not super necessary potentially

eval_interval = 500 # keep frequent because we'll overfit
eval_iters = 10000 // batch_size  # 10000 samples in val/test
log_interval = 10 # don't print too too often

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