File size: 22,296 Bytes
b7d8df8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 |
import logging
import os
from typing import Tuple, Union
import torch
from torch.utils.data import DataLoader
from transformers import PreTrainedTokenizerBase
from .collator import Seq2SeqFinetuningCollator, validate_target_settings
from .tasks import DOWNLOADED_FT_DATASETS_DIRPATH, SUPPORTED_EXTENSIONS, dataset_constructor
from .packing import BinPackCollator, auto_packing_ratio
from .text_data import build_streams, get_tokens_per_batch_func
from .exceptions import MissingHuggingFaceURLSplitError, NotEnoughDatasetSamplesError
log = logging.getLogger(__name__)
_HF_IGNORE_INDEX = -100
_DEFAULT_TARGET_RESPONSES = 'last'
_DEFAULT_TARGET_PROMPTS = 'none'
def build_finetuning_dataloader(cfg: DictConfig, tokenizer: PreTrainedTokenizerBase, device_batch_size: int) -> DataSpec:
"""Builds a finetuning dataloader for training or evaluating.
The underlying dataset can be built through one of two code paths:
1. As a HuggingFace dataset, via `datasets.load_dataset(...)`
2. As a streaming dataset
You will need to set slightly different dataset config fields depending
on which you intend to use, as explained below.
Args:
cfg (DictConfig): An omegaconf dictionary used to configure the loader:
cfg.name (str): The type of dataloader to build. Must = "finetuning".
---
*** HuggingFace dataset config fields ***
cfg.dataset.hf_name (str, optional): The name of the HuggingFace dataset
to use. Can also be a remote http(s) directory or object store bucket
containing the file {split}.jsonl in the format (prompt, response),
in which case the builder will create a HuggingFace dataset.
cfg.dataset.hf_kwargs (DictConfig, optional): Additional kwargs to
pass to `datasets.load_dataset`, which can be used to load
a dataset from local files.
cfg.dataset.preprocessing_fn (str, optional): The name/import path of
the preprocessing function to use for formatting the data examples.
If ``None`` (default), the builder will use the preprocessing function
registered under `hf_name` (see `tasks.py`), if one exists,
otherwise it will skip preprocessing.
If `preprocessing_fn` corresponds to a registered preprocessing
function in `tasks.py`, the builder will use that.
Otherwise, it will interpret `preprocessing_fn` as a
"import.path:function_name" import path; e.g., it will call
`from import.path import function_name` and use the imported
function as the preprocessing function.
*** Streaming dataset config fields ***
cfg.dataset.remote (str, optional): Location of a MDS-formatted
streaming dataset to use. Setting this will tell the builder
to create a streaming dataset rather than a HuggingFace dataset.
cfg.dataset.local (str, optional): Local path where remote data
will be streamed to. Only valid if `cfg.dataset.remote` has
also been set.
*** Shared dataset configs fields ***
cfg.dataset.max_seq_len (int): The maximum length of sequences
in the batch. See :class:`Seq2SeqFinetuningCollator` docstring
for details.
cfg.dataset.decoder_only_format (bool): Whether to format the
examples for a decoder-only model. See :class:`Seq2SeqFinetuningCollator`
docstring for details.
cfg.dataset.target_responses (str): Which responses are used as training targets.
Defaults to "last", meaning only the final response in multi-turn examples
will serve as training targets. See :class:`Seq2SeqFinetuningCollator` docstring for
details.
cfg.dataset.target_prompts (str): Which prompts are used as training targets.
Defaults to "none", meaning prompts are never used as training targets.
See :class:`Seq2SeqFinetuningCollator` docstring for details.
cfg.dataset.allow_pad_trimming (bool, optional): Whether to allow
the collator to trim padding. See :class:`Seq2SeqFinetuningCollator`
docstring for details. Default: ``False``.
cfg.dataset.packing_ratio (Optional[float, Literal['auto']]): If provided, this invokes
a collator wrapper that packs device_batch_size*packing_ratio
raw examples into device_batch_size packed examples. This helps
minimize padding while preserving sequence integrity.
This adds `sequence_id` to the batch, which indicates which unique
sequence each token belongs to.
If set to 'auto', packing_ratio is profiled and the highest observed packing ratio with
zero waste is selected.
In practice, this may result in > 0 waste because profiling is done on only a portion
of the dataset.
Note: Using this feature will not change device_batch_size but it
will determine the number of raw examples consumed by the dataloader
per batch. Some examples may be discarded if they do not fit when
packing.
Select packing_ratio **carefully** based on the dataset
statistics, max_seq_len, and tolerance for discarding samples!
The script `scripts/misc/profile_packing.py` can help
you choose the best packing_ratio.
cfg.dataset.shuffle (bool): Whether to shuffle the dataset.
___
See :class:`StreamingFinetuningDataset` for info on other standard config
options within `cfg.dataset` that will be passed as kwargs if
using the streaming codepath.
---
See :class:`DataLoader` for standard argument options to the pytorch
dataloader, such as `cfg.drop_last`, `cfg.num_workers`, etc.
tokenizer (transformers.PreTrainedTokenizer): The tokenizer used to
prepare the data from raw text. Any missing sentinel tokens will
be added by the collator.
device_batch_size (int): The size of the batches (number of examples)
that the dataloader will produce.
Returns:
A pytorch dataloader
Note:
You can run the script inside `scripts/misc/profile_packing.py` to quickly test the
padding/waste rates for different `cfg.dataset.packing_ratio` choices,
given a starting workload YAML.
"""
_validate_config(cfg.dataset)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
collate_fn, dataloader_batch_size = _build_collate_fn(cfg, tokenizer, device_batch_size)
dataset = None
sampler = None
if cfg.dataset.get('remote') is not None or cfg.dataset.get('streams') is not None:
streams = build_streams(cfg.dataset)
dataset = dataset_constructor.build_from_streaming(tokenizer=tokenizer, streams=streams, local=cfg.dataset.get('local', None), remote=cfg.dataset.get('remote', None), split=cfg.dataset.get('split', None), download_retry=cfg.dataset.get('download_retry', 2), download_timeout=cfg.dataset.get('download_timeout', 60), validate_hash=cfg.dataset.get('validate_hash', None), keep_zip=cfg.dataset.get('keep_zip', False), epoch_size=cfg.dataset.get('epoch_size', None), predownload=cfg.dataset.get('predownload', None), cache_limit=cfg.dataset.get('cache_limit', None), partition_algo=cfg.dataset.get('partition_algo', 'relaxed'), num_canonical_nodes=cfg.dataset.get('num_canonical_nodes', None), batch_size=device_batch_size, shuffle=cfg.dataset.get('shuffle', False), shuffle_algo=cfg.dataset.get('shuffle_algo', 'py1e'), shuffle_seed=cfg.dataset.get('shuffle_seed', 9176), shuffle_block_size=cfg.dataset.get('shuffle_block_size', None), sampling_method=cfg.dataset.get('sampling_method', 'balanced'), sampling_granularity=cfg.dataset.get('sampling_granularity', 1), batching_method=cfg.dataset.get('batching_method', 'random'), max_seq_len=cfg.dataset.max_seq_len)
else:
dataset_name_or_path = cfg.dataset.hf_name
split = cfg.dataset.get('split')
if split is None:
raise MissingHuggingFaceURLSplitError()
backend, _, _ = parse_uri(dataset_name_or_path)
if backend not in ['', None]:
dataset_name_or_path = _download_remote_hf_dataset(remote_path=dataset_name_or_path, split=split)
split = split.replace('-', '_')
proto_preprocessing_fn = cfg.dataset.get('preprocessing_fn')
if isinstance(proto_preprocessing_fn, (dict, DictConfig)):
preprocessing_fn = dataset_constructor.get_preprocessing_fn_from_dict(dict(proto_preprocessing_fn))
else:
preprocessing_fn = dataset_constructor.get_preprocessing_fn_from_str(proto_preprocessing_fn, dataset_name_or_path)
dataset = dataset_constructor.build_from_hf(dataset_name=dataset_name_or_path, split=split, safe_load=cfg.dataset.get('safe_load', False), max_seq_len=cfg.dataset.max_seq_len, preprocessing_fn=preprocessing_fn, tokenizer=tokenizer, target_prompts=cfg.dataset.get('target_prompts', _DEFAULT_TARGET_PROMPTS), target_responses=cfg.dataset.get('target_responses', _DEFAULT_TARGET_RESPONSES), decoder_only_format=cfg.dataset.decoder_only_format, hf_kwargs=cfg.dataset.get('hf_kwargs', {}))
if cfg.drop_last:
world_size = dist.get_world_size()
minimum_dataset_size = world_size * dataloader_batch_size
if hasattr(dataset, '__len__'):
full_dataset_size = len(dataset)
if full_dataset_size < minimum_dataset_size:
raise NotEnoughDatasetSamplesError(dataset_name=cfg.dataset.hf_name, split=split, dataloader_batch_size=dataloader_batch_size, world_size=world_size, full_dataset_size=full_dataset_size, minimum_dataset_size=minimum_dataset_size)
sampler = dist.get_sampler(dataset, drop_last=cfg.drop_last, shuffle=cfg.dataset.shuffle)
assert dataset is not None
dl = DataLoader(dataset, collate_fn=collate_fn, batch_size=dataloader_batch_size, drop_last=cfg.drop_last, sampler=sampler, num_workers=cfg.num_workers, pin_memory=cfg.get('pin_memory', True), prefetch_factor=cfg.get('prefetch_factor', 2), persistent_workers=cfg.get('persistent_workers', True), timeout=cfg.get('timeout', 0))
token_counting_func = get_tokens_per_batch_func()
return DataSpec(dataloader=dl, get_num_tokens_in_batch=token_counting_func)
def _validate_config(dataset_cfg: DictConfig) -> None:
"""Validates the dataset configuration.
Makes sure that the dataset is properly configured for either
a HuggingFace dataset or a streaming dataset. Must be valid for one or
the other.
Args:
dataset_cfg (DictConfig): The dataset configuration to be validated.
Raises:
ValueError: If the dataset configuration does not meet the requirements.
"""
if dataset_cfg.get('hf_name') is not None:
illegal_keys = ['local', 'remote']
discovered_illegal_keys = []
for key in illegal_keys:
if dataset_cfg.get(key) is not None:
discovered_illegal_keys.append('`' + key + '`')
if discovered_illegal_keys:
raise ValueError('The dataset config sets a value for `hf_name` as well as the ' + f"following keys: {', '.join(discovered_illegal_keys)}.\n" + 'Those keys are used when building from a streaming dataset, but ' + 'setting `hf_name` instructs the dataset to build from a HuggingFace dataset.')
elif dataset_cfg.get('remote') is not None:
illegal_keys = ['hf_name', 'hf_kwargs', 'preprocessing_fn', 'safe_load']
discovered_illegal_keys = []
for key in illegal_keys:
if dataset_cfg.get(key) is not None:
discovered_illegal_keys.append('`' + key + '`')
if discovered_illegal_keys:
raise ValueError('The dataset config sets a value for `remote` as well as the ' + f"following keys: {', '.join(discovered_illegal_keys)}.\n" + 'Those keys are used when building from a HuggingFace dataset, but ' + 'setting `remote` instructs the dataset to build from a streaming dataset.')
if dataset_cfg.get('local') is None:
raise ValueError('Using a streaming dataset requires setting both `remote` and `local`, ' + 'but dataset.local is None.')
elif dataset_cfg.get('streams') is not None:
illegal_keys = ['hf_name', 'hf_kwargs', 'preprocessing_fn', 'safe_load']
discovered_illegal_keys = []
for key in illegal_keys:
if dataset_cfg.get(key) is not None:
discovered_illegal_keys.append('`' + key + '`')
if discovered_illegal_keys:
raise ValueError('The dataset config sets a value for `streams` as well as the ' + f"following keys: {', '.join(discovered_illegal_keys)}.\n" + 'Those keys are used when building from a HuggingFace dataset, but ' + 'setting `streams` instructs the dataset to build from a streaming dataset.')
illegal_keys = ['remote', 'local']
discovered_illegal_keys = []
for key in illegal_keys:
if dataset_cfg.get(key) is not None:
discovered_illegal_keys.append('`' + key + '`')
if discovered_illegal_keys:
raise ValueError('The dataset config sets a value for `streams` as well as the ' + f"following keys: {', '.join(discovered_illegal_keys)}.\n" + 'Please either use single stream (set remote/local only) ' + 'or put remote/local under streams')
else:
raise ValueError('In the dataset config, you must set `hf_name` to use a HuggingFace ' + 'dataset, or set `remote` to use a streaming dataset, or set ' + '`streams` to use multiple streaming datasets, but all were None.')
if dataset_cfg.get('max_seq_len') is None:
raise ValueError('In the dataset config, you must set the `max_seq_len`')
target_responses = str(dataset_cfg.get('target_responses', _DEFAULT_TARGET_RESPONSES)).lower()
target_prompts = str(dataset_cfg.get('target_prompts', _DEFAULT_TARGET_PROMPTS)).lower()
decoder_only_format = dataset_cfg.decoder_only_format
validate_target_settings(target_prompts, target_responses, decoder_only_format)
def _download_remote_hf_dataset(remote_path: str, split: str) -> str:
"""Downloads a dataset from a remote object store.
This function supports 'jsonl', 'csv', and 'parquet' file formats for the dataset. It will attempt to download
the dataset, then once it is downloaded, convert it into HuggingFace ``datasets`` format, and then return this
dataset.
The function also ensures synchronicity across multiple processes during the file download. It creates a signal
file that is used to synchronize the start of the download across different processes. Once the download is
completed, the function removes the signal file.
Args:
hf_name (str): The path of the HuggingFace dataset to download.
split (str): The dataset split to download (e.g., 'train', 'validation', 'test').
Returns:
A local directory path where the dataset files are stored.
Raises:
FileNotFoundError: Raised if the dataset file cannot be found with any of the supported extensions.
"""
hf_formatted_split = split.replace('-', '_')
finetune_dir = os.path.join(DOWNLOADED_FT_DATASETS_DIRPATH, hf_formatted_split if hf_formatted_split != 'data' else 'data_not')
os.makedirs(finetune_dir, exist_ok=True)
for extension in SUPPORTED_EXTENSIONS:
name = f"{remote_path.strip('/')}/{split}{extension}"
destination = str(os.path.abspath(os.path.join(finetune_dir, 'data', f'{hf_formatted_split}-00000-of-00001{extension}')))
signal_file_path = os.path.join(finetune_dir, f'.node_{dist.get_node_rank()}_local_rank0_completed')
if dist.get_local_rank() == 0:
try:
get_file(path=name, destination=destination, overwrite=True)
except FileNotFoundError as e:
if extension == SUPPORTED_EXTENSIONS[-1]:
files_searched = [f'{cfg.dataset.hf_name}/{cfg.dataset.split}{ext}' for ext in SUPPORTED_EXTENSIONS]
raise FileNotFoundError(f'Could not find a file with any of ' + f'the supported extensions: {SUPPORTED_EXTENSIONS}\n' + f'at {files_searched}') from e
else:
log.debug(f'Could not find {name}, looking for another extension')
continue
os.makedirs(os.path.dirname(signal_file_path), exist_ok=True)
with open(signal_file_path, 'wb') as f:
f.write(b'local_rank0_completed_download')
with dist.local_rank_zero_download_and_wait(signal_file_path):
dist.barrier()
if dist.get_local_rank() == 0:
os.remove(signal_file_path)
dist.barrier()
break
return finetune_dir
def _build_collate_fn(dataloader_cfg: DictConfig, tokenizer: PreTrainedTokenizerBase, device_batch_size: int) -> Tuple[Union[Seq2SeqFinetuningCollator, BinPackCollator], int]:
dataset_cfg = dataloader_cfg.dataset
max_seq_len = dataset_cfg.max_seq_len
collate_fn = Seq2SeqFinetuningCollator(tokenizer=tokenizer, max_seq_len=max_seq_len, decoder_only_format=dataset_cfg.decoder_only_format, target_responses=dataset_cfg.get('target_responses', _DEFAULT_TARGET_RESPONSES), target_prompts=dataset_cfg.get('target_prompts', _DEFAULT_TARGET_PROMPTS), allow_pad_trimming=dataset_cfg.get('allow_pad_trimming', False))
packing_ratio = dataset_cfg.get('packing_ratio')
max_leftover_bins_to_keep = dataset_cfg.get('max_leftover_bins_to_keep')
if packing_ratio is None:
if max_leftover_bins_to_keep is not None:
raise ValueError('dataset.max_leftover_bins_to_keep has been defined, ' + 'but dataset.packing_ratio has not been set. Please set ' + 'the latter to turn on packing or remove the former from the config.')
return (collate_fn, device_batch_size)
if packing_ratio == 'auto':
packing_ratio = auto_packing_ratio(dataloader_cfg, tokenizer, device_batch_size)
if isinstance(packing_ratio, str):
raise ValueError('dataset.packing_ratio must be a float or "auto", but it was set to ' + f'{packing_ratio}.')
log.info(f'Using packing ratio {packing_ratio}')
if packing_ratio == 1.0:
return (collate_fn, device_batch_size)
elif packing_ratio < 1.0:
raise ValueError('packing_ratio must be >= 1, if supplied')
if not dataset_cfg.decoder_only_format:
raise NotImplementedError('On-the-fly packing is currently only supported for decoder-only formats.')
collate_fn = BinPackCollator(collator=collate_fn, target_batch_size=device_batch_size, max_seq_len=max_seq_len, pad_token_id=tokenizer.pad_token_id, padding_side=tokenizer.padding_side, max_leftover_bins_to_keep=max_leftover_bins_to_keep)
n_examples_to_pack = int(device_batch_size * packing_ratio)
return (collate_fn, n_examples_to_pack)
if __name__ == '__main__':
import torch
from .utils import build_tokenizer
cfg = om.create({'dataset': {'hf_name': 'tatsu-lab/alpaca', 'preprocessing_fn': 'llmfoundry.data.finetuning.tasks:alpaca_preprocessing_function', 'split': 'train', 'packing_ratio': 18.0, 'max_seq_len': 2048, 'decoder_only_format': True, 'allow_pad_trimming': False, 'num_canonical_nodes': 472, 'shuffle': True, 'target_responses': 'last', 'target_prompts': 'none'}, 'drop_last': False, 'num_workers': 0, 'pin_memory': False, 'prefetch_factor': None, 'persistent_workers': False, 'timeout': 0})
tokenizer_name = 'EleutherAI/gpt-neox-20b'
tokenizer_kwargs = {'model_max_length': cfg.dataset.max_seq_len}
tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
device_batch_size = 1
dataloader = build_finetuning_dataloader(cfg, tokenizer, device_batch_size).dataloader
packing = cfg.dataset.get('packing_ratio') is not None
for i, batch in enumerate(dataloader):
if i >= 5:
break
print(f'-----Batch {i}-----')
for k, v in batch.items():
if isinstance(v, torch.Tensor):
print(k, v.shape)
else:
print(k, v)
for j in range(device_batch_size):
print(f'--- Sample {j} ---')
if cfg.dataset.decoder_only_format:
if packing:
for subseq in range(int(batch['sequence_id'][j].max()) + 1):
is_subseq = batch['sequence_id'][j] == subseq
print('\x1b[93m{}\x1b[00m\n'.format('INPUT IDS:'), tokenizer.decode(batch['input_ids'][j, torch.logical_and(is_subseq, batch['attention_mask'][j] == 1)], skip_special_tokens=False, clean_up_tokenization_spaces=True))
print('\x1b[91m{}\x1b[00m\n'.format('TARGET: '), tokenizer.decode(batch['input_ids'][j, torch.logical_and(is_subseq, batch['labels'][j] != _HF_IGNORE_INDEX)], skip_special_tokens=False, clean_up_tokenization_spaces=True))
else:
print('\x1b[93m{}\x1b[00m\n'.format('INPUT IDS:'), tokenizer.decode(batch['input_ids'][j, batch['attention_mask'][j] == 1], skip_special_tokens=False, clean_up_tokenization_spaces=True))
print('\x1b[91m{}\x1b[00m\n'.format('TARGET: '), tokenizer.decode(batch['input_ids'][j, batch['labels'][j] != _HF_IGNORE_INDEX], skip_special_tokens=False, clean_up_tokenization_spaces=True))
else:
print('\x1b[92m{}\x1b[00m\n'.format('CONTEXT: '), tokenizer.decode(batch['input_ids'][j, batch['attention_mask'][j] == 1], skip_special_tokens=False, clean_up_tokenization_spaces=True))
print('\x1b[91m{}\x1b[00m\n'.format('TARGET: '), tokenizer.decode(batch['labels'][j, batch['decoder_attention_mask'][j] == 1], skip_special_tokens=False, clean_up_tokenization_spaces=True))
print(' ') |