Bunch of changes: GPTneoX for more methods, added Wandb, reorganized, compute optimized modify llama

.gitignore

@@ -2,3 +2,5 @@ experiments/**
 outfiles/**
 text_generation/
 **/__pycache__/
+wandb/**
+methods/baselines/topk/bce928f38989812b69c6f8e3a86763e004387d16/**

data/input.txt

@@ -0,0 +1,3 @@
+The United States of America (USA or U.S.A.), commonly known as the United States (US or U.S.) or America, is a country primarily located in North America. 
+The Arsenal Football Club, commonly known as Arsenal, is an English professional football club based in Holloway, North London.
+The three primary colors are

eval_ppl.py

@@ -1,11 +1,23 @@
 from methods import init_tensor_saver
-from configure_model import get_h2o_args, get_topk_args, get_spar_args, get_pca_args, get_save_tensor_args
-from configure_model import get_modifier
+from methods.common.configure_model import get_h2o_args, get_topk_args, get_spar_args, get_pca_args, get_save_tensor_args
+from methods.common.configure_model import get_modifier
+from methods import init_logger, finish_logger
+import methods
 
 import argparse
 import os
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
+#LM_HARNESS_VALID_TASKS = ["hellaswag", "winogrande", "gsm8k", "mmlu", "truthfulqa_mc2", "arc_challenge"]
+LM_HARNESS_TASKS = {
+  "mmlu" : "acc,none",
+  "gsm8k" : "exact_match,strict-match",
+  "hellaswag" : "acc_norm,none",
+  "winogrande" : "acc,none",
+  "truthfulqa_mc2" : "acc,none",
+  "arc_challenge" : "acc_norm,none"
+}
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -14,7 +26,9 @@
     parser.add_argument("--sequence-length", type=int, default=4096, help="sequence length")
     parser.add_argument("--use-axonn", action='store_true', default=False, help="shard a model using AxoNN")
     parser.add_argument("--lm-harness-eval", action='store_true', default=False, help="use lm harness eval")
-    parser.add_argument("--dataset", type=str, default="wikitext-test", help="which dataset to use for ppl eval")
+    parser.add_argument("--dataset", type=str, default="wikitext-test", help="dataset - wikitext, bookcorpus, c4")
+    parser.add_argument("--use-wandb", action='store_true', default=False, help="use wandb")
+    #parser.add_argument("--task", type=str, default="perplexity", help="task - perplexity, <lm_harness_tasks>")
 
     parser = get_h2o_args(parser)
     parser = get_topk_args(parser)
@@ -26,38 +40,49 @@
     if args.save_tensors:
         init_tensor_saver(args.tensors_dir)
 
+    init_logger(args)
+
     modifier_method = get_modifier(args)
     if modifier_method is None:
         raise ValueError("Modifier method not found")
 
     print (modifier_method)
 
     cache = None
-    if args.use_topk:
-        modifier_method(args.top_k)
-    elif args.use_h2o:
-        modifier_method(args.heavy_ratio)
+    if args.use_topk or args.use_h2o or args.use_pca_topk:
+        modifier_method(args)
     elif args.use_sparq or args.use_spark:
         modifier_method(args.top_r, args.top_k)
     elif args.use_spar_hat:
         cache = modifier_method(args.top_r)
     elif args.use_pca:
         modifier_method(args.top_r)
         args.use_axonn = False
-    elif args.use_pca_topk:
-        modifier_method(args.top_r, args.top_k)
-
+
     if args.lm_harness_eval:
         import lm_eval
+        from lm_perplexity_eval import evaluate
+        model = evaluate(model_id=args.model_id,
+                    dataset=args.dataset,
+                    sequence_length=args.sequence_length,
+                    use_axonn=args.use_axonn,
+                    past_key_values=cache,
+                    axonn_low_level_api=True,
+                    return_model=True)
         results = lm_eval.simple_evaluate(
             model = "hf",
-            model_args=f"pretrained={args.model_id}",
-            tasks = ["copa", "rte", "openbookqa", "mathqa", "winogrande", "hellaswag"],
-            #tasks = ["hellaswag"],
+            #model_args=f"pretrained={args.model_id}",
+            #model_args={"pretrained": model, "parallelize": True},
+            model_args={"pretrained": model},
+            tasks = LM_HARNESS_TASKS.keys(),
             log_samples=False,
+            batch_size=16
         )
 
-        print(results["results"])
+        if results is not None:
+            print(results["results"])
+            if methods.LOGGER is not None:
+                methods.LOGGER.log_lm_harness_results(LM_HARNESS_TASKS, results["results"])
     else:
         from lm_perplexity_eval import evaluate
         print(args.use_axonn)
@@ -69,3 +94,7 @@
                     axonn_low_level_api=True)
 
         print(ppl)
+        if methods.LOGGER is not None:
+            methods.LOGGER.log_ppl(ppl)
+
+    finish_logger()

eval_ppl_old.py

@@ -8,8 +8,8 @@
   #make_gemma_attention_top_k
   make_gptneox_attention_top_k
 )
-from configure_model import get_h2o_args, get_topk_args, get_spar_args, get_pca_args, get_save_tensor_args
-from configure_model import get_modifier
+from methods.common.configure_model import get_h2o_args, get_topk_args, get_spar_args, get_pca_args, get_save_tensor_args
+from methods.common.configure_model import get_modifier
 
 from methods import SparHatCache
 import argparse

examples/h2o/h2o-llama.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+set -x
+sbatch examples/h2o/submit_h2o.sh meta-llama/Llama-2-7b-hf llama 4096 0.125
+sbatch examples/h2o/submit_h2o.sh meta-llama/Llama-2-7b-hf llama 4096 0.25
+sbatch examples/h2o/submit_h2o.sh meta-llama/Llama-2-7b-hf llama 4096 0.125 --lm-harness-eval
+sbatch examples/h2o/submit_h2o.sh meta-llama/Llama-2-7b-hf llama 4096 0.25 --lm-harness-eval
+set +x

examples/h2o/h2o-mistral.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+set -x
+sbatch examples/h2o/submit_h2o_1gpu.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.125
+sbatch examples/h2o/submit_h2o_1gpu.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.25
+sbatch examples/h2o/submit_h2o_1gpu.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.125 --lm-harness-eval
+sbatch examples/h2o/submit_h2o_1gpu.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.25 --lm-harness-eval
+set +x

examples/h2o/h2o-pythia.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+set -x
+sbatch examples/h2o/submit_h2o_1gpu.sh EleutherAI/pythia-6.9b gptneox 2048 0.125
+sbatch examples/h2o/submit_h2o_1gpu.sh EleutherAI/pythia-6.9b gptneox 2048 0.25
+sbatch examples/h2o/submit_h2o_1gpu.sh EleutherAI/pythia-6.9b gptneox 2048 0.125 --lm-harness-eval
+sbatch examples/h2o/submit_h2o_1gpu.sh EleutherAI/pythia-6.9b gptneox 2048 0.25 --lm-harness-eval
+set +x

examples/submit_h2o.sh → examples/h2o/submit_h2o.sh

@@ -6,6 +6,9 @@
 #SBATCH --account=m4641_g
 #SBATCH --ntasks-per-node=4
 #SBATCH --time=03:00:00
+#SBATCH -J h2o
+#SBATCH --output=outfiles/%x-%j.out
+
 
 
 # Runs a "10B" parameter model

examples/submit_h2o_1gpu.sh → examples/h2o/submit_h2o_1gpu.sh

@@ -6,6 +6,9 @@
 #SBATCH --account=m4641_g
 #SBATCH --ntasks-per-node=1
 #SBATCH --time=10:00:00
+#SBATCH -J h2o
+#SBATCH --output=outfiles/%x-%j.out
+
 
 
 # Runs a "10B" parameter model

examples/pca/pca-llama.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+set -x
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 128 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 64 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 32 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 128
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 64
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 32
+sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 96
+sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.97 --lm-harness-eval
+sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.95 --lm-harness-eval
+sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.90 --lm-harness-eval
+sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.85 --lm-harness-eval
+sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.75 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.97
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.95
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.90
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.85
+#sbatch examples/pca/submit_pca.sh meta-llama/Llama-2-7b-hf llama 4096 0.75
+set +x

examples/pca/pca-mistral.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+set -x
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 128 --lm-harness-eval
+sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 64 --lm-harness-eval
+sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 32 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 128
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 64
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 32
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 96
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.97 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.95 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.90 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.85 --lm-harness-eval
+sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.70 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.75 --lm-harness-eval
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.97
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.95
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.90
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.85
+#sbatch examples/pca/submit_pca.sh mistralai/Mistral-7B-v0.1 mistral 4096 0.75
+set +x

examples/pca_topk/pca-topk-llama.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+MODEL=$1
+MODEL_TYPE=$2
+SEQLEN=$3
+
+set -x
+sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 64 0.5
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 64 0.25
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 64 0.125
+#
+sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 32 0.5
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 32 0.25
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 32 0.125
+#
+sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 16 0.5
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 16 0.25
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 16 0.125
+#
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 64 0.5 --lm-harness-eval
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 64 0.25 --lm-harness-eval
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 64 0.125 --lm-harness-eval
+#
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 32 0.5 --lm-harness-eval
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 32 0.25 --lm-harness-eval
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 32 0.125 --lm-harness-eval
+#
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 16 0.5 --lm-harness-eval
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 16 0.25 --lm-harness-eval
+#sbatch examples/pca_topk/submit_pca_topk.sh ${MODEL} ${MODEL_TYPE} ${SEQLEN} 16 0.125 --lm-harness-eval
+set +x

examples/pca_topk/pca-topk-mistral.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+set -x
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 64 0.5
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 64 0.25
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 64 0.125
+
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 32 0.5
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 32 0.25
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 32 0.125
+
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 16 0.5
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 16 0.25
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 16 0.125
+
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 64 0.5 --lm-harness-eval
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 64 0.25 --lm-harness-eval
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 64 0.125 --lm-harness-eval
+
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 32 0.5 --lm-harness-eval
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 32 0.25 --lm-harness-eval
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 32 0.125 --lm-harness-eval
+
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 16 0.5 --lm-harness-eval
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 16 0.25 --lm-harness-eval
+sbatch examples/pca_topk/submit_pca_topk.sh mistralai/Mistral-7B-v0.1 mistral 4096 16 0.125 --lm-harness-eval
+set +x