osmarks/nanogpt-experiments
1
0
Fork 0
mirror of https://github.com/osmarks/nanogpt-experiments.git synced 2024-12-18 14:10:28 +00:00
Code Issues Releases Wiki Activity

add device and dtype support to train.py args

Browse Source
This commit is contained in:
Andrej Karpathy 2023-01-08 19:20:38 +00:00
parent e7cd674ce7
commit a855d316fd
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
train.py
View File

@ -12,6 +12,7 @@ $ torchrun --standalone --nproc_per_node=4 train.py
import os import os
import time import time
import math import math
from contextlib import nullcontext
import numpy as np import numpy as np
import torch import torch
@ -56,11 +57,14 @@ min_lr = 6e-5 # minimum learning rate, should be ~= learning_rate/10 per Chinchi
# DDP settings # DDP settings
backend = 'nccl' # 'nccl', 'gloo', etc. backend = 'nccl' # 'nccl', 'gloo', etc.
# system # system
device = 'cuda' device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
dtype = 'bfloat16' # 'float32' or 'bfloat16'
compile = True # use PyTorch 2.0 to compile the model to be faster 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 exec(open('configurator.py').read()) # overrides from command line or config file
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------
# various inits, derived attributes, I/O setup
ddp = int(os.environ.get('LOCAL_RANK', -1)) != -1 # is this a ddp run? ddp = int(os.environ.get('LOCAL_RANK', -1)) != -1 # is this a ddp run?
if ddp: if ddp:
init_process_group(backend=backend) init_process_group(backend=backend)
@ -74,6 +78,10 @@ if gpu_id == 0:
torch.manual_seed(1337 + gpu_id) # note: each worker gets a different seed torch.manual_seed(1337 + gpu_id) # note: each worker gets a different seed
torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn 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
# note: float16 would require us to change the code to use a GradScaler
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16}[dtype]
ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
# poor man's data loader, TODO evaluate need for actual DataLoader # poor man's data loader, TODO evaluate need for actual DataLoader
data_dir = os.path.join('data', dataset) data_dir = os.path.join('data', dataset)
@ -156,7 +164,7 @@ def estimate_loss():
losses = torch.zeros(eval_iters) losses = torch.zeros(eval_iters)
for k in range(eval_iters): for k in range(eval_iters):
X, Y = get_batch(split) X, Y = get_batch(split)
with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): with ctx:
logits, loss = model(X, Y) logits, loss = model(X, Y)
losses[k] = loss.item() losses[k] = loss.item()
out[split] = losses.mean() out[split] = losses.mean()
@ -226,7 +234,7 @@ while True:
break break
X, Y = get_batch('train') X, Y = get_batch('train')
with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): with ctx:
logits, loss = model(X, Y) logits, loss = model(X, Y)
optimizer.zero_grad(set_to_none=True) optimizer.zero_grad(set_to_none=True)