change parallelize a model from AutoModelForCausalLM

xrsrke · Sep 2, 2023 · 9ac8531 · 9ac8531
1 parent c15c80b
commit 9ac8531
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/tests/nn/tensor_parallel/conftest.py b/tests/nn/tensor_parallel/conftest.py
@@ -1,9 +1,9 @@
 import pytest
-from transformers import AutoModel
+from transformers import AutoModelForCausalLM
 
 MODEL_NAME = "bigscience/bloom-560m"
 
 
 @pytest.fixture(scope="session")
 def model():
-    return AutoModel.from_pretrained(MODEL_NAME)
+    return AutoModelForCausalLM.from_pretrained(MODEL_NAME)
diff --git a/tests/nn/tensor_parallel/test_parallelize.py b/tests/nn/tensor_parallel/test_parallelize.py
@@ -85,10 +85,10 @@ def run_parallelize_linear(
 
 
 @pytest.mark.parametrize("tensor_parallel_size, MODULE_NAME, get_module", [
-    (1, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.h[0].mlp.dense_h_to_4h),
-    (2, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.h[0].mlp.dense_h_to_4h),
-    (1, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.h[0].mlp.dense_4h_to_h),
-    (2, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.h[0].mlp.dense_4h_to_h),
+    (1, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.transformer.h[0].mlp.dense_h_to_4h),
+    (2, "transformer.h.0.mlp.dense_h_to_4h", lambda model: model.transformer.h[0].mlp.dense_h_to_4h),
+    (1, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.transformer.h[0].mlp.dense_4h_to_h),
+    (2, "transformer.h.0.mlp.dense_4h_to_h", lambda model: model.transformer.h[0].mlp.dense_4h_to_h),
 ])
 def test_parallelize_linear(model, tensor_parallel_size, MODULE_NAME, get_module):
     PIPELINE_PARALLEL_SIZE = 1