From 7d0042d574368b8ff15ef5b7b32ef683376fda47 Mon Sep 17 00:00:00 2001 From: Andrej Karpathy Date: Wed, 1 May 2024 18:15:10 +0000 Subject: [PATCH] make training data configurable and also dont move the engire dataset to device that is not a good idea --- train_gpt2.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/train_gpt2.py b/train_gpt2.py index abc1eab..f2fa68c 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -390,6 +390,7 @@ if __name__ == "__main__": # if you'd like to e.g. time the forward pass only, call this script as: # python train_gpt2.py --inference_only 1 --write_tensors 0 --sequence_length 1024 parser = argparse.ArgumentParser() + parser.add_argument("--input_bin", type=str, default="data/tiny_shakespeare_val.bin", help="input .bin to train on") parser.add_argument("--write_tensors", type=int, default=1, help="write tensors to disk") parser.add_argument("--inference_only", type=int, default=0, help="only run inference") parser.add_argument("--dtype", type=str, default="float32", help="float32|float16|bfloat16") @@ -477,21 +478,15 @@ if __name__ == "__main__": # data loading related: long but it's just to get a single batch of data # load the tokens - # prefer to use tiny_shakespeare if it's available, otherwise use tiny_stories - # we're using val instead of train split just because it is smaller/faster - shake_tokens_bin = "data/tiny_shakespeare_val.bin" - story_tokens_bin = "data/TinyStories_val.bin" - assert os.path.isfile(shake_tokens_bin) or os.path.isfile(story_tokens_bin), "you must run prepro on some dataset" - tokens_bin = shake_tokens_bin if os.path.isfile(shake_tokens_bin) else story_tokens_bin - assert os.path.isfile(tokens_bin) - print0(f"loading cached tokens in {tokens_bin}") - with open(tokens_bin, "rb") as f: + # note we're using val by default instead of train split just because it is smaller/faster + assert os.path.isfile(args.input_bin) + print0(f"loading cached tokens in {args.input_bin}") + with open(args.input_bin, "rb") as f: tokens = np.frombuffer(f.read(), dtype=np.int32) # np -> tensor, long, on device tokens = torch.tensor(tokens) tokens = tokens.to(torch.long) - tokens = tokens.to(device) # lightweight dataloader def get_batch(): @@ -509,6 +504,8 @@ if __name__ == "__main__": # fetch one batch of data, which we will overfit to data_iter = iter(get_batch()) x, y = next(data_iter) # we'll overfit this batch below + x = x.to(device) + y = y.to(device) # ------------------------------------------------------------------------- # STAGE 1: weights / state logging for C to load later