From 32affd4442fbb4c3e2e8352555276f4143f38f95 Mon Sep 17 00:00:00 2001 From: Rasmus Larsen Date: Fri, 26 Sep 2025 10:32:55 +0200 Subject: [PATCH 1/6] pin triton versions for cuda stable --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 7ee5cc2..1e4bb8b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,7 +131,8 @@ cuda = [ "torch>=2.7", "torchvision", "torchdata", - "pytorch-triton", + "pytorch-triton==3.3.0", + "triton==3.3.1" ] # CUDA 12.8 (nightly) From c3de83b2632cd367d9a2a763814d9b80ae1fc363 Mon Sep 17 00:00:00 2001 From: Peter Schneider-Kamp Date: Sat, 27 Sep 2025 15:02:27 +0200 Subject: [PATCH 2/6] refactored data loading code support for JINX data format in addition to Parquet --- maester/config.py | 7 +- maester/datasets/base.py | 334 +++++++++++++++++++++++ maester/datasets/experimental_otf.py | 383 ++------------------------- maester/datasets/formats/__init__.py | 2 + maester/datasets/formats/jinx.py | 66 +++++ maester/datasets/formats/parquet.py | 73 +++++ maester/datasets/stateful.py | 145 ++++++++++ 7 files changed, 644 insertions(+), 366 deletions(-) create mode 100644 maester/datasets/base.py create mode 100644 maester/datasets/formats/__init__.py create mode 100644 maester/datasets/formats/jinx.py create mode 100644 maester/datasets/formats/parquet.py create mode 100644 maester/datasets/stateful.py diff --git a/maester/config.py b/maester/config.py index 349b068..978d4aa 100644 --- a/maester/config.py +++ b/maester/config.py @@ -6,7 +6,7 @@ TomlConfigSettingsSource, ) from pydantic.fields import FieldInfo -from typing import Callable, Type, Any +from typing import Callable, Type, Any, Optional from pathlib import Path import torch @@ -19,8 +19,9 @@ class DatasetConfig(BaseSettings): data_dirs: list[str] = [ - "data/toy" - ] + "data/toy", + ] + dataset_types: Optional[list[str]] = None dataset_weights: str = "1.0" bos_token: int = 128000 eos_token: int = 128001 diff --git a/maester/datasets/base.py b/maester/datasets/base.py new file mode 100644 index 0000000..25f03f4 --- /dev/null +++ b/maester/datasets/base.py @@ -0,0 +1,334 @@ +import hashlib +import json +import math +import os +import random +import time +import torch.distributed as dist +from typing import Any, Callable, List, Optional, Set + +from .experimental_otf import logger +from .stateful import _Stateful_Dataset + +class BaseDataset(_Stateful_Dataset): + def __init__( + self, + data_dir: str, + data_ext: str, + rank: int, + worldsize: int, + tokenizer, + delimiter_token: Any, + bos_token: Optional[Any] = None, + strip_tokens: Optional[Set[Any]] = set(), + seed: int = 42, + min_length: int = 1, + max_chunksize: int = 1024, + verbose: bool = False, + shuffle: bool = True, + data_column: str = "text", + process_fn: Optional[Callable] = None, + raw_data_mode: bool = False, + ): + super(BaseDataset, self).__init__(rank, worldsize) + self.seed = seed + self.data = data_dir + self.tokenizer = tokenizer + self.min_length = min_length + assert max_chunksize > 0, f"Max chunksize must be a nonzero positive integer" + self.chunksize = max_chunksize + self.eos = delimiter_token + self.bos = bos_token + self.drop = strip_tokens + self.verbose = verbose + self.data_column = data_column + self.raw_data_mode = raw_data_mode + self.process_fn = process_fn or self._default_process + self.docset: List[Any] = [] # map of doc indices to (file_path, min docid, max docid) + self.docs_per_file = {} + + # Guaranteed inconsistent shuffling across workers + random.seed(self.seed + rank) + + # Get all data files in the directory recursively + self.data_files = [os.path.join(root, f) for root, _, files in os.walk(data_dir) for f in files if f.endswith(data_ext)] + self.data_files.sort() # Ensure consistent sharding across machines + assert len(self.data_files) > 0, "No data files found in data directory" + + dataset_hash = self._generate_dataset_hash() + cache_file = os.path.join(data_dir, f"doc_counts_cache_{dataset_hash}.json") + + # Rank 0 handles file I/O, other ranks wait + if dist.get_rank(dist.group.WORLD) == 0: + if not os.path.exists(cache_file): + self._gather_doc_counts() + self._save_cached_doc_counts(cache_file) + + dist.barrier() + + # All ranks load the cache + self._load_cached_doc_counts(cache_file) + + dist.barrier() # ensure all ranks loaded + + # Fragment the files + start_frag = (rank * worldsize * len(self.data_files)) // worldsize + end_frag = ((rank + 1) * worldsize * len(self.data_files)) // worldsize + shardfrags = [ + (self.data_files[i // worldsize], i % worldsize) for i in range(start_frag, end_frag) + ] + + # Read shardfrags, assemble doc list for each file shard (aggregating over fragments): + ndocs = -1 + docset = {} # shardid -> (min docid, max docid) + for i, (shard, frag) in enumerate(shardfrags): + ndocs = self.docs_per_file[shard] + doc_start = (ndocs * frag) // worldsize + doc_end = (ndocs * frag + ndocs) // worldsize - 1 # Inclusive upper bound + if shard not in docset: + docset[shard] = [doc_start, doc_end] + min_d, max_d = docset[shard] + if doc_start < min_d: + docset[shard][0] = doc_start + if doc_end > max_d: + docset[shard][1] = doc_end + + # Add all of this dataset's shard entries to self.docset + doccount = 0 + for shardid in docset: + min_d = docset[shardid][0] + max_d = docset[shardid][1] + self.docset.append((shardid, min_d, max_d)) + doccount += max_d - min_d + 1 + self._len = doccount + + if verbose: + logger.info(f"Worker {rank} responsible for docs: {self.docset}") + logger.info(f"Total docs: {doccount}") + + # Shuffle files + if shuffle: + random.shuffle(self.docset) + + self.docset_index = 0 + self.chunk_index = -1 + self.completed_current_doc = False + + # Stats + self.epochs_seen = -1 + self.tokens_seen = 0 + self.docs_seen = 0 + self.percent_seen = 0 + self.lcg_state = seed + rank + + self.state_params = [ + "docset_index", + "chunk_index", + "completed_current_doc", + "epochs_seen", + "tokens_seen", + "docs_seen", + "percent_seen", + "lcg_state", + ] + + def _default_process(self, data): + """Default processing: tokenize text data.""" + return self.tokenizer.encode(data, add_special_tokens=False, padding=False, truncation=False) + + def _generate_dataset_hash(self): + """Generate a unique hash for the dataset based on file names and sizes.""" + hasher = hashlib.md5() + for file in self.data_files: + hasher.update(file.encode()) + hasher.update(str(os.path.getsize(file)).encode()) + return hasher.hexdigest() + + def _load_cached_doc_counts(self, cache_file): + """Load cached document counts from a file.""" + start = time.time() + with open(cache_file, 'r') as f: + self.docs_per_file = json.load(f) + logger.info(f"Loaded cached document counts in {time.time() - start} seconds") + + def _save_cached_doc_counts(self, cache_file): + """Save document counts to a cache file.""" + with open(cache_file, 'w') as f: + json.dump(self.docs_per_file, f) + logger.info(f"Saved document counts cache to {cache_file}") + + def _gather_doc_counts(self): + """Gather document counts for each Parquet file.""" + start = time.time() + total_rows = 0 + for file in self.data_files: + num_rows = self._gather_doc_count(file) + self.docs_per_file[file] = num_rows + total_rows += num_rows + assert total_rows > 0, "No rows found in parquet files" + logger.info(f"Gathered {total_rows} rows in {time.time() - start} seconds") + + def _get_docid(self, i): + """ + Given a global doc index over the set of docs owned by this worker, + return the corresponding path, num rows + """ + cur = 0 + assert i <= self._len, f"You have requested an illegal doc index {i}, docset length is {self._len}" + for shardid, min_d, max_d in self.docset: + docrange = max_d - min_d + 1 + cur += docrange + if cur > i: + return shardid, docrange, min_d + raise RuntimeError("This should be unreachable") + + def _construct_chunk(self, j, doc, n_chunks): + """ + Construct the jth chunk of doc + """ + start_index = j * self.chunksize + n_pull = self.chunksize + if self.bos is not None: + if j == 0: + n_pull -= 1 + else: + start_index -= 1 + chunk = doc[start_index:start_index + n_pull] + self.tokens_seen += len(chunk) + # Add bos/eos tokens if needed + if self.bos is not None and j == 0: + chunk = [self.bos] + chunk + if j == n_chunks - 1: + chunk = chunk + [self.eos] + return chunk + + def _random_map_docid(self, size): + """ + Given size of document pool, use saved state (prior index) to generate the next index via LCG. + Implements within-shard document shuffling without materializing any large doc lists. + """ + m = 2 ** math.ceil(math.log2(size)) # Round up to nearest power of 2 + a = 5 # A,C values known to work well with powers of 2 (Knuth, 1997, 3.2.1.3) + c = (self.rank + self.seed) * 2 + 1 + state = self.lcg_state + while True: + state = (a * state + c) % m + if state < size: + return state + + def __iter__(self): + docset_offset = self.docset_index + lcg_offset = self.lcg_state + residual_chunks = self.chunk_index + 1 # chunks to skip after restore and create at the end of epoch, 0-indexed + first_doc_mapping = None # Will store the document mapping for the first document + ndocs = self._len + path = "" + reader = None + if self.completed_current_doc: # resuming at the end of a doc + docset_offset = (docset_offset + 1) % ndocs + self.completed_current_doc = False + while True: + for i in range(ndocs): + doc_index = (docset_offset + i) % ndocs + self.completed_current_doc = False # reset + + # Update stats + if doc_index == 0: + self.epochs_seen += 1 + if self.verbose: + logger.info(f"ParquetDataset: entering epoch {self.epochs_seen}") + self.docset_index = doc_index + + # Map docset id to file, owned size and in-doc owned start idx + # This should be the same value many iters in a row, processing each shard + file_path, docrange, mindoc = self._get_docid(doc_index) + + # Map docset ids to consistently shuffled ids + # determine if we need a new document position + if i == 0 and not self.completed_current_doc and self.chunk_index >= 0: + # resuming mid-doc, do not advance lcg + doclcg = self.lcg_state + else: + doclcg = self._random_map_docid(docrange) # shuffled in-doc range + self.lcg_state = doclcg # update lcg state + + # Save the document mapping for the first document for residual processing + if i == 0: + first_doc_mapping = doclcg + + local_row = doclcg + mindoc # map docid to local row + + newpath = file_path + path, reader = self._get_reader(path, newpath, reader) + + row = self._read_specific_row(reader, local_row) + data = row[self.data_column] + + if self.raw_data_mode: + # In raw data mode, yield the data directly without processing + self.docs_seen += 1 + self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) + self.completed_current_doc = True + yield data + else: + # Normal mode: process and chunk the data + doc = self.process_fn(data) + if len(doc) < 2: + logger.warning(f"Empty document detected at {file_path}:{local_row}") + continue + + if doc[0] in self.drop: + doc = doc[1:] + if doc[-1] in self.drop: + doc = doc[:-1] + + doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 + if doclen >= self.min_length: + n_chunks = math.ceil(doclen / self.chunksize) + for j in range(n_chunks): + if i == 0 and not self.completed_current_doc and j < residual_chunks: + pass # skip already processed chunks + # doclcg = self.lcg_state # use saved lcg state when resuming + else: + self.chunk_index = j + # Document complete, update stats + if j == n_chunks - 1: + self.docs_seen += 1 + self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) + self.chunk_index = -1 + self.completed_current_doc = True + out = self._construct_chunk(j, doc, n_chunks) + # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, length {len(out)}, first tokens: {out[:5]}...") + yield out + + # Load any chunks initially skipped in first doc (only in non-raw mode) + if not self.raw_data_mode: + self.docset_index = docset_offset + self.lcg_state = lcg_offset + file_path, docrange, mindoc = self._get_docid(docset_offset) + # Use the saved document mapping from the first document processing + doclcg = first_doc_mapping + local_row = doclcg + mindoc + newpath = file_path + path, reader = self._get_reader(path, newpath, reader) + row = self._read_specific_row(reader, local_row) + data = row[self.data_column] + doc = self.process_fn(data) + + if doc[0] in self.drop: + doc = doc[1:] + if doc[-1] in self.drop: + doc = doc[:-1] + + doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 + if doclen >= self.min_length: + n_chunks = math.ceil(doclen / self.chunksize) + for j in range(residual_chunks): + self.chunk_index = j + out = self._construct_chunk(j, doc, n_chunks) + # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, first tokens: {out[:5]}...") + yield out + + def load_state_dict(self, state_dicts, sharded_input=False): + assert self.load_worldsize == self.worldsize, f"ParquetDataset does not support rescaling: from {self.load_worldsize} to {self.worldsize}" + return super().load_state_dict(state_dicts, sharded_input) diff --git a/maester/datasets/experimental_otf.py b/maester/datasets/experimental_otf.py index f078c36..ad48ec2 100644 --- a/maester/datasets/experimental_otf.py +++ b/maester/datasets/experimental_otf.py @@ -7,7 +7,7 @@ import os import random import time -from typing import Any, Callable, List, Optional, Set +from typing import Any, Callable, List, Optional, Set, Type, Union import pyarrow as pa import pyarrow.parquet as pq @@ -18,6 +18,11 @@ from transformers import AutoTokenizer, PreTrainedTokenizerFast from maester.log_utils import logger +from .formats import ( + JinxDataset, + ParquetDataset, +) + """ The following distributed dataloaders are designed around 3 main principles: @@ -79,16 +84,6 @@ def _shard_partition(itemlist: List[Any], rank: int, worldsize: int) -> List[Any ] -def _shard_inclusive(itemlist: List[Any], rank: int, worldsize: int) -> List[Any]: - """ - In cases where len(itemlist) % worldsize != 0, allow for fractional ownership of items, - and return the span including all owned items, fractional or otherwise. - """ - start = math.floor(len(itemlist) * rank / worldsize) - end = math.ceil(len(itemlist) * (rank + 1) / worldsize) - return itemlist[start:end] - - class _Stateful_Dataset(torch.utils.data.IterableDataset): """ Stub for stateful datasets, extends data.IterableDataset with state_dict methods. @@ -871,353 +866,7 @@ def load_state_dict(self, state_dicts, sharded_input=False): ), f"Dataset mismatch: checkpoint contains {self.dataset}, expected {d}" return out -class ParquetDataset(_Stateful_Dataset): - def __init__( - self, - data_dir: str, - rank: int, - worldsize: int, - tokenizer, - delimiter_token: Any, - bos_token: Optional[Any] = None, - strip_tokens: Optional[Set[Any]] = set(), - seed: int = 42, - min_length: int = 1, - max_chunksize: int = 1024, - verbose: bool = False, - shuffle: bool = True, - data_column: str = "text", - process_fn: Optional[Callable] = None, - raw_data_mode: bool = False, - ): - super(ParquetDataset, self).__init__(rank, worldsize) - self.seed = seed - self.data = data_dir - self.tokenizer = tokenizer - self.min_length = min_length - assert max_chunksize > 0, f"Max chunksize must be a nonzero positive integer" - self.chunksize = max_chunksize - self.eos = delimiter_token - self.bos = bos_token - self.drop = strip_tokens - self.verbose = verbose - self.data_column = data_column - self.raw_data_mode = raw_data_mode - self.process_fn = process_fn or self._default_process - self.docset: List[Any] = [] # map of doc indices to (file_path, min docid, max docid) - self.docs_per_file = {} - - # Guaranteed inconsistent shuffling across workers - random.seed(self.seed + rank) - - # Get all Parquet files in the directory recursively - self.parquet_files = [os.path.join(root, f) for root, _, files in os.walk(data_dir) for f in files if f.endswith('.parquet')] - self.parquet_files.sort() # Ensure consistent sharding across machines - assert len(self.parquet_files) > 0, "No parquet files found in data directory" - - dataset_hash = self._generate_dataset_hash() - cache_file = os.path.join(data_dir, f"doc_counts_cache_{dataset_hash}.json") - - # Rank 0 handles file I/O, other ranks wait - if dist.get_rank(dist.group.WORLD) == 0: - if not os.path.exists(cache_file): - self._gather_doc_counts() - self._save_cached_doc_counts(cache_file) - - dist.barrier() - - # All ranks load the cache - self._load_cached_doc_counts(cache_file) - - dist.barrier() # ensure all ranks loaded - - # Fragment the files - start_frag = (rank * worldsize * len(self.parquet_files)) // worldsize - end_frag = ((rank + 1) * worldsize * len(self.parquet_files)) // worldsize - shardfrags = [ - (self.parquet_files[i // worldsize], i % worldsize) for i in range(start_frag, end_frag) - ] - - # Read shardfrags, assemble doc list for each file shard (aggregating over fragments): - ndocs = -1 - docset = {} # shardid -> (min docid, max docid) - for i, (shard, frag) in enumerate(shardfrags): - ndocs = self.docs_per_file[shard] - doc_start = (ndocs * frag) // worldsize - doc_end = (ndocs * frag + ndocs) // worldsize - 1 # Inclusive upper bound - if shard not in docset: - docset[shard] = [doc_start, doc_end] - min_d, max_d = docset[shard] - if doc_start < min_d: - docset[shard][0] = doc_start - if doc_end > max_d: - docset[shard][1] = doc_end - - # Add all of this dataset's shard entries to self.docset - doccount = 0 - for shardid in docset: - min_d = docset[shardid][0] - max_d = docset[shardid][1] - self.docset.append((shardid, min_d, max_d)) - doccount += max_d - min_d + 1 - self._len = doccount - - if verbose: - logger.info(f"Worker {rank} responsible for docs: {self.docset}") - logger.info(f"Total docs: {doccount}") - - # Shuffle files - if shuffle: - random.shuffle(self.docset) - - self.docset_index = 0 - self.chunk_index = -1 - self.completed_current_doc = False - - # Stats - self.epochs_seen = -1 - self.tokens_seen = 0 - self.docs_seen = 0 - self.percent_seen = 0 - self.lcg_state = seed + rank - - self.state_params = [ - "docset_index", - "chunk_index", - "completed_current_doc", - "epochs_seen", - "tokens_seen", - "docs_seen", - "percent_seen", - "lcg_state", - ] - - def _default_process(self, data): - """Default processing: tokenize text data.""" - return self.tokenizer.encode(data, add_special_tokens=False, padding=False, truncation=False) - - def _generate_dataset_hash(self): - """Generate a unique hash for the dataset based on file names and sizes.""" - hasher = hashlib.md5() - for file in self.parquet_files: - hasher.update(file.encode()) - hasher.update(str(os.path.getsize(file)).encode()) - return hasher.hexdigest() - - def _load_cached_doc_counts(self, cache_file): - """Load cached document counts from a file.""" - start = time.time() - with open(cache_file, 'r') as f: - self.docs_per_file = json.load(f) - logger.info(f"Loaded cached document counts in {time.time() - start} seconds") - - def _save_cached_doc_counts(self, cache_file): - """Save document counts to a cache file.""" - with open(cache_file, 'w') as f: - json.dump(self.docs_per_file, f) - logger.info(f"Saved document counts cache to {cache_file}") - - def _gather_doc_counts(self): - """Gather document counts for each Parquet file.""" - start = time.time() - total_rows = 0 - for file in self.parquet_files: - parquet_file = pq.ParquetFile(file) - num_rows = parquet_file.metadata.num_rows - self.docs_per_file[file] = num_rows - total_rows += num_rows - assert total_rows > 0, "No rows found in parquet files" - logger.info(f"Gathered {total_rows} rows in {time.time() - start} seconds") - - def _get_docid(self, i): - """ - Given a global doc index over the set of docs owned by this worker, - return the corresponding path, num rows - """ - cur = 0 - assert i <= self._len, f"You have requested an illegal doc index {i}, docset length is {self._len}" - for shardid, min_d, max_d in self.docset: - docrange = max_d - min_d + 1 - cur += docrange - if cur > i: - return shardid, docrange, min_d - raise RuntimeError("This should be unreachable") - - def _get_reader(self, path, newpath, reader): - if newpath != path: - del reader - if self.verbose: - logger.info(f"Worker {self.rank} opening new file {newpath}") - reader = pq.ParquetFile(newpath) - path = newpath - return path, reader - - def _construct_chunk(self, j, doc, n_chunks): - """ - Construct the jth chunk of doc - """ - start_index = j * self.chunksize - n_pull = self.chunksize - if self.bos is not None: - if j == 0: - n_pull -= 1 - else: - start_index -= 1 - chunk = doc[start_index:start_index + n_pull] - self.tokens_seen += len(chunk) - # Add bos/eos tokens if needed - if self.bos is not None and j == 0: - chunk = [self.bos] + chunk - if j == n_chunks - 1: - chunk = chunk + [self.eos] - return chunk - def _random_map_docid(self, size): - """ - Given size of document pool, use saved state (prior index) to generate the next index via LCG. - Implements within-shard document shuffling without materializing any large doc lists. - """ - m = 2 ** math.ceil(math.log2(size)) # Round up to nearest power of 2 - a = 5 # A,C values known to work well with powers of 2 (Knuth, 1997, 3.2.1.3) - c = (self.rank + self.seed) * 2 + 1 - state = self.lcg_state - while True: - state = (a * state + c) % m - if state < size: - return state - - def _read_specific_row(self, reader, row_index): - row_group_index = 0 - rows_seen = 0 - for i in range(reader.num_row_groups): - num_rows = reader.metadata.row_group(i).num_rows - if rows_seen + num_rows > row_index: - row_group_index = i - break - rows_seen += num_rows - - row_offset = row_index - rows_seen - table = reader.read_row_group(row_group_index) - row = table.slice(row_offset, 1) - return row - - def __iter__(self): - docset_offset = self.docset_index - lcg_offset = self.lcg_state - residual_chunks = self.chunk_index + 1 # chunks to skip after restore and create at the end of epoch, 0-indexed - first_doc_mapping = None # Will store the document mapping for the first document - ndocs = self._len - path = "" - reader = None - if self.completed_current_doc: # resuming at the end of a doc - docset_offset = (docset_offset + 1) % ndocs - self.completed_current_doc = False - while True: - for i in range(ndocs): - doc_index = (docset_offset + i) % ndocs - self.completed_current_doc = False # reset - - # Update stats - if doc_index == 0: - self.epochs_seen += 1 - if self.verbose: - logger.info(f"ParquetDataset: entering epoch {self.epochs_seen}") - self.docset_index = doc_index - - # Map docset id to file, owned size and in-doc owned start idx - # This should be the same value many iters in a row, processing each shard - file_path, docrange, mindoc = self._get_docid(doc_index) - - # Map docset ids to consistently shuffled ids - # determine if we need a new document position - if i == 0 and not self.completed_current_doc and self.chunk_index >= 0: - # resuming mid-doc, do not advance lcg - doclcg = self.lcg_state - else: - doclcg = self._random_map_docid(docrange) # shuffled in-doc range - self.lcg_state = doclcg # update lcg state - - # Save the document mapping for the first document for residual processing - if i == 0: - first_doc_mapping = doclcg - - local_row = doclcg + mindoc # map docid to local row - - newpath = file_path - path, reader = self._get_reader(path, newpath, reader) - - table = self._read_specific_row(reader, local_row) - data = table[self.data_column][0].as_py() - - if self.raw_data_mode: - # In raw data mode, yield the data directly without processing - self.docs_seen += 1 - self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) - self.completed_current_doc = True - yield data - else: - # Normal mode: process and chunk the data - doc = self.process_fn(data) - if len(doc) < 2: - logger.warning(f"Empty document detected at {file_path}:{local_row}") - continue - - if doc[0] in self.drop: - doc = doc[1:] - if doc[-1] in self.drop: - doc = doc[:-1] - - doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 - if doclen >= self.min_length: - n_chunks = math.ceil(doclen / self.chunksize) - for j in range(n_chunks): - if i == 0 and not self.completed_current_doc and j < residual_chunks: - pass # skip already processed chunks - # doclcg = self.lcg_state # use saved lcg state when resuming - else: - self.chunk_index = j - # Document complete, update stats - if j == n_chunks - 1: - self.docs_seen += 1 - self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) - self.chunk_index = -1 - self.completed_current_doc = True - out = self._construct_chunk(j, doc, n_chunks) - # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, length {len(out)}, first tokens: {out[:5]}...") - yield out - - # Load any chunks initially skipped in first doc (only in non-raw mode) - if not self.raw_data_mode: - self.docset_index = docset_offset - self.lcg_state = lcg_offset - file_path, docrange, mindoc = self._get_docid(docset_offset) - # Use the saved document mapping from the first document processing - doclcg = first_doc_mapping - local_row = doclcg + mindoc - newpath = file_path - path, reader = self._get_reader(path, newpath, reader) - table = self._read_specific_row(reader, local_row) - data = table[self.data_column][0].as_py() - doc = self.process_fn(data) - - if doc[0] in self.drop: - doc = doc[1:] - if doc[-1] in self.drop: - doc = doc[:-1] - - doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 - if doclen >= self.min_length: - n_chunks = math.ceil(doclen / self.chunksize) - for j in range(residual_chunks): - self.chunk_index = j - out = self._construct_chunk(j, doc, n_chunks) - # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, first tokens: {out[:5]}...") - yield out - - def load_state_dict(self, state_dicts, sharded_input=False): - assert self.load_worldsize == self.worldsize, f"ParquetDataset does not support rescaling: from {self.load_worldsize} to {self.worldsize}" - return super().load_state_dict(state_dicts, sharded_input) - class Sampling_Dataset(_Stateful_Dataset): """ A _Stateful_Dataset implementing percentage-based sampling: weights can be floats, and the @@ -1247,6 +896,10 @@ class Sampling_Dataset(_Stateful_Dataset): def __init__( self, data_dirs: list[str], + dataset_types: list[Union[ + Type["ParquetDataset"], + Type["JINXDataset"], + ]], rank: int, worldsize: int, tokenizer, @@ -1272,9 +925,9 @@ def __init__( # Build subdataset iterators self.data = [] - for i, d in enumerate(data_dirs): + for i, (d, dataset_type) in enumerate(zip(data_dirs, dataset_types)): self.data.append( - ParquetDataset( + dataset_type( data_dir=d, rank=rank, worldsize=worldsize, @@ -1494,8 +1147,8 @@ def build_experimental_data_loader(cfg, rank, world_size): Number of distributed workers. Used for handling dataset sharding logic. """ - data_dirs, weights = parse_data_args( - cfg.dataset.data_dirs, cfg.dataset.dataset_weights + data_dirs, dataset_types, weights = parse_data_args( + cfg.dataset.data_dirs, cfg.dataset.dataset_types, cfg.dataset.dataset_weights ) def causal_lm(data_seq, prompt_len=0): @@ -1524,6 +1177,7 @@ def causal_lm(data_seq, prompt_len=0): droplist = droplist + [cfg.dataset.bos_token, cfg.dataset.eos_token] data = Sampling_Dataset( data_dirs, + dataset_types, rank, world_size, tokenizer, @@ -1561,7 +1215,7 @@ def causal_lm(data_seq, prompt_len=0): ) -def parse_data_args(datas, weights): +def parse_data_args(datas, dataset_types, weights): # Convert csv inputs into corresponding lists of values def splitstrip(x): if isinstance(x, str): @@ -1574,5 +1228,8 @@ def splitstrip(x): raise ValueError(f"arg input {x} cannot be parsed.") datas = splitstrip(datas) + if dataset_types is None: + dataset_types = ["ParquetDataset"] * len(datas) + dataset_types = [globals()[x] for x in splitstrip(dataset_types)] weights = [float(x) for x in splitstrip(weights)] - return datas, weights + return datas, dataset_types, weights diff --git a/maester/datasets/formats/__init__.py b/maester/datasets/formats/__init__.py new file mode 100644 index 0000000..ce4ed98 --- /dev/null +++ b/maester/datasets/formats/__init__.py @@ -0,0 +1,2 @@ +from .jinx import JinxDataset +from .parquet import ParquetDataset diff --git a/maester/datasets/formats/jinx.py b/maester/datasets/formats/jinx.py new file mode 100644 index 0000000..b194081 --- /dev/null +++ b/maester/datasets/formats/jinx.py @@ -0,0 +1,66 @@ +from mldataforge.jinx import JinxDatasetReader +from typing import Any, Callable, Optional, Set + +from ..experimental_otf import logger +from ..base import BaseDataset + +class JinxDataset(BaseDataset): + def __init__( + self, + data_dir: str, + rank: int, + worldsize: int, + tokenizer, + delimiter_token: Any, + bos_token: Optional[Any] = None, + strip_tokens: Optional[Set[Any]] = set(), + seed: int = 42, + min_length: int = 1, + max_chunksize: int = 1024, + verbose: bool = False, + shuffle: bool = True, + data_column: str = "text", + process_fn: Optional[Callable] = None, + raw_data_mode: bool = False, + ): + self.readers = {} + super(JinxDataset, self).__init__( + data_dir=data_dir, + data_ext='.jinx', + rank=rank, + worldsize=worldsize, + tokenizer=tokenizer, + delimiter_token=delimiter_token, + bos_token=bos_token, + strip_tokens=strip_tokens, + seed=seed, + min_length=min_length, + max_chunksize=max_chunksize, + verbose=verbose, + shuffle=shuffle, + data_column=data_column, + process_fn=process_fn, + raw_data_mode=raw_data_mode, + ) + + def _gather_doc_count(self, file): + """Count the number of documents in a JINX file.""" + reader = self.readers.get(file, None) + if reader is None: + reader = JinxDatasetReader(file) + self.readers[file] = reader + num_rows = len(reader) + return num_rows + + def _get_reader(self, path, newpath, reader): + # ignore newpath, keep map of files to readers + reader = self.readers.get(newpath, None) + if reader is None: + if self.verbose: + logger.info(f"Worker {self.rank} opening new file {newpath}") + reader = JinxDatasetReader(newpath) + self.readers[newpath] = reader + return newpath, reader + + def _read_specific_row(self, reader, row_index): + return reader[row_index] diff --git a/maester/datasets/formats/parquet.py b/maester/datasets/formats/parquet.py new file mode 100644 index 0000000..e1e0e75 --- /dev/null +++ b/maester/datasets/formats/parquet.py @@ -0,0 +1,73 @@ +import pyarrow.parquet as pq +from typing import Any, Callable, Optional, Set + +from ..experimental_otf import logger +from ..base import BaseDataset + +class ParquetDataset(BaseDataset): + def __init__( + self, + data_dir: str, + rank: int, + worldsize: int, + tokenizer, + delimiter_token: Any, + bos_token: Optional[Any] = None, + strip_tokens: Optional[Set[Any]] = set(), + seed: int = 42, + min_length: int = 1, + max_chunksize: int = 1024, + verbose: bool = False, + shuffle: bool = True, + data_column: str = "text", + process_fn: Optional[Callable] = None, + raw_data_mode: bool = False, + ): + super(ParquetDataset, self).__init__( + data_dir=data_dir, + data_ext='.parquet', + rank=rank, + worldsize=worldsize, + tokenizer=tokenizer, + delimiter_token=delimiter_token, + bos_token=bos_token, + strip_tokens=strip_tokens, + seed=seed, + min_length=min_length, + max_chunksize=max_chunksize, + verbose=verbose, + shuffle=shuffle, + data_column=data_column, + process_fn=process_fn, + raw_data_mode=raw_data_mode, + ) + + def _gather_doc_count(self, file): + """Count the number of documents in a Parquet file.""" + parquet_file = pq.ParquetFile(file) + num_rows = parquet_file.metadata.num_rows + return num_rows + + def _get_reader(self, path, newpath, reader): + if newpath != path: + del reader + if self.verbose: + logger.info(f"Worker {self.rank} opening new file {newpath}") + reader = pq.ParquetFile(newpath) + path = newpath + return path, reader + + def _read_specific_row(self, reader, row_index): + row_group_index = 0 + rows_seen = 0 + for i in range(reader.num_row_groups): + num_rows = reader.metadata.row_group(i).num_rows + if rows_seen + num_rows > row_index: + row_group_index = i + break + rows_seen += num_rows + + row_offset = row_index - rows_seen + table = reader.read_row_group(row_group_index) + row = table.to_struct_array()[row_offset].as_py() + return row diff --git a/maester/datasets/stateful.py b/maester/datasets/stateful.py new file mode 100644 index 0000000..cbfcae8 --- /dev/null +++ b/maester/datasets/stateful.py @@ -0,0 +1,145 @@ +import math +import os +from typing import List +import torch.utils.data +import torch +from typing import Any + +def _shard_inclusive(itemlist: List[Any], rank: int, worldsize: int) -> List[Any]: + """ + In cases where len(itemlist) % worldsize != 0, allow for fractional ownership of items, + and return the span including all owned items, fractional or otherwise. + """ + start = math.floor(len(itemlist) * rank / worldsize) + end = math.ceil(len(itemlist) * (rank + 1) / worldsize) + return itemlist[start:end] + +class _Stateful_Dataset(torch.utils.data.IterableDataset): + """ + Stub for stateful datasets, extends data.IterableDataset with state_dict methods. + All subclasses should specify the params to be considered stateful or reshardable in the + self.state_params and self.reshard_params lists. + """ + + def __init__( + self, + rank: int, + worldsize: int, + ): + assert rank >= 0, f"Rank {rank} must be a positive integer" + assert ( + worldsize > rank + ), f"Worldsize {worldsize} must be greater than rank {rank}" + self.state_params: List[str] = [] + self.reshard_params: List[str] = [] + self.rank = rank + self.worldsize = worldsize + self.load_worldsize = ( + worldsize # Enable calling load_state_dict() directly, assume no rescaling + ) + + def statename(self, x: str): + # Note that this naming convention implicitly disallows repeated layers in the dataset pipeline + return self.__class__.__name__ + "." + x + + def state_dict(self): + """ + Retrieve all state and reshard flags (each worker/process saves its own state dict shard) + """ + return { + self.statename(flag): getattr(self, flag) + for flag in self.state_params + self.reshard_params + } + + def _reshard(self, sharded_list): + """ + Sharded_list is a list of lists, where each "shard" sublist must have the same length. + These shards should tightly span only the partition of data owned by this worker. + (i.e. if global_list is the list of all entries, sharded_list = _shard_inclusive(global_list) ). + Determine fractional ownership of shards, and get the flattened partition owned by this worker. + """ + # How many shards did _shard_inclusive() drop to the left of sharded_list? + shard_offset = math.floor(self.load_worldsize * self.rank / self.worldsize) + # How long are the list shards? + shard_len = len(sharded_list[0]) + for i, shard in enumerate(sharded_list): + assert ( + len(shard) == shard_len + ), f"Shard {i} with length {len(shard)} does not match expected {shard_len}" + # How many list items did _shard_inclusive() drop to the left of the flattened sharded_list? + item_offset = shard_len * shard_offset + # How many list items are there in total? + n_items = self.load_worldsize * shard_len + # The indices of the flattened sharded_list that this worker owns + my_items = range( + int(n_items * self.rank / self.worldsize) - item_offset, + int(n_items * (self.rank + 1) / self.worldsize) - item_offset, + ) + # Pull out owned items + return [sharded_list[i // shard_len][i % shard_len] for i in my_items] + + def load_state_dict(self, state_dicts, sharded_input=False): + """ + Input state_dicts is a list of state_dicts. If sharded_input=False, this is expected to be the + global list of states across all checkpoint shard files. If sharded_input=True, this expects + _shard_inclusive(global_state_list). Handling reduced inputs allows for much more efficient loading. + Workflow: + 1. if sharded_inputs is false, shard the inputs. + 2. If worldsize matches checkpoint, pull state and reshard params from the given checkpoint + shard (state_dicts is a singleton list). + 3. If worldsize does not match checkpoint, toss state params and assemble reshard params from + across given state_dicts. In this case state_dicts may be singleton (for fractional ownership) + or multi-element (for multiple/partitioned ownership). + 4. Return reduced input for use by downstream loading functions + """ + if not sharded_input: + self.load_worldsize = len(state_dicts) + state_dicts = _shard_inclusive(state_dicts, self.rank, self.worldsize) + if self.load_worldsize == self.worldsize: + [ + setattr(self, flag, state_dicts[0][self.statename(flag)]) + for flag in self.state_params + self.reshard_params + ] + else: + for flag in self.reshard_params: + reshard = self._reshard( + [sd[self.statename(flag)] for sd in state_dicts] + ) + setattr(self, flag, reshard) + return state_dicts + + def load_from_path(self, path: str): + """ + Count shard files in the specified checkpoint folder and determine overlap with current + rank and worldsize partition. Load only matching shardfile(s) and pass to load_state_dict. + This is more efficient than sharding the full loaded state. + """ + assert os.path.exists(path), "Specified checkpoint does not exist" + assert not os.path.isfile(path), "Checkpoint should be a folder of shard states" + fileshards = [x for x in os.listdir(path) if "loader" in x] + fileshards = sorted(fileshards, key=lambda x: int(x.split("_")[2][:-4])) + assert ( + len(fileshards) > 0 + ), "Checkpoint directory must contain checkpoint files with 'loader' in the name" + self.load_worldsize = len(fileshards) + # Grab only the shard files holding data we currently own + my_fileshards = _shard_inclusive(fileshards, self.rank, self.worldsize) + states = [torch.load(os.path.join(path, x)) for x in my_fileshards] + self.load_state_dict(states, True) + + def save_to_path(self, path: str): + """ + Grab recursive shard states and save all shard states to the specified checkpoint folder + """ + os.makedirs(path, exist_ok=True) + state = self.state_dict() + + # worker_info = torch.utils.data.get_worker_info() + # if worker_info is None: + # print(f"single process: {worker_info.id}, {worker_info.num_workers}") + # print(state) + # else: + # print(f"worker process: {worker_info.id}, {worker_info.num_workers}") + # print(state) + + torch.save(state, os.path.join(path, f"loader_state_{self.rank}.pth")) From a9d0e1c125e3893476f00f17849ae3ed814e9706 Mon Sep 17 00:00:00 2001 From: Peter Schneider-Kamp Date: Tue, 7 Oct 2025 09:44:49 +0200 Subject: [PATCH 3/6] reworking refactoring --- maester/datasets/__init__.py | 11 - maester/datasets/base.py | 334 --------------------- maester/datasets/experimental_otf.py | 424 +++++++++++++++++++++++++-- maester/datasets/formats/__init__.py | 2 - maester/datasets/formats/jinx.py | 66 ----- maester/datasets/formats/parquet.py | 73 ----- maester/datasets/jinx_dataset.py | 25 ++ maester/datasets/stateful.py | 145 --------- 8 files changed, 428 insertions(+), 652 deletions(-) delete mode 100644 maester/datasets/__init__.py delete mode 100644 maester/datasets/base.py delete mode 100644 maester/datasets/formats/__init__.py delete mode 100644 maester/datasets/formats/jinx.py delete mode 100644 maester/datasets/formats/parquet.py create mode 100644 maester/datasets/jinx_dataset.py delete mode 100644 maester/datasets/stateful.py diff --git a/maester/datasets/__init__.py b/maester/datasets/__init__.py deleted file mode 100644 index 6db8186..0000000 --- a/maester/datasets/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -from .experimental_otf import * - -__all__ = [ - "StreamingDocDataset", - "ScalableShardDataset", - "SamplingDataset", - "PreloadBufferDataset", - "BufferDataset", - "PreprocessDataset", - "build_experimental_data_loader", -] diff --git a/maester/datasets/base.py b/maester/datasets/base.py deleted file mode 100644 index 25f03f4..0000000 --- a/maester/datasets/base.py +++ /dev/null @@ -1,334 +0,0 @@ -import hashlib -import json -import math -import os -import random -import time -import torch.distributed as dist -from typing import Any, Callable, List, Optional, Set - -from .experimental_otf import logger -from .stateful import _Stateful_Dataset - -class BaseDataset(_Stateful_Dataset): - def __init__( - self, - data_dir: str, - data_ext: str, - rank: int, - worldsize: int, - tokenizer, - delimiter_token: Any, - bos_token: Optional[Any] = None, - strip_tokens: Optional[Set[Any]] = set(), - seed: int = 42, - min_length: int = 1, - max_chunksize: int = 1024, - verbose: bool = False, - shuffle: bool = True, - data_column: str = "text", - process_fn: Optional[Callable] = None, - raw_data_mode: bool = False, - ): - super(BaseDataset, self).__init__(rank, worldsize) - self.seed = seed - self.data = data_dir - self.tokenizer = tokenizer - self.min_length = min_length - assert max_chunksize > 0, f"Max chunksize must be a nonzero positive integer" - self.chunksize = max_chunksize - self.eos = delimiter_token - self.bos = bos_token - self.drop = strip_tokens - self.verbose = verbose - self.data_column = data_column - self.raw_data_mode = raw_data_mode - self.process_fn = process_fn or self._default_process - self.docset: List[Any] = [] # map of doc indices to (file_path, min docid, max docid) - self.docs_per_file = {} - - # Guaranteed inconsistent shuffling across workers - random.seed(self.seed + rank) - - # Get all data files in the directory recursively - self.data_files = [os.path.join(root, f) for root, _, files in os.walk(data_dir) for f in files if f.endswith(data_ext)] - self.data_files.sort() # Ensure consistent sharding across machines - assert len(self.data_files) > 0, "No data files found in data directory" - - dataset_hash = self._generate_dataset_hash() - cache_file = os.path.join(data_dir, f"doc_counts_cache_{dataset_hash}.json") - - # Rank 0 handles file I/O, other ranks wait - if dist.get_rank(dist.group.WORLD) == 0: - if not os.path.exists(cache_file): - self._gather_doc_counts() - self._save_cached_doc_counts(cache_file) - - dist.barrier() - - # All ranks load the cache - self._load_cached_doc_counts(cache_file) - - dist.barrier() # ensure all ranks loaded - - # Fragment the files - start_frag = (rank * worldsize * len(self.data_files)) // worldsize - end_frag = ((rank + 1) * worldsize * len(self.data_files)) // worldsize - shardfrags = [ - (self.data_files[i // worldsize], i % worldsize) for i in range(start_frag, end_frag) - ] - - # Read shardfrags, assemble doc list for each file shard (aggregating over fragments): - ndocs = -1 - docset = {} # shardid -> (min docid, max docid) - for i, (shard, frag) in enumerate(shardfrags): - ndocs = self.docs_per_file[shard] - doc_start = (ndocs * frag) // worldsize - doc_end = (ndocs * frag + ndocs) // worldsize - 1 # Inclusive upper bound - if shard not in docset: - docset[shard] = [doc_start, doc_end] - min_d, max_d = docset[shard] - if doc_start < min_d: - docset[shard][0] = doc_start - if doc_end > max_d: - docset[shard][1] = doc_end - - # Add all of this dataset's shard entries to self.docset - doccount = 0 - for shardid in docset: - min_d = docset[shardid][0] - max_d = docset[shardid][1] - self.docset.append((shardid, min_d, max_d)) - doccount += max_d - min_d + 1 - self._len = doccount - - if verbose: - logger.info(f"Worker {rank} responsible for docs: {self.docset}") - logger.info(f"Total docs: {doccount}") - - # Shuffle files - if shuffle: - random.shuffle(self.docset) - - self.docset_index = 0 - self.chunk_index = -1 - self.completed_current_doc = False - - # Stats - self.epochs_seen = -1 - self.tokens_seen = 0 - self.docs_seen = 0 - self.percent_seen = 0 - self.lcg_state = seed + rank - - self.state_params = [ - "docset_index", - "chunk_index", - "completed_current_doc", - "epochs_seen", - "tokens_seen", - "docs_seen", - "percent_seen", - "lcg_state", - ] - - def _default_process(self, data): - """Default processing: tokenize text data.""" - return self.tokenizer.encode(data, add_special_tokens=False, padding=False, truncation=False) - - def _generate_dataset_hash(self): - """Generate a unique hash for the dataset based on file names and sizes.""" - hasher = hashlib.md5() - for file in self.data_files: - hasher.update(file.encode()) - hasher.update(str(os.path.getsize(file)).encode()) - return hasher.hexdigest() - - def _load_cached_doc_counts(self, cache_file): - """Load cached document counts from a file.""" - start = time.time() - with open(cache_file, 'r') as f: - self.docs_per_file = json.load(f) - logger.info(f"Loaded cached document counts in {time.time() - start} seconds") - - def _save_cached_doc_counts(self, cache_file): - """Save document counts to a cache file.""" - with open(cache_file, 'w') as f: - json.dump(self.docs_per_file, f) - logger.info(f"Saved document counts cache to {cache_file}") - - def _gather_doc_counts(self): - """Gather document counts for each Parquet file.""" - start = time.time() - total_rows = 0 - for file in self.data_files: - num_rows = self._gather_doc_count(file) - self.docs_per_file[file] = num_rows - total_rows += num_rows - assert total_rows > 0, "No rows found in parquet files" - logger.info(f"Gathered {total_rows} rows in {time.time() - start} seconds") - - def _get_docid(self, i): - """ - Given a global doc index over the set of docs owned by this worker, - return the corresponding path, num rows - """ - cur = 0 - assert i <= self._len, f"You have requested an illegal doc index {i}, docset length is {self._len}" - for shardid, min_d, max_d in self.docset: - docrange = max_d - min_d + 1 - cur += docrange - if cur > i: - return shardid, docrange, min_d - raise RuntimeError("This should be unreachable") - - def _construct_chunk(self, j, doc, n_chunks): - """ - Construct the jth chunk of doc - """ - start_index = j * self.chunksize - n_pull = self.chunksize - if self.bos is not None: - if j == 0: - n_pull -= 1 - else: - start_index -= 1 - chunk = doc[start_index:start_index + n_pull] - self.tokens_seen += len(chunk) - # Add bos/eos tokens if needed - if self.bos is not None and j == 0: - chunk = [self.bos] + chunk - if j == n_chunks - 1: - chunk = chunk + [self.eos] - return chunk - - def _random_map_docid(self, size): - """ - Given size of document pool, use saved state (prior index) to generate the next index via LCG. - Implements within-shard document shuffling without materializing any large doc lists. - """ - m = 2 ** math.ceil(math.log2(size)) # Round up to nearest power of 2 - a = 5 # A,C values known to work well with powers of 2 (Knuth, 1997, 3.2.1.3) - c = (self.rank + self.seed) * 2 + 1 - state = self.lcg_state - while True: - state = (a * state + c) % m - if state < size: - return state - - def __iter__(self): - docset_offset = self.docset_index - lcg_offset = self.lcg_state - residual_chunks = self.chunk_index + 1 # chunks to skip after restore and create at the end of epoch, 0-indexed - first_doc_mapping = None # Will store the document mapping for the first document - ndocs = self._len - path = "" - reader = None - if self.completed_current_doc: # resuming at the end of a doc - docset_offset = (docset_offset + 1) % ndocs - self.completed_current_doc = False - while True: - for i in range(ndocs): - doc_index = (docset_offset + i) % ndocs - self.completed_current_doc = False # reset - - # Update stats - if doc_index == 0: - self.epochs_seen += 1 - if self.verbose: - logger.info(f"ParquetDataset: entering epoch {self.epochs_seen}") - self.docset_index = doc_index - - # Map docset id to file, owned size and in-doc owned start idx - # This should be the same value many iters in a row, processing each shard - file_path, docrange, mindoc = self._get_docid(doc_index) - - # Map docset ids to consistently shuffled ids - # determine if we need a new document position - if i == 0 and not self.completed_current_doc and self.chunk_index >= 0: - # resuming mid-doc, do not advance lcg - doclcg = self.lcg_state - else: - doclcg = self._random_map_docid(docrange) # shuffled in-doc range - self.lcg_state = doclcg # update lcg state - - # Save the document mapping for the first document for residual processing - if i == 0: - first_doc_mapping = doclcg - - local_row = doclcg + mindoc # map docid to local row - - newpath = file_path - path, reader = self._get_reader(path, newpath, reader) - - row = self._read_specific_row(reader, local_row) - data = row[self.data_column] - - if self.raw_data_mode: - # In raw data mode, yield the data directly without processing - self.docs_seen += 1 - self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) - self.completed_current_doc = True - yield data - else: - # Normal mode: process and chunk the data - doc = self.process_fn(data) - if len(doc) < 2: - logger.warning(f"Empty document detected at {file_path}:{local_row}") - continue - - if doc[0] in self.drop: - doc = doc[1:] - if doc[-1] in self.drop: - doc = doc[:-1] - - doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 - if doclen >= self.min_length: - n_chunks = math.ceil(doclen / self.chunksize) - for j in range(n_chunks): - if i == 0 and not self.completed_current_doc and j < residual_chunks: - pass # skip already processed chunks - # doclcg = self.lcg_state # use saved lcg state when resuming - else: - self.chunk_index = j - # Document complete, update stats - if j == n_chunks - 1: - self.docs_seen += 1 - self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) - self.chunk_index = -1 - self.completed_current_doc = True - out = self._construct_chunk(j, doc, n_chunks) - # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, length {len(out)}, first tokens: {out[:5]}...") - yield out - - # Load any chunks initially skipped in first doc (only in non-raw mode) - if not self.raw_data_mode: - self.docset_index = docset_offset - self.lcg_state = lcg_offset - file_path, docrange, mindoc = self._get_docid(docset_offset) - # Use the saved document mapping from the first document processing - doclcg = first_doc_mapping - local_row = doclcg + mindoc - newpath = file_path - path, reader = self._get_reader(path, newpath, reader) - row = self._read_specific_row(reader, local_row) - data = row[self.data_column] - doc = self.process_fn(data) - - if doc[0] in self.drop: - doc = doc[1:] - if doc[-1] in self.drop: - doc = doc[:-1] - - doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 - if doclen >= self.min_length: - n_chunks = math.ceil(doclen / self.chunksize) - for j in range(residual_chunks): - self.chunk_index = j - out = self._construct_chunk(j, doc, n_chunks) - # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, first tokens: {out[:5]}...") - yield out - - def load_state_dict(self, state_dicts, sharded_input=False): - assert self.load_worldsize == self.worldsize, f"ParquetDataset does not support rescaling: from {self.load_worldsize} to {self.worldsize}" - return super().load_state_dict(state_dicts, sharded_input) diff --git a/maester/datasets/experimental_otf.py b/maester/datasets/experimental_otf.py index 6b3ab25..9bd2864 100644 --- a/maester/datasets/experimental_otf.py +++ b/maester/datasets/experimental_otf.py @@ -9,7 +9,7 @@ import random import time from collections import OrderedDict -from typing import Any, Callable, Dict, List, Optional, Set, Type, Union +from typing import Any, Callable, Dict, List, Optional, Set import pyarrow as pa import pyarrow.parquet as pq @@ -20,11 +20,6 @@ from transformers import AutoTokenizer, PreTrainedTokenizerFast from maester.log_utils import logger -from .formats import ( - JinxDataset, - ParquetDataset, -) - """ The following distributed dataloaders are designed around 3 main principles: @@ -86,6 +81,16 @@ def _shard_partition(itemlist: List[Any], rank: int, worldsize: int) -> List[Any ] +def _shard_inclusive(itemlist: List[Any], rank: int, worldsize: int) -> List[Any]: + """ + In cases where len(itemlist) % worldsize != 0, allow for fractional ownership of items, + and return the span including all owned items, fractional or otherwise. + """ + start = math.floor(len(itemlist) * rank / worldsize) + end = math.ceil(len(itemlist) * (rank + 1) / worldsize) + return itemlist[start:end] + + class _Stateful_Dataset(torch.utils.data.IterableDataset): """ Stub for stateful datasets, extends data.IterableDataset with state_dict methods. @@ -868,7 +873,388 @@ def load_state_dict(self, state_dicts, sharded_input=False): ), f"Dataset mismatch: checkpoint contains {self.dataset}, expected {d}" return out +class ParquetDataset(_Stateful_Dataset): + def __init__( + self, + data_dir: str, + rank: int, + worldsize: int, + tokenizer, + delimiter_token: Any, + bos_token: Optional[Any] = None, + strip_tokens: Optional[Set[Any]] = set(), + seed: int = 42, + min_length: int = 1, + max_chunksize: int = 1024, + verbose: bool = False, + shuffle: bool = True, + data_column: str = "text", + process_fn: Optional[Callable] = None, + raw_data_mode: bool = False, + cache_row_groups: bool = True, + ): + super(ParquetDataset, self).__init__(rank, worldsize) + self.seed = seed + self.data = data_dir + self.tokenizer = tokenizer + self.min_length = min_length + assert max_chunksize > 0, f"Max chunksize must be a nonzero positive integer" + self.chunksize = max_chunksize + self.eos = delimiter_token + self.bos = bos_token + self.drop = strip_tokens + self.verbose = verbose + self.data_column = data_column + self.raw_data_mode = raw_data_mode + self.process_fn = process_fn or self._default_process + self.cache_row_groups = cache_row_groups # toggle to keep parquet row groups in memory + self.docset: List[Any] = [] # map of doc indices to (file_path, min docid, max docid) + self.docs_per_file = {} + # Row-group metadata and cache keyed by file path, used to avoid repeated scans and reads + self._row_group_prefix: Dict[str, List[int]] = {} + self._row_group_cache: Dict[str, OrderedDict[int, pa.Table]] = {} + + # Guaranteed inconsistent shuffling across workers + random.seed(self.seed + rank) + + # Get all Parquet files in the directory recursively + self.parquet_files = [os.path.join(root, f) for root, _, files in os.walk(data_dir) for f in files if f.endswith('.parquet')] + self.parquet_files.sort() # Ensure consistent sharding across machines + assert len(self.parquet_files) > 0, "No parquet files found in data directory" + + dataset_hash = self._generate_dataset_hash() + cache_file = os.path.join(data_dir, f"doc_counts_cache_{dataset_hash}.json") + + # Rank 0 handles file I/O, other ranks wait + if dist.get_rank(dist.group.WORLD) == 0: + if not os.path.exists(cache_file): + self._gather_doc_counts() + self._save_cached_doc_counts(cache_file) + + dist.barrier() + + # All ranks load the cache + self._load_cached_doc_counts(cache_file) + + dist.barrier() # ensure all ranks loaded + + # Fragment the files + start_frag = (rank * worldsize * len(self.parquet_files)) // worldsize + end_frag = ((rank + 1) * worldsize * len(self.parquet_files)) // worldsize + shardfrags = [ + (self.parquet_files[i // worldsize], i % worldsize) for i in range(start_frag, end_frag) + ] + + # Read shardfrags, assemble doc list for each file shard (aggregating over fragments): + ndocs = -1 + docset = {} # shardid -> (min docid, max docid) + for i, (shard, frag) in enumerate(shardfrags): + ndocs = self.docs_per_file[shard] + doc_start = (ndocs * frag) // worldsize + doc_end = (ndocs * frag + ndocs) // worldsize - 1 # Inclusive upper bound + if shard not in docset: + docset[shard] = [doc_start, doc_end] + min_d, max_d = docset[shard] + if doc_start < min_d: + docset[shard][0] = doc_start + if doc_end > max_d: + docset[shard][1] = doc_end + + # Add all of this dataset's shard entries to self.docset + doccount = 0 + for shardid in docset: + min_d = docset[shardid][0] + max_d = docset[shardid][1] + self.docset.append((shardid, min_d, max_d)) + doccount += max_d - min_d + 1 + self._len = doccount + + if verbose: + logger.info(f"Worker {rank} responsible for docs: {self.docset}") + logger.info(f"Total docs: {doccount}") + + # Shuffle files + if shuffle: + random.shuffle(self.docset) + + self.docset_index = 0 + self.chunk_index = -1 + self.completed_current_doc = False + + # Stats + self.epochs_seen = -1 + self.tokens_seen = 0 + self.docs_seen = 0 + self.percent_seen = 0 + self.lcg_state = seed + rank + + self.state_params = [ + "docset_index", + "chunk_index", + "completed_current_doc", + "epochs_seen", + "tokens_seen", + "docs_seen", + "percent_seen", + "lcg_state", + ] + + def _default_process(self, data): + """Default processing: tokenize text data.""" + return self.tokenizer.encode(data, add_special_tokens=False, padding=False, truncation=False) + + def _generate_dataset_hash(self): + """Generate a unique hash for the dataset based on file names and sizes.""" + hasher = hashlib.md5() + for file in self.parquet_files: + hasher.update(file.encode()) + hasher.update(str(os.path.getsize(file)).encode()) + return hasher.hexdigest() + def _load_cached_doc_counts(self, cache_file): + """Load cached document counts from a file.""" + start = time.time() + with open(cache_file, 'r') as f: + self.docs_per_file = json.load(f) + logger.info(f"Loaded cached document counts in {time.time() - start} seconds") + + def _save_cached_doc_counts(self, cache_file): + """Save document counts to a cache file.""" + with open(cache_file, 'w') as f: + json.dump(self.docs_per_file, f) + logger.info(f"Saved document counts cache to {cache_file}") + + def _gather_doc_counts(self): + """Gather document counts for each Parquet file.""" + start = time.time() + total_rows = 0 + for file in self.parquet_files: + parquet_file = pq.ParquetFile(file) + num_rows = parquet_file.metadata.num_rows + self.docs_per_file[file] = num_rows + total_rows += num_rows + assert total_rows > 0, "No rows found in parquet files" + logger.info(f"Gathered {total_rows} rows in {time.time() - start} seconds") + + def _get_docid(self, i): + """ + Given a global doc index over the set of docs owned by this worker, + return the corresponding path, num rows + """ + cur = 0 + assert i <= self._len, f"You have requested an illegal doc index {i}, docset length is {self._len}" + for shardid, min_d, max_d in self.docset: + docrange = max_d - min_d + 1 + cur += docrange + if cur > i: + return shardid, docrange, min_d + raise RuntimeError("This should be unreachable") + + def _prepare_row_group_metadata(self, path: str, reader: pq.ParquetFile): + num_groups = reader.num_row_groups + assert num_groups > 0, f"Parquet file {path} has no row groups" + + if path not in self._row_group_prefix: + prefix = [] + running_total = 0 + for i in range(num_groups): + running_total += reader.metadata.row_group(i).num_rows + prefix.append(running_total) + self._row_group_prefix[path] = prefix + + if self.cache_row_groups and path not in self._row_group_cache: + self._row_group_cache[path] = OrderedDict() + + def _get_reader(self, path, newpath, reader): + if newpath != path: + if self.cache_row_groups and path in self._row_group_cache: + self._row_group_cache.pop(path, None) + if reader is not None: + del reader + if self.verbose: + logger.info(f"Worker {self.rank} opening new file {newpath}") + reader = pq.ParquetFile(newpath) + path = newpath + self._prepare_row_group_metadata(path, reader) + elif path and path not in self._row_group_prefix: + self._prepare_row_group_metadata(path, reader) + return path, reader + + def _construct_chunk(self, j, doc, n_chunks): + """ + Construct the jth chunk of doc + """ + start_index = j * self.chunksize + n_pull = self.chunksize + if self.bos is not None: + if j == 0: + n_pull -= 1 + else: + start_index -= 1 + chunk = doc[start_index:start_index + n_pull] + self.tokens_seen += len(chunk) + # Add bos/eos tokens if needed + if self.bos is not None and j == 0: + chunk = [self.bos] + chunk + if j == n_chunks - 1: + chunk = chunk + [self.eos] + return chunk + + def _random_map_docid(self, size): + """ + Given size of document pool, use saved state (prior index) to generate the next index via LCG. + Implements within-shard document shuffling without materializing any large doc lists. + """ + m = 2 ** math.ceil(math.log2(size)) # Round up to nearest power of 2 + a = 5 # A,C values known to work well with powers of 2 (Knuth, 1997, 3.2.1.3) + c = (self.rank + self.seed) * 2 + 1 + state = self.lcg_state + while True: + state = (a * state + c) % m + if state < size: + return state + + def _read_specific_row(self, path, reader, row_index): + prefix = self._row_group_prefix[path] + row_group_index = bisect.bisect_right(prefix, row_index) + assert ( + row_group_index < len(prefix) + ), f"Row index {row_index} exceeds total rows for {path}" + rows_before_group = prefix[row_group_index - 1] if row_group_index > 0 else 0 + row_offset = row_index - rows_before_group + + if not self.cache_row_groups: + table = reader.read_row_group(row_group_index, columns=[self.data_column]) + return table.slice(row_offset, 1) + + cache = self._row_group_cache[path] + if row_group_index in cache: + table = cache[row_group_index] + cache.move_to_end(row_group_index) + else: + table = reader.read_row_group(row_group_index, columns=[self.data_column]) + cache[row_group_index] = table + cache.move_to_end(row_group_index) + + return table.slice(row_offset, 1) + + def __iter__(self): + docset_offset = self.docset_index + lcg_offset = self.lcg_state + residual_chunks = self.chunk_index + 1 # chunks to skip after restore and create at the end of epoch, 0-indexed + first_doc_mapping = None # Will store the document mapping for the first document + ndocs = self._len + path = "" + reader = None + if self.completed_current_doc: # resuming at the end of a doc + docset_offset = (docset_offset + 1) % ndocs + self.completed_current_doc = False + while True: + for i in range(ndocs): + doc_index = (docset_offset + i) % ndocs + self.completed_current_doc = False # reset + + # Update stats + if doc_index == 0: + self.epochs_seen += 1 + if self.verbose: + logger.info(f"ParquetDataset: entering epoch {self.epochs_seen}") + self.docset_index = doc_index + + # Map docset id to file, owned size and in-doc owned start idx + # This should be the same value many iters in a row, processing each shard + file_path, docrange, mindoc = self._get_docid(doc_index) + + # Map docset ids to consistently shuffled ids + # determine if we need a new document position + if i == 0 and not self.completed_current_doc and self.chunk_index >= 0: + # resuming mid-doc, do not advance lcg + doclcg = self.lcg_state + else: + doclcg = self._random_map_docid(docrange) # shuffled in-doc range + self.lcg_state = doclcg # update lcg state + + # Save the document mapping for the first document for residual processing + if i == 0: + first_doc_mapping = doclcg + + local_row = doclcg + mindoc # map docid to local row + + newpath = file_path + path, reader = self._get_reader(path, newpath, reader) + + table = self._read_specific_row(path, reader, local_row) + data = table[self.data_column][0].as_py() + + if self.raw_data_mode: + # In raw data mode, yield the data directly without processing + self.docs_seen += 1 + self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) + self.completed_current_doc = True + yield data + else: + # Normal mode: process and chunk the data + doc = self.process_fn(data) + if len(doc) < 2: + logger.warning(f"Empty document detected at {file_path}:{local_row}") + continue + + if doc[0] in self.drop: + doc = doc[1:] + if doc[-1] in self.drop: + doc = doc[:-1] + + doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 + if doclen >= self.min_length: + n_chunks = math.ceil(doclen / self.chunksize) + for j in range(n_chunks): + if i == 0 and not self.completed_current_doc and j < residual_chunks: + pass # skip already processed chunks + # doclcg = self.lcg_state # use saved lcg state when resuming + else: + self.chunk_index = j + # Document complete, update stats + if j == n_chunks - 1: + self.docs_seen += 1 + self.percent_seen = (self.docs_seen * 100 / (self._len + 1e-9)) + self.chunk_index = -1 + self.completed_current_doc = True + out = self._construct_chunk(j, doc, n_chunks) + # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, length {len(out)}, first tokens: {out[:5]}...") + yield out + + # Load any chunks initially skipped in first doc (only in non-raw mode) + if not self.raw_data_mode: + self.docset_index = docset_offset + self.lcg_state = lcg_offset + file_path, docrange, mindoc = self._get_docid(docset_offset) + # Use the saved document mapping from the first document processing + doclcg = first_doc_mapping + local_row = doclcg + mindoc + newpath = file_path + path, reader = self._get_reader(path, newpath, reader) + table = self._read_specific_row(path, reader, local_row) + data = table[self.data_column][0].as_py() + doc = self.process_fn(data) + + if doc[0] in self.drop: + doc = doc[1:] + if doc[-1] in self.drop: + doc = doc[:-1] + + doclen = len(doc) + 1 if self.bos is None else len(doc) + 2 + if doclen >= self.min_length: + n_chunks = math.ceil(doclen / self.chunksize) + for j in range(residual_chunks): + self.chunk_index = j + out = self._construct_chunk(j, doc, n_chunks) + # print(f"ParquetDataset: yielding chunk {j}/{n_chunks}, first tokens: {out[:5]}...") + yield out + + def load_state_dict(self, state_dicts, sharded_input=False): + assert self.load_worldsize == self.worldsize, f"ParquetDataset does not support rescaling: from {self.load_worldsize} to {self.worldsize}" + return super().load_state_dict(state_dicts, sharded_input) + class Sampling_Dataset(_Stateful_Dataset): """ A _Stateful_Dataset implementing percentage-based sampling: weights can be floats, and the @@ -898,10 +1284,6 @@ class Sampling_Dataset(_Stateful_Dataset): def __init__( self, data_dirs: list[str], - dataset_types: list[Union[ - Type["ParquetDataset"], - Type["JINXDataset"], - ]], rank: int, worldsize: int, tokenizer, @@ -927,9 +1309,9 @@ def __init__( # Build subdataset iterators self.data = [] - for i, (d, dataset_type) in enumerate(zip(data_dirs, dataset_types)): + for i, d in enumerate(data_dirs): self.data.append( - dataset_type( + ParquetDataset( data_dir=d, rank=rank, worldsize=worldsize, @@ -1149,8 +1531,8 @@ def build_experimental_data_loader(cfg, rank, world_size): Number of distributed workers. Used for handling dataset sharding logic. """ - data_dirs, dataset_types, weights = parse_data_args( - cfg.dataset.data_dirs, cfg.dataset.dataset_types, cfg.dataset.dataset_weights + data_dirs, weights = parse_data_args( + cfg.dataset.data_dirs, cfg.dataset.dataset_weights ) def causal_lm(data_seq, prompt_len=0): @@ -1179,7 +1561,6 @@ def causal_lm(data_seq, prompt_len=0): droplist = droplist + [cfg.dataset.bos_token, cfg.dataset.eos_token] data = Sampling_Dataset( data_dirs, - dataset_types, rank, world_size, tokenizer, @@ -1206,11 +1587,15 @@ def causal_lm(data_seq, prompt_len=0): data = Preprocess_Dataset(data, causal_lm) # Enable auto-saving if cfg.enable_checkpoint: # and not cfg.model_weights_only: # model_weights_only is only for final weight export... + grad_accum_steps = getattr(cfg, "gradient_accumulation_steps", 1) + grad_accum_steps = max(1, grad_accum_steps) + # Track optimizer-scale steps so dataloader checkpoints line up with model checkpoints + steps_per_checkpoint_step = cfg.train_batch_size * grad_accum_steps data = Checkpoint_Dataset( data, os.path.join(cfg.dump_dir, cfg.job_name, cfg.checkpoint_folder, "dataloader"), cfg.checkpoint_interval, - cfg.train_batch_size, + steps_per_checkpoint_step, ) # warning: things *will* break with num_workers > 1 return torch.utils.data.DataLoader( @@ -1218,7 +1603,7 @@ def causal_lm(data_seq, prompt_len=0): ) -def parse_data_args(datas, dataset_types, weights): +def parse_data_args(datas, weights): # Convert csv inputs into corresponding lists of values def splitstrip(x): if isinstance(x, str): @@ -1231,8 +1616,5 @@ def splitstrip(x): raise ValueError(f"arg input {x} cannot be parsed.") datas = splitstrip(datas) - if dataset_types is None: - dataset_types = ["ParquetDataset"] * len(datas) - dataset_types = [globals()[x] for x in splitstrip(dataset_types)] weights = [float(x) for x in splitstrip(weights)] - return datas, dataset_types, weights + return datas, weights diff --git a/maester/datasets/formats/__init__.py b/maester/datasets/formats/__init__.py deleted file mode 100644 index ce4ed98..0000000 --- a/maester/datasets/formats/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .jinx import JinxDataset -from .parquet import ParquetDataset diff --git a/maester/datasets/formats/jinx.py b/maester/datasets/formats/jinx.py deleted file mode 100644 index b194081..0000000 --- a/maester/datasets/formats/jinx.py +++ /dev/null @@ -1,66 +0,0 @@ -from mldataforge.jinx import JinxDatasetReader -from typing import Any, Callable, Optional, Set - -from ..experimental_otf import logger -from ..base import BaseDataset - -class JinxDataset(BaseDataset): - def __init__( - self, - data_dir: str, - rank: int, - worldsize: int, - tokenizer, - delimiter_token: Any, - bos_token: Optional[Any] = None, - strip_tokens: Optional[Set[Any]] = set(), - seed: int = 42, - min_length: int = 1, - max_chunksize: int = 1024, - verbose: bool = False, - shuffle: bool = True, - data_column: str = "text", - process_fn: Optional[Callable] = None, - raw_data_mode: bool = False, - ): - self.readers = {} - super(JinxDataset, self).__init__( - data_dir=data_dir, - data_ext='.jinx', - rank=rank, - worldsize=worldsize, - tokenizer=tokenizer, - delimiter_token=delimiter_token, - bos_token=bos_token, - strip_tokens=strip_tokens, - seed=seed, - min_length=min_length, - max_chunksize=max_chunksize, - verbose=verbose, - shuffle=shuffle, - data_column=data_column, - process_fn=process_fn, - raw_data_mode=raw_data_mode, - ) - - def _gather_doc_count(self, file): - """Count the number of documents in a JINX file.""" - reader = self.readers.get(file, None) - if reader is None: - reader = JinxDatasetReader(file) - self.readers[file] = reader - num_rows = len(reader) - return num_rows - - def _get_reader(self, path, newpath, reader): - # ignore newpath, keep map of files to readers - reader = self.readers.get(newpath, None) - if reader is None: - if self.verbose: - logger.info(f"Worker {self.rank} opening new file {newpath}") - reader = JinxDatasetReader(newpath) - self.readers[newpath] = reader - return newpath, reader - - def _read_specific_row(self, reader, row_index): - return reader[row_index] diff --git a/maester/datasets/formats/parquet.py b/maester/datasets/formats/parquet.py deleted file mode 100644 index e1e0e75..0000000 --- a/maester/datasets/formats/parquet.py +++ /dev/null @@ -1,73 +0,0 @@ -import pyarrow.parquet as pq -from typing import Any, Callable, Optional, Set - -from ..experimental_otf import logger -from ..base import BaseDataset - -class ParquetDataset(BaseDataset): - def __init__( - self, - data_dir: str, - rank: int, - worldsize: int, - tokenizer, - delimiter_token: Any, - bos_token: Optional[Any] = None, - strip_tokens: Optional[Set[Any]] = set(), - seed: int = 42, - min_length: int = 1, - max_chunksize: int = 1024, - verbose: bool = False, - shuffle: bool = True, - data_column: str = "text", - process_fn: Optional[Callable] = None, - raw_data_mode: bool = False, - ): - super(ParquetDataset, self).__init__( - data_dir=data_dir, - data_ext='.parquet', - rank=rank, - worldsize=worldsize, - tokenizer=tokenizer, - delimiter_token=delimiter_token, - bos_token=bos_token, - strip_tokens=strip_tokens, - seed=seed, - min_length=min_length, - max_chunksize=max_chunksize, - verbose=verbose, - shuffle=shuffle, - data_column=data_column, - process_fn=process_fn, - raw_data_mode=raw_data_mode, - ) - - def _gather_doc_count(self, file): - """Count the number of documents in a Parquet file.""" - parquet_file = pq.ParquetFile(file) - num_rows = parquet_file.metadata.num_rows - return num_rows - - def _get_reader(self, path, newpath, reader): - if newpath != path: - del reader - if self.verbose: - logger.info(f"Worker {self.rank} opening new file {newpath}") - reader = pq.ParquetFile(newpath) - path = newpath - return path, reader - - def _read_specific_row(self, reader, row_index): - row_group_index = 0 - rows_seen = 0 - for i in range(reader.num_row_groups): - num_rows = reader.metadata.row_group(i).num_rows - if rows_seen + num_rows > row_index: - row_group_index = i - break - rows_seen += num_rows - - row_offset = row_index - rows_seen - table = reader.read_row_group(row_group_index) - row = table.to_struct_array()[row_offset].as_py() - return row diff --git a/maester/datasets/jinx_dataset.py b/maester/datasets/jinx_dataset.py new file mode 100644 index 0000000..3656997 --- /dev/null +++ b/maester/datasets/jinx_dataset.py @@ -0,0 +1,25 @@ +from mldataforge.jinx import JinxDatasetReader + +from .experimental_otf import logger +from .base import ParquetDataset + +class JinxDataset(ParquetDataset): + def __init__(self, *args, **kwargs): + self._jinx_readers = {} + super(JinxDataset, self).__init__(*args, **kwargs) + + def _gather_doc_count(self, file): + _, reader = self._get_reader(None, file, None) + return len(reader) + + def _get_reader(self, path, newpath, reader): + reader = self._jinx_readers.get(newpath, None) + if reader is None: + if self.verbose: + logger.info(f"Worker {self.rank} opening new file {newpath}") + reader = JinxDatasetReader(newpath) + self._jinx_readers[newpath] = reader + return newpath, reader + + def _read_specific_row(self, reader, row_index): + return reader[row_index] diff --git a/maester/datasets/stateful.py b/maester/datasets/stateful.py deleted file mode 100644 index cbfcae8..0000000 --- a/maester/datasets/stateful.py +++ /dev/null @@ -1,145 +0,0 @@ -import math -import os -from typing import List -import torch.utils.data -import torch -from typing import Any - -def _shard_inclusive(itemlist: List[Any], rank: int, worldsize: int) -> List[Any]: - """ - In cases where len(itemlist) % worldsize != 0, allow for fractional ownership of items, - and return the span including all owned items, fractional or otherwise. - """ - start = math.floor(len(itemlist) * rank / worldsize) - end = math.ceil(len(itemlist) * (rank + 1) / worldsize) - return itemlist[start:end] - -class _Stateful_Dataset(torch.utils.data.IterableDataset): - """ - Stub for stateful datasets, extends data.IterableDataset with state_dict methods. - All subclasses should specify the params to be considered stateful or reshardable in the - self.state_params and self.reshard_params lists. - """ - - def __init__( - self, - rank: int, - worldsize: int, - ): - assert rank >= 0, f"Rank {rank} must be a positive integer" - assert ( - worldsize > rank - ), f"Worldsize {worldsize} must be greater than rank {rank}" - self.state_params: List[str] = [] - self.reshard_params: List[str] = [] - self.rank = rank - self.worldsize = worldsize - self.load_worldsize = ( - worldsize # Enable calling load_state_dict() directly, assume no rescaling - ) - - def statename(self, x: str): - # Note that this naming convention implicitly disallows repeated layers in the dataset pipeline - return self.__class__.__name__ + "." + x - - def state_dict(self): - """ - Retrieve all state and reshard flags (each worker/process saves its own state dict shard) - """ - return { - self.statename(flag): getattr(self, flag) - for flag in self.state_params + self.reshard_params - } - - def _reshard(self, sharded_list): - """ - Sharded_list is a list of lists, where each "shard" sublist must have the same length. - These shards should tightly span only the partition of data owned by this worker. - (i.e. if global_list is the list of all entries, sharded_list = _shard_inclusive(global_list) ). - Determine fractional ownership of shards, and get the flattened partition owned by this worker. - """ - # How many shards did _shard_inclusive() drop to the left of sharded_list? - shard_offset = math.floor(self.load_worldsize * self.rank / self.worldsize) - # How long are the list shards? - shard_len = len(sharded_list[0]) - for i, shard in enumerate(sharded_list): - assert ( - len(shard) == shard_len - ), f"Shard {i} with length {len(shard)} does not match expected {shard_len}" - # How many list items did _shard_inclusive() drop to the left of the flattened sharded_list? - item_offset = shard_len * shard_offset - # How many list items are there in total? - n_items = self.load_worldsize * shard_len - # The indices of the flattened sharded_list that this worker owns - my_items = range( - int(n_items * self.rank / self.worldsize) - item_offset, - int(n_items * (self.rank + 1) / self.worldsize) - item_offset, - ) - # Pull out owned items - return [sharded_list[i // shard_len][i % shard_len] for i in my_items] - - def load_state_dict(self, state_dicts, sharded_input=False): - """ - Input state_dicts is a list of state_dicts. If sharded_input=False, this is expected to be the - global list of states across all checkpoint shard files. If sharded_input=True, this expects - _shard_inclusive(global_state_list). Handling reduced inputs allows for much more efficient loading. - Workflow: - 1. if sharded_inputs is false, shard the inputs. - 2. If worldsize matches checkpoint, pull state and reshard params from the given checkpoint - shard (state_dicts is a singleton list). - 3. If worldsize does not match checkpoint, toss state params and assemble reshard params from - across given state_dicts. In this case state_dicts may be singleton (for fractional ownership) - or multi-element (for multiple/partitioned ownership). - 4. Return reduced input for use by downstream loading functions - """ - if not sharded_input: - self.load_worldsize = len(state_dicts) - state_dicts = _shard_inclusive(state_dicts, self.rank, self.worldsize) - if self.load_worldsize == self.worldsize: - [ - setattr(self, flag, state_dicts[0][self.statename(flag)]) - for flag in self.state_params + self.reshard_params - ] - else: - for flag in self.reshard_params: - reshard = self._reshard( - [sd[self.statename(flag)] for sd in state_dicts] - ) - setattr(self, flag, reshard) - return state_dicts - - def load_from_path(self, path: str): - """ - Count shard files in the specified checkpoint folder and determine overlap with current - rank and worldsize partition. Load only matching shardfile(s) and pass to load_state_dict. - This is more efficient than sharding the full loaded state. - """ - assert os.path.exists(path), "Specified checkpoint does not exist" - assert not os.path.isfile(path), "Checkpoint should be a folder of shard states" - fileshards = [x for x in os.listdir(path) if "loader" in x] - fileshards = sorted(fileshards, key=lambda x: int(x.split("_")[2][:-4])) - assert ( - len(fileshards) > 0 - ), "Checkpoint directory must contain checkpoint files with 'loader' in the name" - self.load_worldsize = len(fileshards) - # Grab only the shard files holding data we currently own - my_fileshards = _shard_inclusive(fileshards, self.rank, self.worldsize) - states = [torch.load(os.path.join(path, x)) for x in my_fileshards] - self.load_state_dict(states, True) - - def save_to_path(self, path: str): - """ - Grab recursive shard states and save all shard states to the specified checkpoint folder - """ - os.makedirs(path, exist_ok=True) - state = self.state_dict() - - # worker_info = torch.utils.data.get_worker_info() - # if worker_info is None: - # print(f"single process: {worker_info.id}, {worker_info.num_workers}") - # print(state) - # else: - # print(f"worker process: {worker_info.id}, {worker_info.num_workers}") - # print(state) - - torch.save(state, os.path.join(path, f"loader_state_{self.rank}.pth")) From 1de91e8634016fb245da938672f6c660f608a01e Mon Sep 17 00:00:00 2001 From: Peter Schneider-Kamp Date: Tue, 7 Oct 2025 09:47:15 +0200 Subject: [PATCH 4/6] back to old config --- maester/config.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/maester/config.py b/maester/config.py index 8a4e25c..5524b80 100644 --- a/maester/config.py +++ b/maester/config.py @@ -6,7 +6,7 @@ TomlConfigSettingsSource, ) from pydantic.fields import FieldInfo -from typing import Callable, Type, Any, Optional +from typing import Callable, Type, Any from pathlib import Path import torch @@ -19,9 +19,8 @@ class DatasetConfig(BaseSettings): data_dirs: list[str] = [ - "data/toy", - ] - dataset_types: Optional[list[str]] = None + "data/toy" + ] dataset_weights: str = "1.0" bos_token: int = 128000 eos_token: int = 128001 @@ -85,6 +84,8 @@ class Config(BaseSettings): data_parallel_replicate_degree: int = 1 tensor_parallel_degree: int = 1 train_batch_size: int = 2 # per device; 2 * 8 gpus * 32 nodes * 8192 seqlen = ~4M tokens per batch + gradient_accumulation_steps: int = 1 + gradient_accumulation_sync_each_step: bool = False train_num_steps: int = 1000 compile: bool = True enable_loss_parallel: bool = True From 3089467a28000dea6b1a915f0cebae7135a23cf2 Mon Sep 17 00:00:00 2001 From: Peter Schneider-Kamp Date: Tue, 7 Oct 2025 09:54:47 +0200 Subject: [PATCH 5/6] and back again --- maester/datasets/__init__.py | 11 +++++++++++ 1 file changed, 11 insertions(+) create mode 100644 maester/datasets/__init__.py diff --git a/maester/datasets/__init__.py b/maester/datasets/__init__.py new file mode 100644 index 0000000..6db8186 --- /dev/null +++ b/maester/datasets/__init__.py @@ -0,0 +1,11 @@ +from .experimental_otf import * + +__all__ = [ + "StreamingDocDataset", + "ScalableShardDataset", + "SamplingDataset", + "PreloadBufferDataset", + "BufferDataset", + "PreprocessDataset", + "build_experimental_data_loader", +] From b394cc1242db688781276d17200f8b9d61b73733 Mon Sep 17 00:00:00 2001 From: Peter Schneider-Kamp Date: Tue, 7 Oct 2025 10:03:12 +0200 Subject: [PATCH 6/6] minimal jinx support --- maester/config.py | 1 + maester/datasets/experimental_otf.py | 13 ++++++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/maester/config.py b/maester/config.py index 5524b80..9385d70 100644 --- a/maester/config.py +++ b/maester/config.py @@ -31,6 +31,7 @@ class DatasetConfig(BaseSettings): num_data_workers: int = 1 # col_name: str = "tokens" # file_type: str = "arrow" + dataset_type: str = "parquet" class SFTConfig(BaseSettings): diff --git a/maester/datasets/experimental_otf.py b/maester/datasets/experimental_otf.py index 9bd2864..a4d5fb7 100644 --- a/maester/datasets/experimental_otf.py +++ b/maester/datasets/experimental_otf.py @@ -1309,9 +1309,20 @@ def __init__( # Build subdataset iterators self.data = [] + match self.cfg.dataset_type: + case "parquet": + dataset_class = ParquetDataset + case "jinx": + try: + from .jinx_dataset import JinxDataset + dataset_class = JinxDataset + except ImportError: + raise ImportError("JinxDataset requires the mldataforge package. Please install it with `pip install mldataforge`.") + case _: + raise ValueError(f"Unsupported dataset type {self.cfg.dataset_type}. Supported types are 'parquet' and 'jinx'.") for i, d in enumerate(data_dirs): self.data.append( - ParquetDataset( + dataset_class( data_dir=d, rank=rank, worldsize=worldsize,