Optimize loading of falcon model checkpoints

tenstorrent · Dec 11, 2024 · 9099753 · 9099753
1 parent 2abd3fe
commit 9099753
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/models/demos/ttnn_falcon7b/tests/test_falcon_attention.py b/models/demos/ttnn_falcon7b/tests/test_falcon_attention.py
@@ -33,7 +33,7 @@ def get_model_prefix(layer_index: int = 0):
 @pytest.fixture(scope="module")
 def torch_model():
     hugging_face_reference_model = transformers.FalconForCausalLM.from_pretrained(
-        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True
+        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True, device_map="auto"
     ).eval()
     state_dict = hugging_face_reference_model.state_dict()
     filtered_state_dict = strip_state_dict_prefix(state_dict, get_model_prefix())

diff --git a/models/demos/ttnn_falcon7b/tests/test_falcon_decoder.py b/models/demos/ttnn_falcon7b/tests/test_falcon_decoder.py
@@ -33,7 +33,7 @@ def get_model_prefix(layer_index: int = 0):
 @pytest.fixture(scope="module")
 def torch_model():
     hugging_face_reference_model = transformers.FalconForCausalLM.from_pretrained(
-        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True
+        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True, device_map="auto"
     ).eval()
     state_dict = hugging_face_reference_model.state_dict()
     mlp_state_dict = strip_state_dict_prefix(state_dict, get_model_prefix())

diff --git a/models/demos/ttnn_falcon7b/tests/test_falcon_mlp.py b/models/demos/ttnn_falcon7b/tests/test_falcon_mlp.py
@@ -25,7 +25,7 @@ def get_model_prefix(layer_index: int = 0):
 @pytest.fixture(scope="module")
 def torch_model():
     hugging_face_reference_model = transformers.FalconForCausalLM.from_pretrained(
-        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True
+        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True, device_map="auto"
     ).eval()
     state_dict = hugging_face_reference_model.state_dict()
     mlp_state_dict = strip_state_dict_prefix(state_dict, get_model_prefix())

diff --git a/models/demos/ttnn_falcon7b/tests/test_falcon_rotary_embedding.py b/models/demos/ttnn_falcon7b/tests/test_falcon_rotary_embedding.py
@@ -29,7 +29,7 @@ def get_model_prefix(layer_index: int = 0):
 @pytest.fixture(scope="module")
 def torch_model():
     hugging_face_reference_model = transformers.FalconForCausalLM.from_pretrained(
-        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True
+        PRETRAINED_MODEL_NAME, low_cpu_mem_usage=True, device_map="auto"
     ).eval()
     state_dict = hugging_face_reference_model.state_dict()
     filtered_state_dict = strip_state_dict_prefix(state_dict, get_model_prefix())