Merge pull request #390 from sxjscience/master

Enable "low_cpu_mem_usage" to reduce the memory usage of HF models

lm_eval/models/gpt2.py

-import transformers
 import torch
+import transformers
 from lm_eval.base import BaseLM
 
 
@@ -9,6 +9,7 @@ class HFLM(BaseLM):
         device="cuda",
         pretrained="gpt2",
         revision="main",
+        low_cpu_mem_usage=None,
         subfolder=None,
         tokenizer=None,
         batch_size=1,
@@ -37,8 +38,7 @@ class HFLM(BaseLM):
         revision = revision + ("/" + subfolder if subfolder is not None else "")
 
         self.gpt2 = transformers.AutoModelForCausalLM.from_pretrained(
-            pretrained,
-            revision=revision,
+            pretrained, revision=revision, low_cpu_mem_usage=low_cpu_mem_usage
         ).to(self.device)
         self.gpt2.eval()
 

lm_eval/utils.py

@@ -7,6 +7,8 @@ import inspect
 import sys
 from typing import List
 
+from omegaconf import OmegaConf
+
 
 class ExitCodeError(Exception):
     pass
@@ -27,10 +29,7 @@ def simple_parse_args_string(args_string):
     if not args_string:
         return {}
     arg_list = args_string.split(",")
-    args_dict = {}
-    for arg in arg_list:
-        k, v = arg.split("=")
-        args_dict[k] = v
+    args_dict = OmegaConf.to_object(OmegaConf.from_dotlist(arg_list))
     return args_dict
 
 

setup.py

@@ -25,6 +25,7 @@ setuptools.setup(
         "jsonlines",
         "numexpr",
         "openai>=0.6.4",
+        "omegaconf>=2.2",
         "pybind11>=2.6.2",
         "pycountry",
         "pytablewriter",