diff --git a/dev/data/data_common.py b/dev/data/data_common.py index b22d2f6..f1a455e 100644 --- a/dev/data/data_common.py +++ b/dev/data/data_common.py @@ -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()) diff --git a/dev/data/fineweb.py b/dev/data/fineweb.py index 72c3129..b13cc8c 100644 --- a/dev/data/fineweb.py +++ b/dev/data/fineweb.py @@ -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])) diff --git a/dev/data/tinystories.py b/dev/data/tinystories.py index fed8bc6..3490004 100644 --- a/dev/data/tinystories.py +++ b/dev/data/tinystories.py @@ -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) diff --git a/train_llama3.py b/train_llama3.py index 31596c3..b523722 100644 --- a/train_llama3.py +++ b/train_llama3.py @@ -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