diff --git a/src/ctx_to_lora/configs.py b/src/ctx_to_lora/configs.py index afe8a06..ad62730 100644 --- a/src/ctx_to_lora/configs.py +++ b/src/ctx_to_lora/configs.py @@ -315,6 +315,13 @@ class CtxTrainingArguments: default=False, metadata={"help": "Whether to use sequence packing."}, ) + max_qas_len: int = field( + default=2**11, + metadata={ + "help": "Maximum question-answering token length of each sample for training. " + "QA pairs that are longer than this value will be split up into multiple samples." + }, + ) max_packed_inp_len: int | None = field( default=2**14, metadata={"help": "Maximum packed input length for training."}, diff --git a/src/ctx_to_lora/data/packing.py b/src/ctx_to_lora/data/packing.py index 8080b28..07870ba 100644 --- a/src/ctx_to_lora/data/packing.py +++ b/src/ctx_to_lora/data/packing.py @@ -4,7 +4,7 @@ import logging import numpy as np -from ctx_to_lora.utils import check_is_iterable, concat_list +from ctx_to_lora.utils import check_is_iterable, concat_list, setup_logging logger = logging.getLogger() @@ -16,44 +16,49 @@ def pack_data_points_by_length( max_packed_inp_len: int, max_packed_ctx_len: int, max_size: int = -1, -) -> list[int]: +) -> tuple[list[int], list[int]]: if not lens: return [] - len_arr = np.array([sum(l) for l in lens], dtype=np.int32) - ctx_len_arr = np.array(ctx_lens, dtype=np.int32) + len_arr = np.array([sum(l) for l in lens], dtype=np.long) + ctx_len_arr = np.array(ctx_lens, dtype=np.long) n = len(len_arr) - assert len(ctx_len_arr) == n, "Length of ctx_lens must match length of lens" + assert len(ctx_len_arr) == n, "Length of ctx_len_arr must match length of lens" assert len(n_queries) == n, "Length of n_queries must match length of lens" if n == 1: return ( - [[0]] - if len_arr[0] <= max_packed_inp_len or ctx_len_arr[0] <= max_packed_ctx_len + [0] + if len_arr[0] <= max_packed_inp_len and ctx_len_arr[0] <= max_packed_ctx_len else [] ) # 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 - start_indices_inp = [0] - start_indices_ctx = [0] + # indices of flattened input_ids + cumsum_n_queries = np.cumsum(n_queries) + + # start_indices_inp = [0] + # start_indices_ctx = [0] + idx_pairs = [] i = 0 while i < n: # Find the maximum j such that sum(lens[i:j+1]) <= max_packed_inp_len start_sum_inp = cumsum_inp_len[i - 1] if i > 0 else 0 - valid_ends_inp = cumsum_inp_len[i:] - start_sum_inp <= max_packed_inp_len + valid_ends_inp = (cumsum_inp_len[i:] - start_sum_inp) <= max_packed_inp_len start_sum_ctx = cumsum_ctx_len[i - 1] if i > 0 else 0 - valid_ends_ctx = cumsum_ctx_len[i:] - start_sum_ctx <= max_packed_ctx_len + valid_ends_ctx = (cumsum_ctx_len[i:] - start_sum_ctx) <= max_packed_ctx_len valid_ends = valid_ends_inp & valid_ends_ctx - # this should never happen? if not np.any(valid_ends): # Single item exceeds max_packed_inp_len, skip it + logging.debug( + f"Skipping item {i} with input length {len_arr[i]} and context length {ctx_len_arr[i]}" + ) i += 1 continue @@ -61,25 +66,28 @@ def pack_data_points_by_length( max_valid_idx = i + np.where(valid_ends)[0][-1] # Apply max_size constraint - if max_size != -1: + if max_size > 0: max_valid_idx = min(max_valid_idx, i + max_size - 1) - start_indices_inp.append(cumsum_n_queries[max_valid_idx]) - start_indices_ctx.append(max_valid_idx + 1) + idx_pairs.append((i, max_valid_idx + 1)) i = max_valid_idx + 1 - return start_indices_inp, start_indices_ctx + return idx_pairs def pack_data_points_FA( batch: dict[str, any], + need_flatten: bool, ) -> dict[str, np.ndarray]: if not batch: raise ValueError("Batch is empty") # Pre-allocate lists with known sizes total_ctx_len = sum(len(x) for x in batch["ctx_ids"]) - total_inp_len = sum(len(x) for x in batch["input_ids"]) + if not need_flatten: + total_inp_len = sum(len(x) for x in batch["input_ids"]) + else: + total_inp_len = sum(len(y) for x in batch["input_ids"] for y in x) ctx_ids = np.empty(total_ctx_len, dtype=np.long) ctx_position_ids = np.empty(total_ctx_len, dtype=np.long) @@ -89,12 +97,26 @@ def pack_data_points_FA( offset = ctx_offset = 0 for input_ids_b, labels_b in zip(batch["input_ids"], batch["labels"]): - inp_len = len(input_ids_b) + inp_start = offset + if need_flatten: + # compute position_ids for each sub-list in input_ids_b + local_start = inp_start + for ids_b in input_ids_b: + position_ids[local_start : local_start + len(ids_b)] = np.arange( + len(ids_b), dtype=np.long + ) + local_start += len(ids_b) - inp_start, inp_end = offset, offset + inp_len + input_ids_b = concat_list(input_ids_b) + labels_b = concat_list(labels_b) + + inp_len = len(input_ids_b) + inp_end = offset + inp_len + + if not need_flatten: + position_ids[inp_start:inp_end] = np.arange(inp_len, dtype=np.long) input_ids[inp_start:inp_end] = input_ids_b - position_ids[inp_start:inp_end] = np.arange(inp_len, dtype=np.long) labels[inp_start:inp_end] = labels_b offset += inp_len @@ -120,17 +142,13 @@ 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 - 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]]) + # 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"]] @@ -142,7 +160,7 @@ def pack_batch( ctx_lens = [len(x) for x in batch["ctx_ids"]] # Group indices - start_indices_inp, start_indices_ctx = pack_data_points_by_length( + idx_pairs = pack_data_points_by_length( inp_lens, ctx_lens, n_queries, @@ -165,35 +183,30 @@ def pack_batch( packing_efficiency_ratios = [] ctx_packing_efficiency_ratios = [] - 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] + for idx_pair in idx_pairs: + start_idx, end_idx = idx_pair[0], idx_pair[1] + # start_idx = start_indices[i] + # end_idx = start_indices[i + 1] group_items = { - "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], + "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], } - packed_item = pack_data_points_FA(group_items) + packed_item = pack_data_points_FA(group_items, need_flatten) packed_batch["ctx_ids"].append(packed_item["ctx_ids"]) packed_batch["ctx_position_ids"].append(packed_item["ctx_position_ids"]) 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]) + packed_batch["n_queries"].append(n_queries[start_idx:end_idx]) - # # Track packing efficiency - # group_inp_tokens = sum(inp_lens[start_idx:end_idx]) - # group_ctx_tokens = sum(ctx_lens[start_idx:end_idx]) + if max_packed_inp_len > 0: + inp_efficiency = len(packed_item["input_ids"]) / max_packed_inp_len + packing_efficiency_ratios.append(inp_efficiency) - # if max_packed_inp_len > 0: - # inp_efficiency = group_inp_tokens / max_packed_inp_len - # packing_efficiency_ratios.append(inp_efficiency) - - # if max_packed_ctx_len > 0: - # ctx_efficiency = group_ctx_tokens / max_packed_ctx_len - # ctx_packing_efficiency_ratios.append(ctx_efficiency) + if max_packed_ctx_len > 0: + ctx_efficiency = len(packed_item["ctx_ids"]) / max_packed_ctx_len + ctx_packing_efficiency_ratios.append(ctx_efficiency) # # Calculate total original tokens # total_original_inp_tokens = sum(inp_lens) @@ -201,33 +214,29 @@ def pack_batch( # total_original_tokens = total_original_inp_tokens + total_original_ctx_tokens # # Calculate length statistics - # packed_inp_lens_arr = np.array([len(x) for x in packed_batch["input_ids"]]) - # packed_ctx_lens_arr = np.array([len(x) for x in packed_batch["ctx_ids"]]) + packed_inp_lens_arr = np.array([len(x) for x in packed_batch["input_ids"]]) + packed_ctx_lens_arr = np.array([len(x) for x in packed_batch["ctx_ids"]]) # # Log performance statistics - # compression_ratio = inp_count / n_samples if n_samples > 0 else 0 - # avg_inp_packing_efficiency = ( - # np.mean(packing_efficiency_ratios) if packing_efficiency_ratios else 0 - # ) - # avg_ctx_packing_efficiency = ( - # np.mean(ctx_packing_efficiency_ratios) if ctx_packing_efficiency_ratios else 0 - # ) + avg_inp_packing_efficiency = ( + np.mean(packing_efficiency_ratios) if packing_efficiency_ratios else 0 + ) + avg_ctx_packing_efficiency = ( + np.mean(ctx_packing_efficiency_ratios) if ctx_packing_efficiency_ratios else 0 + ) - # print( - # f"Packing stats - Original samples: {inp_count}\n" - # f"Packed samples: {n_samples}\n" - # f"Compression ratio (sample-wise): {compression_ratio:.2f}x\n\n" - # f"Total inp tokens: {total_original_inp_tokens}\n" - # f"Total ctx tokens: {total_original_ctx_tokens}\n" - # f"Avg inp packing efficiency: {avg_inp_packing_efficiency:.3f}\n" - # f"Avg ctx packing efficiency: {avg_ctx_packing_efficiency:.3f}\n\n" - # f"Input IDs length stats:\n" - # f" Avg: {np.mean(packed_inp_lens_arr):.1f}, Std: {np.std(packed_inp_lens_arr):.1f}, " - # f"Min: {np.min(packed_inp_lens_arr)}, Max: {np.max(packed_inp_lens_arr)}\n" - # f"Context IDs length stats:\n" - # f" Avg: {np.mean(packed_ctx_lens_arr):.1f}, Std: {np.std(packed_ctx_lens_arr):.1f}, " - # f"Min: {np.min(packed_ctx_lens_arr)}, Max: {np.max(packed_ctx_lens_arr)}" - # ) + logging.debug( + f"Packing stats - Original samples: {inp_count}\n" + f"# Packed samples: {len(idx_pairs)}\n" + f"Avg inp packing efficiency: {avg_inp_packing_efficiency:.3f}\n" + f"Avg ctx packing efficiency: {avg_ctx_packing_efficiency:.3f}\n\n" + f"Input IDs length stats:\n\n" + f" Avg: {np.mean(packed_inp_lens_arr):.1f}, Std: {np.std(packed_inp_lens_arr):.1f}, " + f"Min: {np.min(packed_inp_lens_arr)}, Max: {np.max(packed_inp_lens_arr)}\n" + f"Context IDs length stats:\n\n" + f" Avg: {np.mean(packed_ctx_lens_arr):.1f}, Std: {np.std(packed_ctx_lens_arr):.1f}, " + f"Min: {np.min(packed_ctx_lens_arr)}, Max: {np.max(packed_ctx_lens_arr)}" + ) return packed_batch @@ -240,6 +249,8 @@ if __name__ == "__main__": disable_caching() + setup_logging("tmp/packing_debug.log", debug=True) + logger.info("Starting packing script...") model, tokenizer = get_model_and_tokenizer( "google/gemma-2-2b-it", train=True, @@ -250,8 +261,9 @@ 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_v2_2k_len_level_0_tiny", + ds_name="self_gen/google/gemma-2-2b-it_temp_0.0/fw_qa_v2_2k_len_level_2_tiny", split="train", + max_qas_len=2048, base_model_max_len=model.base_model.config.max_position_embeddings, tokenizer=tokenizer, tokenizer_kwargs={}, @@ -269,13 +281,13 @@ if __name__ == "__main__": packed_ds = ds.map( pack_batch, fn_kwargs={ - "max_packed_inp_len": base_model_max_len * 2, - "max_packed_ctx_len": base_model_max_len * 4, + "max_packed_inp_len": 2**13, # 8k + "max_packed_ctx_len": 2**14, # 16k }, batched=True, - batch_size=1000, + batch_size=100_000, remove_columns=ds.column_names, - # num_proc=0, + num_proc=8, ) print(packed_ds) diff --git a/src/ctx_to_lora/data/processing.py b/src/ctx_to_lora/data/processing.py index 556655f..2284c7e 100644 --- a/src/ctx_to_lora/data/processing.py +++ b/src/ctx_to_lora/data/processing.py @@ -579,6 +579,7 @@ def load_and_process_dataset( def get_tokenized_dataset( ds_name: str, split: str, + max_qas_len: int, base_model_max_len: int, tokenizer: PreTrainedTokenizerBase, tokenizer_kwargs: dict[str, Any], @@ -612,6 +613,7 @@ def get_tokenized_dataset( streaming=streaming, ) tokenize_kwargs = dict( + max_qas_len=max_qas_len, base_model_max_len=base_model_max_len, tokenizer_kwargs=tokenizer_kwargs, ctx_model_max_len=ctx_model_max_len, @@ -627,9 +629,11 @@ def get_tokenized_dataset( all_kwargs = {**load_and_process_kwargs, **tokenize_kwargs} kwargs_str = json.dumps(all_kwargs) - kwargs_str += repr(tokenizer) + repr(ctx_tokenizer) + kwargs_str += tokenizer.name_or_path + ctx_tokenizer.name_or_path logger.debug(f"Tokenizing dataset with kwargs: {kwargs_str}") + kwargs_str += repr(tokenizer) + repr(ctx_tokenizer) ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest() + logger.debug(f"Dataset hash: {ds_hash}") ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}" if path.exists(ds_path) and split == "train" and is_caching_enabled(): @@ -651,6 +655,7 @@ def get_tokenized_dataset( num_proc=num_proc, **tokenize_kwargs, ) + tokenized_ds = tokenized_ds.shuffle() 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 @@ -659,6 +664,7 @@ def get_tokenized_dataset( def construct_and_tokenize_ctx_qa( + max_qas_len, base_model_max_len, tokenizer, tokenizer_kwargs, @@ -793,6 +799,14 @@ def construct_and_tokenize_ctx_qa( num_proc=num_proc, remove_columns=["data", "context"], ) + + tokenized_ds = tokenized_ds.map( + split_too_long_qas, + fn_kwargs={"max_qas_len": max_qas_len}, + batched=True, + batch_size=100_000, + num_proc=num_proc, + ) else: cols_to_remove = ["data", "context"] if "ctx_ids" in tokenized_ds.column_names: @@ -972,6 +986,69 @@ def convert_ctx_prompt_response_to_messages( return dict(data=data) +def split_too_long_qas(samples: dict[str, any], max_qas_len: int): + # samples keys: "input_ids", "attention_mask", "labels", "ctx_ids", "ctx_attn_mask" + # split the qas into multiple samples if they are too long + # e.g., if max_qas_len = 512, and qas is 1024 tokens long, + # we split it such that each sample has at most 512 tokens + # and the ctx_ids and ctx_attn_mask are the same for all samples + input_ids = samples["input_ids"] + attention_mask = samples["attention_mask"] + labels = samples["labels"] + ctx_ids = samples["ctx_ids"] + ctx_attn_mask = samples["ctx_attn_mask"] + out = {k: list() for k in samples} + longest_old_qas_len = 0 + longest_new_qas_len = 0 + for i in range(len(input_ids)): + tot_inp_len = sum([len(x) for x in input_ids[i]]) + longest_old_qas_len = max(longest_old_qas_len, tot_inp_len) + if tot_inp_len <= max_qas_len: + # total length of the qas are shorter than max_qas_len + # no need to split + for k in samples: + out[k].append(samples[k][i]) + continue + + new_qas_len = 0 + new_input_ids = [] + new_attn_mask = [] + new_labels = [] + for inp_ids, attn_mask, label in zip( + input_ids[i], attention_mask[i], labels[i] + ): + if len(inp_ids) > max_qas_len: + # skip + continue + if new_qas_len + len(inp_ids) <= max_qas_len: + new_qas_len += len(inp_ids) + new_input_ids.append(inp_ids) + new_attn_mask.append(attn_mask) + new_labels.append(label) + else: + out["input_ids"].append(new_input_ids) + out["attention_mask"].append(new_attn_mask) + out["labels"].append(new_labels) + out["ctx_ids"].append(ctx_ids[i]) + out["ctx_attn_mask"].append(ctx_attn_mask[i]) + longest_new_qas_len = max(longest_new_qas_len, new_qas_len) + new_qas_len = len(inp_ids) + new_input_ids = [inp_ids] + new_attn_mask = [attn_mask] + new_labels = [label] + if new_input_ids: + longest_new_qas_len = max(longest_new_qas_len, new_qas_len) + out["input_ids"].append(new_input_ids) + out["attention_mask"].append(new_attn_mask) + out["labels"].append(new_labels) + out["ctx_ids"].append(ctx_ids[i]) + out["ctx_attn_mask"].append(ctx_attn_mask[i]) + logger.debug(f"Longest old qas len: {longest_old_qas_len}") + logger.debug(f"Longest new qas len: {longest_new_qas_len}") + + return out + + def unpack_data(sample): return {k: sample["data"][k] for k in sample["data"]} @@ -1239,8 +1316,7 @@ def pack( max_packed_size: int, num_proc: int = 0, ): - # TODO: handle the new data format - # this would generate a giant cache file for the already concat'd + packed ds + # this would generate another cache file for the already concat'd + packed ds kwargs = dict( max_packed_inp_len=max_packed_inp_len, max_packed_ctx_len=max_packed_ctx_len, diff --git a/src/ctx_to_lora/modeling/aggregator.py b/src/ctx_to_lora/modeling/aggregator.py index 118b40c..d33982f 100644 --- a/src/ctx_to_lora/modeling/aggregator.py +++ b/src/ctx_to_lora/modeling/aggregator.py @@ -46,6 +46,9 @@ class AggregatorConfig: lora_r: int = 8 per_rank_gen: bool = False + layer_to_layer_ctx_encoder: bool = False + n_base_queries: int = 208 + def get_aggregator_config( model: PreTrainedModel,