mirror of
https://github.com/osmarks/nanogpt-experiments.git
synced 2024-09-21 03:39:44 +00:00
103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
"""
|
|
A much shorter version of train.py for benchmarking
|
|
"""
|
|
import os
|
|
import numpy as np
|
|
import time
|
|
import torch
|
|
from model import GPTConfig, GPT
|
|
|
|
device = 'cuda'
|
|
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
|
|
dtype = torch.bfloat16
|
|
compile_model = True
|
|
|
|
# data loading init
|
|
real_data = True
|
|
if real_data:
|
|
dataset = 'openwebtext'
|
|
data_dir = os.path.join('data', dataset)
|
|
train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
|
|
def get_batch(split):
|
|
data = train_data # note ignore split in benchmarking script
|
|
ix = torch.randint(len(data) - block_size, (batch_size,))
|
|
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
|
|
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
|
|
x, y = x.to(device), y.to(device)
|
|
return x, y
|
|
else:
|
|
# alternatively, if fixed data is desired to not care about data loading
|
|
x = torch.randint(50257, (batch_size, block_size), device=device)
|
|
y = torch.randint(50257, (batch_size, block_size), device=device)
|
|
get_batch = lambda split: (x, y)
|
|
|
|
# model init
|
|
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)
|
|
|
|
optimizer = model.configure_optimizers(weight_decay=1e-2, learning_rate=1e-4, betas=(0.9, 0.95))
|
|
|
|
if compile_model:
|
|
print("Compiling model...")
|
|
model = torch.compile(model) # pytorch 2.0
|
|
|
|
profile = False # use pytorch profiler, or just simple benchmarking?
|
|
if profile:
|
|
# useful docs on pytorch profiler:
|
|
# - tutorial https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html
|
|
# - api https://pytorch.org/docs/stable/profiler.html#torch.profiler.profile
|
|
wait, warmup, active = 5, 5, 5
|
|
num_steps = wait + warmup + active
|
|
with torch.profiler.profile(
|
|
activities=[torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA],
|
|
schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=1),
|
|
on_trace_ready=torch.profiler.tensorboard_trace_handler('./bench_log'),
|
|
record_shapes=True,
|
|
profile_memory=True,
|
|
with_stack=True, # incurs an additional overhead, disable if not needed
|
|
with_flops=True,
|
|
with_modules=False, # only for torchscript models atm
|
|
) as prof:
|
|
|
|
for k in range(num_steps):
|
|
X, Y = get_batch('train')
|
|
with torch.autocast(device_type='cuda', dtype=dtype):
|
|
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}")
|
|
|
|
prof.step() # notify the profiler at end of each step
|
|
|
|
else:
|
|
|
|
# simple benchmarking
|
|
torch.cuda.synchronize()
|
|
for stage, num_steps in enumerate([10, 20]): # burnin, then benchmark
|
|
t0 = time.time()
|
|
for k in range(num_steps):
|
|
X, Y = get_batch('train')
|
|
with torch.autocast(device_type='cuda', dtype=dtype):
|
|
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()
|
|
if stage == 1:
|
|
print(f"time per iteration: {(t1-t0)/num_steps*1000:.4f}ms")
|