aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

make training data configurable and also dont move the engire dataset to device that is not a good idea

Browse source
This commit is contained in:
Andrej Karpathy 2024-05-01 18:15:10 +00:00
parent 8aa9a9ed53
commit 7d0042d574
1 changed files with 7 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

17
train_gpt2.py
View file

@ -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