Skip to content

Commit ba45256

Browse files
mattdangerwmattdangerw
committed
Remove unused TPU testing for backbones (#1266)
TPU testing would be great, but we should not do it quite like this. - These are not run in any form of CI right now. Do they run at all? I don't think anyone has tried these for many months. - They basically replicate the simplest tests from the main fixture inside a TPU strategy scope. We would be much better off handling this inside some common test setup, rather than replicating test code. - They would not work multi-backend. - They do not test presets, tasks, backprop, and a lot of the actual important functionality for our pretrained model offering. Let's remove for now as I consolidate our preset testing code. We can bring this back later when we have a plan to run these in a more sustainable way for our test suite.
1 parent 4be54b6 commit ba45256

15 files changed

+2
-434
lines changed

‎keras_nlp/conftest.py‎

Lines changed: 2 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -12,30 +12,13 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
import os
16-
1715
import pytest
1816
import tensorflow as tf
1917

2018
from keras_nlp.backend import config as backend_config
2119
from keras_nlp.backend import keras
2220

2321

24-
@pytest.fixture(scope="session")
25-
def tpu_strategy():
26-
tpu_name = os.getenv("KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS")
27-
resolver = tf.distribute.cluster_resolver.TPUClusterResolver.connect(
28-
tpu=tpu_name,
29-
)
30-
return tf.distribute.TPUStrategy(resolver)
31-
32-
33-
@pytest.fixture(scope="class")
34-
def tpu_test_class(request, tpu_strategy):
35-
# set a class attribute on the invoking test context
36-
request.cls.tpu_strategy = tpu_strategy
37-
38-
3922
def pytest_addoption(parser):
4023
parser.addoption(
4124
"--run_large",
@@ -49,18 +32,6 @@ def pytest_addoption(parser):
4932
default=False,
5033
help="run extra_large tests",
5134
)
52-
parser.addoption(
53-
"--run_tpu",
54-
action="store_true",
55-
default=False,
56-
help="run tpu tests",
57-
)
58-
parser.addoption(
59-
"--mixed_precision",
60-
action="store_true",
61-
default=False,
62-
help="run with mixed precision",
63-
)
6435
parser.addoption(
6536
"--docstring_module",
6637
action="store",
@@ -70,18 +41,13 @@ def pytest_addoption(parser):
7041

7142

7243
def pytest_configure(config):
73-
if config.getoption("--mixed_precision"):
74-
keras.mixed_precision.set_global_policy("mixed_float16")
75-
config.addinivalue_line(
76-
"markers", "large: mark test as being slow or requiring a network"
77-
)
7844
config.addinivalue_line(
7945
"markers",
80-
"extra_large: mark test as being too large to run continuously",
46+
"large: mark test as being slow or requiring a network",
8147
)
8248
config.addinivalue_line(
8349
"markers",
84-
"tpu: mark test as tpu test",
50+
"extra_large: mark test as being too large to run continuously",
8551
)
8652
config.addinivalue_line(
8753
"markers",
@@ -93,7 +59,6 @@ def pytest_collection_modifyitems(config, items):
9359
run_extra_large_tests = config.getoption("--run_extra_large")
9460
# Run large tests for --run_extra_large or --run_large.
9561
run_large_tests = config.getoption("--run_large") or run_extra_large_tests
96-
run_tpu = config.getoption("--run_tpu")
9762

9863
# Messages to annotate skipped tests with.
9964
skip_large = pytest.mark.skipif(
@@ -104,10 +69,6 @@ def pytest_collection_modifyitems(config, items):
10469
not run_extra_large_tests,
10570
reason="need --run_extra_large option to run",
10671
)
107-
skip_tpu = pytest.mark.skipif(
108-
not run_tpu,
109-
reason="need --run_tpu option to run",
110-
)
11172
skip_tf_only = pytest.mark.skipif(
11273
not backend_config.backend() == "tensorflow",
11374
reason="tests only run on tf backend",
@@ -117,8 +78,6 @@ def pytest_collection_modifyitems(config, items):
11778
item.add_marker(skip_large)
11879
if "extra_large" in item.keywords:
11980
item.add_marker(skip_extra_large)
120-
if "tpu" in item.keywords:
121-
item.add_marker(skip_tpu)
12281
if "tf_only" in item.keywords:
12382
item.add_marker(skip_tf_only)
12483

‎keras_nlp/models/albert/albert_backbone_test.py‎

Lines changed: 0 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -101,34 +101,3 @@ def test_saved_model(self):
101101
self.assertAllClose(
102102
model_output["pooled_output"], restored_output["pooled_output"]
103103
)
104-
105-
106-
@pytest.mark.tpu
107-
@pytest.mark.usefixtures("tpu_test_class")
108-
class AlbertBackboneTPUTest(TestCase):
109-
def setUp(self):
110-
with self.tpu_strategy.scope():
111-
self.backbone = AlbertBackbone(
112-
vocabulary_size=10,
113-
num_layers=2,
114-
num_heads=2,
115-
num_groups=1,
116-
num_inner_repetitions=1,
117-
embedding_dim=16,
118-
hidden_dim=2,
119-
intermediate_dim=2,
120-
max_sequence_length=4,
121-
)
122-
123-
self.input_batch = {
124-
"token_ids": np.ones((8, 128), dtype="int32"),
125-
"segment_ids": np.ones((8, 128), dtype="int32"),
126-
"padding_mask": np.ones((8, 128), dtype="int32"),
127-
}
128-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
129-
self.input_batch
130-
).batch(2)
131-
132-
def test_predict(self):
133-
self.backbone.compile()
134-
self.backbone.predict(self.input_dataset)

‎keras_nlp/models/bart/bart_backbone_test.py‎

Lines changed: 0 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -91,31 +91,3 @@ def test_saved_model(self):
9191
model_output["decoder_sequence_output"],
9292
restored_output["decoder_sequence_output"],
9393
)
94-
95-
96-
@pytest.mark.tpu
97-
@pytest.mark.usefixtures("tpu_test_class")
98-
class BartBackboneTPUTest(TestCase):
99-
def setUp(self):
100-
with self.tpu_strategy.scope():
101-
self.backbone = BartBackbone(
102-
vocabulary_size=1000,
103-
num_layers=2,
104-
num_heads=2,
105-
hidden_dim=64,
106-
intermediate_dim=128,
107-
max_sequence_length=128,
108-
)
109-
self.input_batch = {
110-
"encoder_token_ids": np.ones((8, 128), dtype="int32"),
111-
"encoder_padding_mask": np.ones((8, 128), dtype="int32"),
112-
"decoder_token_ids": np.ones((8, 128), dtype="int32"),
113-
"decoder_padding_mask": np.ones((8, 128), dtype="int32"),
114-
}
115-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
116-
self.input_batch
117-
).batch(2)
118-
119-
def test_predict(self):
120-
self.backbone.compile()
121-
self.backbone.predict(self.input_dataset)

‎keras_nlp/models/bert/bert_backbone_test.py‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -85,30 +85,3 @@ def test_saved_model(self):
8585
# Check that output matches.
8686
restored_output = restored_model(self.input_batch)
8787
self.assertAllClose(model_output, restored_output)
88-
89-
90-
@pytest.mark.tpu
91-
@pytest.mark.usefixtures("tpu_test_class")
92-
class BertBackboneTPUTest(TestCase):
93-
def setUp(self):
94-
with self.tpu_strategy.scope():
95-
self.backbone = BertBackbone(
96-
vocabulary_size=1000,
97-
num_layers=2,
98-
num_heads=2,
99-
hidden_dim=64,
100-
intermediate_dim=128,
101-
max_sequence_length=128,
102-
)
103-
self.input_batch = {
104-
"token_ids": np.ones((8, 128), dtype="int32"),
105-
"segment_ids": np.ones((8, 128), dtype="int32"),
106-
"padding_mask": np.ones((8, 128), dtype="int32"),
107-
}
108-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
109-
self.input_batch
110-
).batch(2)
111-
112-
def test_predict(self):
113-
self.backbone.compile()
114-
self.backbone.predict(self.input_dataset)

‎keras_nlp/models/deberta_v3/deberta_v3_backbone_test.py‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -90,30 +90,3 @@ def test_saved_model(self):
9090
# Check that output matches.
9191
restored_output = restored_model(self.input_batch)
9292
self.assertAllClose(model_output, restored_output)
93-
94-
95-
@pytest.mark.tpu
96-
@pytest.mark.usefixtures("tpu_test_class")
97-
class DebertaV3BackboneTPUTest(TestCase):
98-
def setUp(self):
99-
with self.tpu_strategy.scope():
100-
self.backbone = DebertaV3Backbone(
101-
vocabulary_size=10,
102-
num_layers=2,
103-
num_heads=2,
104-
hidden_dim=2,
105-
intermediate_dim=4,
106-
max_sequence_length=5,
107-
bucket_size=2,
108-
)
109-
self.input_batch = {
110-
"token_ids": np.ones((2, 5), dtype="int32"),
111-
"padding_mask": np.ones((2, 5), dtype="int32"),
112-
}
113-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
114-
self.input_batch
115-
).batch(2)
116-
117-
def test_predict(self):
118-
self.backbone.compile()
119-
self.backbone.predict(self.input_dataset)

‎keras_nlp/models/distil_bert/distil_bert_backbone_test.py‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -83,29 +83,3 @@ def test_saved_model(self):
8383
# Check that output matches.
8484
restored_output = restored_model(self.input_batch)
8585
self.assertAllClose(model_output, restored_output)
86-
87-
88-
@pytest.mark.tpu
89-
@pytest.mark.usefixtures("tpu_test_class")
90-
class DistilBertTPUTest(TestCase):
91-
def setUp(self):
92-
with self.tpu_strategy.scope():
93-
self.backbone = DistilBertBackbone(
94-
vocabulary_size=1000,
95-
num_layers=2,
96-
num_heads=2,
97-
hidden_dim=64,
98-
intermediate_dim=128,
99-
max_sequence_length=128,
100-
)
101-
self.input_batch = {
102-
"token_ids": np.ones((8, 128), dtype="int32"),
103-
"padding_mask": np.ones((8, 128), dtype="int32"),
104-
}
105-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
106-
self.input_batch
107-
).batch(2)
108-
109-
def test_predict(self):
110-
self.backbone.compile()
111-
self.backbone.predict(self.input_dataset)

‎keras_nlp/models/f_net/f_net_backbone_test.py‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -81,29 +81,3 @@ def test_saved_model(self):
8181
self.assertAllClose(
8282
model_output["pooled_output"], restored_output["pooled_output"]
8383
)
84-
85-
86-
@pytest.mark.tpu
87-
@pytest.mark.usefixtures("tpu_test_class")
88-
class FNetBackboneTPUTest(TestCase):
89-
def setUp(self):
90-
with self.tpu_strategy.scope():
91-
self.backbone = FNetBackbone(
92-
vocabulary_size=100,
93-
num_layers=2,
94-
hidden_dim=16,
95-
intermediate_dim=32,
96-
max_sequence_length=128,
97-
num_segments=4,
98-
)
99-
self.input_batch = {
100-
"token_ids": np.ones((8, 128), dtype="int32"),
101-
"segment_ids": np.ones((8, 128), dtype="int32"),
102-
}
103-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
104-
self.input_batch
105-
).batch(2)
106-
107-
def test_predict(self):
108-
self.backbone.compile()
109-
self.backbone.predict(self.input_dataset)

‎keras_nlp/models/gpt2/gpt2_backbone_test.py‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -101,29 +101,3 @@ def test_create_layout_map(self):
101101
# bridge elsewhere and must disable. See
102102
# https://github.com/keras-team/keras-nlp/issues/1001
103103
tf.config.experimental.disable_mlir_bridge()
104-
105-
106-
@pytest.mark.tpu
107-
@pytest.mark.usefixtures("tpu_test_class")
108-
class GPT2BackboneTPUTest(TestCase):
109-
def setUp(self):
110-
with self.tpu_strategy.scope():
111-
self.model = GPT2Backbone(
112-
vocabulary_size=10,
113-
num_layers=2,
114-
num_heads=2,
115-
hidden_dim=2,
116-
intermediate_dim=4,
117-
max_sequence_length=5,
118-
)
119-
self.input_batch = {
120-
"token_ids": np.ones((2, 5), dtype="int32"),
121-
"padding_mask": np.ones((2, 5), dtype="int32"),
122-
}
123-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
124-
self.input_batch
125-
).batch(2)
126-
127-
def test_predict(self):
128-
self.model.compile()
129-
self.model.predict(self.input_dataset)

‎keras_nlp/models/gpt_neo_x/gpt_neo_x_backbone_test.py‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -83,29 +83,3 @@ def test_saved_model(self):
8383
# Check that output matches.
8484
restored_output = restored_model(self.input_batch)
8585
self.assertAllClose(model_output, restored_output)
86-
87-
88-
@pytest.mark.tpu
89-
@pytest.mark.usefixtures("tpu_test_class")
90-
class GPTNeoXBackboneTPUTest(TestCase):
91-
def setUp(self):
92-
with self.tpu_strategy.scope():
93-
GPTNeoXBackbone(
94-
vocabulary_size=10,
95-
num_layers=4,
96-
num_heads=4,
97-
hidden_dim=64,
98-
intermediate_dim=64,
99-
max_sequence_length=10,
100-
)
101-
self.input_batch = {
102-
"token_ids": np.ones((2, 5), dtype="int32"),
103-
"padding_mask": np.ones((2, 5), dtype="int32"),
104-
}
105-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
106-
self.input_batch
107-
).batch(2)
108-
109-
def test_predict(self):
110-
self.model.compile()
111-
self.model.predict(self.input_dataset)

‎keras_nlp/models/opt/opt_backbone_test.py‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -101,29 +101,3 @@ def test_create_layout_map(self):
101101
# bridge elsewhere and must disable. See
102102
# https://github.com/keras-team/keras-nlp/issues/1001
103103
tf.config.experimental.disable_mlir_bridge()
104-
105-
106-
@pytest.mark.tpu
107-
@pytest.mark.usefixtures("tpu_test_class")
108-
class OPTBackboneTPUTest(TestCase):
109-
def setUp(self):
110-
with self.tpu_strategy.scope():
111-
self.backbone = OPTBackbone(
112-
vocabulary_size=1000,
113-
num_layers=2,
114-
num_heads=2,
115-
hidden_dim=32,
116-
intermediate_dim=128,
117-
max_sequence_length=128,
118-
)
119-
self.input_batch = {
120-
"token_ids": np.ones((8, 128), dtype="int32"),
121-
"padding_mask": np.ones((8, 128), dtype="int32"),
122-
}
123-
self.input_dataset = tf.data.Dataset.from_tensor_slices(
124-
self.input_batch
125-
).batch(2)
126-
127-
def test_predict(self):
128-
self.backbone.compile()
129-
self.backbone.predict(self.input_dataset)

0 commit comments

Comments
 (0)