Tensorflow upgrade on Databricks Runtime ML for GPU enabled clusters

Question

Is it possible to upgrade Tensorflow version from 2.8 (which is preinstalled on Databricks Runtime 10.4 M for GPU clusters) to 2.13 without changing the Databricks Runtime. In other words, can we have Tensorflow 2.13 running on Databricks Runtime 10.4 for ML? If so, how can I adjust the version of Cuda within Databricks Runtime 10.4 for ML?

Answer 1

@Shahrzad Amoozegar - Thanks for the question and using MS Q&A platform.

Based on the documentation, Databricks Runtime 10.4 for ML comes with pre-installed TensorFlow 2.8. However, it is possible to upgrade TensorFlow to version 2.13 without changing the Databricks Runtime. You can install TensorFlow 2.13 using pip install command.

Here is an example command to install TensorFlow 2.13 on Databricks Runtime 10.4 for ML:

%sh 
pip install tensorflow-gpu==2.13

Regarding the version of CUDA within Databricks Runtime 10.4 for ML, it is not possible to adjust the version of CUDA within the runtime. The version of CUDA is pre-installed and cannot be changed.

For additonal details, you may checkout the answer provided on your thread on Databricks community: https://community.databricks.com/t5/machine-learning/upgrading-cudnn-on-databricks-notebook/td-p/61167

Hope this helps.