Added sorting model and modified nanoGPT files

nanoGPT/data/config/train_sort.py

@@ -0,0 +1,43 @@
+"""
+Train a small and simple sorting network
+Requires at least 6 GB of memory
+
+Based on nanoGPT's shakespeare example: https://github.com/karpathy/nanoGPT
+"""
+
+out_dir = 'out-sort-lists'
+eval_interval = 2000
+eval_iters = 200
+log_interval = 200
+verbose_log_interval = 2000
+always_save_checkpoint = False
+
+init_from = 'scratch' # 'scratch' or 'resume'
+
+wandb_log = False
+wandb_project = 'sort_lists'
+wandb_run_name = 'mini-gpt'
+
+dataset = 'sort_lists'
+gradient_accumulation_steps = 1
+batch_size = 12
+block_size = 256 # context window, keep synchronized with training data
+
+# Transformer network parameters
+n_layer = 64
+n_head = 4
+n_embd = 256
+dropout = 0.01
+
+# Training parameters
+learning_rate = 1e-4
+max_iters = 900000
+weight_decay = 1e-1
+beta1 = 0.95
+beta2 = 0.99
+grad_clip = 1.0
+# learning rate decay settings
+decay_lr = True
+warmup_iters = 2000
+lr_decay_iters = max_iters
+min_lr = 0.1*learning_rate

nanoGPT/data/sort_lists/meta.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:749a5d6b08620256f53b531a210014219cea2b542dea0175d5493f6390052f1b
+size 417

nanoGPT/data/sort_lists/prepare.py

@@ -0,0 +1,172 @@
+"""
+Prepare the list-sorting dataset for language modeling.
+Will saves train.bin, val.bin containing the ids, and meta.pkl containing the
+encoder and decoder and some other related info.
+
+Based on nanoGPT: https://github.com/karpathy/nanoGPT
+"""
+import os
+import pickle
+import numpy as np
+import json
+
+random = np.random.default_rng()
+
+context_window_length = 256 # Keep this synchronized with the model definition
+
+num_training_lists = 2e6 # Output file is about 2 GB
+num_val_lists = int(0.1*num_training_lists)
+distribution = "uniform"
+
+def str_array_to_sorted_array(output_string):
+    template = '{{"values": {}}}'
+    array = np.array(json.loads(template.format(output_string))["values"], dtype = np.uint16)
+    array.sort()
+    return array
+
+
+def generate_list_pairs(max_list_length_chars = None,
+                            max_int = 65536,
+                            distribution = "uniform",
+                            num_lists = 1,
+                            data_file = None,
+                            fill_blanks = True):
+    if max_list_length_chars is None:
+        max_list_length_chars = int(np.floor(0.5*context_window_length - 2))
+    
+    if data_file is None:
+        data_file = os.path.join(os.path.dirname(__file__), "train.txt")
+    
+    with open(data_file, 'a') as f:
+        for n_list in range(int(num_lists)):
+
+            max_random_length = int(np.floor(0.5*max_list_length_chars))
+
+            list_length_ints = random.integers(0, max_random_length)
+            if distribution.casefold() == "uniform":
+                random_shuffled = random.integers(0, max_int, size = list_length_ints, endpoint = False)
+            elif distribution.casefold() == "gaussian":
+                random_floats = random.normal(loc = 0.5*max_int,
+                                              scale = max_int,
+                                              size = list_length_ints)
+                random_ints = np.around(random_floats, decimals = 0).astype(np.int64, casting = "unsafe")
+                invalid_integers = np.logical_or(random_ints < 0, random_ints >= max_int)
+                random_ints[invalid_integers]\
+                                = random.integers(0, max_int, size = invalid_integers, endpoint = False)
+            else:
+                raise NotImplementedError("No distribution called '{}'".format(distribution))
+            # Crop based on string length
+            shuffled_str = np.array2string(random_shuffled,
+                                           max_line_width = max_list_length_chars*2,
+                                           separator = ',').replace(" ","")
+            shuffled_str = shuffled_str[:np.min([len(shuffled_str), max_list_length_chars - 1])]
+            shuffled_str = shuffled_str[:-1].strip()
+            
+            if shuffled_str[-1] in [",", "\n"]:
+                shuffled_str = shuffled_str[:-1]
+            
+            shuffled_str += "]"
+            sorted_str = np.array2string(str_array_to_sorted_array(shuffled_str),
+                                            max_line_width = max_list_length_chars*2,
+                                            separator = ',').replace(" ","")
+                     
+            data_line = "(" + shuffled_str[1:-1] + "): [" + sorted_str[1:-1] + "];\n"
+            if fill_blanks:
+                len_align = context_window_length - ((len(data_line) - 5) % context_window_length)
+                filler = "_"*(context_window_length + len_align)
+                f.write(filler)
+            f.write(data_line)
+            if n_list%100 == 0:
+                print("{:.2f}% ({}/{}) -- {}".format(100.*(n_list + 1)/num_lists, n_list, num_lists, len(data_line)))
+    print("Data written to file: {}".format(data_file))
+    return data_file
+
+
+train_file_path = os.path.join(os.path.dirname(__file__), 'train.txt')
+val_file_path = os.path.join(os.path.dirname(__file__), 'val.txt')
+
+print("Generating training data")
+train_data_file = generate_list_pairs(num_lists = num_training_lists,
+                                      fill_blanks = True,
+                                      max_int = 100,
+                                      distribution = distribution,
+                                    data_file = train_file_path)
+print("Generating validation data")
+val_data_file = generate_list_pairs(num_lists = num_val_lists,
+                                    fill_blanks = True,
+                                    max_int = 100,
+                                    distribution = distribution,
+                                    data_file = val_file_path)
+
+
+tokens = ['0','1','2','3','4','5','6','7','8','9',',','(','): [','];\n','_']
+vocab_size = len(tokens)
+print("all the unique characters:", ''.join(tokens))
+print(f"vocab size: {vocab_size:,}")
+print("Still working...")
+
+# create a mapping from characters to integers
+stoi = { ch:i for i,ch in enumerate(tokens) }
+itos = { i:ch for i,ch in enumerate(tokens) }
+char_to_token = {token[0]:token for token in tokens}
+chars_to_skip = {token[0]:len(token)-1 for token in tokens}
+
+
+def encode(s):
+    encoded = []
+    skip = 0
+    for char in s:
+        if skip:
+            skip -= 1
+            continue
+        else:
+            skip = chars_to_skip[char]
+            encoded.append(stoi[char_to_token[char]])
+    return encoded
+
+def decode(l):
+    return ''.join([itos[i] for i in l])
+
+
+# save the meta information as well, to help us encode/decode later
+meta = {
+    'vocab_size': vocab_size,
+    'tokens': tokens,
+    'itos': itos,
+    'stoi': stoi,
+    "char_to_token": char_to_token,
+    "chars_to_skip": chars_to_skip
+}
+with open(os.path.join(os.path.dirname(__file__), 'meta.pkl'), 'wb') as f:
+    pickle.dump(meta, f)
+print("Saved metadata file")
+
+if False:
+    print("Still working...")
+    
+    # encode both to integers
+    with open(train_data_file, 'r') as f:
+        train_data = f.read()
+    train_ids = encode(train_data)
+    del train_data
+    
+    print("Still working...")
+    
+    with open(val_data_file, 'r') as f:
+        val_data = f.read()
+    val_ids = encode(val_data)
+    del val_data
+    print(f"train has {len(train_ids):,} tokens")
+    print(f"val has {len(val_ids):,} tokens")
+    
+    print("Still working...")
+    # export to bin files
+    train_ids = np.array(train_ids, dtype=np.uint16)
+    print("Still working...")
+    val_ids = np.array(val_ids, dtype=np.uint16)
+    print("Still working...")
+    train_ids.tofile(os.path.join(os.path.dirname(__file__), 'train.bin'))
+    print("Still working...")
+    val_ids.tofile(os.path.join(os.path.dirname(__file__), 'val.bin'))
+
+print("Export complete.")

nanoGPT/out-sort-lists/ckpt.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7037aba6fa5003481a8dd53b71a0621d284ad167351d65033cbe62f574df7904
+size 605632927

nanoGPT/sample.py

@@ -0,0 +1,99 @@
+"""
+Sample from a trained model
+"""
+import os
+import pickle
+from contextlib import nullcontext
+import torch
+import tiktoken
+from model import GPTConfig, GPT
+
+# -----------------------------------------------------------------------------
+init_from = 'resume' # either 'resume' (from an out_dir) or a gpt2 variant (e.g. 'gpt2-xl')
+out_dir = 'out' # ignored if init_from is not 'resume'
+start = "\n" # or "<|endoftext|>" or etc. Can also specify a file, use as: "FILE:prompt.txt"
+num_samples = 10 # number of samples to draw
+max_new_tokens = 500 # number of tokens generated in each sample
+temperature = 0.0 # 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 = 'float16' # 'float32' or 'bfloat16' or 'float16'
+compile = False # use PyTorch 2.0 to compile the model to be faster
+exec(open('configurator.py').read()) # overrides from command line or config file
+# -----------------------------------------------------------------------------
+
+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)
+
+# model
+if init_from == 'resume':
+    # init from a model saved in a specific directory
+    ckpt_path = os.path.join(out_dir, 'ckpt.pt')
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    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)
+elif init_from.startswith('gpt2'):
+    # init from a given GPT-2 model
+    model = GPT.from_pretrained(init_from, dict(dropout=0.0))
+
+model.eval()
+model.to(device)
+if compile:
+    model = torch.compile(model) # requires PyTorch 2.0 (optional)
+
+# look for the meta pickle in case it is available in the dataset folder
+load_meta = False
+if init_from == 'resume' and 'config' in checkpoint and 'dataset' in checkpoint['config']: # older checkpoints might not have these...
+    meta_path = os.path.join('data', checkpoint['config']['dataset'], 'meta.pkl')
+    load_meta = os.path.exists(meta_path)
+if load_meta:
+    print(f"Loading meta from {meta_path}...")
+    with open(meta_path, 'rb') as f:
+        meta = pickle.load(f)
+    stoi, itos = meta['stoi'], meta['itos']
+    char_to_token = meta["char_to_token"]
+    chars_to_skip = meta["chars_to_skip"]
+    
+    def encode(s):
+        encoded = []
+        skip = 0
+        for char in s:
+            if skip:
+                skip -= 1
+                continue
+            else:
+                skip = chars_to_skip[char]
+                encoded.append(stoi[char_to_token[char]])
+        return encoded
+    
+    def decode(l):
+        return ''.join([itos[i] for i in l])
+else:
+    raise RuntimeError("No meta.pkl found for sorting! Cannot find token encoder or decoder.")
+
+# encode the beginning of the prompt
+if start.startswith('FILE:'):
+    with open(start[5:], 'r', encoding='utf-8') as f:
+        start = f.read()
+start_ids = encode(start)
+x = (torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...])
+
+# run generation
+with torch.no_grad():
+    with ctx:
+        for k in range(num_samples):
+            y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
+            print(decode(y[0].tolist()))
+            print('---------------')

nanoGPT/train.py

@@ -0,0 +1,520 @@
+"""
+This training script can be run both on a single gpu in debug mode,
+and also in a larger training run with distributed data parallel (ddp).
+
+To run on a single GPU, example:
+$ python train.py config/train_sort.py
+
+Based on nanoGPT by Andrej Karpathy: https://github.com/karpathy/nanoGPT
+Modified for a learn-to-sort experiment by Jacob Bayless
+"""
+
+import os
+import time
+import math
+import pickle
+from contextlib import nullcontext
+
+import numpy as np
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.distributed import init_process_group, destroy_process_group
+
+from model import GPTConfig, GPT
+import json
+
+# -----------------------------------------------------------------------------
+# default config values designed to train a gpt2 (124M) on OpenWebText
+# I/O
+out_dir = 'out'
+eval_interval = 2000
+verbose_log_interval = 250
+log_interval = 1
+eval_iters = 200
+eval_only = False # if True, script exits right after the first eval
+always_save_checkpoint = True # if True, always save a checkpoint after each eval
+init_from = 'scratch' # 'scratch' or 'resume' or 'gpt2*'
+# wandb logging
+wandb_log = False # disabled by default
+wandb_project = 'owt'
+wandb_run_name = 'gpt2' # 'run' + str(time.time())
+# data
+dataset = 'openwebtext'
+gradient_accumulation_steps = 5 * 8 # used to simulate larger batch sizes
+batch_size = 12 # if gradient_accumulation_steps > 1, this is the micro-batch size
+block_size = 1024
+# model
+n_layer = 12
+n_head = 12
+n_embd = 768
+dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+
+bias = False # do we use bias inside LayerNorm and Linear layers?
+# adamw optimizer
+learning_rate = 6e-4 # max learning rate
+max_iters = 600000 # total number of training iterations
+weight_decay = 1e-1
+beta1 = 0.9
+beta2 = 0.95
+grad_clip = 1.0 # clip gradients at this value, or disable if == 0.0
+# learning rate decay settings
+decay_lr = True # whether to decay the learning rate
+warmup_iters = 2000 # how many steps to warm up for
+lr_decay_iters = 600000 # should be ~= max_iters per Chinchilla
+min_lr = 6e-5 # minimum learning rate, should be ~= learning_rate/10 per Chinchilla
+# DDP settings
+backend = 'nccl' # 'nccl', 'gloo', etc.
+# system
+device = 'cuda' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1' etc., or try 'mps' on macbooks
+dtype = 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler
+compile = False # 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
+config = {k: globals()[k] for k in config_keys} # will be useful for logging
+# -----------------------------------------------------------------------------
+
+
+verbose_logfile = os.path.join(out_dir, "verbose_log.txt")
+performance_file = os.path.join(out_dir, "perf_log.txt")
+
+random = np.random.default_rng()
+
+meta_path = os.path.join('data', dataset, 'meta.pkl')
+print(f"Loading meta from {meta_path}...")
+with open(meta_path, 'rb') as f:
+    meta = pickle.load(f)
+stoi, itos = meta['stoi'], meta['itos']
+char_to_token = meta["char_to_token"]
+chars_to_skip = meta["chars_to_skip"]
+
+def encode(s):
+    encoded = []
+    skip = 0
+    for char in s:
+        if skip:
+            skip -= 1
+            continue
+        else:
+            skip = chars_to_skip[char]
+            encoded.append(stoi[char_to_token[char]])
+    return encoded
+
+def decode(l):
+    return ''.join([itos[i] for i in l])
+
+def str_array_to_sorted_array(output_string):
+    template = '{{"values": {}}}'
+    array = np.array(json.loads(template.format(output_string))["values"], dtype = np.uint16)
+    array.sort()
+    return array
+
+def generate_validation_list(max_int = 99,
+                             force_list_length = None,
+                             max_list_length_chars = None):
+    if max_list_length_chars is None:
+        max_list_length_chars = int(np.floor(0.5*block_size - 2))
+    max_random_length = int(np.floor(0.5*max_list_length_chars))
+    if force_list_length is None:
+        list_length_ints = random.integers(0, max_random_length)
+    else:
+        list_length_ints = np.min([force_list_length, max_random_length])
+    random_shuffled = random.integers(0, max_int,
+                                      size = list_length_ints,
+                                      endpoint = False)
+    # Crop based on string length
+    shuffled_str = np.array2string(random_shuffled,
+                                   max_line_width = max_list_length_chars*2,
+                                   separator = ',').replace(" ","")
+    shuffled_str = shuffled_str[:np.min([len(shuffled_str), max_list_length_chars - 1])]
+    shuffled_str = shuffled_str[:-1].strip()
+    
+    if shuffled_str[-1] in [",", "\n"]:
+        shuffled_str = shuffled_str[:-1]
+    shuffled_str += "]"
+    sorted_str = np.array2string(str_array_to_sorted_array(shuffled_str),
+                                    max_line_width = max_list_length_chars*2,
+                                    separator = ',').replace(" ","")
+    input_line = "(" + shuffled_str[1:-1] + "): ["
+    correct_output = sorted_str[1:-1] + "];\n"
+    return (input_line, correct_output)
+    
+
+def score_performance(model_output,
+                      correct_output,
+                      output_terminator = "];\n",
+                      blank_separator = "_",
+                      list_separator = ","):
+    errors = 0
+    
+    model_output, _, _ = model_output.partition(blank_separator)
+    
+    if output_terminator in correct_output and output_terminator not in model_output:
+        errors += 2
+        
+    correct_output, _, _ = correct_output.partition(output_terminator)
+    model_output, _, _ = model_output.partition(output_terminator)
+    
+    correct_output = correct_output.split(list_separator)
+    model_output = model_output.split(list_separator)
+    min_length = np.min([len(correct_output), len(model_output)])
+    
+    length_error = np.abs(len(correct_output) - len(model_output))
+    errors += length_error
+    
+    for entry in range(min_length):
+        if model_output[entry] != correct_output[entry]:
+            errors += 1
+    
+    return errors, errors/float(len(correct_output) + 1), len(correct_output)
+    
+
+def evaluate_performance(model,
+                         force_list_length = None,
+                         max_list_length_chars = None,
+                         output_separator = "): [",
+                         output_terminator = "];\n",
+                         list_separator = ","):
+
+    input_line, correct_output = generate_validation_list(force_list_length = force_list_length,
+                                                    max_list_length_chars = max_list_length_chars)
+    
+    max_new_tokens = len(correct_output) + 10
+    
+    temperature = 0.0001
+    
+    start_ids = encode(input_line)
+    x = (torch.tensor(start_ids, dtype=torch.long, device=device)[None, ...])
+    
+    y = model.generate(x, max_new_tokens, temperature=temperature, top_k=15)
+    model_output = str(decode(y[0].tolist())).partition(output_separator)[-1]
+    
+    error_abs, error_rel, correct_length = score_performance(model_output, correct_output,
+                                               output_terminator = output_terminator,
+                                               list_separator = list_separator)
+    
+    return(input_line, correct_output, model_output, error_abs, error_rel, correct_length)
+
+def log_performance_verbose(model, iter_number,
+                            n = 10,
+                            verbose_log_file = None,
+                            performance_log_file = None):
+    
+    max_list_length_chars = int(np.floor(0.5*block_size - 2))
+    max_random_length = int(np.floor(0.5*max_list_length_chars))
+    force_list_lengths = np.linspace(int(np.floor(0.2*max_random_length)),
+                               max_random_length,
+                               n, dtype = np.int64)
+    
+    
+    errors_abs = []
+    errors_rel = []
+    with open(verbose_log_file, 'a') as log_file:
+        log_file.write("\n\n_____ {} ______\n".format(iter_number))
+    
+    with open(performance_log_file, 'a') as perf_file:
+        perf_file.write("{}:".format(iter_number))
+    
+    list_length_total = 0
+    best_rel_error = np.inf
+    worst_rel_error = -np.inf
+    worst_abs_error = -1
+    worst_list_length = -1
+    best_abs_error = -1
+    best_list_length = -1
+    for n_ind, force_list_length in enumerate(force_list_lengths):
+        input_line, correct_output, model_output,\
+            error_abs, error_rel, correct_length = evaluate_performance(model, force_list_length = force_list_length)
+        errors_abs.append(error_abs)
+        errors_rel.append(error_rel)
+        list_length_total += correct_length
+        
+        if(error_rel > worst_rel_error):
+            worst_rel_error = error_rel
+            worst_abs_error = error_abs
+            worst_list_length = correct_length
+        if(error_rel < best_rel_error):
+            best_rel_error = error_rel
+            best_abs_error = error_abs
+            best_list_length = correct_length
+        
+        with open(verbose_log_file, 'a') as log_file:
+            log_file.write("\n\tINPUT: {}\n\tEXAMPLE: {}\n\t OUTPUT: {}\n\tERRORS:{} / {} ({:.2f}%)\n".format(input_line,
+                                                                                                   correct_output,
+                                                                                                   model_output,
+                                                                                                   error_abs,
+                                                                                                   correct_length,
+                                                                                                   error_rel*100.0))
+        with open(performance_log_file, 'a') as perf_file:
+            if n_ind > 0:
+                perf_file.write(",")
+            perf_file.write(" {} / {} ({:.2f}%)".format(error_abs, correct_length, error_rel*100.0))
+    with open(performance_log_file, 'a') as perf_file:
+        perf_file.write("\n")
+    print("ITER {}: total score is {} errors / {} ({:.2f}%)".format(iter_number,
+                                                     np.sum(errors_abs),
+                                                     np.sum(list_length_total),
+                                                     100.0*np.mean(errors_rel)))
+    print("\t Best: {} errors / {} ({:.2f}%)".format(best_abs_error,
+                                                     best_list_length,
+                                                     100.0*best_rel_error))
+    print("\t Worst: {} errors / {} ({:.2f}%)".format(worst_abs_error,
+                                                     worst_list_length,
+                                                     100.0*worst_rel_error))
+
+# various inits, derived attributes, I/O setup
+ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?
+if ddp:
+    init_process_group(backend=backend)
+    ddp_rank = int(os.environ['RANK'])
+    ddp_local_rank = int(os.environ['LOCAL_RANK'])
+    ddp_world_size = int(os.environ['WORLD_SIZE'])
+    device = f'cuda:{ddp_local_rank}'
+    torch.cuda.set_device(device)
+    master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
+    seed_offset = ddp_rank # each process gets a different seed
+    assert gradient_accumulation_steps % torch.cuda.device_count() == 0
+    gradient_accumulation_steps //= torch.cuda.device_count()
+else:
+    # if not ddp, we are running on a single gpu, and one process
+    master_process = True
+    seed_offset = 0
+    ddp_world_size = 1
+tokens_per_iter = gradient_accumulation_steps * ddp_world_size * batch_size * block_size
+print(f"tokens per iteration will be: {tokens_per_iter:,}")
+
+if master_process:
+    os.makedirs(out_dir, exist_ok=True)
+torch.manual_seed(1337 + seed_offset)
+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
+# note: float16 data type will automatically use a GradScaler
+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)
+
+# poor man's data loader
+data_dir = os.path.join('data', dataset)
+train_data = np.memmap(os.path.join(data_dir, 'train.bin'), dtype=np.uint16, mode='r')
+val_data = np.memmap(os.path.join(data_dir, 'val.bin'), dtype=np.uint16, mode='r')
+def get_batch(split):
+    data = train_data if split == 'train' else val_data
+    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])
+    if device_type == 'cuda':
+        # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
+        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
+    else:
+        x, y = x.to(device), y.to(device)
+    return x, y
+
+# init these up here, can override if init_from='resume' (i.e. from a checkpoint)
+iter_num = 0
+best_val_loss = 1e9
+
+# attempt to derive vocab_size from the dataset
+meta_path = os.path.join(data_dir, 'meta.pkl')
+meta_vocab_size = None
+if os.path.exists(meta_path):
+    with open(meta_path, 'rb') as f:
+        meta = pickle.load(f)
+    meta_vocab_size = meta['vocab_size']
+    print(f"found vocab_size = {meta_vocab_size} (inside {meta_path})")
+
+# model init
+model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,
+                  bias=bias, vocab_size=None, dropout=dropout) # start with model_args from command line
+if init_from == 'scratch':
+    # init a new model from scratch
+    print("Initializing a new model from scratch")
+    # determine the vocab size we'll use for from-scratch training
+    if meta_vocab_size is None:
+        print("defaulting to vocab_size of GPT-2 to 50304 (50257 rounded up for efficiency)")
+    model_args['vocab_size'] = meta_vocab_size if meta_vocab_size is not None else 50304
+    gptconf = GPTConfig(**model_args)
+    model = GPT(gptconf)
+elif init_from == 'resume':
+    print(f"Resuming training from {out_dir}")
+    # resume training from a checkpoint.
+    ckpt_path = os.path.join(out_dir, 'ckpt.pt')
+    checkpoint = torch.load(ckpt_path, map_location=device)
+    checkpoint_model_args = checkpoint['model_args']
+    # force these config attributes to be equal otherwise we can't even resume training
+    # the rest of the attributes (e.g. dropout) can stay as desired from command line
+    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
+        model_args[k] = checkpoint_model_args[k]
+    # create the model
+    gptconf = GPTConfig(**model_args)
+    model = GPT(gptconf)
+    state_dict = checkpoint['model']
+    # fix the keys of the state dictionary :(
+    # honestly no idea how checkpoints sometimes get this prefix, have to debug more
+    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)
+    iter_num = checkpoint['iter_num']
+    best_val_loss = checkpoint['best_val_loss']
+elif init_from.startswith('gpt2'):
+    print(f"Initializing from OpenAI GPT-2 weights: {init_from}")
+    # initialize from OpenAI GPT-2 weights
+    override_args = dict(dropout=dropout)
+    model = GPT.from_pretrained(init_from, override_args)
+    # read off the created config params, so we can store them into checkpoint correctly
+    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:
+        model_args[k] = getattr(model.config, k)
+# crop down the model block size if desired, using model surgery
+if block_size < model.config.block_size:
+    model.crop_block_size(block_size)
+    model_args['block_size'] = block_size # so that the checkpoint will have the right value
+model.to(device)
+
+# initialize a GradScaler. If enabled=False scaler is a no-op
+scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
+
+# optimizer
+optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
+if init_from == 'resume':
+    optimizer.load_state_dict(checkpoint['optimizer'])
+checkpoint = None # free up memory
+
+# compile the model
+if compile:
+    print("compiling the model... (takes a ~minute)")
+    unoptimized_model = model
+    model = torch.compile(model) # requires PyTorch 2.0
+
+# wrap model into DDP container
+if ddp:
+    model = DDP(model, device_ids=[ddp_local_rank])
+
+# helps estimate an arbitrarily accurate loss over either split using many batches
+@torch.no_grad()
+def estimate_loss():
+    out = {}
+    model.eval()
+    for split in ['train', 'val']:
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            X, Y = get_batch(split)
+            with ctx:
+                logits, loss = model(X, Y)
+            losses[k] = loss.item()
+        out[split] = losses.mean()
+    model.train()
+    return out
+
+# learning rate decay scheduler (cosine with warmup)
+def get_lr(it):
+    # 1) linear warmup for warmup_iters steps
+    if it < warmup_iters:
+        return learning_rate * it / warmup_iters
+    # 2) if it > lr_decay_iters, return min learning rate
+    if it > lr_decay_iters:
+        return min_lr
+    # 3) in between, use cosine decay down to min learning rate
+    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
+    assert 0 <= decay_ratio <= 1
+    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1
+    return min_lr + coeff * (learning_rate - min_lr)
+
+# logging
+if wandb_log and master_process:
+    import wandb
+    wandb.init(project=wandb_project, name=wandb_run_name, config=config)
+
+# training loop
+X, Y = get_batch('train') # fetch the very first batch
+t0 = time.time()
+local_iter_num = 0 # number of iterations in the lifetime of this process
+raw_model = model.module if ddp else model # unwrap DDP container if needed
+running_mfu = -1.0
+while True:
+
+    # determine and set the learning rate for this iteration
+    lr = get_lr(iter_num) if decay_lr else learning_rate
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+    # evaluate the loss on train/val sets and write checkpoints
+    if iter_num % eval_interval == 0 and master_process:
+        losses = estimate_loss()
+        print(f"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}")
+        if wandb_log:
+            wandb.log({
+                "iter": iter_num,
+                "train/loss": losses['train'],
+                "val/loss": losses['val'],
+                "lr": lr,
+                "mfu": running_mfu*100, # convert to percentage
+            })
+        if losses['val'] < best_val_loss or always_save_checkpoint:
+            best_val_loss = losses['val']
+            if iter_num > 0:
+                checkpoint = {
+                    'model': raw_model.state_dict(),
+                    'optimizer': optimizer.state_dict(),
+                    'model_args': model_args,
+                    'iter_num': iter_num,
+                    'best_val_loss': best_val_loss,
+                    'config': config,
+                }
+                print(f"saving checkpoint to {out_dir}")
+                torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
+    if iter_num == 0 and eval_only:
+        break
+
+    # forward backward update, with optional gradient accumulation to simulate larger batch size
+    # and using the GradScaler if data type is float16
+    for micro_step in range(gradient_accumulation_steps):
+        if ddp:
+            # in DDP training we only need to sync gradients at the last micro step.
+            # the official way to do this is with model.no_sync() context manager, but
+            # I really dislike that this bloats the code and forces us to repeat code
+            # looking at the source of that context manager, it just toggles this variable
+            model.require_backward_grad_sync = (micro_step == gradient_accumulation_steps - 1)
+        with ctx:
+            logits, loss = model(X, Y)
+            loss = loss / gradient_accumulation_steps # scale the loss to account for gradient accumulation
+        # immediately async prefetch next batch while model is doing the forward pass on the GPU
+        X, Y = get_batch('train')
+        # backward pass, with gradient scaling if training in fp16
+        scaler.scale(loss).backward()
+    # clip the gradient
+    if grad_clip != 0.0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+    # step the optimizer and scaler if training in fp16
+    scaler.step(optimizer)
+    scaler.update()
+    # flush the gradients as soon as we can, no need for this memory anymore
+    optimizer.zero_grad(set_to_none=True)
+
+    # timing and logging
+    t1 = time.time()
+    dt = t1 - t0
+    t0 = t1
+    if iter_num % log_interval == 0 and master_process:
+        # get loss as float. note: this is a CPU-GPU sync point
+        # scale up to undo the division above, approximating the true total loss (exact would have been a sum)
+        lossf = loss.item() * gradient_accumulation_steps
+        if local_iter_num >= 5: # let the training loop settle a bit
+            mfu = raw_model.estimate_mfu(batch_size * gradient_accumulation_steps, dt)
+            running_mfu = mfu if running_mfu == -1.0 else 0.9*running_mfu + 0.1*mfu
+        print(f"iter {iter_num}: loss {lossf:.4f}, time {dt*1000:.2f}ms, mfu {running_mfu*100:.2f}%")
+    
+    if (iter_num % verbose_log_interval == 0) and master_process and (local_iter_num > 0):
+        log_performance_verbose(raw_model, iter_num,
+                                verbose_log_file = verbose_logfile,
+                                performance_log_file = performance_file)
+    
+    iter_num += 1
+    local_iter_num += 1
+
+    # termination conditions
+    if iter_num > max_iters:
+        break
+
+if ddp:
+    destroy_process_group()