Cannot use XLA accelerator Tensorflow in M1 silicon

Question

DanielOh OP

Created Jan ’23

Replies 3

Boosts 0

Views 2.7k

Participants 6

I'm trying to use XLA in TensorFlow using jit_compile=True, currently. However, the error looks like it cannot.

First of all, I tried simple code as below,

import tensorflow as tf

@tf.function(jit_compile=True)
def operation(a, b):
    return a * b

And, I got an error as below,

(base) ***@MacBook-Pro folder % /Users/***/miniforge3/envs/testtf/bin/python /Users/***/workplace/research/folder/subfolder/test/test_jit_compile.py
Metal device set to: Apple M1 Pro

systemMemory: 32.00 GB
maxCacheSize: 10.67 GB

2023-01-11 12:45:18.794439: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:306] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.
2023-01-11 12:45:18.794689: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:272] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)
2023-01-11 12:45:18.808796: W tensorflow/core/framework/op_kernel.cc:1830] OP_REQUIRES failed at xla_ops.cc:418 : NOT_FOUND: could not find registered platform with id: 0x10c7ba7e0
Traceback (most recent call last):
  File "/Users/***/workplace/research/folder/subfolder/test/test_jit_compile.py", line 25, in <module>
    print(operation(3, 4))
  File "/Users/***/miniforge3/envs/testtf/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/Users/***/miniforge3/envs/testtf/lib/python3.8/site-packages/tensorflow/python/eager/execute.py", line 52, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.NotFoundError: could not find registered platform with id: 0x10c7ba7e0 [Op:__inference_operation_5]

This seems like XLA did not register. Would you guide how to use jit_compile=True in this environment? I attached my environment as below also.

Package                 Version
----------------------- ---------
absl-py                 1.3.0
astunparse              1.6.3
cachetools              5.2.1
certifi                 2022.12.7
charset-normalizer      2.1.1
flatbuffers             23.1.4
gast                    0.4.0
google-auth             2.16.0
google-auth-oauthlib    0.4.6
google-pasta            0.2.0
grpcio                  1.42.0
h5py                    3.6.0
idna                    3.4
importlib-metadata      6.0.0
keras                   2.11.0
libclang                15.0.6.1
Markdown                3.4.1
MarkupSafe              2.1.1
numpy                   1.22.3
oauthlib                3.2.2
opt-einsum              3.3.0
packaging               23.0
pip                     22.3.1
protobuf                3.19.6
pyasn1                  0.4.8
pyasn1-modules          0.2.8
requests                2.28.1
requests-oauthlib       1.3.1
rsa                     4.9
setuptools              65.6.3
six                     1.16.0
tensorboard             2.11.0
tensorboard-data-server 0.6.1
tensorboard-plugin-wit  1.8.1
tensorflow-estimator    2.11.0
tensorflow-macos        2.11.0
tensorflow-metal        0.7.0
termcolor               2.2.0
typing_extensions       4.4.0
urllib3                 1.26.13
Werkzeug                2.2.2
wheel                   0.37.1
wrapt                   1.14.1
zipp                    3.11.0

Thanks.

Boost

Answer 1

Frameworks Engineer OP

Apple

Jan ’23

Hi @DanielOh, the Tensorflow metal backend uses Pluggable architecture to provide GPU acceleration. XLA compiler path is not supported through the Pluggable design and fails in this expected way.

0

Answer 2

jvikoren OP

Jan ’23

Okay, so how do we circumvent the error if we are not using JIT? I installed tensorflow following the instructions here and got the error shown above when running the sample code proposed on the site. The only other install in the conda env is python==3.10.4

import tensorflow as tf

cifar = tf.keras.datasets.cifar100 (x_train, y_train), (x_test, y_test) = cifar.load_data() model = tf.keras.applications.ResNet50( include_top=True, weights=None, input_shape=(32, 32, 3), classes=100,)

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True) model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"]) model.fit(x_train, y_train, epochs=5, batch_size=64)

1

Answer 3

ladams OP

Jan ’23

@jvikoren using the legacy optimizer should work in this case, see here: https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/legacy/Optimizer

import tensorflow as tf
from tensorflow.keras.optimizers.legacy import Adam

cifar = tf.keras.datasets.cifar100
(x_train, y_train), (x_test, y_test) = cifar.load_data()
model = tf.keras.applications.ResNet50(
    include_top=True,
    weights=None,
    input_shape=(32, 32, 3),
    classes=100,)

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
# Instantiate legacy Adam optimizer
model.compile(optimizer=Adam(), loss=loss_fn, metrics=["accuracy"])
model.fit(x_train, y_train, epochs=5, batch_size=64)

Works for me on M1 Air with a fresh version of miniconda.

2