From cce43deb64ee5c91c9ea8fe9f330a5d4486b0756 Mon Sep 17 00:00:00 2001 From: 51616 Date: Fri, 17 Jan 2025 17:41:10 +0000 Subject: [PATCH] restructure --- configs/fw_qa.yaml | 4 +- configs/fw_qa_ctx_qa.yaml | 4 +- configs/fw_qa_large_ctx_pwc_hotpot.yaml | 55 +++++++++++++++++++ configs/fw_qa_pwc_hotpot.yaml | 4 +- configs/fw_qa_tiny.yaml | 4 +- configs/hotpot_qa.yaml | 4 +- configs/pwc_hotpot_qa.yaml | 4 +- install.sh | 1 + hyperlora/intx_sft.py => intx_sft.py | 20 ++++--- requirements.txt | 5 +- setup.py | 12 ++++ src/ctx_to_lora/__init__.py | 0 {hyperlora => src/ctx_to_lora}/configs.py | 0 {hyperlora => src/ctx_to_lora}/data_utils.py | 35 +++++++++++- {hyperlora => src/ctx_to_lora}/eval.py | 8 +-- {hyperlora => src/ctx_to_lora}/hooks.py | 2 +- .../ctx_to_lora}/model_loading.py | 0 .../ctx_to_lora}/modeling_utils.py | 25 ++++++--- {hyperlora => src/ctx_to_lora}/pooling.py | 0 .../ctx_to_lora}/training_utils.py | 0 {hyperlora => src/ctx_to_lora}/utils.py | 0 hyperlora/watcher.py => watcher.py | 2 +- 22 files changed, 154 insertions(+), 35 deletions(-) create mode 100644 configs/fw_qa_large_ctx_pwc_hotpot.yaml create mode 100644 install.sh rename hyperlora/intx_sft.py => intx_sft.py (97%) create mode 100644 setup.py create mode 100644 src/ctx_to_lora/__init__.py rename {hyperlora => src/ctx_to_lora}/configs.py (100%) rename {hyperlora => src/ctx_to_lora}/data_utils.py (95%) rename {hyperlora => src/ctx_to_lora}/eval.py (98%) rename {hyperlora => src/ctx_to_lora}/hooks.py (99%) rename {hyperlora => src/ctx_to_lora}/model_loading.py (100%) rename {hyperlora => src/ctx_to_lora}/modeling_utils.py (97%) rename {hyperlora => src/ctx_to_lora}/pooling.py (100%) rename {hyperlora => src/ctx_to_lora}/training_utils.py (100%) rename {hyperlora => src/ctx_to_lora}/utils.py (100%) rename hyperlora/watcher.py => watcher.py (99%) diff --git a/configs/fw_qa.yaml b/configs/fw_qa.yaml index 6ad3122..cc7dcb8 100644 --- a/configs/fw_qa.yaml +++ b/configs/fw_qa.yaml @@ -27,8 +27,8 @@ neftune_noise_alpha: 5 weight_decay: 0.01 warmup_ratio: 0.1 -dataloader_prefetch_factor: 16 -dataloader_num_workers: 16 +dataloader_prefetch_factor: 8 +dataloader_num_workers: 8 # LoRA lora_r: 8 lora_dropout: 0.05 diff --git a/configs/fw_qa_ctx_qa.yaml b/configs/fw_qa_ctx_qa.yaml index 9575033..0263f92 100644 --- a/configs/fw_qa_ctx_qa.yaml +++ b/configs/fw_qa_ctx_qa.yaml @@ -27,8 +27,8 @@ neftune_noise_alpha: 5 weight_decay: 0.01 warmup_ratio: 0.1 -dataloader_prefetch_factor: 16 -dataloader_num_workers: 16 +dataloader_prefetch_factor: 8 +dataloader_num_workers: 8 # LoRA lora_r: 8 lora_dropout: 0.05 diff --git a/configs/fw_qa_large_ctx_pwc_hotpot.yaml b/configs/fw_qa_large_ctx_pwc_hotpot.yaml new file mode 100644 index 0000000..0efd743 --- /dev/null +++ b/configs/fw_qa_large_ctx_pwc_hotpot.yaml @@ -0,0 +1,55 @@ +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: 8 +per_device_eval_batch_size: 8 +max_val_samples_per_ds: 1000 +# optim: schedule_free_adamw + +learning_rate: 0.00002 +# lr_scheduler_type: "constant_with_warmup" +neftune_noise_alpha: 5 +weight_decay: 0.01 +warmup_ratio: 0.1 + +dataloader_prefetch_factor: 8 +dataloader_num_workers: 8 +# LoRA +lora_r: 8 +lora_dropout: 0.05 +target_modules: + - down_proj + - up_proj + +# data +train_ds_names: +- fw_qa +- fw_qa_large +- ctx_qa +- pwc +- hotpot_qa + +val_ds_names: +- fw_qa +- fw_qa_large +- ctx_qa +- pwc +- hotpot_qa + +load_best_model_at_end: true +metric_for_best_model: eval_pwc_loss diff --git a/configs/fw_qa_pwc_hotpot.yaml b/configs/fw_qa_pwc_hotpot.yaml index f29b20b..1188076 100644 --- a/configs/fw_qa_pwc_hotpot.yaml +++ b/configs/fw_qa_pwc_hotpot.yaml @@ -27,8 +27,8 @@ neftune_noise_alpha: 5 weight_decay: 0.01 warmup_ratio: 0.1 -dataloader_prefetch_factor: 16 -dataloader_num_workers: 16 +dataloader_prefetch_factor: 8 +dataloader_num_workers: 8 # LoRA lora_r: 8 lora_dropout: 0.05 diff --git a/configs/fw_qa_tiny.yaml b/configs/fw_qa_tiny.yaml index 2c314f4..50598c0 100644 --- a/configs/fw_qa_tiny.yaml +++ b/configs/fw_qa_tiny.yaml @@ -27,8 +27,8 @@ neftune_noise_alpha: 5 weight_decay: 0.01 warmup_ratio: 0.1 -dataloader_prefetch_factor: 16 -dataloader_num_workers: 16 +dataloader_prefetch_factor: 8 +dataloader_num_workers: 8 # LoRA lora_r: 8 lora_dropout: 0.05 diff --git a/configs/hotpot_qa.yaml b/configs/hotpot_qa.yaml index a0c1416..1ad924d 100644 --- a/configs/hotpot_qa.yaml +++ b/configs/hotpot_qa.yaml @@ -16,8 +16,8 @@ label_names: ["labels"] # w/o this the trainer stores logits of all sample in memory... # batch_eval_metrics: true -per_device_train_batch_size: 16 -per_device_eval_batch_size: 16 +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 max_val_samples_per_ds: 1000 # optim: schedule_free_adamw learning_rate: 0.00001 diff --git a/configs/pwc_hotpot_qa.yaml b/configs/pwc_hotpot_qa.yaml index 3da3f64..e59e494 100644 --- a/configs/pwc_hotpot_qa.yaml +++ b/configs/pwc_hotpot_qa.yaml @@ -16,8 +16,8 @@ label_names: ["labels"] # w/o this the trainer stores logits of all sample in memory... # batch_eval_metrics: true -per_device_train_batch_size: 16 -per_device_eval_batch_size: 16 +per_device_train_batch_size: 8 +per_device_eval_batch_size: 8 max_val_samples_per_ds: 1000 # optim: schedule_free_adamw diff --git a/install.sh b/install.sh new file mode 100644 index 0000000..41e3228 --- /dev/null +++ b/install.sh @@ -0,0 +1 @@ +pip install -e . \ No newline at end of file diff --git a/hyperlora/intx_sft.py b/intx_sft.py similarity index 97% rename from hyperlora/intx_sft.py rename to intx_sft.py index 624f5fd..8ef5ef9 100755 --- a/hyperlora/intx_sft.py +++ b/intx_sft.py @@ -14,7 +14,7 @@ import torch import wandb import yaml -from data_utils import ( +from ctx_to_lora.data_utils import ( convert_ctx_prompt_response_to_messages, get_preprocessing_fn, get_sft_prompt_formatting_fn, @@ -29,15 +29,19 @@ from datasets import ( load_dataset, IterableDataset, ) -from model_loading import get_lora_config, get_model_and_tokenizer, get_tokenizer -from modeling_utils import ( +from ctx_to_lora.model_loading import ( + get_lora_config, + get_model_and_tokenizer, + get_tokenizer, +) +from ctx_to_lora.modeling_utils import ( EarlyExit, HyperLoRA, ModulatedPretrainedModel, get_hypernet_config, ) from rouge_score import rouge_scorer -from training_utils import TRAINING_TASK, train_model +from ctx_to_lora.training_utils import TRAINING_TASK, train_model from transformers import ( AutoModelForCausalLM, AutoTokenizer, @@ -50,7 +54,7 @@ from transformers import ( from peft import PeftModel from transformers.utils import is_liger_kernel_available from torch.utils.data import DataLoader -from utils import ( +from ctx_to_lora.utils import ( extract_cli_args, get_base_model, get_run_name, @@ -60,7 +64,7 @@ from utils import ( validate_args, ) -from configs import ( +from ctx_to_lora.configs import ( ArgumentParser, CtxTrainingArguments, TrainingArguments, @@ -483,8 +487,10 @@ def main(): # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer # TODO: use packing with SFTTrainer - # HACK: see transformers/trainer_seq2seq.py for supressing + # HACK [local patch]: see transformers/trainer_seq2seq.py for supressing # "Trainer.tokenizer is now deprecated. You should use Trainer.processing_class instead." + # HACK [local patch]: deepspeed model loading problem (for resume training) + # see https://github.com/microsoft/DeepSpeed/pull/6626/files if training_args.use_liger_kernel and is_liger_kernel_available(): from liger_kernel.transformers import _apply_liger_kernel_to_instance diff --git a/requirements.txt b/requirements.txt index b66eb78..4e0ffcd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,9 @@ einops jaxtyping -schedulefree liger-kernel - +transformers==4.46.3 +deepspeed==0.15.4 # curretly 0.16.2 +accelerate==1.2.1 # visualization tensorboard flask diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..f6c53cc --- /dev/null +++ b/setup.py @@ -0,0 +1,12 @@ +# read the contents of the README file +from pathlib import Path + +from setuptools import find_packages, setup + +this_directory = Path(__file__).parent +long_description = (this_directory / "README.md").read_text() + +setup( + long_description=long_description, + long_description_content_type="text/markdown", +) diff --git a/src/ctx_to_lora/__init__.py b/src/ctx_to_lora/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/hyperlora/configs.py b/src/ctx_to_lora/configs.py similarity index 100% rename from hyperlora/configs.py rename to src/ctx_to_lora/configs.py diff --git a/hyperlora/data_utils.py b/src/ctx_to_lora/data_utils.py similarity index 95% rename from hyperlora/data_utils.py rename to src/ctx_to_lora/data_utils.py index 0ed5bad..731510c 100644 --- a/hyperlora/data_utils.py +++ b/src/ctx_to_lora/data_utils.py @@ -4,9 +4,10 @@ from glob import glob from typing import Any, Callable, Iterator, Optional from datasets import load_dataset, IterableDataset -from training_utils import TRAINING_TASK from transformers import PreTrainedTokenizerBase +from ctx_to_lora.training_utils import TRAINING_TASK + IGNORE_INDEX = -100 logger = logging.getLogger() @@ -32,6 +33,10 @@ DS_KWARGS = { train=dict(path="sggetao/PwC", split="train[900:2000]"), validation=dict(path="sggetao/PwC", split="train[:900]"), ), + "squad": dict( + train=dict(path="rajpurkar/squad", split="train[900:]"), + validation=dict(path="rajpurkar/squad", split="train[:900]"), + ), "fw_qa_tiny": dict( train=dict( path="parquet", @@ -56,6 +61,18 @@ DS_KWARGS = { split="train", ), ), + "fw_qa_large": dict( + train=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_large/*res.parquet"), + split="train", + ), + validation=dict( + path="parquet", + data_files=glob("data/raw_datasets/fw_qa_large/*val.parquet"), + split="train", + ), + ), "ctx_qa": dict( train=dict( path="parquet", @@ -98,6 +115,10 @@ def filter_long_samples(samples): return [len(ctx) < 10000 for ctx in samples["context"]] +def filter_long_chat(samples): + return [len(chat) < 2**13 for chat in samples["chat"]] + + def add_repeat_prompt_fn(samples): unique_contexts = set() ctxs, prompts, responses = [], [], [] @@ -144,6 +165,8 @@ def add_negative_prompt_fn(samples): 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 @@ -258,6 +281,7 @@ def construct_and_tokenize_ctx_qa( get_sft_prompt_formatting_fn(TRAINING_TASK.COMPLETION, tokenizer), num_proc=16 ) # tokenize the chat + mask the assistant inputs + ds = ds.filter(filter_long_chat, batched=True, num_proc=16) # add "input_ids", "attention_mask", "labels" tokenized_ds = ds.map( @@ -440,6 +464,15 @@ def get_preprocessing_fn(ds_name: str) -> Callable[[dict[str, Any]], dict[str, A "response": sample["answer"], } + elif "squad" in ds_name: + + def f(sample): + return { + "context": sample["context"], + "prompt": sample["question"], + "response": sample["answer"]["text"][0], + } + return f diff --git a/hyperlora/eval.py b/src/ctx_to_lora/eval.py similarity index 98% rename from hyperlora/eval.py rename to src/ctx_to_lora/eval.py index 0c380e1..7da43c9 100644 --- a/hyperlora/eval.py +++ b/src/ctx_to_lora/eval.py @@ -22,9 +22,9 @@ from transformers import ( set_seed, ) -from data_utils import get_tokenized_dataset -from modeling_utils import ModulatedPretrainedModel -from model_loading import get_tokenizer +from ctx_to_lora.data_utils import get_tokenized_dataset +from ctx_to_lora.modeling_utils import ModulatedPretrainedModel +from ctx_to_lora.model_loading import get_tokenizer def clear_gpu(): @@ -386,7 +386,7 @@ def evaluate(checkpoint_path, args, split, generative): eval_trainer_args["eval_strategy"] = "no" eval_trainer_args["save_strategy"] = "no" eval_trainer_args["overwrite_output_dir"] = True - eval_trainer_args["per_device_eval_batch_size"] = 64 + eval_trainer_args["per_device_eval_batch_size"] = 16 eval_trainer_args = Seq2SeqTrainingArguments( **eval_trainer_args, diff --git a/hyperlora/hooks.py b/src/ctx_to_lora/hooks.py similarity index 99% rename from hyperlora/hooks.py rename to src/ctx_to_lora/hooks.py index 8945811..5ee312b 100644 --- a/hyperlora/hooks.py +++ b/src/ctx_to_lora/hooks.py @@ -8,7 +8,7 @@ from einops import einsum from jaxtyping import Float from torch import Tensor from torch.utils.hooks import RemovableHandle -from utils import get_layers +from ctx_to_lora.utils import get_layers logger = logging.getLogger() diff --git a/hyperlora/model_loading.py b/src/ctx_to_lora/model_loading.py similarity index 100% rename from hyperlora/model_loading.py rename to src/ctx_to_lora/model_loading.py diff --git a/hyperlora/modeling_utils.py b/src/ctx_to_lora/modeling_utils.py similarity index 97% rename from hyperlora/modeling_utils.py rename to src/ctx_to_lora/modeling_utils.py index 27345c8..0f60b8d 100644 --- a/hyperlora/modeling_utils.py +++ b/src/ctx_to_lora/modeling_utils.py @@ -8,14 +8,10 @@ from typing import Any, Iterable, Optional, Tuple, Union import torch import torch.nn.functional as F - -from configs import AggregatorArguments, HypernetArguments, CtxEncoderArguments from einops import rearrange, repeat, unpack, einsum from einops.layers.torch import EinMix as Mix from einops.layers.torch import Reduce -from hooks import add_generated_lora_hook, remove_hook_handles from jaxtyping import Float, Integer -from model_loading import get_lora_config, get_model, get_model_and_tokenizer from peft import ( get_peft_config, load_peft_weights, @@ -29,7 +25,6 @@ from peft import ( from peft.utils import PeftType, TaskType, ModulesToSaveWrapper from peft.tuners._buffer_dict import BufferDict from peft.tuners.tuners_utils import BaseTunerLayer, check_target_module_exists -from pooling import POOL_FN, get_pooling_fn from torch import Tensor, nn from transformers import ( PerceiverConfig, @@ -41,7 +36,12 @@ from transformers.models.perceiver.modeling_perceiver import ( PerceiverBasicDecoder, ) from transformers.modeling_outputs import ModelOutput -from utils import ( + +from ctx_to_lora.configs import AggregatorArguments, HypernetArguments, CtxEncoderArguments +from ctx_to_lora.hooks import add_generated_lora_hook, remove_hook_handles +from ctx_to_lora.model_loading import get_lora_config, get_model, get_model_and_tokenizer +from ctx_to_lora.pooling import POOL_FN, get_pooling_fn +from ctx_to_lora.utils import ( get_lora_module_names, get_num_layers, get_peft_in_out_features, @@ -100,7 +100,7 @@ class HypernetConfig: lora_config: LoraConfig module_names: dict[str, list[str]] - trainable_base_modules: Optional[list[str]] + # trainable_base_modules: Optional[list[str]] layer_indices: Iterable[int] feature_sizes: tuple[dict[str, int], dict[str, int]] aggregator_config: AggregatorConfig @@ -761,6 +761,17 @@ class ModulatedPretrainedModel(nn.Module): # out["ctx_features"] = features # return out + def generate_loras( + self, + ctx_ids: Integer[Tensor, "bs ctx_len"], + ctx_attn_mask: Integer[Tensor, "bs ctx_len"], + ): + with torch.no_grad(): + ctx_features = self.ctx_encoder( + input_ids=ctx_ids, attention_mask=ctx_attn_mask + ) + return self.hypernet.generate_loras(ctx_features, ctx_attn_mask) + def forward( self, # ctx_features: Optional[Float[Tensor, "bs ctx_length feature_dim"]] = None, diff --git a/hyperlora/pooling.py b/src/ctx_to_lora/pooling.py similarity index 100% rename from hyperlora/pooling.py rename to src/ctx_to_lora/pooling.py diff --git a/hyperlora/training_utils.py b/src/ctx_to_lora/training_utils.py similarity index 100% rename from hyperlora/training_utils.py rename to src/ctx_to_lora/training_utils.py diff --git a/hyperlora/utils.py b/src/ctx_to_lora/utils.py similarity index 100% rename from hyperlora/utils.py rename to src/ctx_to_lora/utils.py diff --git a/hyperlora/watcher.py b/watcher.py similarity index 99% rename from hyperlora/watcher.py rename to watcher.py index 0f362b5..179a963 100644 --- a/hyperlora/watcher.py +++ b/watcher.py @@ -87,7 +87,7 @@ if __name__ == "__main__": new_files = watcher.watch() for file in new_files: # workaround to prevent loading incomplete checkpoints - time.sleep(10) + time.sleep(20) if not os.path.exists(file): # cp is delete before we can read it continue