all the llamas in canary (#1803)

Summary:
All of these OOM on a single device but want to make them all available for the distributed tests H-Huang is working on

Pull Request resolved: #1803

Reviewed By: H-Huang, xuzhao9

Differential Revision: D47994029

Pulled By: msaroufim

fbshipit-source-id: ad335348b911ddb2379a87a4ec9ef3b0dcc91ea0

torchbenchmark/canary_models/llama_v2_13b/__init__.py

@@ -0,0 +1,15 @@
+from torchbenchmark.tasks import NLP
+from torchbenchmark.util.framework.huggingface.model_factory import HuggingFaceModel, HuggingFaceAuthMixin
+
+class Model(HuggingFaceModel, HuggingFaceAuthMixin):
+    task = NLP.LANGUAGE_MODELING
+    DEFAULT_TRAIN_BSIZE = 1
+    DEFAULT_EVAL_BSIZE = 1
+    DEEPCOPY = False 
+
+    def __init__(self, test, device, batch_size=None, extra_args=[]):
+        HuggingFaceAuthMixin.__init__(self)
+        super().__init__(name="llama_v2_13b", test=test, device=device, batch_size=batch_size, extra_args=extra_args)
+
+    def train(self):
+        return NotImplementedError("FSDP should implement a training loop")

torchbenchmark/canary_models/llama_v2_13b/install.py

@@ -0,0 +1,9 @@
+
+import subprocess
+import sys
+import os
+from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
+
+if __name__ == '__main__':
+    model_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))
+    cache_model(model_name)

torchbenchmark/canary_models/llama_v2_13b/metadata.yaml

@@ -0,0 +1,12 @@
+devices:
+  NVIDIA A100-SXM4-40GB:
+    eval_batch_size: 1
+eval_benchmark: false
+eval_deterministic: false
+eval_nograd: true
+not_implemented:
+- device: cpu
+- device: cuda
+  test: train
+train_benchmark: false
+train_deterministic: false

torchbenchmark/canary_models/llama_v2_70b/__init__.py

@@ -0,0 +1,16 @@
+from torchbenchmark.tasks import NLP
+from torchbenchmark.util.framework.huggingface.model_factory import HuggingFaceModel, HuggingFaceAuthMixin
+
+class Model(HuggingFaceModel, HuggingFaceAuthMixin):
+    task = NLP.LANGUAGE_MODELING
+    DEFAULT_TRAIN_BSIZE = 1
+    DEFAULT_EVAL_BSIZE = 1
+    DEEPCOPY = False 
+
+    def __init__(self, test, device, batch_size=None, extra_args=[]):
+        HuggingFaceAuthMixin.__init__(self)
+        super().__init__(name="llama_v2_70b", test=test, device=device, batch_size=batch_size, extra_args=extra_args)
+
+
+    def train(self):
+        return NotImplementedError("FSDP should implement a training loop")

torchbenchmark/canary_models/llama_v2_70b/install.py

@@ -0,0 +1,9 @@
+
+import subprocess
+import sys
+import os
+from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
+
+if __name__ == '__main__':
+    model_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))
+    cache_model(model_name)

torchbenchmark/canary_models/llama_v2_70b/metadata.yaml

@@ -0,0 +1,12 @@
+devices:
+  NVIDIA A100-SXM4-40GB:
+    eval_batch_size: 1
+eval_benchmark: false
+eval_deterministic: false
+eval_nograd: true
+not_implemented:
+- device: cpu
+- device: cuda
+  test: train
+train_benchmark: false
+train_deterministic: false

torchbenchmark/canary_models/llama_v2_7b/__init__.py

@@ -0,0 +1,16 @@
+from torchbenchmark.tasks import NLP
+from torchbenchmark.util.framework.huggingface.model_factory import HuggingFaceModel, HuggingFaceAuthMixin
+
+class Model(HuggingFaceModel, HuggingFaceAuthMixin):
+    task = NLP.LANGUAGE_MODELING
+    DEFAULT_TRAIN_BSIZE = 1
+    DEFAULT_EVAL_BSIZE = 1
+    DEEPCOPY = False 
+
+    def __init__(self, test, device, batch_size=None, extra_args=[]):
+        HuggingFaceAuthMixin.__init__(self)
+        super().__init__(name="llama_v2_7b", test=test, device=device, batch_size=batch_size, extra_args=extra_args)
+
+
+    def train(self):
+        return NotImplementedError("FSDP should implement a training loop")

torchbenchmark/canary_models/llama_v2_7b/install.py

@@ -0,0 +1,9 @@
+
+import subprocess
+import sys
+import os
+from torchbenchmark.util.framework.huggingface.patch_hf import patch_transformers, cache_model
+
+if __name__ == '__main__':
+    model_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))
+    cache_model(model_name)

torchbenchmark/canary_models/llama_v2_7b/metadata.yaml

@@ -0,0 +1,12 @@
+devices:
+  NVIDIA A100-SXM4-40GB:
+    eval_batch_size: 1
+eval_benchmark: false
+eval_deterministic: false
+eval_nograd: true
+not_implemented:
+- device: cpu
+- device: cuda
+  test: train
+train_benchmark: false
+train_deterministic: false

torchbenchmark/util/framework/huggingface/model_factory.py

@@ -30,6 +30,9 @@
     'hf_Whisper': (1024, 1024, 'WhisperConfig()', 'AutoModelForAudioClassification'),
     # default num_hidden_layers=32 but that OOMs, feel free to change this config to something more real
     'llama_v2_7b_16h' : (512,512, 'LlamaConfig(num_hidden_layers=16)', 'AutoModelForCausalLM'),
+    'llama_v2_7b' : (512,512, 'AutoConfig.from_pretrained("meta-llama/Llama-2-7b-hf")', 'AutoModelForCausalLM'),
+    'llama_v2_13b' : (512,512, 'AutoConfig.from_pretrained("meta-llama/Llama-2-13b-hf")', 'AutoModelForCausalLM'),
+    'llama_v2_70b' : (512, 512, 'AutoConfig.from_pretrained("meta-llama/Llama-2-70b-hf")', 'AutoModelForMaskedLM'),
 }
 
 cpu_input_slice = {