canonical / canonical/notebook-operators

Tensorflow and Pytorch Images are not compatible with Blackwell GPUs

Open
#538 1 comment 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
6
Forks
12
Avg merge
3d 5h
Merged PRs (30d)
12

Description

### Bug Description

Default jupyter images provided with the charmed kubeflow is not compatible in Blackwell GPUs.
I have tried with both pytorch `jupyter-pytorch-cuda-full:v1.10.0-rc.1` and tensorflow `jupyter-tensorflow-cuda-full:v1.10.0-rc.1` images as explained below.

From tensorflow notebook:
```
W0000 00:00:1770648523.797505 1654 gpu_device.cc:2342] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
```

From pytorch notebook:
```
NVIDIA B200 with CUDA capability sm_100 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_50 sm_60 sm_70 sm_75 sm_80 sm_86 sm_90.
If you want to use the NVIDIA B200 GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/
```

### To Reproduce

I have a kubeflow deployment which gives me below list of default jupyter-ui notebook images.
```
- charmedkubeflow/jupyter-scipy:1.10.0-0be57a5
- charmedkubeflow/jupyter-pytorch-full:v1.10.0-ef1fc67
- kubeflownotebookswg/jupyter-pytorch-cuda-full:v1.10.0-rc.1
- kubeflownotebookswg/jupyter-pytorch-gaudi-full:v1.10.0-rc.1
- charmedkubeflow/jupyter-tensorflow-full:1.10.0-448e378
- kubeflownotebookswg/jupyter-tensorflow-cuda-full:v1.10.0-rc.1
```

**Scenario 1:**
Notebooks with jupyter-tensorflow-cuda-full:v1.10.0-rc.1 image on H200 and B200.
Below is a snippet of the code which returns error in B200.
```
import tensorflow as tf
from tensorflow.keras import layers, models

gpus = tf.config.list_physical_devices('GPU')
if gpus:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
mnist = tf.keras.datasets.mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
print("Data loaded successfully!")
```
As per below trace of logs, GPU device registration is skipped. With that, it will fail the exact same notebook which is successfully executed in H200.

```
2026-02-09 14:48:38.531363: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2026-02-09 14:48:38.622841: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2026-02-09 14:48:40.659586: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1770648523.797505 1654 gpu_device.cc:2342] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://www.tensorflow.org/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
```

**Scenario 2:**
Notebooks with jupyter-pytorch-cuda-full:v1.10.0-rc.1 image on H200 and B200.
Tested with below simple code.
```
import torch
print(torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.cuda.device(0))
print(torch.cuda.get_device_name(0))
```

This is the output in the B200 notebook.

```
True
2

NVIDIA B200

/opt/conda/lib/python3.11/site-packages/torch/cuda/__init__.py:235: UserWarning:
NVIDIA B200 with CUDA capability sm_100 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_50 sm_60 sm_70 sm_75 sm_80 sm_86 sm_90.
If you want to use the NVIDIA B200 GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

warnings.warn(
```

This is the output in H200.
```
True
2

NVIDIA H200
```

### Environment

Canonical K8s - 1.32 LTS
Charmed Kubeflow - 1.10
gpu-operator - v25.10.1
network-operator - 25.10.0
Nvidia driver.version - 580.105.08

### Relevant Log Output

```shell
I have added the relevant logs while explaining the steps above in "To Reproduce" section.
```

### Additional Context

_No response_

Contributor guide

Open the contributing guide

Research direction

No source files or tests are named. Start by reproducing the TensorFlow and PyTorch notebook snippets with the listed CUDA images on H200 and B200, then trace how those images are built and their framework dependencies are selected. Done means the default images recognize B200 GPUs with sm_100 while continuing to work on H200.

Written by the indexing model from the issue text.

Assessment

Tech stack
kubernetes, python, pytorch, tensorflow
Domain
devops, infrastructure, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.