From 889bf08ba840ff8fe843a021d900ee073b7f8cce Mon Sep 17 00:00:00 2001 From: 51616 Date: Tue, 14 Jan 2025 18:14:00 +0000 Subject: [PATCH] add watcher --- hyperlora/eval.py | 16 +++++-- hyperlora/watcher.py | 109 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 121 insertions(+), 4 deletions(-) create mode 100644 hyperlora/watcher.py diff --git a/hyperlora/eval.py b/hyperlora/eval.py index 4f22195..aad198a 100644 --- a/hyperlora/eval.py +++ b/hyperlora/eval.py @@ -153,7 +153,7 @@ def decode_test_result(test_dataset, test_result, tokenizer): def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs): if not isinstance(datasets, dict): datasets = {"": datasets} - + out = {} for ds_name, ds in datasets.items(): split_name = f"{split}_{ds_name}" if ds_name else split eval_result = eval_trainer.predict( @@ -173,22 +173,26 @@ def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs): split=split_name, output_dir=eval_trainer.args.output_dir, ) + out[split_name] = eval_result.metrics eval_trainer.log_metrics(split_name, eval_result.metrics) eval_trainer.save_metrics(split_name, eval_result.metrics) clear_gpu() + return out @torch.no_grad() def eval_teacher_forcing(eval_trainer, datasets, split): if not isinstance(datasets, dict): datasets = {"": datasets} - + out = {} for ds_name, ds in datasets.items(): split_name = f"{split}_{ds_name}" if ds_name else split metrics = eval_trainer.evaluate(ds, metric_key_prefix=split_name) + out[split_name] = metrics eval_trainer.log_metrics(split_name, metrics) eval_trainer.save_metrics(split_name, metrics) clear_gpu() + return out def train_collator(inp_list, tokenizer): @@ -376,6 +380,7 @@ def evaluate(checkpoint_path, args, split, generative): "args": eval_trainer_args, "data_collator": partial(collator, tokenizer=tokenizer), } + out = {} if not generative: trainer_kwargs["compute_metrics"] = partial( @@ -390,12 +395,15 @@ def evaluate(checkpoint_path, args, split, generative): # spents a few hours on this but couldn't find the reason for ds_name, ds in datasets.items(): eval_trainer = Trainer(**trainer_kwargs) - eval_teacher_forcing(eval_trainer, {ds_name: ds}, split) + metrics = eval_teacher_forcing(eval_trainer, {ds_name: ds}, split) + out.update(metrics) else: eval_trainer = Seq2SeqTrainer(**trainer_kwargs) - eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs) + metrics = eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs) + out.update(metrics) clear_gpu() + return out if __name__ == "__main__": diff --git a/hyperlora/watcher.py b/hyperlora/watcher.py new file mode 100644 index 0000000..2985f34 --- /dev/null +++ b/hyperlora/watcher.py @@ -0,0 +1,109 @@ +# 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 itertools +import shutil +import time +import os +import argparse +from glob import glob + +import numpy as np +import pandas as pd +import wandb +import yaml + +from eval import evaluate + + +CP_PATTERN = "train_outputs/runs/*/checkpoint*/pytorch_model.bin" + + +def flatten(l): + return itertools.chain.from_iterable(l) + + +class Watcher: + def __init__(self, patterns): + self.patterns = patterns + self.files = self.get_files() + self.last_files = self.files + + def get_files(self): + return set(flatten(glob(pattern) for pattern in self.patterns)) + + def watch(self): + self.files = self.get_files() + new_files = self.files - self.last_files + self.last_files = self.files + return new_files + + def save_state(self): + with open("watcher_state.yaml", "w") as f: + yaml.dump({"last_files": self.last_files}, f) + + def load_state(self): + if not os.path.exists("watcher_state.yaml"): + return + with open("watcher_state.yaml", "r") as f: + state = yaml.safe_load(f) + self.last_files = state["last_files"] + + +if __name__ == "__main__": + os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" + os.environ["TOKENIZERS_PARALLELISM"] = "true" + 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() + print("Watching for new files...") + while True: + time.sleep(10) + new_files = watcher.watch() + for file in new_files: + # workaround to prevent loading incomplete checkpoints + time.sleep(10) + if not os.path.exists(file): + # cp is delete before we can read it + continue + run_dir = file.split("/checkpoint")[0] + checkpoint_dir = os.path.dirname(file) + args = argparse.Namespace( + **yaml.safe_load(open(f"{run_dir}/args.yaml", "r")) + ) + args.output_dir = checkpoint_dir + args.logging_dir = checkpoint_dir + args.run_name = run_dir.split("/")[-1] + print("Evaluating...") + print(f"checkpoint_dir: {checkpoint_dir}") + 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", + "resume": "allow", + } + wandb.init(**wandb_kwargs) + metrics = evaluate(file, args, split="validation", generative=False) + gen_metrics = evaluate(file, args, split="validation", generative=True) + metrics.update(gen_metrics) + wandb.log(metrics, step=curstep) + wandb.finish() + + watcher.save_state()