code/length_gen.py

""" Length generalization: testing how well RoPE and PoPE generalize to 
extrapolated sequence lengths compared to that used for training. 
This script loads a ckpt and runs inference on sequences of greater length 
and plots perplexity vs sequence length. """

import os
from contextlib import nullcontext
import numpy as np
import torch
from data.pg19.prepare import PG19DataLoader
from model import GPTConfig, GPT

init_from = 'resume' # either 'resume' (from an out_dir) or a gpt2 variant (e.g. 'gpt2-xl')
base_dir = ''
ckpt_dir = 'final-owt-ckpts/' # ignored if init_from is not 'resume'
ckpt_fname = 'gpt2-124M-rope-ckpt.pt' # ignored if init_from is not 'resume'
dataset = 'pg19' # 'openwebtext', 'pg19'
split = 'val'  # for pg19 we use the test split
# model
n_layer = 12
n_head = 12
n_embd = 768
block_size = 1024
batch_size = 1
dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+
norm_type = 'rmsnorm' # 'layernorm' or 'rmsnorm'
pos_type = 'rope' # 'rope' or 'pope'
base_freq = 10000 # base frequency for rotary positional encoding
rotate_fraction = 1.0
thetab_init = 'zero'
bias = False # do we use bias inside LayerNorm and Linear layers?
complex_flash = False
eval_iters = 1500 # how many iterations to run the evaluation loop
# sampling
num_samples = 1 # number of samples to draw
max_new_tokens = 500 # number of tokens generated in each sample
temperature = 0.8 # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions
top_k = 200 # retain only the top_k most likely tokens, clamp others to have 0 probability
seed = 1337
device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16'
compile = True # use PyTorch 2.0 to compile the model to be faster
exec(open('configurator.py').read()) # overrides from command line or config file

np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)

data_dir = os.path.join(base_dir, 'data', dataset)
# init from a model saved in a specific directory
def load_ckpt():
    ckpt_path = os.path.join(os.path.join(base_dir, ckpt_dir, ckpt_fname))
    print(f"Loading checkpoint from {ckpt_path}...")
    checkpoint = torch.load(ckpt_path, map_location=device)
    # updating block size in model args
    gptconf = GPTConfig(**checkpoint['model_args'])
    model = GPT(gptconf)
    state_dict = checkpoint['model']
    unwanted_prefix = '_orig_mod.'
    for k,v in list(state_dict.items()):
        if k.startswith(unwanted_prefix):
            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
    model.load_state_dict(state_dict, strict=False)
    return model

def get_batch(split, chunk_size):
    target_mask = None
    data = np.memmap(os.path.join(data_dir, split + '.bin'), dtype=np.uint16, mode='r')
    ix = torch.randint(len(data) - chunk_size, (batch_size,))
    x = torch.stack([torch.from_numpy((data[i:i+chunk_size]).astype(np.int64)) for i in ix])
    y = torch.stack([torch.from_numpy((data[i+1:i+1+chunk_size]).astype(np.int64)) for i in ix])
    x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
    return x, y, target_mask

extra_lengths = np.arange(block_size, 10*block_size+1, block_size)
loss_vs_length = {}
# load model
model = load_ckpt()
model.eval()
model.to(device)
if compile:
    model = torch.compile(model)

# run evals at various lengths
for l in extra_lengths:
    print(f"Eval with block size {l}...")
    # init dataloader for PG-19 dataset
    ds = PG19DataLoader(
        seq_len=l+1,
        batch_size=batch_size,
        streaming=False,
        max_books=None,
        num_workers=0,
        shuffle=False,
        pin_memory=True
    )
    ds_itr = iter(ds)
    if dataset == 'pg19':
        eval_iters = len(ds_itr)  # eval over full test split of PG-19
    losses = torch.zeros(eval_iters)
    # run inference over dataset
    for k in range(eval_iters):
        if dataset == 'openwebtext':
            x, y, target_mask = get_batch(split, l)
        elif dataset == 'pg19':
            batch = next(ds_itr)
            x = batch['input_ids'].to(device, non_blocking=True)
            y = x.clone()[:, 1:]
            x = x[:, :-1]
        with torch.no_grad():
            logits, loss = model(x, y, target_mask=None)
        losses[k] = loss.item()
    loss_vs_length[l] = losses.mean().item()
    print(f"Done with block size {l}.")

if not os.path.exists("length_gen"):
    os.makedirs("length_gen")
torch.save(loss_vs_length, f'length_gen/{ckpt_fname[:-8]}-loss-vs-length.pt')