apunkt/doc-to-lora
1
0
Fork 0
mirror of https://github.com/SakanaAI/doc-to-lora.git synced 2026-07-23 17:01:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

long qas splitting + packing properly skip long qas

Browse source
This commit is contained in:
51616 2025-07-03 11:03:56 +00:00
parent 1a03b01a52
commit 2efba2bb32
4 changed files with 178 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

7
src/ctx_to_lora/configs.py
View file

@ -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."},

166
src/ctx_to_lora/data/packing.py
View file

@ -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)

82
src/ctx_to_lora/data/processing.py
View file

@ -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,

3
src/ctx_to_lora/modeling/aggregator.py
View file

@ -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,