[BUG][all_reduce] INVALID_ARGUMENT: You must feed a value for placeholder tensor

[edwardyehuang]

  File "/usr/local/lib/python3.11/dist-packages/keras/src/utils/traceback_utils.py", line 118, in error_handler

  File "/usr/local/lib/python3.11/dist-packages/keras/src/backend/tensorflow/trainer.py", line 323, in fit

  File "/usr/local/lib/python3.11/dist-packages/keras/src/backend/tensorflow/trainer.py", line 117, in one_step_on_iterator

3 root error(s) found.
  (0) INVALID_ARGUMENT:  You must feed a value for placeholder tensor 'StatefulPartitionedCall/cond/else/_1258/cond/Placeholder_831' with dtype int32
	 [[{{node cond/Placeholder_831}}]]
	 [[StatefulPartitionedCall/cond/then/_1257/cond/Placeholder_3/_2616]]
	 [[StatefulPartitionedCall/cond/pivot_f/_1260/_985]]
  (1) INVALID_ARGUMENT:  You must feed a value for placeholder tensor 'StatefulPartitionedCall/cond/else/_1258/cond/Placeholder_831' with dtype int32
	 [[{{node cond/Placeholder_831}}]]
	 [[StatefulPartitionedCall/cond/then/_1257/cond/Placeholder_3/_2616]]
  (2) INVALID_ARGUMENT:  You must feed a value for placeholder tensor 'StatefulPartitionedCall/cond/else/_1258/cond/Placeholder_831' with dtype int32
	 [[{{node cond/Placeholder_831}}]]

Well, currently, I have no idea how to debug this issue because there is no useful information.

[edwardyehuang]

The current situation is:

1. tf.distribute.ReplicaContext.all_reduce will directly fail in graph mode of keras 3 + TensorFlow 2.15
2. Any other case (e.g. eager with keras 3.0 or graph with keras 2.15) with the same TensorFlow environment is ok.

[SuryanarayanaY]

Hi @edwardyehuang ,

If you want to use Keras with TF2.x backend, you need to install tf-keras package using pip install tf-keras and set the environment variable TF_USE_LEGACY_KERAS=1.

Also could you please confirm the behaviour with TF2.16v which is compatible with Keras3?

Thanks!

[edwardyehuang]

Hi @edwardyehuang ,

If you want to use Keras with TF2.x backend, you need to install tf-keras package using pip install tf-keras and set the environment variable TF_USE_LEGACY_KERAS=1.

Also could you please confirm the behaviour with TF2.16v which is compatible with Keras3?

Thanks!

Thanks for the reply. I am discussing solving the bug of keras 3, and my code is working fine in keras 2. This issue also happened in TensorFlow 2.16

[edwardyehuang]

Update: After removing replica_ctx.all_reduce, the error is gone. So now the problem is, why replica_ctx.all_reduce works fine in Keras 2.x but has errors in Keras3.x ? Still investigating

@SuryanarayanaY Note that, this is a bug, not a support

[SuryanarayanaY]

Hi @edwardyehuang ,If possible could we have a reproducible code snippet for this?

If this happened with TF2.16v with keras3 then it may need investigation. If you working on this I will leave it as it is for now.

[edwardyehuang]

@fchollet @SuryanarayanaY

A reproducible code snippet is presented below. Make sure you test it on at least 2 GPUs (and set batch_size >= num_gpus).

Remove either conv or all_reduce, and the problem will disappear.

The code below is working fine with Keras 2.15, it only has the error in Keras 3

It looks like this is caused by the all_reduce and tf.placeholder (for building the model in keras.Layer.__call__l)

However, given my limited knowledge and time, I'm unable to provide a quick fix. Thus, I need help

import keras
import tensorflow as tf

BATCH_SIZE = 4

tf.get_logger().setLevel('INFO')
strategy = tf.distribute.MirroredStrategy()

# Make model ##########################################################################
with strategy.scope():


    class SimpleModel (keras.Model):

        def __init__(self, name=None):
            super().__init__(name=name)


        def build(self, input_shape):

            self.l = keras.layers.Conv2D(3, (1, 1), padding='same')

            super().build(input_shape)


        def call (self, inputs, training=False):

            x = inputs

            if training:
                x = tf.distribute.get_replica_context().all_reduce(tf.distribute.ReduceOp.SUM, x)
                
            x = self.l(x)

            return x


    m = SimpleModel()
    m.compile(
        optimizer=keras.optimizers.SGD(learning_rate=1e-3),
        loss=keras.losses.MeanSquaredError(),
    )


# Make dataset ##########################################################################
def simple_data_generator(num_samples=-1, size=(17, 17)):

    counter = 0

    while True:
        
        random_data = tf.random.uniform(
            shape=size, 
            minval=-2, 
            maxval=2, 
            dtype=tf.float32
        )
        
        yield random_data, random_data

        counter += 1

        if num_samples > 0 and counter >= num_samples:
            break


io_shape = (17, 17, 3)

train_dataset = tf.data.Dataset.from_generator(
    simple_data_generator,
    args=(-1, io_shape),
    output_signature=(
        tf.TensorSpec(shape=io_shape, dtype=tf.float32),
        tf.TensorSpec(shape=io_shape, dtype=tf.float32),
    )
)
train_dataset = train_dataset.batch(BATCH_SIZE)

options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
train_dataset = train_dataset.with_options(options)

train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)


# Train #################################################################################

m.fit(train_dataset, epochs=1, verbose=1, steps_per_epoch=1000)

Note that I found Keras 3 results 2 more placeholders in graph.capture than Keras 2.15, still under investigation.