diff --git a/configs/hotpot_qa.yaml b/configs/hotpot_qa.yaml new file mode 100644 index 0000000..cffffd4 --- /dev/null +++ b/configs/hotpot_qa.yaml @@ -0,0 +1,45 @@ +output_dir: "" # just a placeholder +bf16: true +model_name_or_path: meta-llama/Llama-3.2-1B-Instruct +label_names: ["labels"] +# eval_on_start: True +# eval_strategy: "steps" +# eval_steps: 500 +# save_strategy: "no" +# # save_steps: 500 +# logging_strategy: "steps" +# logging_steps: 100 +# use_liger_kernel: true +# remove_unused_columns: false + +# needed to avoid OOM by compute the metrics batch by batch +# w/o this the trainer stores logits of all sample in memory... +# batch_eval_metrics: true + +per_device_train_batch_size: 32 +per_device_eval_batch_size: 32 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw +learning_rate: 0.00001 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 1 +weight_decay: 0.1 +warmup_ratio: 0.05 + +# LoRA +lora_r: 16 +lora_dropout: 0.05 +target_modules: + - down_proj + - up_proj + - gate_proj + +# data +train_ds_names: +- hotpotqa/hotpot_qa + +val_ds_names: +- hotpotqa/hotpot_qa + +test_ds_names: +- hotpotqa/hotpot_qa diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py index 5159d6e..68f30c8 100644 --- a/hyperlora/data_utils.py +++ b/hyperlora/data_utils.py @@ -13,6 +13,29 @@ IGNORE_INDEX = -100 logger = logging.getLogger() +DS_KWARGS = { + "hotpotqa/hotpot_qa": dict( + train=dict(name="fullwiki", split="train[1000:2000]"), + validation=dict(name="fullwiki", split="train[:1000]"), + test=dict(name="fullwiki", split="validation"), + ), + "sggetao/PwC": dict( + train=dict(split="train[1000:]"), + validation=dict(split="train[:1000]"), + test=dict(split="test"), + ), +} + + +def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]: + if (ds_name not in DS_KWARGS) or (split not in DS_KWARGS[ds_name]): + logger.warning( + f"No dataset kwargs found for '{ds_name}' with split '{split}'.\n" + f"Using default kwargs: path={ds_name}, split={split}" + ) + return dict(split=split) + return DS_KWARGS[ds_name][split] + def validate_columns(tokenized_ds): cols = ["input_ids", "attention_mask", "labels"] @@ -45,6 +68,9 @@ def add_repeat_prompt_fn(samples): responses.append(ctx) prompts.append("Repeat the text above.") + logger.debug(f"Adding repeat prompt...") + logger.debug(f"# unique contexts: {len(unique_contexts)}") + return dict( context=ctxs + samples["context"], response=responses + samples["response"], @@ -79,6 +105,9 @@ def add_negative_prompt_fn(samples): prompts.append(prompt) responses.append(response) + logger.debug(f"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() @@ -104,6 +133,15 @@ def add_negative_prompt_fn(samples): ) +def filter_none(samples): + out = [True] * len(samples["context"]) + for k in samples: + for i, sample in enumerate(samples[k]): + if not bool(sample): + out[i] = False + return out + + def get_tokenized_dataset( ds_name: str, split: str, @@ -115,30 +153,23 @@ def get_tokenized_dataset( use_kl_loss: bool, ) -> 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(ds_name, split=split) + ds = load_dataset(ds_name, **get_ds_kwargs(ds_name, split)) except ValueError as e: logger.info( f"Failed to load dataset {ds_name} with split {split}. Error: {e}\nSkipping..." ) return None - ds = ds.map(get_preprocessing_fn(ds_name)) + ds = ds.map(get_preprocessing_fn(ds_name), remove_columns=ds.column_names) + ds = ds.filter(filter_none, batched=True) ds = ds.filter(filter_long_samples, batched=True) - if add_negative_prompt: - cols_to_remove = [ - col - for col in ds.column_names - if col not in ["context", "prompt", "response"] - ] - ds = ds.map(add_negative_prompt_fn, batched=True, remove_columns=cols_to_remove) - if add_repeat_prompt and "context_numbers" not in ds_name: - cols_to_remove = [ - col - for col in ds.column_names - if col not in ["context", "prompt", "response"] - ] - ds = ds.map(add_repeat_prompt_fn, batched=True, remove_columns=cols_to_remove) + if split != "test": + if add_negative_prompt: + ds = ds.map(add_negative_prompt_fn, batched=True, batch_size=None) + if add_repeat_prompt and "context_numbers" not in ds_name: + ds = ds.map(add_repeat_prompt_fn, batched=True, batch_size=None) # for sft + chat_model, we need to convert the dataset to chat format # add "messages" field @@ -195,9 +226,7 @@ def get_tokenized_dataset( ) tokenized_ds = tokenized_ds.remove_columns(other_cols) - tokenized_ds.set_format(type="pt") - validate_columns(tokenized_ds) return tokenized_ds @@ -310,54 +339,25 @@ def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, A if ds_name == "sggetao/PwC": - def f(example): - example["context"] = example["input"] - example["response"] = example["answer"] - return example + def f(sample): + return { + "context": sample["context"], + "prompt": sample["prompt"], + "response": sample["answer"], + } - # if ds_name == "context_numbers": + elif ds_name == "hotpotqa/hotpot_qa": - # def f(example): - # example["messages"] = [ - # {"role": "user", "content": example["prompt"]}, - # {"role": "assistant", "content": example["response"]}, - # ] - # return example + def f(sample): + txt = "" + for p in sample["context"]["sentences"]: + txt += " " + "".join(p) - # if ds_name.startswith("lol_"): - - # def f(example): - # txt = example["input"] - # task_def = txt.split("Definition: ")[1].split("\n\nPositive Example")[0] - # task_def += " Please complete the task without any explanation." - # if len(example["output"]) > 1: - # task_def += "\nThe answer should be a comma-separated list of possible completions." - # problem = txt.split("Now complete the following example -")[1].split("Input: ")[1].split("\nOutput:")[0] - # answer = ", ".join(example["output"]) - # return dict(task_def=task_def, problem=problem, answer=answer) - - # if ds_name.startswith("arc_"): - # ABCD = ["A", "B", "C", "D"] - - # def f(example): - # choices = example["choices"] - # assert len(choices["text"]) == len(choices["label"]) - # n_to_fill = 4 - len(choices["text"]) - # if len(choices["text"]) < 4: - # choices["text"] += ["N/A"] * n_to_fill - # if len(choices["label"]) < 4: - # if choices["label"][0].isdigit(): - # choices["label"] += [str(len(choices["label"]) + i + 1) for i in range(n_to_fill)] - # else: - # choices["label"] += [ABCD[len(choices["label"]) + i] for i in range(n_to_fill)] - # example["choices"] = choices - # return example - - # if ds_name.startswith("mbpp"): - # # for training an oracle lora on mbpp - # def f(example): - # example["assertions"] = "\n".join(example["test_list"]) - # return example + return { + "context": txt, + "prompt": sample["question"], + "response": sample["answer"], + } return f diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index a79946c..5df4173 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -87,6 +87,7 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks): perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf) return {"prefix_matching": perf.mean().item()} + @torch.no_grad() def compute_perplexity(shift_logits, shift_labels, valid_masks): loss_fct = torch.nn.CrossEntropyLoss(reduction="none") @@ -488,5 +489,5 @@ if __name__ == "__main__": os.environ["WANDB_WATCH"] = "" # "all" os.environ["WANDB_CONSOLE"] = "off" os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" - # disable_caching() + disable_caching() main()