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

Add llama 3 tokenization and data loader

Browse source
This commit is contained in:
Aleksa Gordic 2024-08-10 21:25:57 +02:00
parent 6e6a528111
commit d3151c4ace
4 changed files with 65 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

27
dev/data/data_common.py
View file

@ -23,27 +23,30 @@ def download_file(url: str, fname: str, chunk_size=1024):
bar.update(size)
def write_datafile(filename, toks):
HEADERS_INFO = {
"gpt-2": (20240520, 1),
"llama": (20240801, 7),
}
def write_datafile(filename, toks, model="gpt-2"):
"""
Saves token data as a .bin file, for reading in C.
- First comes a header with 256 int32s
- The tokens follow, each as a uint16
- The tokens follow, each as uint16 (gpt-2) or uint32 (llama)
"""
assert len(toks) < 2**31, "token count too large" # ~2.1B tokens
assert model in ["gpt-2", "llama"], f"unknown model {model}"
# construct the header
header = np.zeros(256, dtype=np.int32)
header[0] = 20240520 # magic
header[1] = 1 # version
header[2] = len(toks) # number of tokens after the 256*4 bytes of header (each 2 bytes as uint16)
# construct the tokens numpy array, if not already
if not isinstance(toks, np.ndarray) or not toks.dtype == np.uint16:
# validate that no token exceeds a uint16
maxtok = 2**16
assert all(0 <= t < maxtok for t in toks), "token dictionary too large for uint16"
header[0] = HEADERS_INFO[model][0] # magic
header[1] = HEADERS_INFO[model][1] # version
header[2] = len(toks) # number of tokens after the 256*4 bytes of header
if model == "gpt-2":
toks_np = np.array(toks, dtype=np.uint16)
elif model == "llama":
toks_np = np.array(toks, dtype=np.uint32)
else:
toks_np = toks
# write to file
raise ValueError(f"unknown model {model}")
print(f"writing {len(toks):,} tokens to {filename}")
with open(filename, "wb") as f:
f.write(header.tobytes())

4
dev/data/fineweb.py
View file

@ -99,7 +99,7 @@ with mp.Pool(nprocs) as pool:
remainder = args.shard_size - token_count
progress_bar.update(remainder)
all_tokens_np[token_count:token_count+remainder] = tokens[:remainder]
write_datafile(filename, all_tokens_np)
write_datafile(filename, list(all_tokens_np))
shard_index += 1
progress_bar = None
# populate the next shard with the leftovers of the current doc
@ -110,4 +110,4 @@ with mp.Pool(nprocs) as pool:
if token_count != 0:
split = "val" if shard_index == 0 else "train"
filename = os.path.join(DATA_CACHE_DIR, f"{name}_{split}_{shard_index:06d}.bin")
write_datafile(filename, all_tokens_np[:token_count])
write_datafile(filename, list(all_tokens_np[:token_count]))

50
dev/data/tinystories.py
View file

@ -1,17 +1,25 @@
"""
Downloads and tokenizes the TinyStories dataset.
- The download is from HuggingFace datasets.
- The tokenization is GPT-2 tokenizer with tiktoken
- The tokenization is using either GPT-2 or LLaMA 3 tokenizer.
The output is written to a newly created tinystories/ folder.
The script prints:
For GPT-2:
Tokenizing val split...
Saved 19043638 tokens to tinystories/TinyStories_val.bin
Tokenizing train split...
Saved 925653391 tokens to tinystories/TinyStories_train.bin
And runs in 1-2 minutes two depending on your internet
For LLaMA 3:
Number of shards: 50
Tokenizing val split...
writing 18,660,516 tokens to tinystories/TinyStories_val.bin
Tokenizing train split...
writing 907,021,844 tokens to tinystories/TinyStories_train.bin
And runs in few minutes two depending on your internet
connection and computer. The .bin files are raw byte
streams of int32 numbers indicating the token ids.
"""
@ -20,19 +28,16 @@ import os
import glob
import json
import random
import requests
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor, as_completed
import tiktoken
import numpy as np
from data_common import download_file, write_datafile
from transformers import AutoTokenizer
# -----------------------------------------------------------------------------
DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), "tinystories")
enc = tiktoken.get_encoding("gpt2")
encode = lambda s: enc.encode_ordinary(s)
def download():
"""Downloads the TinyStories dataset to DATA_CACHE_DIR"""
os.makedirs(DATA_CACHE_DIR, exist_ok=True)
@ -63,10 +68,21 @@ def download():
# data = json.load(f)
# print(f"Example story:\n{data[0]}")
def process_shard(shard_index, shard_filename):
def process_shard(shard_index, shard_filename, model):
if model == "gpt-2":
enc = tiktoken.get_encoding("gpt2")
encode = lambda s: enc.encode_ordinary(s)
eot = enc._special_tokens['<|endoftext|>'] # end of text token
elif model == "llama":
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
def encode(x):
return tokenizer(x).input_ids
eot = None
else:
raise ValueError(f"unknown model {model}")
with open(shard_filename, "r") as f:
data = json.load(f)
eot = enc._special_tokens['<|endoftext|>'] # end of text token
rng = random.Random(1337 + shard_index)
rng.shuffle(data)
all_tokens = []
@ -74,11 +90,12 @@ def process_shard(shard_index, shard_filename):
text = example["story"]
text = text.strip() # get rid of leading/trailing whitespace
tokens = encode(text)
all_tokens.append(eot)
if eot is not None:
all_tokens.append(eot)
all_tokens.extend(tokens)
return all_tokens
def tokenize():
def tokenize(model):
# shard 0 will be the val split, rest is train
data_dir = os.path.join(DATA_CACHE_DIR, "TinyStories_all_data")
shard_filenames = sorted(glob.glob(os.path.join(data_dir, "*.json")))
@ -89,20 +106,15 @@ def tokenize():
print(f"Tokenizing {split_name} split...")
all_tokens = []
with ProcessPoolExecutor() as executor:
futures = [executor.submit(process_shard, shard_index, shard_filename)
futures = [executor.submit(process_shard, shard_index, shard_filename, model)
for shard_index, shard_filename in enumerate(split_shards)]
for future in as_completed(futures):
all_tokens.extend(future.result())
split_filename = os.path.join(DATA_CACHE_DIR, f"TinyStories_{split_name}.bin")
write_datafile(split_filename, all_tokens)
write_datafile(split_filename, all_tokens, model)
if __name__ == "__main__":
model = "gpt-2" # gpt-2 or llama
download()
tokenize()
# Prints:
# Tokenizing val split...
# Saved 19043638 tokens to data/TinyStories_val.bin
# Tokenizing train split...
# Saved 925653391 tokens to data/TinyStories_train.bin
tokenize(model)

31
train_llama3.py
View file

@ -773,35 +773,38 @@ class Tokenizer:
# Our own simple Distributed Data Loader
def _peek_data_shard(filename):
raise NotImplementedError("_peek_data_shard not yet implemented for llama 3")
# only reads the header, returns header data
with open(filename, "rb") as f:
# first read the header, which is 256 int32 integers (4 bytes each)
header = np.frombuffer(f.read(256*4), dtype=np.int32)
if header[0] != 20240520:
if header[0] != 20240801:
print("ERROR: magic number mismatch in the data .bin file!")
print("---> HINT: Are you passing in a correct file with --input_bin?")
print("---> HINT: Dataset encoding changed recently, re-run data prepro or refer again to README")
print("---> HINT: For example re-run: `python dev/data/tinyshakespeare.py`, then re-try")
exit(1)
assert header[1] == 1, "unsupported version"
assert header[1] == 7, "unsupported version"
ntok = header[2] # number of tokens (claimed)
return ntok # for now just return the number of tokens
def _load_data_shard(filename):
raise NotImplementedError("_load_data_shard not yet implemented for llama 3")
with open(filename, "rb") as f:
# first read the header, which is 256 int32 integers (4 bytes each)
header = np.frombuffer(f.read(256*4), dtype=np.int32)
assert header[0] == 20240520, "magic number mismatch in the data .bin file"
assert header[1] == 1, "unsupported version"
assert header[0] == 20240801, "magic number mismatch in the data .bin file"
assert header[1] == 7, "unsupported version"
ntok = header[2] # number of tokens (claimed)
# the rest of it are tokens, stored as uint16
tokens = np.frombuffer(f.read(), dtype=np.uint16)
tokens = np.frombuffer(f.read(), dtype=np.uint32)
assert len(tokens) == ntok, "number of tokens read does not match header?"
return tokens
class DistributedDataLoader:
class DistributedShardedDataLoader:
"""
This DataLoader is both:
- distributed (works correctly in case of multiple processes in DDP)
- sharded (supports datasets that are broken up into multiple data shards)
It is not *permuted*, meaning that it itearates over the data in the order
of the dataset on disk, so the user should make sure to shuffle their examples
during the creation of their data shards for best performance.
"""
def __init__(self, filename_pattern, B, T, process_rank, num_processes):
self.process_rank = process_rank
self.num_processes = num_processes
@ -842,7 +845,7 @@ class DistributedDataLoader:
B = self.B
T = self.T
buf = self.tokens[self.current_position : self.current_position+B*T+1]
buf = torch.tensor(buf.astype(np.int32), dtype=torch.long)
buf = torch.tensor(buf, dtype=torch.long)
x = (buf[:-1]).view(B, T) # inputs
y = (buf[1:]).view(B, T) # targets
# advance the start pointer in current shard
@ -1097,10 +1100,10 @@ if __name__ == "__main__":
# Our own version of a simple DistributedDataLoader
# load tokens
train_loader = DistributedDataLoader(args.input_bin, B, T, ddp_rank, ddp_world_size)
train_loader = DistributedShardedDataLoader(args.input_bin, B, T, ddp_rank, ddp_world_size)
val_loader = None
if args.input_val_bin:
val_loader = DistributedDataLoader(args.input_val_bin, B, T, ddp_rank, ddp_world_size)
val_loader = DistributedShardedDataLoader(args.input_val_bin, B, T, ddp_rank, ddp_world_size)
# -------------------------------------------------------------------------
# PyTorch -> C bridge: save some weights and state for C to load later as reference