load_state_dict correctly + eval script (crazy trainer eval bug...)

hyperlora/eval.py

@@ -0,0 +1,416 @@
+import yaml
+import sys
+import gc
+import json
+import os
+from argparse import Namespace
+from dataclasses import fields
+from functools import partial
+from collections import defaultdict
+from typing import Callable
+
+import numpy as np
+import torch
+from rouge_score import rouge_scorer
+from transformers import (
+    GenerationConfig,
+    Trainer,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+    EvalPrediction,
+)
+
+from data_utils import get_tokenized_dataset
+from modeling_utils import ModulatedPretrainedModel
+from model_loading import get_tokenizer
+
+
+def clear_gpu():
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_max_memory_allocated()
+    torch.cuda.reset_max_memory_cached()
+
+
+def compute_rouge(pred_texts, label_texts):
+    out = defaultdict(list)
+    scorer = rouge_scorer.RougeScorer(["rouge1", "rougeL"], use_stemmer=False)
+    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}
+
+
+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()}
+
+
+@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}
+
+
+class Evaluator:
+    def __init__(self, metric_fns: list[Callable]):
+        self.metric_fns = metric_fns
+        self.reset()
+
+    def reset(self):
+        self.accum_metrics = 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)
+
+    def compute(self):
+        # Get result across entire eval set
+        result = {k: np.mean(v) for k, v in self.accum_metrics.items()}
+        # Reset batch statistics
+        self.reset()
+        return result
+
+
+@torch.no_grad()
+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):
+    with open(f"{output_dir}/{split}_generated_text.jsonl", "w") as f:
+        for sample in samples:
+            f.write(json.dumps(sample) + "\n")
+
+
+def decode_test_result(test_dataset, test_result, tokenizer):
+    out = []
+    for sample, pred_toks in zip(test_dataset, test_result.predictions):
+        # sample and labels are not padded
+        # pred_toks are padded though
+        d = dict()
+        if "labels" in sample:
+            start_idx = np.argmax(sample["labels"] != -100)
+            label_toks = sample["labels"][start_idx:]
+            # labels are padded with -100, so we need to
+            # replace them with the pad token id
+            label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks)
+            label_text = tokenizer.decode(label_toks, skip_special_tokens=True)
+            d["label"] = label_text
+
+        # remove the label part
+        input_toks = sample["input_ids"][:start_idx]
+
+        gen_toks = pred_toks[np.argmax(pred_toks != tokenizer.pad_token_id) :]
+        gen_toks = gen_toks[start_idx:]
+        gen_toks = np.where(gen_toks == -100, tokenizer.pad_token_id, gen_toks)
+
+        d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True)
+        d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True)
+        if "ctx_ids" in sample:
+            d["context"] = tokenizer.decode(sample["ctx_ids"], skip_special_tokens=True)
+        out.append(d)
+
+    return out
+
+
+@torch.no_grad()
+def eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs):
+    if not isinstance(datasets, dict):
+        datasets = {"": datasets}
+
+    for ds_name, ds in datasets.items():
+        split_name = f"{split}_{ds_name}" if ds_name else split
+        eval_result = eval_trainer.predict(
+            ds,
+            metric_key_prefix=split_name,
+            **gen_kwargs,
+        )
+        decoded_txts = decode_test_result(ds, eval_result, tokenizer)
+        rouge_metrics = compute_rouge(
+            [txt["generated"] for txt in decoded_txts],
+            [txt["label"] for txt in decoded_txts],
+        )
+        for k, v in rouge_metrics.items():
+            eval_result.metrics[f"{split_name}_{k}"] = v
+        save_generated_text(
+            decoded_txts,
+            split=split_name,
+            output_dir=eval_trainer.args.output_dir,
+        )
+        eval_trainer.log_metrics(split_name, eval_result.metrics)
+        eval_trainer.save_metrics(split_name, eval_result.metrics)
+        clear_gpu()
+
+
+@torch.no_grad()
+def eval_teacher_forcing(eval_trainer, datasets, split):
+    if not isinstance(datasets, dict):
+        datasets = {"": datasets}
+
+    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)
+        eval_trainer.log_metrics(split_name, metrics)
+        eval_trainer.save_metrics(split_name, metrics)
+        clear_gpu()
+
+
+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
+    )
+
+    out["labels"] = labels
+
+    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, args, split, generative):
+    assert split in ["validation", "test"]
+    model = ModulatedPretrainedModel.from_state_dict(
+        torch.load(open(checkpoint_path, "rb")), train=False
+    )
+
+    tokenizer = get_tokenizer(args.model_name_or_path, train=False)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    model.base_model.config.pad_token_id = tokenizer.pad_token_id
+    model.base_model.generation_config.pad_token_id = tokenizer.pad_token_id
+
+    ctx_tokenizer = tokenizer
+    if args.ctx_encoder_model_name_or_path:
+        ctx_tokenizer = get_tokenizer(args.ctx_encoder_model_name_or_path, train=False)
+        if ctx_tokenizer.pad_token_id is None:
+            ctx_tokenizer.pad_token_id = ctx_tokenizer.eos_token_id
+
+    tokenizer_kwargs = {"max_length": args.max_base_len}
+    ctx_tokenizer_kwargs = {"max_length": args.max_ctx_len}
+    add_ctx_to_chat = not isinstance(model, ModulatedPretrainedModel)
+
+    _get_tokenized_dataset = partial(
+        get_tokenized_dataset,
+        tokenizer=tokenizer,
+        tokenizer_kwargs=tokenizer_kwargs,
+        ctx_tokenizer=ctx_tokenizer,
+        ctx_tokenizer_kwargs=ctx_tokenizer_kwargs,
+        add_ctx_to_chat=add_ctx_to_chat,
+        add_repeat_prompt=False,
+        add_negative_prompt=False,
+        use_kl_loss=False,
+    )
+
+    datasets = dict()
+    ds_names = args.val_ds_names if split == "validation" else args.test_ds_names
+    for ds_name in ds_names:
+        datasets[ds_name] = _get_tokenized_dataset(ds_name, split)
+    print(datasets)
+
+    gen_kwargs = dict(do_sample=False, max_new_tokens=args.max_new_tokens)
+
+    eval_trainer_args = {}
+
+    # Copy only necessary attributes from training_args to eval_trainer_args
+    seq2seq_training_args_fields = {f.name for f in fields(Seq2SeqTrainingArguments)}
+    for attr, value in dict(**vars(args)).items():
+        if attr in seq2seq_training_args_fields and not attr.startswith("_"):
+            eval_trainer_args[attr] = value
+
+    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 = Seq2SeqTrainingArguments(
+        **eval_trainer_args,
+        predict_with_generate=generative,
+        generation_config=GenerationConfig(**gen_kwargs),
+    )
+
+    # Seq2SeqTrainer is actually just the same as Trainer
+    # (although it uses a different data collator, i.e., explicit prompt/answer separation)
+    # it just allows `predict_with_generate`
+    # allowing us to compute metrics on the generated outputs
+    # no clue why they call this seq2seq...
+
+    print("=" * 80 + "\n" + "Evaluating model..." + "\n" + "=" * 80)
+    print(f"checkpoint_path: {checkpoint_path}")
+
+    model.eval()
+
+    collator = generation_collator if generative else train_collator
+
+    trainer_kwargs = {
+        "model": model,
+        "args": eval_trainer_args,
+        "data_collator": partial(collator, tokenizer=tokenizer),
+    }
+    if not generative:
+
+        trainer_kwargs["compute_metrics"] = partial(
+            compute_metrics,
+            evaluator=Evaluator(
+                [compute_per_token_acc, compute_prefix_matching, compute_perplexity]
+            ),
+        )
+        # this is insane
+        # but i don't know why calling trainer.evaluate() on different datasets
+        # always gives the same numbers across datasets...
+        # 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)
+    else:
+        eval_trainer = Seq2SeqTrainer(**trainer_kwargs)
+        eval_generation(eval_trainer, tokenizer, datasets, split, gen_kwargs)
+
+    clear_gpu()
+
+
+if __name__ == "__main__":
+    os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+    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}")
+
+    args.output_dir = checkpoint_dir
+    args.logging_dir = checkpoint_dir
+    args.run_name = run_dir.split("/")[-1]
+
+    evaluate(checkpoint_path, args, split="validation", generative=False)
+    evaluate(checkpoint_path, args, split="validation", generative=True)

hyperlora/intx_sft.py

@@ -80,7 +80,7 @@ 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}
+    return {"perplexity": acc}
 
 
 def compute_prefix_matching(shift_logits, shift_labels, valid_masks):
@@ -222,18 +222,6 @@ def main():
         ]
     )
 
-    args = {
-        **vars(deepcopy(data_args)),
-        **vars(deepcopy(ctx_args)),
-        **vars(deepcopy(model_args)),
-        **vars(deepcopy(lora_args)),
-        **vars(deepcopy(training_args)),
-        **vars(deepcopy(hypernet_args)),
-        **vars(deepcopy(aggregator_args)),
-        **vars(deepcopy(ctx_encoder_args)),
-    }
-    args["deepspeed_plugin"] = None
-
     checkpoint_dir = training_args.resume_from_checkpoint
     if checkpoint_dir and not os.path.isdir(checkpoint_dir):
         raise NotADirectoryError(f"Checkpoint{checkpoint_dir} is not a directory")
@@ -242,12 +230,12 @@ def main():
     # still possible to have a name crash though
     # logging_dir is just "runs/DATE_TIME_HOSTNAME"
     run_name = (
-        get_run_name(seed_str=training_args.logging_dir)
+        get_run_name(seed_str=training_args.logging_dir.strip("runs/"))
         if not checkpoint_dir
         else checkpoint_dir.split("/")[-2]
     )
 
-    output_dir = f"train_outputs/{run_name}"
+    output_dir = f"train_outputs/runs/{run_name}"
     setup_logging(output_dir, debug=os.getenv("DEBUG", False))
     logger.debug(f"CMD: {' '.join(os.sys.argv)}")
     cli_args = extract_cli_args(os.sys.argv)
@@ -266,6 +254,17 @@ def main():
         and training_args.lr_scheduler_kwargs is None
     ):
         training_args.lr_scheduler_kwargs = {"min_lr": 1e-7}
+    args = {
+        **vars(deepcopy(data_args)),
+        **vars(deepcopy(ctx_args)),
+        **vars(deepcopy(model_args)),
+        **vars(deepcopy(lora_args)),
+        **vars(deepcopy(training_args)),
+        **vars(deepcopy(hypernet_args)),
+        **vars(deepcopy(aggregator_args)),
+        **vars(deepcopy(ctx_encoder_args)),
+    }
+    args["deepspeed_plugin"] = None
     logger.debug(f"args: {args}")
     save_yaml(args, f"{output_dir}/args.yaml")
     set_seed(training_args.seed)
@@ -455,42 +454,6 @@ def main():
             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
-        )
-        out["labels"] = labels
-
-        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
-
     # TODO: use SFTTrainer instead? https://huggingface.co/docs/trl/en/sft_trainer
     # TODO: use packing with SFTTrainer
 

hyperlora/modeling_utils.py

@@ -575,9 +575,11 @@ class ModulatedPretrainedModel(nn.Module):
             requires_grad=False,
             peft_config=lora_config,
         )
-        hypernet_config = state_dict.pop("hypernet_config")
-        ctx_encoder_args = state_dict.pop("ctx_encoder_args")
-        return cls(base_model, hypernet_config, ctx_encoder_args)
+        hypernet_config = state_dict["hypernet_config"]
+        ctx_encoder_args = state_dict["ctx_encoder_args"]
+        model = cls(base_model, hypernet_config, ctx_encoder_args)
+        model.load_state_dict(state_dict)
+        return model
 
     def _init_model(self):
         self.hypernet = HyperLoRA(self.hypernet_config).to(self.device)
@@ -862,7 +864,7 @@ if __name__ == "__main__":
 
     # set torch randomness seed
     torch.manual_seed(42)
-    model_name = "meta-llama/Llama-3.2-1B-Instruct"
+    model_name = "meta-llama/Llama-3.1-8B-Instruct"
     base_model, tokenizer = get_model_and_tokenizer(
         model_name,
         train=True,
@@ -871,13 +873,16 @@ if __name__ == "__main__":
     )
     # TODO: base model shoul be init with target lora config
     print(base_model)
+    ctx_model = base_model
     device = base_model.device
     # lora_config = base_model.peft_config["default"]
     # d_in, d_out = get_peft_in_out_features(base_model, peft_config=lora_config)
 
     hypernet_args = HypernetArguments(latent_size=512)
     aggregator_args = AggregatorArguments(aggregator_type=AGGREGATOR_TYPE.PERCEIVER)
-    hypernet_config = get_hypernet_config(base_model, hypernet_args, aggregator_args)
+    hypernet_config = get_hypernet_config(
+        base_model, ctx_model, hypernet_args, aggregator_args
+    )
     ctx_encoder_args = CtxEncoderArguments(layer_idx=4)
     # ctx_encoder = EarlyExit(get_base_model(base_model), 4)
     # hypernet = HyperLoRA(
@@ -919,7 +924,10 @@ if __name__ == "__main__":
 
     # model.load_state_dict(
     #     torch.load(
-    #         open("train_outputs/20250108-193257_BB1p30QQ/pytorch_model.bin", "rb")
+    #         open(
+    #             "train_outputs/runs/Jan14_11-10-56_slurm0-a3nodeset-1_30c9e93d/checkpoint-1720/pytorch_model.bin",
+    #             "rb",
+    #         )
     #     )
     # )
     # modelout = model(ctx_ids, ctx_attn_mask, **prompt_inputs)

hyperlora/training_utils.py

@@ -28,84 +28,6 @@ def clear_gpu():
     torch.cuda.reset_max_memory_cached()
 
 
-def compute_rouge(pred_texts, label_texts):
-    out = defaultdict(list)
-    scorer = rouge_scorer.RougeScorer(["rouge1", "rougeL"], use_stemmer=False)
-    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 save_generated_text(samples, output_dir, split):
-    with open(f"{output_dir}/{split}_generated_text.jsonl", "w") as f:
-        for sample in samples:
-            f.write(json.dumps(sample) + "\n")
-
-
-def decode_test_result(test_dataset, test_result, tokenizer):
-    out = []
-    for sample, pred_toks in zip(test_dataset, test_result.predictions):
-        d = dict()
-        if "labels" in sample:
-            start_idx = np.argmax(sample["labels"] != -100)
-            label_toks = sample["labels"][start_idx:]
-            # labels are padded with -100, so we need to
-            # replace them with the pad token id
-            label_toks = np.where(label_toks == -100, tokenizer.pad_token_id, label_toks)
-            label_text = tokenizer.decode(label_toks, skip_special_tokens=True)
-            d["label"] = label_text
-
-        # remove the label part
-        input_toks = sample["input_ids"][:start_idx]
-        gen_toks = pred_toks[len(input_toks) :]
-        gen_toks = np.where(gen_toks == -100, tokenizer.pad_token_id, gen_toks)
-
-        d["input"] = tokenizer.decode(input_toks, skip_special_tokens=True)
-        d["generated"] = tokenizer.decode(gen_toks, skip_special_tokens=True)
-        if "ctx_ids" in sample:
-            d["context"] = tokenizer.decode(sample["ctx_ids"], skip_special_tokens=True)
-        out.append(d)
-
-    return out
-
-
-def eval_generation(eval_trainer, tokenizer, dataset, split, gen_kwargs):
-    if not isinstance(dataset, dict):
-        dataset = {"": dataset}
-
-    for ds_name, ds in dataset.items():
-        split_name = f"{split}_{ds_name}" if ds_name else split
-        eval_result = eval_trainer.predict(
-            ds,
-            metric_key_prefix=split_name,
-            **gen_kwargs,
-        )
-        decoded_txts = decode_test_result(ds, eval_result, tokenizer)
-        rouge_metrics = compute_rouge(
-            [txt["generated"] for txt in decoded_txts],
-            [txt["label"] for txt in decoded_txts],
-        )
-        for k, v in rouge_metrics.items():
-            eval_result.metrics[f"{split}_{k}"] = v
-
-        save_generated_text(
-            decoded_txts,
-            split=split_name,
-            output_dir=eval_trainer.args.output_dir,
-        )
-        eval_trainer.log_metrics(split_name, eval_result.metrics)
-        eval_trainer.save_metrics(split_name, eval_result.metrics)
-        clear_gpu()
-
-
-# def per_sample_loss_avg_fn(outputs, labels, num_items_in_batch):
-#     ...
-
-
 def train_model(
     model,
     # tokenizer,
@@ -138,8 +60,8 @@ def train_model(
     # Trainer loads the best model after training
     # is done when load_best_model_at_end=True
     train_result = trainer.train(resume_from_checkpoint=checkpoint)
-    trainer.log_metrics("train", train_result.metrics)
     trainer.save_model()  # just in case OOM when run trainer.evaluate()
+    trainer.log_metrics("train", train_result.metrics)
     clear_gpu()
     metrics = trainer.evaluate(dict(**val_dataset, test=test_dataset))
     trainer.log_metrics("eval", metrics)