adding configs + loading + peft training

hyperlora/configs.py

@@ -0,0 +1,54 @@
+import dataclasses
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Literal, NewType, Optional, Tuple
+from enum import Enum, auto
+
+from transformers import MODEL_FOR_CAUSAL_LM_MAPPING, HfArgumentParser
+
+
+class ExperimentSetup(Enum):
+    LORA = "lora"
+    HYPER_LORA = "hyper_lora"
+    FULL_FINETUNE = "full_finetune"
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments for the base model.
+    """
+
+    model_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={"help": ("Base model name or path.")},
+    )
+    # use_peft: bool = field(
+    #     default=False,
+    #     metadata={"help": ("Whether to use PEFT or not for training.")},
+    # )
+
+
+@dataclass
+class LoRAArguments:
+    r: Optional[int] = field(
+        default=8,
+        metadata={"help": ("LoRA R value.")},
+    )
+    lora_dropout: Optional[float] = field(
+        default=0.05,
+        metadata={"help": ("LoRA dropout.")},
+    )
+    target_modules: Optional[List[str]] = field(
+        default=None,
+        metadata={"help": ("LoRA target modules.")},
+    )
+
+
+@dataclass
+class CtxTrainingArguments:
+    exp_setup: ExperimentSetup = field(
+        default=ExperimentSetup.LORA,
+        metadata={"help": "Experiment setup - LoRA, HyperLoRA, or full finetuning"},
+    )

hyperlora/intx_sft.py

@@ -1,3 +1,4 @@
+import logging
 import numpy as np
 import torch
 
@@ -12,6 +13,9 @@ from transformers import (
 )
 
 
+from configs import CtxTrainingArguments, LoRAArguments, ModelArguments, ExperimentSetup
+from utils import log_num_train_params
+from model_loading import get_model_and_tokenizer, get_lora_config
 from modeling_utils import ModulatedPretrainedModel
 from data_utils import (
     convert_ctx_prompt_response_to_messages,
@@ -21,6 +25,8 @@ from data_utils import (
 )
 from training_utils import TRAINING_TASK, train_model
 
+logger = logging.getLogger(__name__)
+
 
 def compute_metrics(eval_pred: EvalPrediction) -> dict:
     """
@@ -41,9 +47,15 @@ def compute_metrics(eval_pred: EvalPrediction) -> dict:
 
 
 def main():
-    parser = HfArgumentParser((TrainingArguments,))
-    training_args, *_ = parser.parse_args_into_dataclasses()
+    # Set logging verbosity to INFO
+    logging.basicConfig(level=logging.INFO)
 
+    parser = HfArgumentParser(
+        (CtxTrainingArguments, ModelArguments, LoRAArguments, TrainingArguments)
+    )
+    ctx_args, model_args, lora_args, training_args = parser.parse_args_into_dataclasses()
+
+    training_args.label_names = ["labels"]
     training_args.eval_on_start = True
     training_args.eval_strategy = "steps"
     training_args.eval_steps = 500
@@ -60,61 +72,35 @@ def main():
         "use_reentrant": False
     }  # manually add this argument in the code
 
-    model = AutoModelForCausalLM.from_pretrained(
-        # "meta-llama/Llama-3.1-8B-Instruct",
-        "meta-llama/Llama-3.2-1B-Instruct",
-        torch_dtype=torch.bfloat16,
-        attn_implementation="flash_attention_2",
+    # "meta-llama/Llama-3.1-8B-Instruct",
+
+    # model = AutoModelForCausalLM.from_pretrained(
+    #     base_model_name,
+    #     torch_dtype=torch.bfloat16,
+    #     attn_implementation="flash_attention_2",
+    # )
+    # tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+    # tokenizer.pad_token_id = tokenizer.eos_token_id
+    # tokenizer.padding_side = "right"
+
+    model_name = model_args.model_name_or_path
+
+    model, tokenizer = get_model_and_tokenizer(
+        **vars(model_args),
+        train=True,
+        requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
+        peft_config=get_lora_config(model_name, **vars(lora_args)),
     )
-    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-    tokenizer.padding_side = "right"
+    if ctx_args.exp_setup == ExperimentSetup.HYPER_LORA:
+        hypernet = ...
+        model = ModulatedPretrainedModel(model, hypernet)
+    else:
+        # activate LoRA
+        model.set_adapter("default")
 
-    max_seq_len = 1024
+    log_num_train_params(model)
 
-    # if isinstance(model, ModulatedPretrainedModel):
-
-    #     # TODO: how to tokenize properly?
-    #     def tokenize(example):
-    #         model_inputs = tokenizer(
-    #             example["prompt"],
-    #             truncation=True,
-    #             padding=False,
-    #         )
-    #         model_inputs["ctx_ids"] = tokenizer(example["context"]).input_ids
-    #         model_inputs["ctx_attention_mask"] = tokenizer(
-    #             example["context"]
-    #         ).attention_mask
-    #         model_inputs["labels"] = ...
-    #         return model_inputs
-
-    # else:
-
-    #     def tokenize(example):
-    #         inp = [
-    #             ctx + "\n" + prompt
-    #             for ctx, prompt in zip(example["context"], example["prompt"])
-    #         ]
-    #         model_inputs = tokenizer(
-    #             inp,
-    #             example["answer"],
-    #             add_special_tokens=False,
-    #             truncation=True,
-    #             padding=False,
-    #             max_length=max_seq_len,
-    #         )
-
-    #         input_ids = model_inputs["input_ids"]
-    #         labels = [None] * len(input_ids)
-
-    #         for i in range(len(input_ids)):
-    #             sequence_ids = model_inputs.sequence_ids(i)
-    #             labels[i] = [
-    #                 -100 if sequence_id == 0 else label
-    #                 for sequence_id, label in zip(sequence_ids, input_ids[i])
-    #             ]
-    #         model_inputs["labels"] = labels
-    #         return model_inputs
+    # max_seq_len = 1024
 
     print("Loading dataset...")
     train_file = "../data/raw_datasets/context_numbers/train.jsonl"
@@ -138,7 +124,7 @@ def main():
             "tokenizer": tokenizer,
             "mask_assistant_inputs": True,
             "tokenizer_kwargs": {
-                "max_length": max_seq_len,
+                "max_length": None,
             },
         },
         remove_columns=ds["train"].column_names,
@@ -158,8 +144,12 @@ def main():
     # DataCollatorForSeq2Seq also pads the `labels`
     # useful when we're computing the labels manually
     # or masking the loss only on completion
+    # TODO: change to a faster collator? e.g.,
+    # https://huggingface.co/blog/packing-with-FA2
     data_collator = DataCollatorForSeq2Seq(tokenizer, model, pad_to_multiple_of=8)
 
+    # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer
+    # TODO: use packing with SFTTrainer
     train_model(
         model,
         train_ds,

hyperlora/model_loading.py

@@ -0,0 +1,168 @@
+import logging
+import os
+
+import torch
+from peft import PeftModel, LoraConfig, VeraConfig, PeftConfig
+from peft import get_peft_config as _get_peft_config
+from peft.utils import PeftType
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel
+
+from hyper_llm_modulator.utils.pooling import get_pooling_fn
+from hyper_llm_modulator.utils.preprocessing import add_full_stop, apply_sfr_template
+
+logger = logging.getLogger()
+
+
+def get_model_and_tokenizer(
+    model_name_or_path,
+    train,
+    requires_grad,
+    use_flash_attn=True,
+    peft_config=None,
+    model_kwargs=None,
+    tokenizer_kwargs=None,
+    device="cuda:0",
+    dtype=torch.bfloat16,
+):
+    model = get_model(
+        model_name_or_path,
+        train,
+        requires_grad,
+        use_flash_attn,
+        peft_config,
+        model_kwargs,
+        device,
+        dtype,
+    )
+    tokenizer = get_tokenizer(model_name_or_path, tokenizer_kwargs, peft_config, train)
+    return model, tokenizer
+
+
+def get_tokenizer(
+    model_name_or_path, tokenizer_kwargs=None, peft_config=None, train=False
+):
+    # tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, padding_side="left")
+    # NOTE: lora models don't have tokenizer config in the folder
+
+    # left pad for generation, right pad for training (why?)
+    padding_side = "left" if not train else "right"
+    truncation_side = "left"
+
+    # tokenizer = AutoTokenizer.from_pretrained(
+    #         "models/Mistral-7B-v0.1/", padding_side=padding_side
+    #     )
+    if peft_config:
+        model_name_or_path = peft_config.base_model_name_or_path
+
+    if tokenizer_kwargs is None:
+        tokenizer_kwargs = {}
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name_or_path,
+        padding_side=padding_side,
+        truncation_side=truncation_side,
+        **tokenizer_kwargs,
+    )
+
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        # tokenizer.pad_token = tokenizer.eos_token
+
+    # tokenizer.pad_token = tokenizer.eos_token
+    # tokenizer.add_bos_token = True
+
+    # assert os.path.exists(f"{model_name_or_path}/chat_template.jinja"), (
+    #     f"Chat template not found in {model_name_or_path}\n\n"
+    #     "We assume a specfic form of chat template for consistency between models. Please use the templates provided",
+    # )
+    # if os.path.exists(f"{model_name_or_path}/chat_template.jinja"):
+    #     chat_template = open(f"{model_name_or_path}/chat_template.jinja").read()
+    #     chat_template = chat_template.replace("    ", "").replace("\n", "")
+    #     tokenizer.chat_template = chat_template
+
+    # tokenizer.add_eos_token = False
+    # if train:
+    #     # NOTE: this correctly add an eos token that has attention = 1
+    #     # while padding eos tokens have attention = 0
+    #     tokenizer.add_eos_token = True
+
+    #     # shouldn't be needed as we manually compute the labels now
+    #     # NOTE: this seems oddly important
+    #     # and is not necessarily in the alignment handbook scripts
+    #     # tokenizer.pad_token = tokenizer.unk_token  # fix model not generating eos
+
+    return tokenizer
+
+
+def get_model(
+    model_name_or_path,
+    train,
+    requires_grad,
+    use_flash_attn=True,
+    peft_config=None,
+    model_kwargs=None,
+    device="cuda:0",
+    dtype=torch.bfloat16,
+):
+    # TODO: we could use lower precision during finetuning
+    model_init_kwargs = dict(
+        pretrained_model_name_or_path=model_name_or_path,
+        device_map=device,
+        torch_dtype=dtype,
+        trust_remote_code=True,
+        # attn_implementation="flash_attention_2",
+        # load_in_4bit=True,
+        # load_in_8bit=True,
+    )
+    if model_kwargs is not None:
+        model_init_kwargs.update(model_kwargs)
+    if use_flash_attn:
+        model_init_kwargs["attn_implementation"] = "flash_attention_2"
+        # for training disable cache
+    if train:
+        model_init_kwargs["use_cache"] = False
+    logger.debug(f"Model init kwargs: {model_init_kwargs}")
+    model = AutoModelForCausalLM.from_pretrained(**model_init_kwargs)
+    if peft_config is not None:
+        model = PeftModel(model, peft_config)
+    model.train(train)
+    for param in model.parameters():
+        param.requires_grad = requires_grad
+    return model
+
+
+def get_lora_config(model_dir, **kwargs):
+    r = kwargs.get("r", 8)
+    peft_conf_kwargs = dict(
+        r=r,
+        peft_type=PeftType.LORA,
+        base_model_name_or_path=model_dir,
+        task_type="CAUSAL_LM",
+        lora_dropout=0.05,
+        lora_alpha=r ** (3 / 2) * 2,
+    )
+
+    peft_conf_kwargs.update(kwargs)
+    peft_config = _get_peft_config(peft_conf_kwargs)
+    return peft_config
+
+
+# def get_emb_model_and_fns(emb_model_name, device):
+#     emb_model = AutoModel.from_pretrained(
+#         emb_model_name,
+#         device_map=device,
+#         torch_dtype=torch.float32 if "gte" in emb_model_name else torch.bfloat16,
+#         trust_remote_code=True,
+#     ).eval()
+#     emb_tokenizer = AutoTokenizer.from_pretrained(emb_model_name)
+#     if emb_tokenizer.pad_token_id is None:
+#         emb_tokenizer.pad_token_id = emb_tokenizer.eos_token_id
+#         # assert bool(args.query_intx), "query_intx must be provided for the emb_model_name"
+#         # partial(get_detailed_instruct, task_description=args.query_intx)
+#     task_desc_format_fn = add_full_stop
+#     if "SFR" in emb_model_name:
+#         task_desc_format_fn = apply_sfr_template
+#         pooling_fn = get_pooling_fn("last_token")
+#     elif "gte" in emb_model_name:
+#         pooling_fn = get_pooling_fn("cls")
+#     return emb_model, emb_tokenizer, task_desc_format_fn, pooling_fn

hyperlora/utils.py

@@ -0,0 +1,28 @@
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def get_num_params(model):
+    total_params = 0
+    trainable_params = 0
+    for p in model.parameters():
+        total_params += p.numel()
+        if p.requires_grad:
+            trainable_params += p.numel()
+
+    return total_params, trainable_params
+
+
+def log_num_train_params(model):
+    logger.debug("Trainable model parameters:")
+    for name, p in model.named_parameters():
+        if p.requires_grad:
+            logger.debug(f"{name}, dtype:{p.dtype}")
+
+    num_total_params, num_trainable_params = get_num_params(model)
+    logger.info(
+        f"trainable params: {num_trainable_params:,d} "
+        f"|| all params: {num_total_params:,d} "
+        f"|| trainable%: {100 * num_trainable_params / num_total_params:.4f}"
+    )