Running ops in specific device-(id) with global distribution setting (multi-gpu/tpu)

[innat]

1. For a specific reason, I need to run some operations into specific device while keeping other operations runable with global distribution. In torch, we can do .to(device_id) to place operations to specific device. Is it possible in keras yet? Maybe using keras.distribution.initialize, no?

2. Also, current distribution api only support jax. Any update on supporting tensorflow and torch?

3. Is there any known bug or limitation in the Keras distributed API when using the JAX backend? I realize this might sound a bit vague, but in one of my setups, training works fine with a single GPU. However, when I switch to multi-GPU or TPU training with jax, things go wrong. Interestingly, using TensorFlow's distribution strategy for multi-GPU training runs without issues. In short, training with JAX across multiple devices seems to corrupt the model weights. While reproducing the issue would take some time, I wanted to check if there are any known issues related to this.

devices = jax.devices("tpu")
devices
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),
 TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),
 TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),
 TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),
 TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),
 TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),
 TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),
 TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]

data_parallel = keras.distribution.DataParallel(devices=devices)
keras.distribution.set_distribution(data_parallel)

[hertschuh]

@innat

1. We have the device scope API that works across backends.

with keras.device("cpu"):
   # some operations here

Although I'm not sure if it's what you're looking for. Let me know.

2. keras.distribution.DataParallel and keras.distribution.ModelParallel work with the TensorFlow backend. Torch is not supported at this time.

3. Not that we know of. You will have to provide more details. What tells you the weights are corrupted? Also, is it a multi-host setup? Or a single-host multiple accelerator setup?

[innat]

@hertschuh
Thanks for the response.

1. Having this keras.device("cpu") will be usefull. But for my current case, the training is running on multiple device. But I want to perform inference on single gpu. Now, both of my training and inference code running on single notebook. So, I was looking for something like tf.device('/device:GPU:0') .
2. Its good to know that tf now supports keras DataParallel API. I got some error, I need to recheck if I overlooked something. Also, this guide may nee some update.
3. Its hard to point out the exact reason, i.e. weight damage or something else. I am trying to reproduce the issue - will report.

[hertschuh]

1. Yes, keras.device is supposed to work like tf.device. The syntax of the string you pass may depend on the backend though. What's your backend? This might work with keras.device("gpu:0"):.

[innat]

Currently I'm using jax backend.

[hertschuh]

So the syntax "gpu:0" should work:

import os
os.environ["KERAS_BACKEND"] = "jax"
import keras

keras.distribution.list_devices()

['gpu:0']

with keras.device("cpu:0"):
  x = keras.ops.convert_to_tensor(1.0)
x.device

CpuDevice(id=0)

with keras.device("gpu:0"):
  y = keras.ops.square(x)
y.device

CudaDevice(id=0)

[innat]

@hertschuh

I didn't try multi-gpu yet, will do.

In TPU-VM (kaggle),

keras backend:  tensorflow 
keras version:  3.9.2 
tensorflow version:  2.18.0 

# with tf backend
devices = keras.distribution.list_devices()
devices
['tpu:0',
 'tpu:1',
 'tpu:2',
 'tpu:3',
 'tpu:4',
 'tpu:5',
 'tpu:6',
 'tpu:7',
 'tpu_system:0']

data_parallel = keras.distribution.DataParallel(devices=devices)
keras.distribution.set_distribution(data_parallel)

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[5], line 11
      5 keras.utils.set_random_seed(101)
      7 # strategy = tf.distribute.MirroredStrategy()
      8 # devices = strategy.num_replicas_in_sync
      9 # devices
---> 11 data_parallel = keras.distribution.DataParallel(devices=devices)
     12 keras.distribution.set_distribution(data_parallel)

File /usr/local/lib/python3.10/site-packages/keras/src/distribution/distribution_lib.py:404, in DataParallel.__init__(self, device_mesh, devices, auto_shard_dataset)
    401     self._initialize_mesh_from_list_devices()
    403 # Those following attributes might get convert to public methods.
--> 404 self._num_process = distribution_lib.num_processes()
    405 self._process_id = distribution_lib.process_id()
    406 self._is_multi_process = self._num_process > 1

AttributeError: module 'keras.src.backend.tensorflow.distribution_lib' has no attribute 'num_processes'

update

I've just tried multi-gpu with keras.distribution.DataParallel, but same error as above. Seems like, it is still in progress.

[innat]

@hertschuh

Is there any known bug or limitation in the Keras distributed API when using the JAX backend? I realize this might sound a bit vague, but in one of my setups, training works fine with a single GPU. However, when I switch to multi-GPU or TPU training with jax, things go wrong. Interestingly, using TensorFlow's distribution strategy for multi-GPU training runs without issues. In short, training with JAX across multiple devices seems to corrupt the model weights. While reproducing the issue would take some time, I wanted to check if there are any known issues related to this.

Here is the reproduced issue I talked about. I might overlook something. The thing is, when evaluating the model within a callback. It happens both multi-gpu and tpu. Look at the validation score received from callback.

os.environ["KERAS_BACKEND"] = "jax"
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.9"

import jax
import keras
from keras import layers
import tensorflow as tf

tf.__version__, keras.version() # ('2.19.0', '3.9.2')

from keras import mixed_precision

policy = mixed_precision.Policy('mixed_bfloat16')  # Best for TPUs
mixed_precision.set_global_policy(policy)

devices = keras.distribution.list_devices()
devices
['tpu:0', 'tpu:1', 'tpu:2', 'tpu:3', 'tpu:4', 'tpu:5', 'tpu:6', 'tpu:7']

data_parallel = keras.distribution.DataParallel(devices=devices)
keras.distribution.set_distribution(data_parallel)

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

x_train = np.expand_dims(x_train, axis=-1)
x_train = np.repeat(x_train, 3, axis=-1)
x_train = x_train.astype('float32') / 255
x_train = tf.image.resize(x_train, [100,100])
y_train = tf.one_hot(y_train , depth=10) 

x_test = np.expand_dims(x_test, axis=-1)
x_test = np.repeat(x_test, 3, axis=-1)
x_test = x_test.astype('float32') / 255
x_test = tf.image.resize(x_test, [100,100])
y_test = tf.one_hot(y_test , depth=10)

x_train.shape, y_train.shape, x_test.shape, y_test.shape
(TensorShape([60000, 100, 100, 3]),
 TensorShape([60000, 10]),
 TensorShape([10000, 100, 100, 3]),
 TensorShape([10000, 10]))

class EvaluateWithMetricsCallback(keras.callbacks.Callback):
    def __init__(self, validation_data, interval=1):
        super().__init__()
        self.x_val, self.y_val = validation_data
        self.interval = interval
        # Instantiate metrics
        self.top1 = keras.metrics.TopKCategoricalAccuracy(k=1, name='val_acc_top1')
        self.loss_fn = keras.losses.CategoricalCrossentropy()

    def on_epoch_end(self, epoch, logs=None):
        if (epoch + 1) % self.interval == 0:
            y_pred = self.model.predict(
                self.x_val, batch_size=6 * len(devices), verbose=0
            )

            # Reset state before calculation
            self.top1.reset_state()

            # Update metrics
            self.top1.update_state(self.y_val, y_pred)
            val_loss = self.loss_fn(self.y_val, y_pred)

            print(
                f"\nEpoch {epoch + 1}: Validation loss = {val_loss:.4f}, "
                f"Top-1 Acc = {self.top1.result():.4f}"
            )


custom_model.fit(
    x_train,
    y_train,
    batch_size=6 * len(devices),
    epochs=10,
    verbose=1,
    validation_data=(x_test, y_test),
    validation_batch_size=6 * len(devices),
    callbacks=[
        EvaluateWithMetricsCallback(
            validation_data=(x_test, y_test), 
            interval=2
        )
    ]
)

Epoch 1/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 81s 65ms/step - acc_top1: 0.9959 - loss: 0.0134 - val_acc_top1: 0.9913 - val_loss: 0.0269
Epoch 2/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 0s 45ms/step - acc_top1: 0.9971 - loss: 0.0088
Epoch 2: Validation loss = 14.1312, Top-1 Acc = 0.0887
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 63s 50ms/step - acc_top1: 0.9971 - loss: 0.0088 - val_acc_top1: 0.9941 - val_loss: 0.0211
Epoch 3/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 60s 48ms/step - acc_top1: 0.9964 - loss: 0.0108 - val_acc_top1: 0.9948 - val_loss: 0.0187
Epoch 4/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 0s 44ms/step - acc_top1: 0.9975 - loss: 0.0090
Epoch 4: Validation loss = 14.1378, Top-1 Acc = 0.0886
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 62s 50ms/step - acc_top1: 0.9975 - loss: 0.0090 - val_acc_top1: 0.9907 - val_loss: 0.0445
Epoch 5/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 59s 47ms/step - acc_top1: 0.9978 - loss: 0.0073 - val_acc_top1: 0.9913 - val_loss: 0.0379
Epoch 6/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 0s 45ms/step - acc_top1: 0.9974 - loss: 0.0081
Epoch 6: Validation loss = 14.3566, Top-1 Acc = 0.0883
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 63s 50ms/step - acc_top1: 0.9974 - loss: 0.0081 - val_acc_top1: 0.9943 - val_loss: 0.0193
Epoch 7/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 60s 48ms/step - acc_top1: 0.9980 - loss: 0.0068 - val_acc_top1: 0.9926 - val_loss: 0.0312
Epoch 8/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 0s 44ms/step - acc_top1: 0.9982 - loss: 0.0056
Epoch 8: Validation loss = 14.3617, Top-1 Acc = 0.0886
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 62s 50ms/step - acc_top1: 0.9982 - loss: 0.0056 - val_acc_top1: 0.9942 - val_loss: 0.0186
Epoch 9/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 57s 46ms/step - acc_top1: 0.9988 - loss: 0.0047 - val_acc_top1: 0.9934 - val_loss: 0.0235
Epoch 10/10
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 0s 43ms/step - acc_top1: 0.9980 - loss: 0.0064
Epoch 10: Validation loss = 14.3836, Top-1 Acc = 0.0883
1250/1250 ━━━━━━━━━━━━━━━━━━━━ 62s 49ms/step - acc_top1: 0.9980 - loss: 0.0064 - val_acc_top1: 0.9938 - val_loss: 0.0236
<keras.src.callbacks.history.History at 0x7cc34c1cc520>

update

If not mistaken, it happens with model.predict...(and model.evaluate might have issue as well), regardless using it within callback. + .predict_on_batch doesn't work either.