diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index 5fb8ca0..d2e9f14 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -154,7 +154,7 @@ class TrainingArguments(TrainingArguments): }, ) dataloader_prefetch_factor: int = field( - default=8, + default=16, metadata={"help": "Number of batches loaded in advance by each worker."}, ) dataloader_num_workers: int = field( diff --git a/src/ctx_to_lora/data/collator.py b/src/ctx_to_lora/data/collator.py index 6d3e006..f3cde31 100644 --- a/src/ctx_to_lora/data/collator.py +++ b/src/ctx_to_lora/data/collator.py @@ -1,6 +1,3 @@ -import itertools -from collections.abc import Iterable - import numpy as np import torch from transformers.data import ( @@ -8,36 +5,40 @@ from transformers.data import ( default_data_collator, ) -from ctx_to_lora.utils import concat_list +from ctx_to_lora.utils import check_is_iterable, concat_list flattener = DataCollatorWithFlattening() -def concat_batch(inp_list): - sample = inp_list[0] - out = [ - { - k: list(itertools.chain.from_iterable([x[k] for x in inp_list])) - for k in sample - } - ] - return out - - def flatten_if_not_packed(inp_list): - # TODO: handle new data format # no padding sample = inp_list[0] - n = len(sample) + n = len(inp_list) if "position_ids" in sample: if n == 1: - return default_data_collator(inp_list, return_tensors="pt") + n_queries = sample.pop("n_queries") + batch = default_data_collator(inp_list, return_tensors="pt") + batch["n_queries"] = torch.tensor(n_queries) + return batch elif n > 1: - # needed for when batch_size > 1 (never used?) - return default_data_collator(concat_batch(inp_list), return_tensors="pt") + raise NotImplementedError("Please use batch_size=1 when using packed data") + # when batch_size > 1 (never used?) + # return default_data_collator(concat_batch(inp_list), return_tensors="pt") - # for eval data during training - packed_inputs = flattener(inp_list, return_tensors="pt") + # for eval data (not packed) during training + need_flatten = check_is_iterable(sample["input_ids"][0]) + if need_flatten: + n_queries = torch.tensor([len(x["input_ids"]) for x in inp_list]) + tokens = dict() + for k in ["input_ids", "attention_mask", "labels"]: + tokens[k] = concat_list([x[k] for x in inp_list]) + + packed_inputs = flattener(tokens, return_tensors="pt") + else: + n_queries = torch.ones(len(inp_list), dtype=torch.int32) + packed_inputs = flattener(inp_list, return_tensors="pt") + + packed_inputs["n_queries"] = n_queries if "ctx_ids" in sample: ctx_ids = [{"input_ids": example["ctx_ids"]} for example in inp_list] packed_ctx = flattener(ctx_ids, return_tensors="pt") @@ -81,7 +82,7 @@ def train_collator(inp_list, tokenizer): padding_value=0, ) - need_flatten = isinstance(sample["input_ids"][0], Iterable) + need_flatten = check_is_iterable(sample["input_ids"][0]) if need_flatten: # needed for one-lora-multi-q @@ -89,17 +90,16 @@ def train_collator(inp_list, tokenizer): tokens = dict() for k in ["input_ids", "attention_mask", "labels"]: tokens[k] = concat_list([x[k] for x in inp_list]) - labels = tokens.pop("labels") padded_seq = tokenizer.pad(tokens, **padding_kwargs) - # hacky explicit padding since the labels are not padded by default - labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"] else: # needed for eval data during training labels = [x.pop("labels") for x in inp_list] padded_seq = tokenizer.pad(inp_list, **padding_kwargs) n_queries = torch.ones(len(inp_list), dtype=torch.int32) + # hacky explicit padding since the labels are not padded by default + labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"] labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels) out = {**padded_seq, "labels": labels, "n_queries": n_queries} if ctx_ids is not None: diff --git a/src/ctx_to_lora/data/packing.py b/src/ctx_to_lora/data/packing.py index bb1a2a5..97e81dd 100644 --- a/src/ctx_to_lora/data/packing.py +++ b/src/ctx_to_lora/data/packing.py @@ -4,27 +4,27 @@ import logging import numpy as np +from ctx_to_lora.utils import check_is_iterable, concat_list + logger = logging.getLogger() def pack_data_points_by_length( - lens: list[int], + lens: list[list[int]], ctx_lens: list[int], + n_queries: list[int], max_packed_inp_len: int, max_packed_ctx_len: int, max_size: int = -1, ) -> list[int]: - """Fully vectorized version using NumPy operations for maximum performance. - - This is an alternative implementation that's more complex but potentially faster - for very large datasets. - """ if not lens: return [] - len_arr = np.array(lens, dtype=np.int32) + len_arr = np.array([sum(l) for l in lens], dtype=np.int32) ctx_len_arr = np.array(ctx_lens, dtype=np.int32) n = len(len_arr) + assert len(ctx_len_arr) == n, "Length of ctx_lens must match length of lens" + assert len(n_queries) == n, "Length of n_queries must match length of lens" if n == 1: return ( @@ -36,8 +36,10 @@ def pack_data_points_by_length( # Create cumulative sum arrays for efficient range queries cumsum_inp_len = np.cumsum(len_arr) cumsum_ctx_len = np.cumsum(ctx_len_arr) + cumsum_n_queries = np.cumsum(n_queries) # the last one is not used - boundaries = [0] + start_indices_inp = [0] + start_indices_ctx = [0] i = 0 while i < n: @@ -62,16 +64,16 @@ def pack_data_points_by_length( if max_size != -1: max_valid_idx = min(max_valid_idx, i + max_size - 1) - boundaries.append(max_valid_idx + 1) + start_indices_inp.append(cumsum_n_queries[max_valid_idx]) + start_indices_ctx.append(max_valid_idx + 1) i = max_valid_idx + 1 - return boundaries + return start_indices_inp, start_indices_ctx def pack_data_points_FA( batch: dict[str, any], ) -> dict[str, np.ndarray]: - """Optimized packing function using torch operations.""" if not batch: raise ValueError("Batch is empty") @@ -86,10 +88,7 @@ def pack_data_points_FA( labels = np.empty(total_inp_len, dtype=np.long) offset = ctx_offset = 0 - for ctx_ids_b, input_ids_b, labels_b in zip( - batch["ctx_ids"], batch["input_ids"], batch["labels"] - ): - ctx_len = len(ctx_ids_b) + for input_ids_b, labels_b in zip(batch["input_ids"], batch["labels"]): inp_len = len(input_ids_b) inp_start, inp_end = offset, offset + inp_len @@ -99,6 +98,8 @@ def pack_data_points_FA( labels[inp_start:inp_end] = labels_b offset += inp_len + for ctx_ids_b in batch["ctx_ids"]: + ctx_len = len(ctx_ids_b) ctx_start, ctx_end = ctx_offset, ctx_offset + ctx_len ctx_ids[ctx_start:ctx_end] = ctx_ids_b ctx_position_ids[ctx_start:ctx_end] = np.arange(ctx_len, dtype=np.long) @@ -119,9 +120,20 @@ def pack_batch( max_packed_ctx_len: int, max_packed_size: int = -1, ) -> dict[str, any]: + # TODO: handle new data format # Extract lengths # we do not pad so we can just take the length of the tokens - inp_lens = [len(x) for x in batch["input_ids"]] + keys = list(batch.keys()) + need_flatten = check_is_iterable(batch["input_ids"][0][0]) + if need_flatten: + n_queries = [len(x) for x in batch["input_ids"]] + inp_lens = [[len(y) for y in x] for x in batch["input_ids"]] + # batch = {k: concat_list([x for x in batch[k]]) for k in keys} + for k in ["input_ids", "attention_mask", "labels"]: + batch[k] = concat_list([x for x in batch[k]]) + else: + n_queries = [1] * len(batch["input_ids"]) + inp_lens = [[len(x)] for x in batch["input_ids"]] inp_count = len(inp_lens) # total_original_tokens = sum(inp_lens) if "ctx_ids" not in batch: @@ -130,34 +142,38 @@ def pack_batch( ctx_lens = [len(x) for x in batch["ctx_ids"]] # Group indices - boundaries = pack_data_points_by_length( + start_indices_inp, start_indices_ctx = pack_data_points_by_length( inp_lens, ctx_lens, + n_queries, max_packed_inp_len, max_packed_ctx_len, max_packed_size, ) # Pack groups - n_samples = len(boundaries) - 1 + # n_samples = len(boundaries) - 1 packed_batch = { "ctx_ids": [], "ctx_position_ids": [], "input_ids": [], "position_ids": [], "labels": [], + "n_queries": [], } packing_efficiency_ratios = [] ctx_packing_efficiency_ratios = [] - for i in range(len(boundaries) - 1): - start_idx = boundaries[i] - end_idx = boundaries[i + 1] + for i in range(len(start_indices_inp) - 1): + start_idx_inp = start_indices_inp[i] + end_idx_inp = start_indices_inp[i + 1] + start_idx_ctx = start_indices_ctx[i] + end_idx_ctx = start_indices_ctx[i + 1] group_items = { - "ctx_ids": batch["ctx_ids"][start_idx:end_idx], - "input_ids": batch["input_ids"][start_idx:end_idx], - "labels": batch["labels"][start_idx:end_idx], + "ctx_ids": batch["ctx_ids"][start_idx_ctx:end_idx_ctx], + "input_ids": batch["input_ids"][start_idx_inp:end_idx_inp], + "labels": batch["labels"][start_idx_inp:end_idx_inp], } packed_item = pack_data_points_FA(group_items) packed_batch["ctx_ids"].append(packed_item["ctx_ids"]) @@ -165,6 +181,7 @@ def pack_batch( packed_batch["input_ids"].append(packed_item["input_ids"]) packed_batch["position_ids"].append(packed_item["position_ids"]) packed_batch["labels"].append(packed_item["labels"]) + packed_batch["n_queries"].append(n_queries[start_idx_ctx:end_idx_ctx]) # # Track packing efficiency # group_inp_tokens = sum(inp_lens[start_idx:end_idx]) @@ -216,9 +233,13 @@ def pack_batch( if __name__ == "__main__": + from datasets import disable_caching + from ctx_to_lora.data.processing import get_tokenized_dataset from ctx_to_lora.model_loading import get_model_and_tokenizer + disable_caching() + model, tokenizer = get_model_and_tokenizer( "google/gemma-2-2b-it", train=True, @@ -229,7 +250,7 @@ if __name__ == "__main__": tokenizer_kwargs = {"max_length": base_model_max_len} # not used ctx_tokenizer_kwargs = {"max_length": base_model_max_len} # not used for now ds = get_tokenized_dataset( - ds_name="fw_qa_3_mini_pretrain", + ds_name="squad", split="train", base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, @@ -254,7 +275,7 @@ if __name__ == "__main__": batched=True, batch_size=1000, remove_columns=ds.column_names, - num_proc=8, + # num_proc=0, ) print(packed_ds) diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 84703a3..c723382 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -626,7 +626,7 @@ def get_tokenized_dataset( num_proc=num_proc, **tokenize_kwargs, ) - if split == "train": + if split == "train" and is_caching_enabled(): tokenized_ds.save_to_disk(ds_path, num_proc=num_proc) # force reload from disk for fingerprint consistency tokenized_ds = datasets.load_from_disk(ds_path) @@ -789,7 +789,7 @@ def construct_and_tokenize_ctx_qa( add_length_info, fn_kwargs={"columns": ["input_ids"]}, ) - if "ctx_ids" in tokenized_ds: + if "ctx_ids" in tokenized_ds.column_names: tokenized_ds = tokenized_ds.map( truncate_middle_if_too_long, fn_kwargs={ @@ -952,6 +952,7 @@ def unpack_data(sample): def unpack_data_eval(samples): + # n_queries always == 1 for eval data = samples["data"] out = dict(input_ids=[], attention_mask=[], labels=[]) for d in data: @@ -1221,7 +1222,7 @@ def pack( logger.info( f"Packing dataset {ds_hash} with max_packed_inp_len={max_packed_inp_len} and max_packed_ctx_len={max_packed_ctx_len}" ) - if path.exists(ds_path): + if path.exists(ds_path) and is_caching_enabled(): logger.info(f"Loading a cached packed dataset for {ds_path}") return datasets.load_from_disk(ds_path) else: diff --git a/src/ctx_to_lora/utils.py b/src/ctx_to_lora/utils.py index c59add9..86ea369 100644 --- a/src/ctx_to_lora/utils.py +++ b/src/ctx_to_lora/utils.py @@ -265,3 +265,11 @@ def concat_list(l): for x in l: out += x return out + + +def check_is_iterable(x): + try: + iter(x) + except TypeError: + return False + return True