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

watcher + partial refactor eval + debug train w/ pwc and hotpot

Browse source
This commit is contained in:
51616 2025-06-02 09:02:20 +00:00
parent 31de0779da
commit dcda29bae6
13 changed files with 409 additions and 444 deletions
Download patch file Download diff file Expand all files Collapse all files

4
.gitignore vendored
View file

@ -24,4 +24,6 @@ watcher_state.yaml
ds_config.json
eval_results/
.github/
*.code-workspace
*.code-workspace
.wandb/
.ruff_cache/

16
README.md
View file

@ -33,6 +33,22 @@ print(model.decode(outputs))
```bash
WANDB_MODE=disabled uv run python intx_sft.py configs/context_numbers_10.yaml --model_name_or_path=meta-llama/Llama-3.2-1B-Instruct --num_train_epochs=100 --per_device_train_batch_size=64 --per_device_eval_batch_size=64 --exp_setup=hyper_lora --aggregator_type=perceiver --target_modules=down_proj --per_rank_gen=True --per_layer_processing=True --decoder_depth=2 --seed=1 --gen_lora_l1_reg_coef=0 --use_token_mixing=False --lora_r=8 --bf16=True --tf32=True --dataloader_num_workers=8 --dataloader_prefetch_factor=8
```
### PwC + Hotpot training (for testing/debugging)
```bash
uv run intx_sft.py configs/pwc_hotpot_qa.yaml \
--model_name_or_path=google/gemma-2-2b-it --num_train_epochs=1 --per_device_train_batch_size=32 \
--gradient_accumulation_steps=1 --per_device_eval_batch_size=32 --exp_setup=hyper_lora --aggregator_type=perceiver \
--target_modules=down_proj \
--num_self_attends_per_block=4 --num_latent_factor=2 \
--lora_r=8 \
--eval_steps=1000 --save_steps=1000 --learning_rate=4e-5 --lora_dropout=0.0 \
--neftune_noise_alpha=5 --use_light_weight_lora=False \
--load_best_model_at_end=True --metric_for_best_model=pwc_loss --add_negative_prompt=False \
--add_repeat_prompt=False \
--use_sequence_packing=True --per_rank_gen=True \
--per_layer_processing=True \
--gen_lora_l1_reg_coef=0.1
```
<!-- ### HyperLoRA w/ context_numbers_128
```bash

24
configs/pwc_hotpot_qa.yaml
View file

@ -1,6 +1,6 @@
output_dir: "" # just a placeholder
bf16: true
model_name_or_path: meta-llama/Llama-3.2-1B-Instruct
model_name_or_path: google/gemma-2-2b-it
label_names: ["labels"]
# eval_on_start: True
# eval_strategy: "steps"
@ -21,29 +21,31 @@ per_device_eval_batch_size: 8
max_val_samples_per_ds: 1000
# optim: schedule_free_adamw
learning_rate: 0.00002
learning_rate: 0.00004
# lr_scheduler_type: "constant_with_warmup"
neftune_noise_alpha: 1
weight_decay: 0.01
warmup_steps: 100
dataloader_prefetch_factor: 8
dataloader_num_workers: 8
# LoRA
lora_r: 8
lora_dropout: 0.05
lora_dropout: 0.02
target_modules:
- down_proj
- up_proj
- gate_proj
# data
train_ds_names:
- pwc
- hotpot_qa
- pwc
- hotpot_qa
val_ds_names:
- pwc
- hotpot_qa
- pwc
- hotpot_qa
test_ds_names:
- pwc
- hotpot_qa
- pwc
- hotpot_qa

73
intx_sft.py
View file

@ -6,6 +6,7 @@ from math import ceil
import numpy as np
import torch
import wandb
from datasets import (
disable_caching,
interleave_datasets,
@ -17,8 +18,6 @@ from transformers import (
)
from transformers.utils import is_liger_kernel_available
import wandb
from ctx_to_lora.configs import (
AggregatorArguments,
ArgumentParser,
@ -32,7 +31,6 @@ from ctx_to_lora.configs import (
TrainingArguments,
)
from ctx_to_lora.data.collator import train_collator, train_packed_collator
from ctx_to_lora.data.definitions import DS_LEN
from ctx_to_lora.data.processing import get_tokenized_dataset
from ctx_to_lora.metrics import (
Evaluator,
@ -65,6 +63,21 @@ logger = logging.getLogger()
LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0"))
def get_ds_prob(train_ds_len: list[int], total_len: int):
# if a dataset is smaller than 1%, make it 1%
probs = [0 for _ in train_ds_len]
for i, ds_len in enumerate(train_ds_len):
if ds_len / total_len <= 0.01:
probs[i] = 0.01
res_probs = 1 - sum(probs)
res_total_len = sum([l for l in train_ds_len if l / total_len > 0.01])
for i, ds_len in enumerate(train_ds_len):
if ds_len / total_len > 0.01:
probs[i] = ds_len / res_total_len * res_probs
assert sum(probs) == 1
return probs
def main():
############ Argument parsing
parser = ArgumentParser(
@ -283,44 +296,44 @@ def main():
val_ds[ds_name] = val_ds[ds_name].select(val_indices)
# if data_args.streaming:
# max_steps = ceil(
# sum(DS_LEN[ds] for ds in train_ds)
# * training_args.num_train_epochs
# / training_args.per_device_train_batch_size
# / training_args.gradient_accumulation_steps
# / training_args.world_size
# )
# training_args.max_steps = max_steps
# # interleaving streaming datasets
# # simplify the probs for smaller datasets
# # slightly upsample those datasets
# probs = [
# 0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
# for ds_name in train_ds
# ]
# train_ds = interleave_datasets(
# list(train_ds.values()),
# probabilities=probs,
# stopping_strategy="all_exhausted",
# seed=training_args.seed,
# )
# else:
train_ds_len = [len(ds) for ds in train_ds.values()]
total_len = sum(train_ds_len)
max_steps = ceil(
sum(DS_LEN[ds] for ds in train_ds)
total_len
* training_args.num_train_epochs
/ training_args.per_device_train_batch_size
/ training_args.gradient_accumulation_steps
/ training_args.world_size
)
training_args.max_steps = max_steps
# interleaving streaming datasets
# simplify the probs for smaller datasets
# slightly upsample those datasets
probs = [
0.01 if "fw_qa" not in ds_name else 1 + 0.01 - 0.01 * len(train_ds)
for ds_name in train_ds
]
train_ds = interleave_datasets(
list(train_ds.values()),
probabilities=probs,
stopping_strategy="all_exhausted",
probabilities=get_ds_prob(train_ds_len, total_len),
seed=training_args.seed,
)
# else:
# train_ds_len = [len(ds) for ds in train_ds.values()]
# total_len = sum(train_ds_len)
# max_steps = ceil(
# total_len
# * training_args.num_train_epochs
# / training_args.per_device_train_batch_size
# / training_args.gradient_accumulation_steps
# / training_args.world_size
# )
# training_args.max_steps = max_steps
# train_ds = interleave_datasets(
# list(train_ds.values()),
# probabilities=[l / total_len for l in train_ds_len],
# seed=training_args.seed,
# )
logger.info(f"train_ds: {train_ds}")
logger.info(f"val_ds: {val_ds}")

5
pyproject.toml
View file

@ -59,6 +59,9 @@ exclude = [
"plots",
"tmp",
"wandb",
".wandb",
".ruff_cache",
"assets",
]
typeCheckingMode = "off"
@ -69,6 +72,4 @@ ignore = ["E", "F"]
[tool.isort]
profile = "black"
known_first_party = ["wandb"]
known_local_folder = ["ctx_to_lora"]
known_third_party = ["wandb"]

84
run_eval.py Normal file
View file

@ -0,0 +1,84 @@
import logging
from ctx_to_lora.eval_utils import (
run_eval,
)
logger = logging.getLogger()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Evaluate a checkpoint")
parser.add_argument(
"--model_name_or_path",
type=str,
default=None,
help="Evaluate a base model from HuggingFace Hub, without loading checkpoint",
)
parser.add_argument(
"--checkpoint_path",
type=str,
default=None,
help="Path to the checkpoint to evaluate",
)
parser.add_argument(
"--split",
type=str,
choices=["validation", "test"],
default="validation",
help="Which split to evaluate on",
)
parser.add_argument(
"--datasets",
type=str,
nargs="+",
help=(
"Specific datasets to evaluate on."
"If not provided, uses default from args.yaml"
),
)
parser.add_argument(
"--eval_batch_size",
type=int,
default=8,
help="Eval batch size for teacher forcing",
)
parser.add_argument(
"--eval_batch_size_gen",
type=int,
default=32,
help="Eval batch size for generation",
)
parser.add_argument(
"--remove_context",
action="store_true",
help="Remove context when evaluating the base model.",
)
cli_args = vars(parser.parse_args())
# setup_logging(output_dir, debug=os.getenv("DEBUG", False))
eval_batch_size_gen = cli_args.pop("eval_batch_size_gen")
eval_batch_size = cli_args.pop("eval_batch_size")
run_eval(
**cli_args,
# cli_args.checkpoint_path,
# cli_args.model_name_or_path,
# cli_args.eval_batch_size,
# args,
# split=cli_args.split,
eval_batch_size=eval_batch_size,
generative=False,
)
run_eval(
**cli_args,
# cli_args.checkpoint_path,
# cli_args.model_name_or_path,
# cli_args.eval_batch_size_gen,
# args,
# split=cli_args.split,
eval_batch_size=eval_batch_size_gen,
generative=True,
)

108
src/ctx_to_lora/data/collator.py
View file

@ -1,3 +1,4 @@
import numpy as np
import torch
from transformers.data import DataCollatorWithFlattening
@ -56,3 +57,110 @@ def train_collator(inp_list, tokenizer):
out["ctx_attn_mask"] = ctx_attn_mask
return out
def eval_collator(inp_list, tokenizer):
# input is a list of tokenized sequences
padding_kwargs = dict(
padding=True,
padding_side="right",
pad_to_multiple_of=8,
return_tensors="pt",
)
ctx_ids = None
if "ctx_ids" in inp_list[0]:
# have to be manual since it has [ctx_len, features] shape
# pad to the longest ctx_len in the batch
# which can have a different length from the input_ids, attn_mask, labels
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
ctx_ids = torch.nn.utils.rnn.pad_sequence(
ctx_ids,
batch_first=True,
padding_value=0,
)
# exotic keys won't be padded, so we need to pad them as well
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
ctx_attn_mask,
batch_first=True,
padding_value=0,
)
chat_ids = None
if "chat_ids" in inp_list[0]:
chat_ids = [x.pop("chat_ids") for x in inp_list]
chat_ids = torch.nn.utils.rnn.pad_sequence(
chat_ids,
batch_first=True,
padding_value=0,
)
chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list]
chat_attn_mask = torch.nn.utils.rnn.pad_sequence(
chat_attn_mask,
batch_first=True,
padding_value=0,
)
chat_labels = [x.pop("chat_labels") for x in inp_list]
chat_labels = torch.nn.utils.rnn.pad_sequence(
chat_labels,
batch_first=True,
padding_value=-100,
)
chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels)
labels = [x.pop("labels") for x in inp_list]
padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
# hacky explicit padding since the labels are not padded by default
labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
out = {**padded_seq, "labels": labels}
if ctx_ids is not None:
out["ctx_ids"] = ctx_ids
out["ctx_attn_mask"] = ctx_attn_mask
if chat_ids is not None:
out["chat_ids"] = chat_ids
out["chat_attn_mask"] = chat_attn_mask
out["chat_labels"] = chat_labels
return out
def generation_collator(inp_list, tokenizer):
padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
input_ids = [x.pop("input_ids") for x in inp_list]
attn_mask = [x.pop("attention_mask") for x in inp_list]
labels = [x.pop("labels") for x in inp_list]
for i, label in enumerate(labels):
# remove the response tokens
idx = np.argmax(label != -100)
input_ids[i] = input_ids[i][:idx]
attn_mask[i] = attn_mask[i][:idx]
out = tokenizer.pad(
{"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs
)
# we don't include the labels in the output
# since we don't want to compute the loss on the labels
# during generation
if "ctx_ids" in inp_list[0]:
# have to be manual since it has [ctx_len, features] shape
# pad to the longest ctx_len in the batch
# which can have a different length from the input_ids, attn_mask, labels
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
ctx_ids = torch.nn.utils.rnn.pad_sequence(
ctx_ids,
batch_first=True,
padding_value=0,
)
# exotic keys won't be padded, so we need to pad them as well
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
ctx_attn_mask,
batch_first=True,
padding_value=0,
)
out["ctx_ids"] = ctx_ids
out["ctx_attn_mask"] = ctx_attn_mask
return out

20
src/ctx_to_lora/data/definitions.py
View file

@ -333,6 +333,26 @@ for ds_name in LONGBENCH_TASKS + LONGBENCH_E_TASKS:
)
)
CLOSED_QA_DATASETS = {
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"hotpot_qa",
"squad",
"triviaqa_retrieved",
"negative_nq",
}
for ds_name in list(CLOSED_QA_DATASETS):
if ds_name.startswith("longbench/"):
CLOSED_QA_DATASETS.add(f"{ds_name}_e")
# for training closed qa datasets, e.g., hotpot_qa, squad, etc.
CLOSED_QA_INTX_TEMPLATES = [
"Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}",

32
src/ctx_to_lora/data/processing.py
View file

@ -254,18 +254,18 @@ def get_ds_kwargs(ds_name: str, split: str) -> dict[str, Any]:
# return DS_KWARGS[ds_name][split]
kwargs = DS_KWARGS[ds_name][split]
# custom logic for slicing iterable datasets
take, skip = None, None
kw_split = kwargs["split"]
if ("[" in kw_split) and kw_split.endswith("]"):
kwargs["split"], slice = kw_split.split("[")
slice = slice.strip("]")
skip = slice.split(":")[0]
if skip:
kwargs["skip"] = int(skip)
take = slice.split(":")[1]
if take:
kwargs["take"] = int(take)
# # custom logic for slicing iterable datasets
# take, skip = None, None
# kw_split = kwargs["split"]
# if ("[" in kw_split) and kw_split.endswith("]"):
# kwargs["split"], slice = kw_split.split("[")
# slice = slice.strip("]")
# skip = slice.split(":")[0]
# if skip:
# kwargs["skip"] = int(skip)
# take = slice.split(":")[1]
# if take:
# kwargs["take"] = int(take)
return kwargs
@ -439,7 +439,9 @@ def _load_and_process_dataset(
# TODO: drop samples or split into multiple contexts
ds = ds.filter(filter_long_samples, batched=True, num_proc=num_proc)
if add_negative_prompt:
# TODO: make explicit dataset for this
# TODO: add explicit dataset for this
# w/ qs that are answerable 0-shot
# here we add the context to the prompt
ds = ds.map(
add_negative_prompt_fn,
batched=True,
@ -484,8 +486,6 @@ def get_tokenized_dataset(
load_and_process_kwargs = dict(
ds_name=ds_name,
split=split,
base_model_max_len=base_model_max_len,
ctx_model_max_len=ctx_model_max_len,
add_negative_prompt=add_negative_prompt,
add_repeat_prompt=add_repeat_prompt,
repeat_prob=repeat_prob,
@ -544,8 +544,10 @@ def construct_and_tokenize_ctx_qa(
num_proc=None,
):
kwargs = dict(
base_model_max_len=base_model_max_len,
tokenizer=repr(tokenizer),
tokenizer_kwargs=json.dumps(tokenizer_kwargs),
ctx_model_max_len=ctx_model_max_len,
ctx_tokenizer=repr(ctx_tokenizer),
ctx_tokenizer_kwargs=json.dumps(ctx_tokenizer_kwargs),
add_ctx_to_chat=add_ctx_to_chat,

0
src/ctx_to_lora/eval.py
View file

397
eval.py → src/ctx_to_lora/eval_utils.py
View file

@ -1,4 +1,3 @@
import gc
import json
import logging
import os
@ -7,10 +6,8 @@ import string
import sys
from argparse import Namespace
from collections import Counter, defaultdict
from collections.abc import Callable
from dataclasses import fields
from functools import partial
from typing import Any
import numpy as np
import pandas as pd
@ -18,9 +15,7 @@ import torch
import yaml
from datasets import disable_caching
from peft import PeftModel
from rouge_score import rouge_scorer
from transformers import (
EvalPrediction,
GenerationConfig,
PreTrainedModel,
Seq2SeqTrainer,
@ -30,20 +25,31 @@ from transformers import (
)
from transformers.utils import is_liger_kernel_available
from ctx_to_lora.data.definitions import LONGBENCH_E_TASKS, LONGBENCH_TASKS
from ctx_to_lora.data.collator import eval_collator, generation_collator
from ctx_to_lora.data.definitions import (
CLOSED_QA_DATASETS,
LONGBENCH_E_TASKS,
LONGBENCH_TASKS,
)
from ctx_to_lora.data.processing import get_tokenized_dataset
from ctx_to_lora.metrics import (
Evaluator,
compute_metrics,
compute_per_token_acc,
compute_perplexity,
compute_prefix_matching,
compute_rouge,
)
from ctx_to_lora.model_loading import get_model, get_tokenizer
from ctx_to_lora.modeling import hypernet
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
# bandaid for loading old models (before restructure)
sys.modules["ctx_to_lora.modeling_utils"] = hypernet
from ctx_to_lora.utils import clear_gpu, setup_logging
logger = logging.getLogger()
# TODO: add negative samples eval (per-sample?)
# TODO: add length info for more fine grain eval
# bandaid for loading old models (before restructure)
sys.modules["ctx_to_lora.modeling_utils"] = hypernet
# longbench metrics
@ -93,30 +99,6 @@ def compute_qa_f1_score(pred_texts: list[str], label_texts: list[str]):
return dict(qa_f1=np.mean(res))
CLOSED_QA_DATASETS = {
"longbench/narrativeqa",
"longbench/qasper",
"longbench/multifieldqa_en",
"longbench/hotpotqa",
"longbench/2wikimqa",
"longbench/musique",
"hotpot_qa",
"squad",
"triviaqa_retrieved",
"negative_nq",
}
for ds_name in list(CLOSED_QA_DATASETS):
CLOSED_QA_DATASETS.add(f"{ds_name}_e")
def clear_gpu():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_max_memory_cached()
def add_longbench_tasks(ds_names: list[str]):
if "longbench" in ds_names:
ds_names.remove("longbench")
@ -126,106 +108,6 @@ def add_longbench_tasks(ds_names: list[str]):
ds_names += LONGBENCH_E_TASKS
def compute_rouge(pred_texts, label_texts):
out = defaultdict(list)
scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
for pred_text, label_text in zip(pred_texts, label_texts):
scores = scorer.score(pred_text, label_text)
for k, v in scores.items():
out[f"{k}.f1"].append(v.fmeasure)
for k in out:
out[k] = np.mean(out[k])
return out
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()
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):
lengths = valid_masks.sum(dim=1)
is_wrong = (shift_logits.argmax(-1) != shift_labels) * valid_masks
is_correct = (shift_logits.argmax(-1) == shift_labels) * valid_masks
# NOTE: not reliable for multi-turn conversations
# ie, all tokens in the following user's turn will be correct
# still monotonically correlate with perf though
wrong_pos = torch.argmax(is_wrong, dim=1) - torch.argmax(valid_masks, dim=1)
perf = wrong_pos / lengths
# 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_matchings": perf.tolist(),
"n_prefix_matchings": valid_masks.shape[0],
}
@torch.no_grad()
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}
preplexities = torch.exp(loss)
return {
"perplexities": preplexities.tolist(),
"n_perplexities": valid_masks.shape[0],
}
class Evaluator:
def __init__(self, metric_fns: list[Callable]):
self.metric_fns = metric_fns
self.reset()
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():
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.sum(v) / np.sum(self.count[k]) for k, v in self.accum_metrics.items()
}
# Reset batch statistics
self.reset()
return result
@torch.inference_mode()
def compute_metrics(
eval_pred: EvalPrediction,
compute_result: bool,
evaluator: Evaluator,
):
logits, labels = eval_pred.predictions, eval_pred.label_ids
shift_logits = logits[..., :-1, :]
shift_labels = labels[..., 1:]
valid_masks = torch.where(shift_labels != -100, 1, 0)
evaluator.update(shift_logits, shift_labels, valid_masks)
if compute_result:
return evaluator.compute()
def save_generated_text(samples, output_dir, split):
os.makedirs(output_dir, exist_ok=True)
# Create any necessary subdirectories if split contains path separators
@ -239,7 +121,7 @@ def save_generated_text(samples, output_dir, split):
def create_metrics_csv(
metrics_dict: dict[str, dict[str, Any]],
metrics_dict: dict[str, dict[str, any]],
output_dir: str,
model_name: str,
is_hypernet_model: bool = False,
@ -261,12 +143,10 @@ def create_metrics_csv(
# Collect all unique metric names, length groups, and tasks
all_metrics = set()
all_length_groups = set()
all_tasks = set()
all_splits = set()
# FIXME: fix name split (currently wrong)
for split_name, metrics in metrics_dict.items():
task_name = split_name.split("_")[1] if "_" in split_name else split_name
all_tasks.add(task_name)
all_splits.add(split_name)
for metric_key in metrics.keys():
if not metric_key.startswith(split_name):
@ -310,7 +190,7 @@ def create_metrics_csv(
return (2, 0, 0) # Put any malformed strings last
all_length_groups = sorted(list(all_length_groups), key=sort_length_groups)
all_tasks = sorted(all_tasks)
all_splits = sorted(all_splits)
# Create single row with hierarchical columns
row_data = {
@ -319,16 +199,9 @@ def create_metrics_csv(
}
# Create hierarchical column structure: task_metric_lengthgroup
for task in all_tasks:
for task in all_splits:
# Find the corresponding split in metrics_dict
split_key = None
for split_name in metrics_dict.keys():
if task in split_name:
split_key = split_name
break
if split_key is None:
continue
split_key = split_name
metrics = metrics_dict[split_key]
@ -553,113 +426,6 @@ def eval_teacher_forcing(eval_trainer, datasets, split, remove_context):
return out
def train_collator(inp_list, tokenizer):
# input is a list of tokenized sequences
padding_kwargs = dict(
padding=True,
padding_side="right",
pad_to_multiple_of=8,
return_tensors="pt",
)
ctx_ids = None
if "ctx_ids" in inp_list[0]:
# have to be manual since it has [ctx_len, features] shape
# pad to the longest ctx_len in the batch
# which can have a different length from the input_ids, attn_mask, labels
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
ctx_ids = torch.nn.utils.rnn.pad_sequence(
ctx_ids,
batch_first=True,
padding_value=0,
)
# exotic keys won't be padded, so we need to pad them as well
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
ctx_attn_mask,
batch_first=True,
padding_value=0,
)
chat_ids = None
if "chat_ids" in inp_list[0]:
chat_ids = [x.pop("chat_ids") for x in inp_list]
chat_ids = torch.nn.utils.rnn.pad_sequence(
chat_ids,
batch_first=True,
padding_value=0,
)
chat_attn_mask = [x.pop("chat_attn_mask") for x in inp_list]
chat_attn_mask = torch.nn.utils.rnn.pad_sequence(
chat_attn_mask,
batch_first=True,
padding_value=0,
)
chat_labels = [x.pop("chat_labels") for x in inp_list]
chat_labels = torch.nn.utils.rnn.pad_sequence(
chat_labels,
batch_first=True,
padding_value=-100,
)
chat_labels = torch.where(chat_attn_mask == 0, -100, chat_labels)
labels = [x.pop("labels") for x in inp_list]
padded_seq = tokenizer.pad(inp_list, **padding_kwargs)
# hacky explicit padding since the labels are not padded by default
labels = tokenizer.pad({"input_ids": labels}, **padding_kwargs)["input_ids"]
labels = torch.where(padded_seq["attention_mask"] == 0, -100, labels)
out = {**padded_seq, "labels": labels}
if ctx_ids is not None:
out["ctx_ids"] = ctx_ids
out["ctx_attn_mask"] = ctx_attn_mask
if chat_ids is not None:
out["chat_ids"] = chat_ids
out["chat_attn_mask"] = chat_attn_mask
out["chat_labels"] = chat_labels
return out
def generation_collator(inp_list, tokenizer):
padding_kwargs = dict(padding=True, padding_side="left", return_tensors="pt")
input_ids = [x.pop("input_ids") for x in inp_list]
attn_mask = [x.pop("attention_mask") for x in inp_list]
labels = [x.pop("labels") for x in inp_list]
for i, label in enumerate(labels):
# remove the response tokens
idx = np.argmax(label != -100)
input_ids[i] = input_ids[i][:idx]
attn_mask[i] = attn_mask[i][:idx]
out = tokenizer.pad(
{"input_ids": input_ids, "attention_mask": attn_mask}, **padding_kwargs
)
# we don't include the labels in the output
# since we don't want to compute the loss on the labels
# during generation
if "ctx_ids" in inp_list[0]:
# have to be manual since it has [ctx_len, features] shape
# pad to the longest ctx_len in the batch
# which can have a different length from the input_ids, attn_mask, labels
ctx_ids = [example.pop("ctx_ids") for example in inp_list]
ctx_ids = torch.nn.utils.rnn.pad_sequence(
ctx_ids,
batch_first=True,
padding_value=0,
)
# exotic keys won't be padded, so we need to pad them as well
ctx_attn_mask = [example.pop("ctx_attn_mask") for example in inp_list]
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
ctx_attn_mask,
batch_first=True,
padding_value=0,
)
out["ctx_ids"] = ctx_ids
out["ctx_attn_mask"] = ctx_attn_mask
return out
def evaluate(
checkpoint_path,
model_name_or_path,
@ -793,7 +559,7 @@ def evaluate(
model.eval()
collator = generation_collator if generative else train_collator
collator = generation_collator if generative else eval_collator
trainer_kwargs = {
"model": model,
@ -835,67 +601,25 @@ def evaluate(
return out
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Evaluate a checkpoint")
parser.add_argument(
"--model_name_or_path",
type=str,
default=None,
help="Evaluate a base model from HuggingFace Hub, without loading checkpoint",
)
parser.add_argument(
"--checkpoint_path",
type=str,
default=None,
help="Path to the checkpoint to evaluate",
)
parser.add_argument(
"--split",
type=str,
choices=["validation", "test"],
default="validation",
help="Which split to evaluate on",
)
parser.add_argument(
"--datasets",
type=str,
nargs="+",
help=(
"Specific datasets to evaluate on."
"If not provided, uses default from args.yaml"
),
)
parser.add_argument(
"--eval_batch_size",
type=int,
default=8,
help="Eval batch size for teacher forcing",
)
parser.add_argument(
"--eval_batch_size_gen",
type=int,
default=32,
help="Eval batch size for generation",
)
parser.add_argument(
"--remove_context",
action="store_true",
help="Remove context when evaluating the base model.",
)
cli_args = parser.parse_args()
def run_eval(
checkpoint_path: str = None,
model_name_or_path: str = None,
datasets: list[str] = None,
split: str = "validation",
eval_batch_size: int = 8,
remove_context: bool = False,
generative: bool = False,
):
# setup_logging(output_dir, debug=os.getenv("DEBUG", False))
assert bool(cli_args.model_name_or_path) ^ bool(cli_args.checkpoint_path), (
assert bool(model_name_or_path) ^ bool(checkpoint_path), (
"Either --model_name_or_path or --checkpoint_path must be provided"
)
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"
# os.environ["WANDB_MODE"] = "disabled"
disable_caching()
set_seed(42)
torch.use_deterministic_algorithms(True, warn_only=True)
@ -905,51 +629,46 @@ if __name__ == "__main__":
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
if cli_args.checkpoint_path:
checkpoint_dir = "/".join(cli_args.checkpoint_path.split("/")[:-1])
run_dir = "/".join(cli_args.checkpoint_path.split("/")[:-2])
cur_it = int(cli_args.checkpoint_path.split("checkpoint-")[1].split("/")[0])
if checkpoint_path:
checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])
run_dir = "/".join(checkpoint_path.split("/")[:-2])
cur_it = int(checkpoint_path.split("checkpoint-")[1].split("/")[0])
args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
print(f"checkpoint_path: {cli_args.checkpoint_path}")
print(f"checkpoint_path: {checkpoint_path}")
print(f"run_dir: {run_dir}")
args.output_dir = f"{run_dir}/eval-results-{cur_it}"
args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
args.run_name = run_dir.split("/")[-1]
args.remove_context = False # modulated model doesn't see ctx by default
# modulated model doesn't see ctx by default
# but remove_context has to be false for correct file naming
args.remove_context = False
else:
args = Namespace(
model_name_or_path=cli_args.model_name_or_path,
output_dir=f"eval_results/{cli_args.model_name_or_path}",
logging_dir=f"eval_results/{cli_args.model_name_or_path}",
run_name=f"eval_results/{cli_args.model_name_or_path}",
model_name_or_path=model_name_or_path,
output_dir=f"eval_results/{model_name_or_path}",
logging_dir=f"eval_results/{model_name_or_path}",
run_name=f"eval_results/{model_name_or_path}",
# max_base_len=2**13,
# max_ctx_len=-1, # not used
val_ds_names=[],
test_ds_names=[],
remove_context=cli_args.remove_context,
remove_context=remove_context,
)
setup_logging(args.logging_dir)
# Override dataset names if provided via CLI
if cli_args.datasets:
if cli_args.split == "validation":
args.val_ds_names = cli_args.datasets
if datasets:
if split == "validation":
args.val_ds_names = datasets
else:
args.test_ds_names = cli_args.datasets
args.test_ds_names = datasets
# evaluate(
# cli_args.checkpoint_path,
# cli_args.model_name_or_path,
# cli_args.eval_batch_size,
# args,
# split=cli_args.split,
# generative=False,
# )
evaluate(
cli_args.checkpoint_path,
cli_args.model_name_or_path,
cli_args.eval_batch_size_gen,
checkpoint_path,
model_name_or_path,
eval_batch_size,
args,
split=cli_args.split,
generative=True,
split,
generative=generative,
)

17
src/ctx_to_lora/metrics.py
View file

@ -3,9 +3,23 @@ from collections.abc import Callable
import numpy as np
import torch
from rouge_score import rouge_scorer
from transformers import EvalPrediction
def compute_rouge(pred_texts, label_texts):
out = defaultdict(list)
scorer = rouge_scorer.RougeScorer(["rougeL"], use_stemmer=True)
for pred_text, label_text in zip(pred_texts, label_texts):
scores = scorer.score(pred_text, label_text)
for k, v in scores.items():
out[f"{k}.f1"].append(v.fmeasure)
for k in out:
out[k] = np.mean(out[k])
return out
@torch.inference_mode()
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()
@ -15,6 +29,7 @@ def compute_per_token_acc(shift_logits, shift_labels, valid_masks):
}
@torch.inference_mode()
def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
lengths = valid_masks.sum(dim=1)
@ -34,7 +49,7 @@ def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
}
@torch.no_grad()
@torch.inference_mode()
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)

73
watcher.py
View file

@ -1,19 +1,13 @@
# handmade file watcher using glob
# not using watchdog because there are too many saved files
# but we want to just watch when these files are created
# */checkpoints/it_*/hypermod.pt (for HyperLoRA)
import argparse
import gc
import itertools
import os
import time
from glob import glob
import torch
import wandb
import yaml
import wandb
from eval import evaluate
from ctx_to_lora.eval_utils import run_eval
from ctx_to_lora.utils import clear_gpu
CP_PATTERN = "train_outputs/runs/*/checkpoint*/pytorch_model.bin"
@ -22,13 +16,9 @@ def flatten(l):
return itertools.chain.from_iterable(l)
def clear_gpu():
gc.collect()
torch.cuda.empty_cache()
torch.cuda.reset_max_memory_allocated()
torch.cuda.reset_max_memory_cached()
# handmade file watcher using glob
# not using watchdog because there are too many saved files
# but we want to just watch CP_PATTERN files
class Watcher:
def __init__(self, patterns):
self.patterns = patterns
@ -42,7 +32,7 @@ class Watcher:
self.files = self.get_files()
new_files = self.files - self.last_files
self.last_files = self.files
return new_files
return sorted(list(new_files))
def save_state(self):
with open("watcher_state.yaml", "w") as f:
@ -56,26 +46,11 @@ class Watcher:
self.last_files = state["last_files"]
# TODO: run eval on "validation" split during training
# using function from `src/ctx-to-lora/eval.py`
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"
# checkpoint_path = sys.argv[1]
# checkpoint_dir = "/".join(checkpoint_path.split("/")[:-1])
# run_dir = "/".join(checkpoint_path.split("/")[:-2])
# args = Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml", "r")))
# print(f"checkpoint_path: {checkpoint_path}")
# print(f"run_dir: {run_dir}")
# # print(f"args: {args}")
# evaluate(checkpoint_path, args, split="validation", generative=False)
# evaluate(checkpoint_path, args, split="validation", generative=True)
watcher = Watcher([CP_PATTERN])
watcher.load_state()
@ -90,25 +65,33 @@ if __name__ == "__main__":
# cp is delete before we can read it
continue
run_dir = file.split("/checkpoint")[0]
cur_it = int(file.split("checkpoint-")[1].split("/")[0])
checkpoint_dir = os.path.dirname(file)
args = argparse.Namespace(**yaml.unsafe_load(open(f"{run_dir}/args.yaml")))
args.output_dir = f"{run_dir}/eval-results-{cur_it}"
args.logging_dir = f"{run_dir}/eval-results-{cur_it}"
args.run_name = run_dir.split("/")[-1]
print("Evaluating...")
print(f"checkpoint_dir: {checkpoint_dir}")
run_name = run_dir.split("/")[-1]
print(f"Evaluating {file}")
curstep = int(file.split("checkpoint-")[1].split("/")[0])
wandb_kwargs = {
"project": os.getenv("WANDB_PROJECT"),
"group": args.run_name,
"name": f"{args.run_name}-eval",
"id": f"{args.run_name}-eval",
"group": run_name,
"name": f"{run_name}-eval",
"id": f"{run_name}-eval",
"resume": "allow",
}
wandb.init(**wandb_kwargs)
metrics = evaluate(file, args, split="validation", generative=False)
gen_metrics = evaluate(file, args, split="validation", generative=True)
# TODO: have to change this for bigger models
eval_batch_size = 8
eval_batch_size_gen = 32
metrics = run_eval(
checkpoint_path=file,
eval_batch_size=eval_batch_size,
split="validation",
generative=False,
)
gen_metrics = run_eval(
checkpoint_path=file,
eval_batch_size=eval_batch_size_gen,
split="validation",
generative=True,
)
metrics.update(gen_metrics)
wandb.log(metrics, step=curstep)
wandb.finish()