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

deprecate trainable_base_modules + fix vision model load + fix metric

Browse source
This commit is contained in:
51616 2025-01-16 11:46:06 +00:00
parent 9dcc6bdb97
commit 9035b1834c
6 changed files with 139 additions and 51 deletions
Download patch file Download diff file Expand all files Collapse all files

4
hyperlora/configs.py
View file

@ -332,6 +332,10 @@ class HypernetArguments:
default=128,
metadata={"help": "Latent size for light-weight LoRA."},
)
# trainable_base_modules: Optional[list[str]] = field(
# default=None,
# metadata={"help": ("Modules to train of the base model.")},
# )
@dataclass

49
hyperlora/eval.py
View file

@ -3,6 +3,7 @@ import sys
import gc
import json
import os
import random
from argparse import Namespace
from dataclasses import fields
from functools import partial
@ -18,6 +19,7 @@ from transformers import (
Seq2SeqTrainer,
Seq2SeqTrainingArguments,
EvalPrediction,
set_seed,
)
from data_utils import get_tokenized_dataset
@ -46,8 +48,13 @@ def compute_rouge(pred_texts, label_texts):
def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
indices = torch.where(valid_masks)
acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
return {"per_token_acc": acc}
# acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
# return {"per_token_acc": acc}
acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
return {
"per_token_accs": acc.flatten().tolist(),
"n_per_token_accs": valid_masks.sum().item(),
}
def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
@ -63,7 +70,11 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
# if all tokens are correct, set to 1
perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf)
return {"prefix_matching": perf.mean().item()}
# return {"prefix_matching": perf.mean().item()}
return {
"prefix_matchings": perf.tolist(),
"n_prefix_matchings": valid_masks.shape[0],
}
@torch.no_grad()
@ -71,8 +82,13 @@ def compute_perplexity(shift_logits, shift_labels, valid_masks):
loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1)
perplexity = torch.exp(loss).mean().item()
return {"perplexity": perplexity}
# perplexity = torch.exp(loss).mean().item()
# return {"perplexity": perplexity}
preplexities = torch.exp(loss)
return {
"perplexities": preplexities.tolist(),
"n_perplexities": valid_masks.shape[0],
}
class Evaluator:
@ -82,16 +98,22 @@ class Evaluator:
def reset(self):
self.accum_metrics = defaultdict(list)
self.count = defaultdict(list)
def update(self, shift_logits, shift_labels, valid_masks):
for metric_fn in self.metric_fns:
metric = metric_fn(shift_logits, shift_labels, valid_masks)
for k, v in metric.items():
self.accum_metrics[k].append(v)
if k.startswith("n_"):
self.count[k[2:]].append(v)
else:
self.accum_metrics[k] += v
def compute(self):
# Get result across entire eval set
result = {k: np.mean(v) for k, v in self.accum_metrics.items()}
result = {
k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items()
}
# Reset batch statistics
self.reset()
return result
@ -304,9 +326,18 @@ def generation_collator(inp_list, tokenizer):
def evaluate(checkpoint_path, args, split, generative):
assert split in ["validation", "test"]
set_seed(42)
torch.use_deterministic_algorithms(True, warn_only=True)
torch.backends.cudnn.benchmark = False
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
model = ModulatedPretrainedModel.from_state_dict(
torch.load(open(checkpoint_path, "rb")), train=False
torch.load(open(checkpoint_path, "rb")),
train=False,
use_flash_attn=False,
)
# NOTE: there is still some randomness in the eval result
# despite all the deterministic settings
tokenizer = get_tokenizer(args.model_name_or_path, train=False)
if tokenizer.pad_token_id is None:
@ -408,7 +439,9 @@ def evaluate(checkpoint_path, args, split, generative):
if __name__ == "__main__":
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
os.environ["WANDB_MODE"] = "disabled"
checkpoint_path = sys.argv[1]
checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])

48
hyperlora/intx_sft.py
View file

@ -80,8 +80,13 @@ LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
indices = torch.where(valid_masks)
acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
return {"per_token_acc": acc}
# acc = (shift_logits.argmax(-1) == shift_labels)[indices].float().mean().item()
# return {"per_token_acc": acc}
acc = (shift_logits.argmax(-1) == shift_labels)[indices].float()
return {
"per_token_accs": acc.flatten().tolist(),
"n_per_token_accs": valid_masks.sum().item(),
}
def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
@ -97,7 +102,11 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
# if all tokens are correct, set to 1
perf = torch.where(is_correct.sum(dim=1) == lengths, 1, perf)
return {"prefix_matching": perf.mean().item()}
# return {"prefix_matching": perf.mean().item()}
return {
"prefix_matchings": perf.tolist(),
"n_prefix_matchings": valid_masks.shape[0],
}
@torch.no_grad()
@ -105,8 +114,13 @@ def compute_perplexity(shift_logits, shift_labels, valid_masks):
loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
loss = loss_fct(shift_logits.transpose(1, 2), shift_labels)
loss = (loss * valid_masks).sum(dim=1) / valid_masks.sum(dim=1)
perplexity = torch.exp(loss).mean().item()
return {"perplexity": perplexity}
# perplexity = torch.exp(loss).mean().item()
# return {"perplexity": perplexity}
preplexities = torch.exp(loss)
return {
"perplexities": preplexities.tolist(),
"n_perplexities": valid_masks.shape[0],
}
class Evaluator:
@ -116,16 +130,22 @@ class Evaluator:
def reset(self):
self.accum_metrics = defaultdict(list)
self.count = defaultdict(list)
def update(self, shift_logits, shift_labels, valid_masks):
for metric_fn in self.metric_fns:
metric = metric_fn(shift_logits, shift_labels, valid_masks)
for k, v in metric.items():
self.accum_metrics[k].append(v)
if k.startswith("n_"):
self.count[k[2:]].append(v)
else:
self.accum_metrics[k] += v
def compute(self):
# Get result across entire eval set
result = {k: np.mean(v) for k, v in self.accum_metrics.items()}
result = {
k: np.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items()
}
# Reset batch statistics
self.reset()
return result
@ -279,11 +299,12 @@ def main():
requires_grad=ctx_args.exp_setup == ExperimentSetup.FULL_FINETUNE,
peft_config=get_lora_config(model_name, **vars(lora_args)),
)
if ctx_encoder_args.ctx_encoder_model_name_or_path is not None:
ctx_encoder_model_config = AutoConfig.from_pretrained(
ctx_encoder_args.ctx_encoder_model_name_or_path
)
ctx_tokenizer = get_tokenizer(ctx_encoder_args.ctx_encoder_model_name_or_path)
ctx_name = ctx_encoder_args.ctx_encoder_model_name_or_path
if ctx_name is not None:
ctx_encoder_model_config = AutoConfig.from_pretrained(ctx_name)
if "Llama" in ctx_name and "Vision" in ctx_name:
ctx_encoder_model_config = ctx_encoder_model_config.text_config
ctx_tokenizer = get_tokenizer(ctx_name)
else:
ctx_encoder_model_config = model.config
ctx_tokenizer = tokenizer
@ -371,6 +392,7 @@ def main():
train_ds = interleave_datasets(
list(train_ds.values()),
probabilities=[l / total_len for l in train_ds_len],
seed=training_args.seed,
)
# val_train_indices = np.random.permutation(len(train_ds))[:500]
@ -469,7 +491,7 @@ def main():
if isinstance(model, ModulatedPretrainedModel):
logger.info("Applying liger-kernel to ModulatedPretrainedModel")
_apply_liger_kernel_to_instance(model=model.base_model.base_model.model)
if ctx_encoder_args.ctx_encoder_model_name_or_path is not None:
if ctx_name is not None:
logger.info("Applying liger-kernel to ctx_encoder_model")
_apply_liger_kernel_to_instance(model=model.ctx_encoder)
elif isinstance(model, PeftModel):

21
hyperlora/model_loading.py
View file

@ -143,16 +143,17 @@ def get_model(
model = PeftModel(model, peft_config)
model.train(train)
for name, param in model.named_parameters():
if "modules_to_save" not in name:
param.requires_grad = requires_grad
else:
logging.debug(f"modules_to_save: {name} to be trained")
# always train "modules_to_save"
if not "layernorm" in name:
raise NotImplementedError(
f"modules_to_save should only be layernorm, got {name}"
)
param.requires_grad = True
# if "modules_to_save" not in name:
# param.requires_grad = requires_grad
# else:
# logging.debug(f"modules_to_save: {name} to be trained")
# # always train "modules_to_save"
# if not "layernorm" in name:
# raise NotImplementedError(
# f"modules_to_save should only be layernorm, got {name}"
# )
# param.requires_grad = True
param.requires_grad = requires_grad
return model

66
hyperlora/modeling_utils.py
View file

@ -26,7 +26,7 @@ from peft import (
get_peft_model_state_dict,
set_peft_model_state_dict,
)
from peft.utils import PeftType, TaskType
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
@ -100,6 +100,7 @@ class HypernetConfig:
lora_config: LoraConfig
module_names: dict[str, list[str]]
trainable_base_modules: Optional[list[str]]
layer_indices: Iterable[int]
feature_sizes: tuple[dict[str, int], dict[str, int]]
aggregator_config: AggregatorConfig
@ -574,7 +575,12 @@ class ModulatedPretrainedModel(nn.Module):
# self.register_module("ctx_encoder", ctx_encoder)
@classmethod
def from_state_dict(cls, state_dict: dict, train: bool = True):
def from_state_dict(
cls,
state_dict: dict,
train: bool = True,
use_flash_attn: bool = True,
):
lora_config = state_dict["hypernet_config"].lora_config
model_name_or_path = lora_config.base_model_name_or_path
base_model = get_model(
@ -582,6 +588,7 @@ class ModulatedPretrainedModel(nn.Module):
train=train,
requires_grad=False,
peft_config=lora_config,
use_flash_attn=use_flash_attn,
)
hypernet_config = state_dict["hypernet_config"]
ctx_encoder_args = state_dict["ctx_encoder_args"]
@ -590,6 +597,17 @@ class ModulatedPretrainedModel(nn.Module):
return model
def _init_model(self):
# for name, module in self.base_model.named_modules():
# if isinstance(module, ModulesToSaveWrapper):
# module.set_adapter("default")
# trainable_base_modules = getattr(
# self.hypernet_config, "trainable_base_modules", None
# )
# if trainable_base_modules:
# for name, param in self.base_model.named_parameters():
# if any(module_name in name for module_name in trainable_base_modules):
# param.requires_grad = True
self.hypernet = HyperLoRA(self.hypernet_config).to(self.device)
# ctx_encoder_name =
# if self.ctx_encoder_args.ctx_encoder_model_name_or_path is not None:
@ -675,11 +693,18 @@ class ModulatedPretrainedModel(nn.Module):
state_dict = self.hypernet.state_dict(*args, **kwargs)
# base_state_dict = dict()
# if self.base_model.modules_to_save:
# for k, v in get_peft_model_state_dict(self.base_model).items():
# if any(module in k for module in self.base_model.modules_to_save):
# # save only "modules_to_save"
# base_state_dict[k] = v
# if self.hypernet_config.trainable_base_modules:
# # for k, v in get_peft_model_state_dict(self.base_model).items():
# # if any(module in k for module in self.hypernet_config.trainable_base_modules):
# # # save only "modules_to_save"
# # base_state_dict[k] = v
# for name, param in self.base_model.named_parameters():
# if param.requires_grad and any(
# module in name
# for module in self.hypernet_config.trainable_base_modules
# ):
# base_state_dict[name] = param.data
# state_dict["base_model"] = base_state_dict
state_dict["hypernet_config"] = self.hypernet_config
state_dict["ctx_encoder_args"] = self.ctx_encoder_args
@ -699,16 +724,19 @@ class ModulatedPretrainedModel(nn.Module):
f"but the hypernet config is for: {self.hypernet_config.lora_config.base_model_name_or_path}"
)
self._init_model()
state_dict.pop("base_model", None)
# base_state_dict = state_dict.pop("base_model", None)
# if base_state_dict:
# logging.info("Loading base model state dict")
# peft_load_result = set_peft_model_state_dict(
# self.base_model,
# base_state_dict,
# *args,
# **kwargs,
# )
# logging.info(f"peft_load_result: {peft_load_result}")
# print("Loading base model state dict")
# print(base_state_dict.keys())
# # peft_load_result = set_peft_model_state_dict(
# # self.base_model,
# # base_state_dict,
# # *args,
# # **kwargs,
# # )
# load_result = self.base_model.load_state_dict(base_state_dict, strict=False)
# print(f"base_load_result: {load_result}")
return self.hypernet.load_state_dict(state_dict, *args, **kwargs)
# @torch.no_grad()
@ -910,9 +938,7 @@ if __name__ == "__main__":
model_name,
train=True,
requires_grad=False,
peft_config=get_lora_config(
model_name, modules_to_save=["input_layernorm", "post_attention_layernorm"]
),
peft_config=get_lora_config(model_name),
)
print(base_model.modules_to_save)
ctx_model_config = base_model.config
@ -966,11 +992,11 @@ if __name__ == "__main__":
state_dict = torch.load(
open(
"train_outputs/runs/Jan15_17-40-28_slurm0-a3nodeset-0_fd03d041/checkpoint-1000/pytorch_model.bin",
"train_outputs/runs/Jan15_19-10-21_slurm0-a3nodeset-11_d1842c41/checkpoint-55000/pytorch_model.bin",
"rb",
)
)
breakpoint()
model.load_state_dict(state_dict)
modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs)
print(modelout)

2
hyperlora/watcher.py
View file

@ -62,6 +62,8 @@ class Watcher:
if __name__ == "__main__":
os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
os.environ["FLASH_ATTENTION_DETERMINISTIC"] = "1"
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
os.environ["WANDB_PROJECT"] = "ctx_to_lora"
os.environ["WANDB_WATCH"] = "" # "all"
os.environ["WANDB_CONSOLE"] = "off"