restructure

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

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

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

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

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

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

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
 

install.sh

@@ -0,0 +1 @@
+pip install -e . 

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
 

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

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",
+)

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
 
 

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,

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

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,

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