From 37050b84d07f7cfa731fc85ecb9de3faef1d282e Mon Sep 17 00:00:00 2001
From: Conglong Li <conglong.li@gmail.com>
Date: Tue, 14 Nov 2023 23:48:46 -0800
Subject: [PATCH] Fix lm_eval_harness for GPT models (#292)

---
 examples_deepspeed/MoE/ds_evalharness.sh      |  3 ++-
 examples_deepspeed/MoE/readme_evalharness.md  |  8 +++----
 .../compression/ds_evalharness.sh             |  5 +++--
 .../gpt/eval/ds_evalharness_1gpu.sh           |  3 ++-
 .../gpt/eval/ds_evalharness_parallel_run.sh   |  1 +
 .../ds_evalharness_parallel_run_10shot.sh     |  1 +
 tasks/eval_harness/evaluate.py                | 22 +++++++++++--------
 7 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/examples_deepspeed/MoE/ds_evalharness.sh b/examples_deepspeed/MoE/ds_evalharness.sh
index f989b1f37..3496ada20 100644
--- a/examples_deepspeed/MoE/ds_evalharness.sh
+++ b/examples_deepspeed/MoE/ds_evalharness.sh
@@ -28,7 +28,7 @@ TASKS="lambada"
 VOCAB_FILE=/data/Megatron-LM/data/gpt2-vocab.json
 MERGE_FILE=/data/Megatron-LM/data/gpt2-merges.txt
 
-export HF_DATASETS_OFFLINE=1
+# export HF_DATASETS_OFFLINE=1
 
 # Dummy arguments to make megatron happy. No need to configure them.
 # The reason we don't need to configure them and many other arguments is
@@ -53,6 +53,7 @@ CMD="../../tasks/eval_harness/evaluate.py \
     --no-load-rng \
     --inference \
     --disable-moe-token-dropping \
+    --tokenizer-type GPT2BPETokenizer \
     --adaptive_seq_len\
     --eval_fp32\
     --task_list $TASKS\
diff --git a/examples_deepspeed/MoE/readme_evalharness.md b/examples_deepspeed/MoE/readme_evalharness.md
index 426d63735..d30075e2f 100644
--- a/examples_deepspeed/MoE/readme_evalharness.md
+++ b/examples_deepspeed/MoE/readme_evalharness.md
@@ -11,11 +11,10 @@ This particular setup uses the normal deepspeed checkpoint and requires no conve
 On login console with external network
 
 Get lm-eval harness (https://github.com/EleutherAI/lm-evaluation-harness) and `best-download==0.0.7` needed to download some tasks.
+Below package version numbers are what we tested that work.
 ```
 (maybe need pip install --upgrade pip)
-pip install best-download==0.0.7
-pip install lm-eval
-(previously we used "pip install git+https://github.com/EleutherAI/lm-evaluation-harness" to install, but later found the command above has less dependency issues)
+pip install best-download==0.0.7 lm-eval==0.2.0 datasets==1.15.1 transformers==4.20.1 huggingface-hub==0.8.1
 ```
 
 2. Pre-download needed datasets
@@ -33,7 +32,8 @@ Then install datasets for the tasks:
 ```
 python ../../tasks/eval_harness/download.py --task_list hellaswag,lambada,triviaqa,webqs,winogrande,piqa,arc_challenge,arc_easy,openbookqa,race,boolq,cb,copa,rte,wic,wsc,multirc,record,anli_r1,anli_r2,anli_r3,wikitext,logiqa,mathqa,mc_taco,mrpc,prost,pubmedqa,qnli,qqp,sciq,sst,wnli
 ```
-and make sure that `export HF_DATASETS_OFFLINE=1`
+
+Previously we set `export HF_DATASETS_OFFLINE=1` to make the dataset offline after the above manual download. But somehow now this could trigger error on some kind of online verification for some of the datasets, so it's recommended to only set offline mode when necessary.
 
 <!-- If there are things like custom tokenizers, pre-download those too, e.g.:
 
diff --git a/examples_deepspeed/compression/ds_evalharness.sh b/examples_deepspeed/compression/ds_evalharness.sh
index a1ac63ce2..0922dc033 100644
--- a/examples_deepspeed/compression/ds_evalharness.sh
+++ b/examples_deepspeed/compression/ds_evalharness.sh
@@ -1,4 +1,4 @@
-# This is an example zero-shot eval script. Please first read the readme_evalharness.md under the same directory.
+# This is an example zero-shot eval script. Please first read the readme_evalharness.md under the ../MoE directory.
 
 # CHECKPOINT_PATH=/blob/users/minjiaz/compression_library/checkpoint/125M10L_Compression_Test_INT8_64gpu_lr6e-5_tokens5.25B_nocl_alpha-no_pp/global_step2000/
 # CHECKPOINT_PATH=/blob/users/conglli/project/gpt3_with_pile/checkpoint/gpt3-with-pile-0.125B-lr-2.4e-3-minlr-6.0e-5-bs-2048-gpus-64-zero-0-mp-1-pp-1-no_pp-cl-startseqlen-72-step-27638-token-60B/global_step71000/
@@ -31,7 +31,7 @@ TASKS="lambada,wikitext"
 VOCAB_FILE=/blob/data/the_pile_public_merged_nopreprocessing/gpt2-vocab.json
 MERGE_FILE=/blob/data/the_pile_public_merged_nopreprocessing/gpt2-merges.txt
 
-export HF_DATASETS_OFFLINE=1
+# export HF_DATASETS_OFFLINE=1
 
 # Dummy arguments to make megatron happy. No need to configure them.
 # The reason we don't need to configure them and many other arguments is
@@ -56,6 +56,7 @@ CMD="../../tasks/eval_harness/evaluate.py \
     --no-load-rng \
     --inference \
     --disable-moe-token-dropping \
+    --tokenizer-type GPT2BPETokenizer \
     --adaptive_seq_len\
     --eval_fp32\
     --task_list $TASKS\
diff --git a/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_1gpu.sh b/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_1gpu.sh
index 28992f71a..32ade4917 100644
--- a/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_1gpu.sh
+++ b/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_1gpu.sh
@@ -27,7 +27,7 @@ if [ ! -f "$merge_file" ]; then
     wget https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt
 fi
 
-export HF_DATASETS_OFFLINE=1
+# export HF_DATASETS_OFFLINE=1
 
 dir2=$(dirname "$checkpoint_path")
 dirname=$(basename "$dir2")/$(basename "$checkpoint_path")
@@ -58,6 +58,7 @@ command="../../../../tasks/eval_harness/evaluate.py \
     --no-load-rng \
     --inference \
     --disable-moe-token-dropping \
+    --tokenizer-type GPT2BPETokenizer \
     --adaptive_seq_len \
     --eval_fp32 \
     --num_fewshot ${num_fewshot} \
diff --git a/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run.sh b/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run.sh
index a7fd6318e..2bfbec3a1 100644
--- a/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run.sh
+++ b/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run.sh
@@ -48,6 +48,7 @@ num_fewshot=0
 num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
 cuda_id=-1
 total_mem=$(nvidia-smi --query-gpu=memory.total --format=csv -i 0 | grep -Eo [0-9]+)
+total_mem=$(( ${total_mem}*99/100 )) # somehow there could exist tiny (4MB or so) gpu memory leak
 
 ## Code below only works when you run each evalharness task on a single GPU.
 ## For multi-GPU evalharness, check Megatron-DeepSpeed/blob/main/examples_deepspeed/MoE/ds_evalharness.sh
diff --git a/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run_10shot.sh b/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run_10shot.sh
index ca7d24f80..8e6406477 100644
--- a/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run_10shot.sh
+++ b/examples_deepspeed/data_efficiency/gpt/eval/ds_evalharness_parallel_run_10shot.sh
@@ -43,6 +43,7 @@ batch_size=16
 num_gpus=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
 cuda_id=-1
 total_mem=$(nvidia-smi --query-gpu=memory.total --format=csv -i 0 | grep -Eo [0-9]+)
+total_mem=$(( ${total_mem}*99/100 )) # somehow there could exist tiny (4MB or so) gpu memory leak
 
 ## Code below only works when you run each evalharness task on a single GPU.
 ## For multi-GPU evalharness, check Megatron-DeepSpeed/blob/main/examples_deepspeed/MoE/ds_evalharness.sh
diff --git a/tasks/eval_harness/evaluate.py b/tasks/eval_harness/evaluate.py
index 0392bd4f6..860d3cf01 100644
--- a/tasks/eval_harness/evaluate.py
+++ b/tasks/eval_harness/evaluate.py
@@ -23,9 +23,10 @@
 from megatron import get_args
 from megatron import print_rank_0
 from megatron import get_tokenizer
+from megatron.core.enums import ModelType
 from megatron.core import mpu
 from megatron.training import setup_model_and_optimizer, get_model
-from megatron.mpu.mappings import gather_from_tensor_model_parallel_region
+from megatron.core.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
 
 from megatron.utils import get_ltor_masks_and_position_ids, unwrap_model
 from megatron.p2p_communication import recv_forward, send_forward
@@ -222,8 +223,7 @@ def _model_call(self, inps):
                     a_output, *other_losses = self.model(tokens,
                         position_ids,
                         attention_mask,
-                        tokentype_ids=None,
-                        forward_method_parallel_output=False)
+                        tokentype_ids=None)
                     output.append(a_output)
 
                 if output is not None:
@@ -320,7 +320,7 @@ def load_ds_checkpoint_and_setup_megatron(extra_args_provider):
     # avoid printing the arguments, since they will later be overridden.
     _print_args = megatron.arguments._print_args
     megatron.arguments._print_args = lambda *_args, **kwarg: None
-    args = _parse_args(extra_args_provider)
+    args = parse_args(extra_args_provider=extra_args_provider)
 
     ds_checkpoint = DeepSpeedCheckpoint(args.load,
                                         tp_degree=args.tensor_model_parallel_size,
@@ -340,20 +340,24 @@ def load_ds_checkpoint_and_setup_megatron(extra_args_provider):
         cp_args.bf16 = False
         cp_args.params_dtype = torch.float32
 
+    cp_args.tokenizer_type = 'GPT2BPETokenizer'
+
     override_args(args, cp_args, skip_keys, skip_if_specified)
 
     # stop megatron from reparsing the arguments.
-    megatron.global_vars._parse_args = lambda *_args, **kwarg: args
+    megatron.arguments.parse_args = lambda *_args, **kwarg: args
+    megatron.global_vars._ensure_var_is_not_initialized = lambda *_args, **kwarg: None
     megatron.global_vars._GLOBAL_ARGS = args
 
-    initialize_megatron()
+    initialize_megatron(extra_args_provider=extra_args_provider)
+    megatron.global_vars._GLOBAL_ARGS = args
     torch.distributed.barrier()
 
     # Initializing megatron will update eg. tokenizer size. Override again.
     override_args(args, cp_args, skip_keys, skip_if_specified)
 
     # print final arguments.
-    _print_args(args)
+    _print_args("eval_harness arguments", args)
     if args.deepspeed:
 
         # Hack #3:
@@ -369,7 +373,7 @@ def load_ds_checkpoint_and_setup_megatron(extra_args_provider):
 
         cp_path = args.load
         args.load = None
-        model, _, _ = setup_model_and_optimizer(model_provider)
+        model, _, _ = setup_model_and_optimizer(model_provider, ModelType.encoder_or_decoder)
         model = model[0]
         zero_enabled = model._config.zero_enabled
         model._config.zero_enabled = False
@@ -399,7 +403,7 @@ def tasks_args(parser):
     group.add_argument('--eval_fp32',  default = False, action='store_true', help='Should the evaluation run in fp32')
     return parser
 
-from megatron.global_vars import _parse_args
+from megatron.arguments import parse_args
 
 def main():
     start = time.time()