diff --git a/configs/pretrain_all_xl_3.yaml b/configs/pretrain_all_xl_3.yaml index e19f84a..be40251 100644 --- a/configs/pretrain_all_xl_3.yaml +++ b/configs/pretrain_all_xl_3.yaml @@ -37,18 +37,19 @@ target_modules: - down_proj # data train_ds_names: - - fw_qa_3_mini - - ctx_qa - - pwc - - hotpot_qa - - squad - - drop - - narrativeqa - - quoref - - ropes - - synthetic_convqa + - fw_qa_3 # ~ 267M + - ctx_qa # 300k + - pwc # 240k + - hotpot_qa # 90k + - squad # 90k + - drop # 77k + - narrativeqa # 40k + - quoref # 11k + - ropes # 11k + - synthetic_convqa # 40k val_ds_names: + - fw_qa_3 - fw_qa_xl - ctx_qa - pwc diff --git a/intx_sft.py b/intx_sft.py index 7a65c16..ede5cca 100755 --- a/intx_sft.py +++ b/intx_sft.py @@ -3,6 +3,7 @@ import os import random import string import time +from math import ceil from collections import defaultdict from copy import copy, deepcopy from functools import partial @@ -15,6 +16,7 @@ import wandb import yaml from ctx_to_lora.data_utils import ( + DS_LEN, convert_ctx_prompt_response_to_messages, get_preprocessing_fn, get_sft_prompt_formatting_fn, @@ -282,7 +284,7 @@ def main(): training_args.lr_scheduler_type == "cosine_with_min_lr" and training_args.lr_scheduler_kwargs is None ): - training_args.lr_scheduler_kwargs = {"min_lr": 1e-7} + training_args.lr_scheduler_kwargs = {"min_lr": 1e-8} args = { **vars(deepcopy(data_args)), **vars(deepcopy(ctx_args)), @@ -380,6 +382,7 @@ def main(): add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel) tokenizer_kwargs = {"max_length": ctx_args.max_base_len} ctx_tokenizer_kwargs = {"max_length": ctx_args.max_ctx_len} # not used for now + _get_tokenized_dataset = partial( get_tokenized_dataset, tokenizer=tokenizer, @@ -400,12 +403,16 @@ def main(): ): if not ds_names: continue + streaming = (split == "train") and data_args.streaming tokenized_ds[split] = { - os.path.basename(ds_name): _get_tokenized_dataset(ds_name, split) + os.path.basename(ds_name): _get_tokenized_dataset( + ds_name, split, streaming=streaming + ) for ds_name in ds_names } train_ds = tokenized_ds["train"] + logging.info(f"train_ds: {train_ds}") val_ds = dict() if "validation" in tokenized_ds: n_val_samples = data_args.max_val_samples_per_ds @@ -419,35 +426,31 @@ def main(): val_indices = np.random.permutation(len(ds))[:n_val_samples] val_ds[ds_name] = val_ds[ds_name].select(val_indices) - # train_ds = concatenate_datasets(list(train_ds.values())) + max_steps = ceil( + sum([DS_LEN[ds] for ds in train_ds]) + * training_args.num_train_epochs + / training_args.per_device_train_batch_size + / training_args.gradient_accumulation_steps + / training_args.world_size + ) + training_args.max_steps = max_steps - # total_len = sum(len(ds) for ds in train_ds.values()) - train_ds_len = [len(ds) for ds in train_ds.values()] - total_len = sum(train_ds_len) + # interleaving streaming datasets + # simplify the probs for smaller datasets + # slightly upsample those datasets + probs = [ + 0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds) + for ds_name in train_ds + ] train_ds = interleave_datasets( list(train_ds.values()), - probabilities=[l / total_len for l in train_ds_len], + probabilities=probs, + stopping_strategy="all_exhausted", seed=training_args.seed, ) - # val_train_indices = np.random.permutation(len(train_ds))[:500] - # val_ds["train"] = train_ds.select(val_train_indices) - - test_ds = dict() - if "test" in tokenized_ds: - n_test_samples = data_args.max_test_samples_per_ds - for ds_name, ds in tokenized_ds["test"].items(): - test_ds[ds_name] = ds - test_indices = np.random.permutation(len(ds))[:n_test_samples] - test_ds[ds_name] = test_ds[ds_name].select(test_indices) - logger.info(f"train_ds: {train_ds}") logger.info(f"val_ds: {val_ds}") - logger.info(f"test_ds: {test_ds}") - - # TODO: change to a faster collator? e.g., - # https://huggingface.co/blog/packing-with-FA2 - # data_collator = DataCollatorForSeq2Seq(tokenizer, model, pad_to_multiple_of=8) flattener = DataCollatorWithFlattening() def train_packed_collator(inp_list): @@ -549,7 +552,6 @@ def main(): training_args, train_ds, val_ds, - test_ds, train_collator, # partial(generation_collator, tokenizer=tokenizer), compute_metrics=partial( @@ -571,6 +573,8 @@ if __name__ == "__main__": os.environ["WANDB_WATCH"] = "" # "all" os.environ["WANDB_CONSOLE"] = "off" os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + os.environ["KMP_AFFINITY"] = "disabled" # fixing iterable dataset stuck + os.environ["OMP_NUM_THREADS"] = "16" if os.getenv("DEBUG", False): disable_caching() # randomly sleep to avoid run_name collision diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 2617346..8245097 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -338,6 +338,10 @@ class DataArguments: default=None, metadata={"help": "Training dataset names."}, ) + streaming: bool = field( + default=False, + metadata={"help": "Whether to use streaming dataset for training."}, + ) val_ds_names: Optional[list[str]] = field( default=None, metadata={"help": "Validation dataset names."}, diff --git a/src/ctx_to_lora/data_utils.py b/src/ctx_to_lora/data_utils.py index 8dc7b50..5157f83 100644 --- a/src/ctx_to_lora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -1,4 +1,5 @@ import logging +from os import path import numpy as np from glob import glob from typing import Any, Callable, Iterator, Optional @@ -14,6 +15,24 @@ logger = logging.getLogger() FW_QA_PATHS = [ f"data/raw_datasets/fw_qa/{i:05d}.parquet" for i in [0, 1, 6, 7, 8, 10, 22, 30, 35] ] + +# approximate length of the datasets +# needed for streaming datasets +DS_LEN = { + "fw_qa_3_mini": 100_000, + "fw_qa_3": 270_000_000, + "fw_qa_xl": 27_000_000, + "ctx_qa": 300_000, + "pwc": 240_000, + "hotpot_qa": 100_000, + "squad": 90_000, + "drop": 77_000, + "narrativeqa": 40_000, + "quoref": 11_000, + "ropes": 11_000, + "synthetic_convqa": 40_000, +} + DS_KWARGS = { "hotpot_qa": dict( train=dict(path="hotpotqa/hotpot_qa", name="fullwiki", split="train[900:]"), @@ -101,10 +120,22 @@ DS_KWARGS = { "fw_qa_3_mini": dict( train=dict( path="parquet", - data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"), + data_files=glob("data/raw_datasets/fw_qa_3/000_00001.parquet"), split="train[:100000]", ), ), + "fw_qa_3": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_3/*[!val].parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_3/*val.parquet"), + split="train", + ), + ), "ctx_qa": dict( train=dict( path="parquet", @@ -225,8 +256,24 @@ def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]: f"No dataset kwargs found for '{ds_name}' with split '{split}'.\n" f"Using default kwargs: {kwargs}" ) - return kwargs - return DS_KWARGS[ds_name][split] + # return kwargs + else: + # return DS_KWARGS[ds_name][split] + kwargs = DS_KWARGS[ds_name][split] + + # custom logic for slicing iterable datasets + take, skip = None, None + kw_split = kwargs["split"] + if ("[" in kw_split) and kw_split.endswith("]"): + kwargs["split"], slice = kw_split.split("[") + slice = slice.strip("]") + skip = slice.split(":")[0] + if skip: + kwargs["skip"] = int(skip) + take = slice.split(":")[1] + if take: + kwargs["take"] = int(take) + return kwargs def validate_columns(tokenized_ds): @@ -354,12 +401,24 @@ def get_tokenized_dataset( add_negative_prompt: bool, use_kl_loss: bool, set_format: Optional[str] = None, + streaming: bool = False, ) -> dict[str, Any]: logger.debug(f"Loading dataset {ds_name} with split {split}...") need_ctx_ids = not add_ctx_to_chat try: - ds = load_dataset(**get_ds_kwargs(ds_name, split), trust_remote_code=True) + ds_kwargs = get_ds_kwargs(ds_name, split) + skip = ds_kwargs.pop("skip", None) + take = ds_kwargs.pop("take", None) + ds = load_dataset( + **get_ds_kwargs(ds_name, split), + trust_remote_code=True, + streaming=streaming, + ) + if skip is not None: + ds = ds.skip(skip) + if take is not None: + ds = ds.take(take) except ValueError as e: logger.info( f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..." @@ -368,10 +427,10 @@ def get_tokenized_dataset( cols_to_remove = [ col for col in ds.column_names if col not in ["context", "prompt", "response"] ] - ds = ds.map(get_preprocessing_fn(ds_name), num_proc=16) + ds = ds.map(get_preprocessing_fn(ds_name)) ds = ds.remove_columns(cols_to_remove) - ds = ds.filter(filter_none, batched=True, num_proc=16) - ds = ds.filter(filter_long_samples, batched=True, num_proc=16) + ds = ds.filter(filter_none, batched=True) + ds = ds.filter(filter_long_samples, batched=True) if split == "train": if add_negative_prompt: ds = ds.map(add_negative_prompt_fn, batched=True) @@ -407,14 +466,11 @@ def construct_and_tokenize_ctx_qa( ds = ds.map( convert_ctx_prompt_response_to_messages, fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat}, - num_proc=16, ) # add "chat" field - ds = ds.map( - get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), num_proc=16 - ) + ds = ds.map(get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer)) # tokenize the chat + mask the assistant inputs - ds = ds.filter(filter_long_chat, batched=True, num_proc=16) + ds = ds.filter(filter_long_chat, batched=True) # add "input_ids", "attention_mask", "labels" tokenized_ds = ds.map( @@ -424,7 +480,6 @@ def construct_and_tokenize_ctx_qa( "mask_assistant_inputs": True, "tokenizer_kwargs": tokenizer_kwargs, }, - num_proc=16, ) # for use_kl_loss, we need "chat_ids" and "chat_attn_mask" @@ -447,25 +502,25 @@ def construct_and_tokenize_ctx_qa( "tokenizer_kwargs": tokenizer_kwargs, "for_kl_loss": True, }, - num_proc=16, ) if need_ctx_ids: - # TODO: can we batch this? - # TODO: can we cache this? # tokenize the ctx_text to get ctx_ids and ctx_attn_mask tokenized_ds = tokenized_ds.map( tokenize_ctx_text, fn_kwargs={"tokenizer": ctx_tokenizer}, batched=True, - num_proc=16, ) tokenized_ds = tokenized_ds.remove_columns( ["messages", "chat", "context", "prompt", "response"] ) - tokenized_ds.set_format(type=set_format) - validate_columns(tokenized_ds) + if set_format: + tokenized_ds.set_format(type=set_format) + # if isinstance(tokenized_ds, IterableDataset): + # # the columns are unknown when using streaming dataset + # tokenized_ds = tokenized_ds._resolve_features() + # validate_columns(tokenized_ds) return tokenized_ds diff --git a/src/ctx_to_lora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py index 55e22ab..8449aed 100644 --- a/src/ctx_to_lora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -1212,6 +1212,9 @@ class ModulatedPretrainedModel(nn.Module): A = A[0:1] B = B[:, 0:1] biases = [A, B.T] + # TODO: scale the biases by the output size? + # e.g., rank=16 gives bigger (2x?) gradient magnitudes at initialization + # compared to rank=8 # bias-hyperinit # init weights to zeros and bias to the base weights diff --git a/src/ctx_to_lora/training_utils.py b/src/ctx_to_lora/training_utils.py index 2153d04..2a11845 100644 --- a/src/ctx_to_lora/training_utils.py +++ b/src/ctx_to_lora/training_utils.py @@ -99,7 +99,7 @@ def train_model( training_args, train_dataset=None, val_dataset=None, - test_dataset=None, + # test_dataset=None, train_collator=None, # generation_collator=None, compute_metrics=None,