Add revision

benchmarks/llava/main.py

@@ -62,8 +62,12 @@ def main():
         "llava-hf/llava-1.5-7b-hf",
         torch_dtype=torch.bfloat16,
         device_map=compat.device_type,
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
+    )
+    processor = AutoProcessor.from_pretrained(
+        "llava-hf/llava-1.5-7b-hf",
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
     )
-    processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
 
     # Load dataset and create DataLoader
     dataset = load_dataset("HuggingFaceM4/the_cauldron", "aokvqa")["train"]

benchmarks/llava/prepare.py

@@ -11,8 +11,12 @@ def main():
         "llava-hf/llava-1.5-7b-hf",
         torch_dtype=torch.float32,  # Change to float32
         device_map="auto",
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
+    )
+    _ = AutoProcessor.from_pretrained(
+        "llava-hf/llava-1.5-7b-hf",
+        revision="a272c74b2481d8aff3aa6fc2c4bf891fe57334fb"
     )
-    _ = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
 
     # Load dataset and create DataLoader
     _ = load_dataset("HuggingFaceM4/the_cauldron", "aokvqa")["train"]

milabench/sizer.py

@@ -53,9 +53,10 @@ def to_octet(value: str) -> float:
 class Sizer:
     """Automatically scale the batch size to match GPU spec"""
 
-    def __init__(self, scaling_config=None):
+    def __init__(self, sizer=None, scaling_config=None):
         self.path = scaling_config
-
+        self.sizer_override = sizer
+
         if scaling_config is None:
             scaling_config = default_scaling_config
 
@@ -64,6 +65,8 @@ def __init__(self, scaling_config=None):
 
     @property
     def options(self):
+        if self.sizer_override:
+            return self.sizer_override
         return SizerOptions()
 
     def benchscaling(self, benchmark):