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

guess the config from globals() and log all of it with wandb

Browse Source
This commit is contained in:
Andrej Karpathy 2023-01-11 01:00:22 +00:00
parent 8b2e622b27
commit c2a402f7f7
1 changed files with 6 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
train.py
View File

@ -48,7 +48,8 @@ dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+
learning_rate = 6e-4 # max learning rate learning_rate = 6e-4 # max learning rate
max_iters = 600000 # total number of training iterations max_iters = 600000 # total number of training iterations
weight_decay = 1e-2 weight_decay = 1e-2
betas = (0.9, 0.95) beta1 = 0.9
beta2 = 0.95
# learning rate decay settings # learning rate decay settings
decay_lr = True # whether to decay the learning rate decay_lr = True # whether to decay the learning rate
warmup_iters = 2000 # how many steps to warm up for warmup_iters = 2000 # how many steps to warm up for
@ -61,7 +62,9 @@ device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
dtype = 'bfloat16' # 'float32' or 'bfloat16' 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
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------
config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
exec(open('configurator.py').read()) # overrides from command line or config file exec(open('configurator.py').read()) # overrides from command line or config file
config = {k: globals()[k] for k in config_keys} # will be useful for logging
# ----------------------------------------------------------------------------- # -----------------------------------------------------------------------------
# various inits, derived attributes, I/O setup # various inits, derived attributes, I/O setup
@ -142,7 +145,7 @@ if block_size < model.config.block_size:
model.to(device) model.to(device)
# optimizer # optimizer
optimizer = model.configure_optimizers(weight_decay, learning_rate, betas) optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2))
if init_from == 'resume': if init_from == 'resume':
optimizer.load_state_dict(checkpoint['optimizer']) optimizer.load_state_dict(checkpoint['optimizer'])
@ -188,12 +191,7 @@ def get_lr(iter):
# logging # logging
if wandb_log and gpu_id == 0: if wandb_log and gpu_id == 0:
import wandb import wandb
wandb.init(project=wandb_project, name=wandb_run_name) wandb.init(project=wandb_project, name=wandb_run_name, config=config)
wandb.config = {
"batch_size": batch_size,
"block_size": block_size,
"learning_rate": learning_rate, # TODO log everything else too
}
# training loop # training loop
t0 = time.time() t0 = time.time()