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

new trainng data format working non-packing

Browse source
This commit is contained in:
51616 2025-06-23 23:21:22 +09:00
parent a37998a015
commit 0d8e50d218
5 changed files with 267 additions and 211 deletions
Download patch file Download diff file Expand all files Collapse all files

11
configs/context_numbers_10.yaml
View file

@ -31,11 +31,9 @@ weight_decay: 0.01
# LoRA
lora_r: 16
lora_dropout: 0.05
lora_dropout: 0.0
target_modules:
- down_proj
- up_proj
- gate_proj
# data
train_ds_names:
@ -59,10 +57,3 @@ val_ds_names:
- data/raw_datasets/context_numbers_8
- data/raw_datasets/context_numbers_9
- data/raw_datasets/context_numbers_10
test_ds_names:
- data/raw_datasets/context_numbers_11
- data/raw_datasets/context_numbers_12
- data/raw_datasets/context_numbers_13
- data/raw_datasets/context_numbers_14
- data/raw_datasets/context_numbers_15

4
intx_sft.py
View file

@ -470,9 +470,11 @@ if __name__ == "__main__":
# os.environ["KMP_AFFINITY"] = "disabled" # fixing iterable dataset stuck
os.environ["OMP_NUM_THREADS"] = "23"
# os.environ["HF_DATASETS_IN_MEMORY_MAX_SIZE"] = "137438953472" # 128 TB
torch._dynamo.config.capture_scalar_outputs = True
# DOES THIS CAUSE ACCURACY DEGRADATION?
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch._dynamo.config.capture_scalar_outputs = True
if os.getenv("DEBUG", False):
disable_caching()
main()

51
src/ctx_to_lora/data/collator.py
View file

@ -1,4 +1,5 @@
import itertools
from collections.abc import Iterable
import numpy as np
import torch
@ -7,6 +8,8 @@ from transformers.data import (
default_data_collator,
)
from ctx_to_lora.utils import concat_list
flattener = DataCollatorWithFlattening()
@ -21,24 +24,8 @@ def concat_batch(inp_list):
return out
# def train_packed_collator(inp_list):
# # no padding
# packed_inputs = flattener(inp_list, return_tensors="pt")
# if "ctx_ids" in inp_list[0]:
# ctx_ids = [{"input_ids": example["ctx_ids"]} for example in inp_list]
# packed_ctx = flattener(ctx_ids, return_tensors="pt")
# packed_inputs["ctx_ids"] = packed_ctx["input_ids"]
# packed_inputs["ctx_position_ids"] = packed_ctx["position_ids"]
# # for eval info
# if "ctx_ids_len" in inp_list[0]:
# packed_inputs["ctx_ids_len"] = [
# example["ctx_ids_len"] for example in inp_list
# ]
# return packed_inputs
def flatten_if_not_packed(inp_list):
# TODO: handle new data format
# no padding
sample = inp_list[0]
n = len(sample)
@ -46,8 +33,10 @@ def flatten_if_not_packed(inp_list):
if n == 1:
return default_data_collator(inp_list, return_tensors="pt")
elif n > 1:
# needed for 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")
if "ctx_ids" in sample:
ctx_ids = [{"input_ids": example["ctx_ids"]} for example in inp_list]
@ -71,9 +60,9 @@ def train_collator(inp_list, tokenizer):
pad_to_multiple_of=8,
return_tensors="pt",
)
sample = inp_list[0]
ctx_ids = None
if "ctx_ids" in inp_list[0]:
if "ctx_ids" in sample:
# have to be manual since it has [ctx_len, features] shape
# pad to the longest ctx_len in the batch
# which can have a different length from the input_ids, attn_mask, labels
@ -92,13 +81,27 @@ def train_collator(inp_list, tokenizer):
padding_value=0,
)
labels = [x.pop("labels") for x in inp_list]
padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
need_flatten = isinstance(sample["input_ids"][0], Iterable)
if need_flatten:
# needed for one-lora-multi-q
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])
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}
out = {**padded_seq, "labels": labels, "n_queries": n_queries}
if ctx_ids is not None:
out["ctx_ids"] = ctx_ids
out["ctx_attn_mask"] = ctx_attn_mask

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

@ -8,8 +8,7 @@ from os import path
from typing import Any
import datasets
import numpy as np
from datasets import Dataset, load_dataset
from datasets import Dataset, is_caching_enabled, load_dataset
from transformers import PreTrainedTokenizerBase
from ctx_to_lora.data.definitions import (
@ -22,7 +21,7 @@ from ctx_to_lora.data.definitions import (
TRANSFORMED_DATA_DIR,
)
from ctx_to_lora.data.packing import pack_batch
from ctx_to_lora.utils import TRAINING_TASK
from ctx_to_lora.utils import TRAINING_TASK, concat_list
logger = logging.getLogger()
@ -69,7 +68,30 @@ def get_preprocessing_fn(
"""
f = lambda x: x
if ds_name.startswith("longbench"):
if "fw_qa_v2" in ds_name and "level" in ds_name:
def f_train(sample):
# get questions/answers from all levels in the ds
q_cols = [col for col in sample.keys() if col.startswith("prompts_level")]
a_cols = [col for col in sample.keys() if col.startswith("answers_level")]
questions = concat_list([sample[col] for col in q_cols])
answers = concat_list([sample[col] for col in a_cols])
return {
"context": sample["context"],
"prompts": questions,
"responses": answers,
}
def f_eval(sample):
return {
"context": sample["context"],
"prompt": sample["prompt"],
"response": sample["response"],
}
f = f_eval if is_eval else f_train
elif ds_name.startswith("longbench"):
def f(sample):
return {
@ -142,7 +164,7 @@ def get_preprocessing_fn(
return {
"context": sample["passage"],
"prompt": prompt,
"response": ", ".join(sample["answers_spans"]["spans"]),
"response": ", ".join(set(sample["answers_spans"]["spans"])),
}
elif ds_name in ["narrativeqa", "quoref", "ropes"]:
@ -247,6 +269,21 @@ def get_preprocessing_fn(
f = eval_intx_decorator(f)
if not is_pretrain:
def maybe_convert_to_list(f):
def g(sample):
sample = f(sample)
if "prompt" in sample:
sample["prompts"] = [sample.pop("prompt")]
if "response" in sample:
sample["responses"] = [sample.pop("response")]
return sample
return g
f = maybe_convert_to_list(f)
return f
@ -305,6 +342,19 @@ def len_filter(sample, max_length: int, keys: list[str]):
return sum(m) == len(keys)
def len_filter_data(sample, max_length: int, keys: list[str]):
data = sample["data"]
indices = dict()
for k in keys:
indices[k] = {i for i, arr in enumerate(data[k]) if len(arr) <= max_length}
all_indices = set.intersection(*indices.values())
for k in data:
data[k] = [data[k][i] for i in all_indices]
return dict(data=data)
# def len_filter(ds, max_length: int, keys: list[str]):
# for k in keys:
# ds = ds.select([i for i, x in enumerate(ds[k]) if len(x) <= max_length])
@ -322,9 +372,7 @@ def len_filter(sample, max_length: int, keys: list[str]):
def add_repeat_prompt_fn(samples: dict[str, list[str]], prob: float):
unique_contexts = set()
ctxs, prompts, responses = [], [], []
for ctx, prompt, response in zip(
samples["context"], samples["prompt"], samples["response"]
):
for ctx in samples["context"]:
if random.random() > prob:
continue
# Only process if the context is not already in the set
@ -335,74 +383,74 @@ def add_repeat_prompt_fn(samples: dict[str, list[str]], prob: float):
# continue
unique_contexts.add(ctx)
ctxs.append(ctx)
responses.append(ctx)
prompts.append(get_repeat_prompt())
responses.append([ctx])
prompts.append([get_repeat_prompt()])
logger.debug("Adding repeat prompt...")
logger.debug(f"# unique contexts: {len(unique_contexts)}")
return dict(
context=ctxs + samples["context"],
response=responses + samples["response"],
prompt=prompts + samples["prompt"],
responses=responses + samples["responses"],
prompts=prompts + samples["prompts"],
)
def add_negative_prompt_fn(samples):
unique_contexts = set()
ctxs, prompts, responses = [], [], []
keywords = [
"repeat",
"rephrase",
"summarize",
"rewrite",
"title",
"keyword",
"continuation",
]
# def add_negative_prompt_fn(samples):
# unique_contexts = set()
# ctxs, prompts, responses = [], [], []
# keywords = [
# "repeat",
# "rephrase",
# "summarize",
# "rewrite",
# "title",
# "keyword",
# "continuation",
# ]
for ctx, prompt, response in zip(
samples["context"], samples["prompt"], samples["response"]
):
if ctx in unique_contexts:
continue
unique_contexts.add(ctx)
if len(ctx) > 3000:
continue
if any(keyword in prompt for keyword in keywords):
# Skip samples where the prompt contains any of the specified keywords
continue
# for ctx, prompt, response in zip(
# samples["context"], samples["prompt"], samples["response"]
# ):
# if ctx in unique_contexts:
# continue
# unique_contexts.add(ctx)
# if len(ctx) > 3000:
# continue
# if any(keyword in prompt for keyword in keywords):
# # Skip samples where the prompt contains any of the specified keywords
# continue
ctxs.append(ctx)
prompts.append(prompt)
responses.append(response)
# ctxs.append(ctx)
# prompts.append(prompt)
# responses.append(response)
logger.debug("Adding negative prompt...")
logger.debug(f"# unique contexts: {len(unique_contexts)}")
# logger.debug("Adding negative prompt...")
# logger.debug(f"# unique contexts: {len(unique_contexts)}")
# remove one last sample if the number of samples is odd
if len(ctxs) % 2 != 0:
ctxs.pop()
prompts.pop()
responses.pop()
# # remove one last sample if the number of samples is odd
# if len(ctxs) % 2 != 0:
# ctxs.pop()
# prompts.pop()
# responses.pop()
# to make sure that the negative prompt/response is not the same as the original
indices = list(np.random.permutation(len(ctxs))) + list(
np.random.permutation(len(ctxs))
)
neg_ctxs, neg_prompts, neg_responses = [], [], []
for idx in range(0, len(indices), 2):
i = indices[idx]
j = indices[idx + 1]
neg_ctxs.append(ctxs[i])
neg_prompts.append(ctxs[j] + "\n\n" + prompts[j])
neg_responses.append(responses[j])
# # to make sure that the negative prompt/response is not the same as the original
# indices = list(np.random.permutation(len(ctxs))) + list(
# np.random.permutation(len(ctxs))
# )
# neg_ctxs, neg_prompts, neg_responses = [], [], []
# for idx in range(0, len(indices), 2):
# i = indices[idx]
# j = indices[idx + 1]
# neg_ctxs.append(ctxs[i])
# neg_prompts.append(ctxs[j] + "\n\n" + prompts[j])
# neg_responses.append(responses[j])
return dict(
context=neg_ctxs + samples["context"],
prompt=neg_prompts + samples["prompt"],
response=neg_responses + samples["response"],
)
# return dict(
# context=neg_ctxs + samples["context"],
# prompt=neg_prompts + samples["prompt"],
# response=neg_responses + samples["response"],
# )
def filter_none(sample):
@ -460,7 +508,7 @@ def load_and_process_dataset(
cols_to_remove = [
col
for col in ds.column_names
if col not in ["context", "prompt", "response", "qas", "variation"]
if col not in ["context", "prompts", "responses", "qas", "variation"]
]
is_eval = split != "train"
ds = ds.map(
@ -559,7 +607,7 @@ def get_tokenized_dataset(
ds_hash = hashlib.sha256(kwargs_str.encode()).hexdigest()
ds_path = f"{TRANSFORMED_DATA_DIR}/{ds_hash}"
if path.exists(ds_path) and split == "train":
if path.exists(ds_path) and split == "train" and is_caching_enabled():
# load the cached ds
logger.info(f"Loaded tokenized dataset from {ds_path}")
tokenized_ds = datasets.load_from_disk(ds_path)
@ -631,79 +679,47 @@ def construct_and_tokenize_ctx_qa(
)
else:
# for sft + chat_model, we need to convert the dataset to chat format
# add "messages" field
# convert everything into a new intermediate single column format
# (required for tracking n_queries and
# ejecting entire (ctx,questions,answers) triplet after filtering)
# then convert back to columns after filtering
# for sft + chat_model, we need to convert the dataset to chat format
# add "data" field
ds = ds.map(
convert_ctx_prompt_response_to_messages,
fn_kwargs={"add_ctx_to_chat": add_ctx_to_chat},
num_proc=num_proc,
remove_columns=[col for col in ds.column_names if col != "context"],
)
# add "chat" field
# add input_ids, attention_mask, labels to "data"
tokenized_ds = ds.map(
get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer),
batched=True,
batch_size=100_000,
# # num_proc=num_proc,
)
# tokenize the chat + mask the assistant inputs
# TODO: drop samples or split into multiple contexts
# if split == "train":
# ds = ds.filter(
# filter_long_chat,
# batched=True,
# batch_size=100_000,
# num_proc=num_proc,
# )
# add "input_ids", "attention_mask", "labels"
# TODO: this is slow!
# replace manual assistant mask with `return_assistant_tokens_mask`
# when using apply_chat_template (function above)
# tokenized_ds = ds.map(
# tokenize_chat_messages,
# fn_kwargs={
# "tokenizer": tokenizer,
# "mask_assistant_inputs": True,
# "tokenizer_kwargs": tokenizer_kwargs,
# },
# # # num_proc=num_proc,
# )
if split == "train":
# base model cant handle these samples naturally
# should skip
tokenized_ds = tokenized_ds.filter(
len_filter,
fn_kwargs={
"max_length": base_model_max_len,
"keys": ["input_ids", "labels"],
},
# batched=False,
# batch_size=0,
num_proc=num_proc,
)
# tokenized_ds = len_filter(
# tokenized_ds, base_model_max_len, ["input_ids", "labels"]
# tokenized_ds = tokenized_ds.filter(
# len_filter,
# fn_kwargs={
# "max_length": base_model_max_len,
# "keys": ["input_ids", "labels"],
# },
# # batched=False,
# # batch_size=0,
# num_proc=num_proc,
# )
else:
tokenized_ds = tokenized_ds.map(
truncate_middle_if_too_long,
len_filter_data,
fn_kwargs={
"max_length": base_model_max_len,
"columns": ["input_ids", "attention_mask", "labels"],
"max_new_tokens": 256,
"keys": ["input_ids"],
},
# batched=True,
# batch_size=100_000,
# num_proc=num_proc,
)
tokenized_ds = tokenized_ds.map(
add_length_info,
fn_kwargs={"columns": ["input_ids"]},
# batched=True,
# batch_size=100_000,
# num_proc=num_proc,
num_proc=num_proc,
)
# for use_kl_loss, we need "chat_ids" and "chat_attn_mask"
@ -743,47 +759,61 @@ def construct_and_tokenize_ctx_qa(
tokenized_ds = tokenized_ds.filter(
len_filter,
fn_kwargs={"max_length": ctx_model_max_len, "keys": ["ctx_ids"]},
# batched=False,
# batch_size=0,
num_proc=num_proc,
)
# tokenized_ds = len_filter(tokenized_ds, ctx_model_max_len, ["ctx_ids"])
if not is_pretrain:
if split == "train":
tokenized_ds = tokenized_ds.map(
unpack_data,
num_proc=num_proc,
remove_columns=["data"],
)
else:
# truncate in the middle
tokenized_ds = tokenized_ds.map(
unpack_data_eval,
num_proc=num_proc,
remove_columns=["data"],
batched=True,
batch_size=100_000,
)
tokenized_ds = tokenized_ds.map(
truncate_middle_if_too_long,
fn_kwargs={
"max_length": ctx_model_max_len,
"columns": ["ctx_ids", "ctx_attn_mask"],
# cxt encoder doesnt need to add new_tokens
"max_new_tokens": 0,
"max_length": base_model_max_len,
"columns": ["input_ids", "attention_mask", "labels"],
"max_new_tokens": 256,
},
# batched=True,
# batch_size=100_000,
# num_proc=num_proc,
)
tokenized_ds = tokenized_ds.map(
add_length_info,
fn_kwargs={"columns": ["ctx_ids"]},
# batched=True,
# batch_size=100_000,
# num_proc=num_proc,
fn_kwargs={"columns": ["input_ids"]},
)
if "ctx_ids" in tokenized_ds:
tokenized_ds = tokenized_ds.map(
truncate_middle_if_too_long,
fn_kwargs={
"max_length": ctx_model_max_len,
"columns": ["ctx_ids", "ctx_attn_mask"],
# cxt encoder doesnt need to add new_tokens
"max_new_tokens": 0,
},
)
tokenized_ds = tokenized_ds.map(
add_length_info,
fn_kwargs={"columns": ["ctx_ids"]},
)
if is_pretrain:
tokenized_ds = tokenized_ds.remove_columns(["context", "text"])
if not is_paraphrased:
tokenized_ds = tokenized_ds.remove_columns(["qas"])
else:
tokenized_ds = tokenized_ds.remove_columns(
["messages", "context", "prompt", "response"]
)
tokenized_ds = tokenized_ds.remove_columns(["context"])
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
@ -829,33 +859,17 @@ def get_sft_prompt_formatting_fn(
"Only chat models + SFT are supported. "
"Training with pre-training data is not supported yet."
)
# TODO: add support for causal_lm training (for pre-training data)
# TODO: add support for recon training (for pre-training data)
# TODO: add support for non-chat models
# def f(example):
# output_texts = (
# dict(text=[]) if sft_mode == "causal_lm" else dict(prompt=[], response=[])
# )
# df = pd.DataFrame(dict(example))
# for i, inp_txt in df.iterrows():
# if sft_mode == "causal_lm":
# text = metadata["text_template"].format(**inp_txt)
# output_texts["text"].append(text)
# elif sft_mode == "completion":
# prompt = metadata["user_prompt_template"].format(**inp_txt)
# output_texts["prompt"].append(prompt)
# output_texts["response"].append(str(inp_txt[metadata["response_field"]]))
# return output_texts
def f_intx(samples):
# flatten all the messages into a list
# tokenize, the pack back correctly
messages_list = [x for x in samples["data"]]
# def f_intx(example):
# chat_text = tokenizer.apply_chat_template(
# example["messages"], tokenize=False, add_generation_prompt=False
# )
# return dict(chat=chat_text)
def f_intx(example):
n_queries = [len(x) for x in messages_list]
messages = concat_list(messages_list)
logger.info(f"Tokenizing {len(messages)} messages...")
tokens = tokenizer.apply_chat_template(
example["messages"],
messages,
tokenize=True,
add_special_token=False,
truncation=False,
@ -869,7 +883,19 @@ def get_sft_prompt_formatting_fn(
labels.append(o)
del tokens["assistant_masks"]
tokens["labels"] = labels
return tokens
per_ctx_tokens = []
i = 0
for n in n_queries:
per_ctx_tokens.append(
{
"input_ids": tokens["input_ids"][i : i + n],
"attention_mask": tokens["attention_mask"][i : i + n],
"labels": tokens["labels"][i : i + n],
}
)
i += n
return dict(data=per_ctx_tokens)
# return f if not apply_chat_template_fn is not None else f_intx
return f_intx
@ -877,7 +903,7 @@ def get_sft_prompt_formatting_fn(
def convert_ctx_prompt_response_to_messages(
example: dict[str, Any],
add_ctx_to_chat: bool = True,
add_ctx_to_chat: bool,
) -> dict[str, Any]:
"""
Convert context/prompt/response format to chat messages format.
@ -892,24 +918,52 @@ def convert_ctx_prompt_response_to_messages(
Raises:
ValueError: If 'prompt' or 'response' keys are missing
"""
if "prompt" not in example or "response" not in example:
raise ValueError(f"'prompt' and 'response' are required. Got: {example}")
prompt_field = "prompts"
res_field = "responses"
if prompt_field not in example or res_field not in example:
raise ValueError(
f"'{prompt_field}' and '{res_field}' are required. Got: {example}"
)
system_msg = ""
if "system_message" in example:
system_msg = example["system_message"].strip()
user_msg = example["prompt"].strip()
if add_ctx_to_chat:
user_msg = example["context"].strip() + "\n\n" + user_msg.strip()
data = []
for prompt, response in zip(example[prompt_field], example[res_field]):
user_msg = prompt.strip()
if add_ctx_to_chat:
user_msg = example["context"].strip() + "\n\n" + user_msg
messages = [
{"role": "system", "content": system_msg.strip()},
{"role": "user", "content": user_msg.strip()},
{"role": "assistant", "content": example["response"].strip()},
]
data.append(
[
{"role": "system", "content": system_msg.strip()},
{"role": "user", "content": user_msg.strip()},
{"role": "assistant", "content": response.strip()},
]
)
return dict(messages=messages)
return dict(data=data)
def unpack_data(sample):
return {k: sample["data"][k] for k in sample["data"]}
def unpack_data_eval(samples):
data = samples["data"]
out = dict(input_ids=[], attention_mask=[], labels=[])
for d in data:
for tokens in zip(
d["input_ids"],
d["attention_mask"],
d["labels"],
):
out["input_ids"].append(tokens[0])
out["attention_mask"].append(tokens[1])
out["labels"].append(tokens[2])
return out
# adapted from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547
@ -1153,9 +1207,8 @@ def pack(
max_packed_size: int,
num_proc: int = 0,
):
# TODO: packing has to happen after concat'ing all the ds
# might have to do this on the fly...
# or have a giant cached file for the already concat'd + packed ds
# TODO: handle the new data format
# this would generate a giant 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,

7
src/ctx_to_lora/utils.py
View file

@ -258,3 +258,10 @@ def clear_gpu():
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_max_memory_cached()
def concat_list(l):
out = []
for x in l:
out += x
return out