diff --git a/README.md b/README.md index b017587..31c13b5 100644 --- a/README.md +++ b/README.md @@ -57,3 +57,7 @@ and observe the following losses on train and val: | gpt2-xl | 1558M | 2.56 | 2.54 | I briefly tried finetuning gpt2 a bit more on our OWT and didn't notice dramatic improvements, suggesting that OWT is not much much different from WT in terms of the data distribution, but this needs a bit more thorough attempt once the code is in a better place. + +## benchmarking + +For model benchmarking `bench.py` might be useful. It's identical what happens in `train.py` except we're measuring just the fwd+bwd+update time of the model on a fixed random batch of data. diff --git a/bench.py b/bench.py new file mode 100644 index 0000000..8356d21 --- /dev/null +++ b/bench.py @@ -0,0 +1,48 @@ +""" +A much shorter version of train.py for benchmarking the model +""" + +import time +import torch +from model import GPTConfig, GPT + +device = 'cuda:3' +torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul +torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn +torch.manual_seed(1337) + +batch_size = 8 +block_size = 1024 + +gptconf = GPTConfig( + block_size = block_size, # how far back does the model look? i.e. context size + n_layer = 12, n_head = 12, n_embd = 768, # size of the model + dropout = 0, # for determinism +) +model = GPT(gptconf) +model.to(device) + +x = torch.randint(50257, (batch_size, block_size), device=device) +y = torch.randint(50257, (batch_size, block_size), device=device) + +optimizer = model.configure_optimizers(weight_decay=1e-2, learning_rate=1e-4, betas=(0.9, 0.95)) + +burn_in = 10 # number of burn in steps where we don't measure time +num_steps = 30 +for k in range(num_steps): + + if k == burn_in: + t0 = time.time() # start the timer + + with torch.autocast(device_type='cuda', dtype=torch.bfloat16): + logits, loss = model(x, y) + + optimizer.zero_grad(set_to_none=True) + loss.backward() + optimizer.step() + lossf = loss.item() + print(f"{k}/{num_steps} loss: {lossf:.4f}") + +torch.cuda.synchronize() +t1 = time.time() +print("time in ms per iteration: %.2f" % ((t1 - t0) / (num_steps - burn_in) * 1000))