# -*- coding: utf-8 -*- from __future__ import annotations import copy import pickle from copy import deepcopy from dataclasses import dataclass from typing import Any, Callable, Dict, Iterable, List, Optional, Union import datasets import numpy as np import torch from datasets import Dataset, IterableDataset, interleave_datasets, load_dataset from datasets.iterable_dataset import ShufflingConfig from torch.distributed.checkpoint.stateful import Stateful from torchdata.stateful_dataloader import StatefulDataLoader from transformers import PreTrainedTokenizer from torchtitan.tools import utils from torchtitan.tools.logging import logger datasets.logging.set_verbosity_info() class BufferShuffledIterableDataset(IterableDataset): def __init__( self, dataset: Dataset, tokenizer: PreTrainedTokenizer, seq_len: int = 2048, rank: int = 0, world_size: int = 1, buffer_size: int = 1024, ) -> BufferShuffledIterableDataset: self.dataset = dataset self.tokenizer = tokenizer self.data = dataset.shard(world_size, rank) self.seq_len = seq_len self.rank = rank self.world_size = world_size self.buffer_size = buffer_size if tokenizer.vocab_size < torch.iinfo(torch.uint16).max: self.dtype = torch.uint16 elif tokenizer.vocab_size < torch.iinfo(torch.uint32).max: self.dtype = torch.uint32 else: self.dtype = torch.uint64 self.states = None self.buffer = torch.tensor([], dtype=self.dtype) self.tokens = [] self.rand_id = 0 self.token_id = 0 self.rng_state = None self._epoch = 0 def __iter__(self): g = torch.Generator() g.manual_seed(self._epoch + self.rank) if self.rng_state is not None: g.set_state(self.rng_state) rand_it = self.randint(0, self.buffer_size, g=g) if self.states is not None: self.data.load_state_dict(self.states) # max number of tokens allowed in the chunk buffer n_tokens = self.buffer_size * self.seq_len while True: for sample in self.tokenize(self.data): # keep appending the samples to the token buffer self.tokens += sample # if the token buffer is full, start sampling # NOTE: we first convert the token ids to a tensor of shape [n_chunks, seq_len] for efficiency if len(self.buffer) == 0 and len(self.tokens) >= n_tokens: self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=self.dtype).view(self.buffer_size, -1) self.tokens = self.tokens[n_tokens:] if len(self.buffer) == self.buffer_size: yield from self.sample(rand_it) n_chunks = len(self.tokens) // self.seq_len # handle the left tokens in the buffer if n_chunks > 0: n_tokens = n_chunks * self.seq_len indices = torch.randperm(n_chunks, generator=g).tolist() self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=torch.long).view(n_chunks, -1) self.tokens = self.tokens[n_tokens:] for i in indices: yield {'input_ids': self.buffer[i]} def tokenize(self, data, batch_size: int = 64): texts, states = [], [] for sample in data: texts.append(sample['text']) states.append(self.data.state_dict()) if len(texts) == batch_size: for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']): self.states = s yield tokenized texts, states = [], [] if len(texts) > 0: for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']): self.states = s yield tokenized def sample(self, indices): n_tokens = (len(self.tokens) // self.seq_len) * self.seq_len while self.token_id < n_tokens: i = next(indices) start, end = self.token_id, self.token_id + self.seq_len self.token_id += self.seq_len yield {'input_ids': self.buffer[i].to(torch.long)} self.buffer[i] = torch.tensor(self.tokens[start:end], dtype=self.dtype) self.token_id = 0 self.tokens = self.tokens[n_tokens:] def randint(self, low: int, high: int, buffer_size: int = 1024, g: torch.Generator = torch.Generator()) -> Iterable[int]: indices = torch.empty(buffer_size, dtype=torch.long) while True: # record the generator states before sampling self.rng_state = g.get_state() indices = torch.randint(low, high, (buffer_size,), out=indices, generator=g) for i in indices[self.rand_id:].tolist(): self.rand_id += 1 yield i self.rand_id = 0 def set_epoch(self, epoch): self._epoch = epoch if hasattr(self.dataset, 'set_epoch'): self.dataset.set_epoch(epoch) def state_dict(self): return { 'states': self.states, 'buffer': self.buffer.clone(), 'tokens': deepcopy(self.tokens), 'rand_id': self.rand_id, 'token_id': self.token_id, 'rng_state': self.rng_state, 'epoch': self._epoch, } def load_state_dict(self, state_dict): self.states = state_dict['states'] self.buffer = state_dict['buffer'].clone() self.tokens = deepcopy(state_dict['tokens']) self.rand_id = state_dict['rand_id'] self.token_id = state_dict['token_id'] self.rng_state = state_dict['rng_state'].clone() if state_dict['rng_state'] is not None else None self._epoch = state_dict['epoch'] class OnlineTokenizedIterableDataset(IterableDataset): def __init__( self, dataset: Dataset, tokenizer: PreTrainedTokenizer, seq_len: int = 2048, rank: int = 0, world_size: int = 1 ) -> OnlineTokenizedIterableDataset: self.dataset = dataset self.tokenizer = tokenizer self.data = dataset.shard(world_size, rank) self.seq_len = seq_len self.rank = rank self.world_size = world_size self.states = None self.tokens = [] def __iter__(self): if self.states is not None: self.data.load_state_dict(self.states) while True: for sample in self.tokenize(self.data): # keep appending the samples to the token buffer self.tokens += sample while len(self.tokens) >= self.seq_len: input_ids = torch.tensor(self.tokens[:self.seq_len], dtype=torch.long) self.tokens = self.tokens[self.seq_len:] yield {'input_ids': input_ids} def tokenize(self, data, buffer_size: int = 64): buffer, states = [], [] for sample in data: if sample.get('text', None) is not None: buffer.append(sample['text']) elif sample.get('content', None) is not None: buffer.append(sample['content']) else: raise ValueError(f"No 'text' or 'content' field found in sample:\n{sample}") states.append(self.data.state_dict()) if len(buffer) == buffer_size: for s, tokenized in zip(states, self.tokenizer(buffer, return_attention_mask=False)['input_ids']): self.states = s yield tokenized buffer, states = [], [] if len(buffer) > 0: for s, tokenized in zip(states, self.tokenizer(buffer, return_attention_mask=False)['input_ids']): self.states = s yield tokenized def state_dict(self): return {'states': self.states, 'tokens': deepcopy(self.tokens)} def load_state_dict(self, state_dict): self.states = state_dict['states'] self.tokens = deepcopy(state_dict['tokens']) class BufferShuffledExamplesIterable(datasets.iterable_dataset.BufferShuffledExamplesIterable): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def _init_state_dict(self) -> dict: self._state_dict = self.ex_iterable._init_state_dict() self._state_dict['mem_buffer'] = ([],) self._state_dict['bit_generator_state'] = self.generator.bit_generator.state self._state_dict['bit_generator_index_offset'] = 0 self._state_dict['bit_generator_index_offset_shuffle'] = 0 return self._state_dict def __iter__(self): buffer_size = self.buffer_size rng = deepcopy(self.generator) # this is the shuffle buffer that we keep in memory mem_buffer = self._state_dict['mem_buffer'][0] # this is an infinite iterator that randomly samples the index of the source to pick examples from index_offset = self._state_dict['bit_generator_index_offset'] if self._state_dict else 0 if self._state_dict: rng.bit_generator.state = self._state_dict['bit_generator_state'] indices_iterator = self._iter_random_indices(rng, buffer_size, random_batch_size=buffer_size) # skip already consumed ones for _ in range(index_offset): i = next(indices_iterator) for x in self.ex_iterable: if len(mem_buffer) < buffer_size: # if the buffer is not full, keep filling the buffer mem_buffer.append(x) else: # otherwise, pick an example from it i = next(indices_iterator) index_offset = (index_offset + 1) % buffer_size if self._state_dict: self._state_dict['bit_generator_index_offset'] = index_offset if index_offset == 0: self._state_dict['bit_generator_state'] = rng.bit_generator.state selected = mem_buffer[i] mem_buffer[i] = x # replace the picked example by a new one yield selected index_offset = self._state_dict['bit_generator_index_offset_shuffle'] if self._state_dict else 0 if self._state_dict: rng.bit_generator.state = self._state_dict['bit_generator_state'] # when we run out of examples, we shuffle the remaining examples in the buffer and yield them for i in rng.permutation(len(mem_buffer))[index_offset:].tolist(): index_offset = index_offset + 1 if self._state_dict: self._state_dict['bit_generator_index_offset_shuffle'] = index_offset yield mem_buffer[i] def shuffle_data_sources(self, generator: np.random.Generator) -> BufferShuffledExamplesIterable: """Shuffle the wrapped examples iterable as well as the shuffling buffer.""" return BufferShuffledExamplesIterable( self.ex_iterable.shuffle_data_sources(generator), buffer_size=self.buffer_size, generator=generator ) def shard_data_sources(self, num_shards: int, index: int, contiguous=True) -> BufferShuffledExamplesIterable: """Keep only the requested shard.""" return BufferShuffledExamplesIterable( self.ex_iterable.shard_data_sources(num_shards, index, contiguous=contiguous), buffer_size=self.buffer_size, generator=self.generator, ) def load_state_dict(self, state_dict: dict) -> dict: def _inner_load_state_dict(state, new_state): if new_state is not None and isinstance(state, dict): for key in new_state: state[key] = _inner_load_state_dict(state[key], new_state[key]) return state elif new_state is not None and isinstance(state, list): for i in range(len(state)): state[i] = _inner_load_state_dict(state[i], new_state[i]) return state return new_state return _inner_load_state_dict(self._state_dict, state_dict) def shuffle( dataset: IterableDataset, seed: int = 42, generator: np.random.Generator = None, buffer_size: int = 1024, ): generator = np.random.default_rng(seed) if generator is None else deepcopy(generator) return IterableDataset( ex_iterable=BufferShuffledExamplesIterable(dataset._ex_iterable, buffer_size=buffer_size, generator=generator), info=dataset._info.copy(), split=dataset._split, formatting=dataset._formatting, shuffling=ShufflingConfig(generator=generator, _original_seed=seed), distributed=copy.deepcopy(dataset._distributed), token_per_repo_id=dataset._token_per_repo_id, ) @dataclass class DataCollatorForLanguageModeling: """ Data collator used for language modeling. Inputs are dynamically padded if `varlen=False`. If `varlen=True`, sequences are expected to be concatenated, and labels match inputs. Args: tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]): The tokenizer used for encoding the data. context_len (`int`, optional): When `varlen=True`, sequences longer than this length within a document (as determined by `cu_seqlens`) will be further chunked. varlen (`bool`): Whether to handle variable length concatenated sequences (`True`) or padded batches (`False`). Returns: A dictionary with the following keys: - `input_ids`: Tensor of input IDs. Shape `[batch_size, seq_len]` if `varlen=False`, `[1, total_len]` if `varlen=True`. - `labels`: Tensor of labels. Shape matches `input_ids`. Padding positions are masked with -100 if `varlen=False`. - `attention_mask`: Tensor indicating non-padding tokens (only if `varlen=False`). Shape matches `input_ids`. - `cu_seqlens`: Tensor of cumulative sequence lengths (only if `varlen=True`). Shape `[1, num_sequences + 1]`. NOTE: When `varlen=True`, the `batch_size` must be 1. """ tokenizer: PreTrainedTokenizer context_len: Optional[int] = None varlen: bool = False def __call__(self, examples: List[Union[List[int], Dict[str, Any]]]) -> Dict[str, Any]: if not isinstance(examples[0], Dict): examples = [{'input_ids': example} for example in examples] def tensorize(example: Dict[str, Any]) -> Dict[str, Any]: tensorized = {} for key in ['input_ids', 'cu_seqlens']: if key not in example: continue if isinstance(example[key], List): tensorized[key] = torch.tensor(example[key], dtype=torch.long) elif isinstance(example[key], np.ndarray): tensorized[key] = torch.from_numpy(example[key]) else: tensorized[key] = example[key] return tensorized examples = list(map(tensorize, examples)) if not self.varlen: # --- Handling for varlen=False (Batch Padding) --- length_of_first = examples[0]['input_ids'].size(0) needs_padding = not all(example['input_ids'].size(0) == length_of_first for example in examples) if needs_padding: # Check for pad token if padding is actually required if self.tokenizer.pad_token_id is None: raise ValueError( f'You are attempting to pad samples but the tokenizer you are using ' f'({self.tokenizer.__class__.__name__}) does not have a pad token.' ) # Pad using the tokenizer, ensuring attention_mask is returned batch = self.tokenizer.pad(examples, return_tensors='pt', return_attention_mask=True) else: # No padding needed, stack directly and create a full attention mask input_ids = torch.stack([example['input_ids'] for example in examples], dim=0) batch = { 'input_ids': input_ids, # Create attention mask of all ones 'attention_mask': torch.ones_like(input_ids), } # Create labels by cloning input_ids labels = batch['input_ids'].clone() # Mask labels only where attention_mask is 0 (padding positions) if 'attention_mask' in batch: labels[batch['attention_mask'] == 0] = -100 batch['labels'] = labels else: # --- Handling for varlen=True (Concatenated Sequences) --- if len(examples) > 1: raise ValueError('The batch size must be 1 for inputs with variable lengths (varlen=True).') batch = {'input_ids': torch.cat([example['input_ids'] for example in examples], dim=0).unsqueeze(0)} # --- cu_seqlens calculation logic remains the same --- if 'cu_seqlens' in examples[0]: batch['cu_seqlens'] = ( torch.cat([example['cu_seqlens'] for example in examples], dim=0).unsqueeze(0).to(dtype=torch.int32) ) # Ensure int32 else: # determine boundaries by bos/eos positions # Check for bos_token_id first if self.tokenizer.bos_token_id is not None: cu_seqlens = [] # Handle case where the sequence doesn't start with BOS if batch['input_ids'][0, 0] != self.tokenizer.bos_token_id: cu_seqlens.append(torch.tensor([0], device=batch['input_ids'].device)) # Match device # Find all BOS token positions bos_positions = torch.where(batch['input_ids'].eq(self.tokenizer.bos_token_id))[1] # Ensure bos_positions is on the correct device if empty if bos_positions.numel() == 0 and len(cu_seqlens) > 0: cu_seqlens.append(bos_positions.to(cu_seqlens[0].device)) elif bos_positions.numel() > 0: cu_seqlens.append(bos_positions) # Add the end of the entire batch cu_seqlens.append( torch.tensor([batch['input_ids'].size(1)], device=batch['input_ids'].device) ) # Match device and use size(1) # Filter out empty tensors before cat cu_seqlens = [t for t in cu_seqlens if t.numel() > 0] if not cu_seqlens: # Handle case where input is empty or has no BOS batch['cu_seqlens'] = torch.tensor( [0, batch['input_ids'].size(1)], dtype=torch.int32, device=batch['input_ids'].device ) else: batch['cu_seqlens'] = torch.cat(cu_seqlens, dim=0).to(dtype=torch.int32) # Else, check for eos_token_id elif self.tokenizer.eos_token_id is not None: cu_seqlens = [torch.tensor([0], device=batch['input_ids'].device)] # Match device # Find positions *after* EOS tokens eos_positions = torch.where(batch['input_ids'].eq(self.tokenizer.eos_token_id))[1] + 1 # Ensure eos_positions is on the correct device if empty if eos_positions.numel() > 0: cu_seqlens.append(eos_positions) # Handle case where the sequence doesn't end with EOS if batch['input_ids'][0, -1] != self.tokenizer.eos_token_id: # Only add the final length if the last found EOS wasn't already the end if eos_positions.numel() == 0 or eos_positions[-1] != batch['input_ids'].size(1): cu_seqlens.append( torch.tensor([batch['input_ids'].size(1)], device=batch['input_ids'].device) ) # Match device and use size(1) # Filter out empty tensors before cat cu_seqlens = [t for t in cu_seqlens if t.numel() > 0] if not cu_seqlens: # Handle case where input is empty or has no EOS batch['cu_seqlens'] = torch.tensor( [0, batch['input_ids'].size(1)], dtype=torch.int32, device=batch['input_ids'].device ) else: batch['cu_seqlens'] = torch.cat(cu_seqlens, dim=0).to(dtype=torch.int32) # Else, neither BOS nor EOS is usable else: raise ValueError( 'For varlen=True without precomputed cu_seqlens, the tokenizer must have either a bos_token_id ' 'or an eos_token_id defined to act as sequence separators.' ) # --- cu_seqlens validation checks remain the same --- if batch['cu_seqlens'].numel() < 2: raise ValueError(f'Calculated cu_seqlens must have at least start and end: {batch["cu_seqlens"]}') if not torch.all(batch['cu_seqlens'][1:] >= batch['cu_seqlens'][:-1]): raise ValueError(f'Calculated cu_seqlens are not monotonically increasing: {batch["cu_seqlens"]}') if batch['cu_seqlens'][0] != 0: raise ValueError(f'Calculated cu_seqlens do not start at 0: {batch["cu_seqlens"]}') if batch['cu_seqlens'][-1] != batch['input_ids'].size(1): # Allow empty sequence case where cu_seqlens=[0, 0] and input_ids.size(1)=0 if not (batch['cu_seqlens'].tolist() == [0, 0] and batch['input_ids'].size(1) == 0): raise ValueError( f'Calculated cu_seqlens do not end at total length {batch["input_ids"].size(1)}: ' f'{batch["cu_seqlens"]}' ) # --- context_len splitting logic remains the same --- if self.context_len is not None: # This logic splits sequences based on context_len *after* initial boundaries are found bos = batch['cu_seqlens'][:-1].tolist() eos = batch['cu_seqlens'][1:].tolist() # Handle empty sequences between boundaries split_boundaries = [] for i, j in zip(bos, eos): if i < j: # Only process non-empty sequences split_boundaries.append(torch.arange(i, j, self.context_len, device=batch['input_ids'].device)) # Add the final end point if it wasn't included by arange final_end_point = torch.tensor([batch['input_ids'].size(1)], device=batch['input_ids'].device) # Concatenate all boundaries if not split_boundaries: # Handle case of completely empty input batch['cu_seqlens'] = torch.tensor([0, 0], dtype=torch.int32, device=batch['input_ids'].device) else: batch['cu_seqlens'] = torch.cat(split_boundaries + [final_end_point]).to(dtype=torch.int32) # Ensure uniqueness and sort, as arange might duplicate the endpoint batch['cu_seqlens'] = torch.unique(batch['cu_seqlens']) # Create labels directly from input_ids, NO padding mask needed for varlen labels = batch['input_ids'].clone() batch['labels'] = labels return batch class ParallelAwareDataLoader(StatefulDataLoader, Stateful): """ A wrapper around the StatefulDataLoader that ensures that the state is stored only once per DP rank. """ def __init__( self, rank: int, dataset: IterableDataset, batch_size: int, collate_fn: Callable, num_workers: int = 0, pin_memory: bool = False, prefetch_factor: int = 2, persistent_workers: bool = False, snapshot_every_n_steps: Optional[int] = 1, ): super().__init__( dataset=dataset, batch_size=batch_size, collate_fn=collate_fn, num_workers=num_workers, pin_memory=pin_memory, prefetch_factor=prefetch_factor, persistent_workers=persistent_workers, snapshot_every_n_steps=snapshot_every_n_steps, ) self.rank = rank def state_dict(self) -> Dict[str, Any]: # Store state only for dp rank to avoid replicating the same state across other dimensions return {f'rank_{self.rank}': pickle.dumps(super().state_dict())} def load_state_dict(self, state_dict: Dict[str, Any]) -> None: # State being empty is valid if not state_dict: return if f'rank_{self.rank}' not in state_dict: logger.warning(f'DataLoader state is empty for dp rank {self.rank}, expected key rank_{self.rank}') return super().load_state_dict(pickle.loads(state_dict[f'rank_{self.rank}'])) def build_dataset( dataset: str, dataset_name: str = None, dataset_split: str = 'train', data_dir: str = None, data_files: str = None, data_probs: List[float] = None, streaming: bool = False, dp_degree: Optional[int] = None, num_workers: int = 32, seed: Optional[int] = None, ) -> IterableDataset: color = utils.Color min_num_shards = dp_degree * num_workers if dp_degree else None if len(dataset.split(',')) == 1: dataset = load_dataset( path=dataset, name=dataset_name, split=dataset_split, data_dir=data_dir, data_files=data_files, # trust_remote_code=True, streaming=streaming, num_proc=num_workers if not streaming else None, ) logger.info(f"Shuffling the dataset with seed {seed}") if not streaming: # the states of map-style dataset is recoverable after shuffling if seed is not None: dataset = dataset.shuffle(seed=seed) if min_num_shards is not None: dataset = dataset.to_iterable_dataset(num_shards=min_num_shards) else: if min_num_shards is not None and dataset.num_shards < min_num_shards: logger.warning( f"{color.red}" f"Dataset {dataset} has insufficient shards ({dataset.num_shards}). " f"Need {min_num_shards} shards minimum for {dp_degree} data parallel workers × " f"{num_workers} dataloader workers. " f"Disabling the streaming mode and resharding dataset to {min_num_shards} shards." f"{color.reset}" ) dataset = load_dataset( path=dataset, name=dataset_name, split=dataset_split, data_dir=data_dir, data_files=data_files, # trust_remote_code=True, streaming=False, num_proc=num_workers, ) if seed is not None: dataset = dataset.shuffle(seed=seed) dataset = dataset.to_iterable_dataset(num_shards=min_num_shards) else: if seed is not None: dataset = shuffle(dataset, seed=seed) else: datasets = dataset.split(",") if dataset_name is not None: dataset_names = [ name or None for name in dataset_name.split(",") ] assert len(dataset_names) == len(datasets), ( "The number of dataset names must match the number of datasets" ) else: dataset_names = [None] * len(datasets) if dataset_split is not None: dataset_splits = [split or "train"for split in dataset_split.split(",")] assert len(dataset_splits) == len(datasets), ( "The number of dataset splits must match the number of datasets" ) else: dataset_splits = ["train"] * len(datasets) if data_dir is not None: data_dirs = [ data_dir or None for data_dir in data_dir.split(",") ] assert len(data_dirs) == len(datasets), ( "The number of data dirs must match the number of datasets" ) else: data_dirs = [None] * len(datasets) if data_files is not None: data_files = data_files.split(",") assert len(data_files) == len(datasets), ( "The number of data files must match the number of datasets" ) else: data_files = [None] * len(datasets) if data_probs is not None: data_probs = [float(p) for p in data_probs.split(",")] assert len(data_probs) == len(datasets), ( "The number of data probabilities must match the number of datasets" ) else: raise ValueError( "Data sampling probabilities are required if using multiple datasets" ) subsets = [] for i, prob in enumerate(data_probs): subset = load_dataset( path=datasets[i], name=dataset_names[i], split=dataset_splits[i], data_dir=data_dirs[i], data_files=data_files[i], # trust_remote_code=True, streaming=streaming, num_proc=( num_workers if not streaming else None ), ) logger.info( f"Subset {color.cyan}{datasets[i]}" + (f":{dataset_names[i]} " if dataset_names[i] else " ") + f"(p = {prob:.3f}){color.reset}:\n" + f"{subset}" ) logger.info(f"Shuffling the dataset with seed {seed}") if not streaming: # the states of map-style dataset is recoverable after shuffling if seed is not None: subset = subset.shuffle(seed=seed) if min_num_shards is not None: subset = subset.to_iterable_dataset(num_shards=min_num_shards) else: if min_num_shards is not None and subset.num_shards < min_num_shards: logger.warning( f"{color.red}" f"Dataset {datasets[i]} has insufficient shards ({subset.num_shards}). " f"Need {min_num_shards} shards minimum for desired data parallel workers × " f"{num_workers} dataloader workers. " f"Resharding dataset to {min_num_shards} shards and disabling streaming mode." f"{color.reset}" ) # again, it's ok to directly shuffle the map-style dataset # we expect an error raised if the map-style dataset still has not enough data shards subset = load_dataset( path=datasets[i], name=dataset_names[i], split=dataset_splits[i], data_dir=data_dirs[i], data_files=data_files[i], # trust_remote_code=True, streaming=False, num_proc=num_workers, ) if seed is not None: subset = subset.shuffle(seed=seed) subset = subset.to_iterable_dataset(num_shards=min_num_shards) else: # we set relatively small buffer size here as interleaving could provide some randomness if seed is not None: subset = shuffle(subset, seed=seed, buffer_size=max(128, 1024 // len(datasets))) if "text" in subset.column_names: subset = subset.select_columns("text") elif "content" in subset.column_names: subset = subset.select_columns("content") else: raise ValueError( f"Subset {datasets[i]} has no 'text' or 'content' column" ) subsets.append(subset) logger.info( f"Interleaving {len(subsets)} datasets with probabilities {data_probs}" ) dataset = interleave_datasets( datasets=subsets, probabilities=data_probs, stopping_strategy="all_exhausted", seed=seed, ) logger.info(f"{dataset}") return dataset def build_dataloader( dataset: IterableDataset, tokenizer: PreTrainedTokenizer, rank: int, world_size: int, batch_size: int, seq_len: int, context_len: Optional[int] = None, varlen: bool = False, num_workers: int = 0, pin_memory: bool = False, persistent_workers: bool = False, snapshot_every_n_steps: Optional[int] = 1, ): dataset = OnlineTokenizedIterableDataset( dataset=dataset, tokenizer=tokenizer, seq_len=seq_len, rank=rank, world_size=world_size ) return ParallelAwareDataLoader( rank=rank, dataset=dataset, batch_size=batch_size, collate_fn=DataCollatorForLanguageModeling(tokenizer=tokenizer, context_len=context_len, varlen=varlen), num_workers=num_workers, pin_memory=pin_memory, persistent_workers=persistent_workers, snapshot_every_n_steps=snapshot_every_n_steps, )