From 15e7a22a413aadc4ef2da861486dbc1faa690d5e Mon Sep 17 00:00:00 2001 From: Shizhe Diao Date: Sun, 19 Oct 2025 07:53:44 -0700 Subject: [PATCH] support custom training data --- nanochat/dataloader.py | 15 ++++++++++++--- nanochat/dataset.py | 21 ++++++++++++++------- scripts/base_loss.py | 4 +++- 3 files changed, 29 insertions(+), 11 deletions(-) diff --git a/nanochat/dataloader.py b/nanochat/dataloader.py index 6c864d3d..1878a062 100644 --- a/nanochat/dataloader.py +++ b/nanochat/dataloader.py @@ -6,8 +6,17 @@ from nanochat.common import get_dist_info from nanochat.dataset import parquets_iter_batched from nanochat.tokenizer import get_tokenizer -def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokenizer_batch_size=128, device="cuda"): - """Stream pretraining text from parquet files, tokenize, yield training batches.""" +def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokenizer_batch_size=128, device="cuda", data_dir=None): + """Stream pretraining text from parquet files, tokenize, yield training batches. + + Args: + B: batch size + T: sequence length + split: "train" or "val" + tokenizer_threads: number of threads for tokenization + tokenizer_batch_size: batch size for tokenization + data_dir: optional custom directory containing parquet files (None = use default) + """ assert split in ["train", "val"], "split must be 'train' or 'val'" ddp, ddp_rank, ddp_local_rank, ddp_world_size = get_dist_info() needed_tokens = B * T + 1 # +1 is because we also need the target at the last token @@ -21,7 +30,7 @@ def tokenizing_distributed_data_loader(B, T, split, tokenizer_threads=4, tokeniz def document_batches(): while True: # batch will iterate in group size of the parquet files, usually e.g. 1024 rows - for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size): + for batch in parquets_iter_batched(split=split, start=ddp_rank, step=ddp_world_size, data_dir=data_dir): # for the tokenizer we might want to go in usually smaller batches, e.g. 128 rows for i in range(0, len(batch), tokenizer_batch_size): yield batch[i:i+tokenizer_batch_size] diff --git a/nanochat/dataset.py b/nanochat/dataset.py index 602daedc..62eb8168 100644 --- a/nanochat/dataset.py +++ b/nanochat/dataset.py @@ -31,7 +31,11 @@ os.makedirs(DATA_DIR, exist_ok=True) # These functions are useful utilities to other modules, can/should be imported def list_parquet_files(data_dir=None): - """ Looks into a data dir and returns full paths to all parquet files. """ + """Looks into a data dir and returns full paths to all parquet files. + + Args: + data_dir: optional custom directory containing parquet files (None = use default DATA_DIR) + """ data_dir = DATA_DIR if data_dir is None else data_dir parquet_files = sorted([ f for f in os.listdir(data_dir) @@ -40,14 +44,17 @@ def list_parquet_files(data_dir=None): parquet_paths = [os.path.join(data_dir, f) for f in parquet_files] return parquet_paths -def parquets_iter_batched(split, start=0, step=1): - """ - Iterate through the dataset, in batches of underlying row_groups for efficiency. - - split can be "train" or "val". the last parquet file will be val. - - start/step are useful for skipping rows in DDP. e.g. start=rank, step=world_size +def parquets_iter_batched(split, start=0, step=1, data_dir=None): + """Iterate through the dataset, in batches of underlying row_groups for efficiency. + + Args: + split: "train" or "val". the last parquet file will be val. + start: starting row group index (useful for DDP, e.g. start=rank) + step: step size for row groups (useful for DDP, e.g. step=world_size) + data_dir: optional custom directory containing parquet files (None = use default) """ assert split in ["train", "val"], "split must be 'train' or 'val'" - parquet_paths = list_parquet_files() + parquet_paths = list_parquet_files(data_dir=data_dir) parquet_paths = parquet_paths[:-1] if split == "train" else parquet_paths[-1:] for filepath in parquet_paths: pf = pq.ParquetFile(filepath) diff --git a/scripts/base_loss.py b/scripts/base_loss.py index abcde5f9..9391ec35 100644 --- a/scripts/base_loss.py +++ b/scripts/base_loss.py @@ -22,6 +22,7 @@ split_tokens = 20*524288 # number of tokens to evaluate per split model_tag = None # optional model tag for the output directory name model_step = None # optional model step for the output directory name device_type = "" # cuda|cpu|mps (empty => autodetect) +data_dir = "" # path to directory containing parquet files with 'text' column (empty string = use default) exec(open(os.path.join('nanochat', 'configurator.py')).read()) # overrides from command line or config file # Load the base model and the tokenizer @@ -37,8 +38,9 @@ assert split_tokens % tokens_per_step == 0, "split_tokens must be divisible by t steps = split_tokens // tokens_per_step token_bytes = get_token_bytes(device=device) bpb_results = {} +custom_data_dir = data_dir if data_dir else None for split_name in ["train", "val"]: - loader = tokenizing_distributed_data_loader(device_batch_size, sequence_len, split_name, device=device) + loader = tokenizing_distributed_data_loader(device_batch_size, sequence_len, split_name, device=device, data_dir=custom_data_dir) with autocast_ctx: bpb = evaluate_bpb(model, loader, steps, token_bytes) print0(f"{split_name} bpb: {bpb:.4f}")