Fix #41630: include max_seq_length in cudnn descriptor cache key

[lissyx]

No description provided.

[googlebot]

We found a Contributor License Agreement for you (the sender of this pull request), but were unable to find agreements for all the commit author(s) or Co-authors. If you authored these, maybe you used a different email address in the git commits than was used to sign the CLA (login here to double check)? If these were authored by someone else, then they will need to sign a CLA as well, and confirm that they're okay with these being contributed to Google.
In order to pass this check, please resolve this problem and then comment @googlebot I fixed it.. If the bot doesn't comment, it means it doesn't think anything has changed.

ℹ️ Googlers: Go here for more info.

[googlebot]

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

[sanjoy]

Looks fine to me, but I'll wait for @kaixih to LGTM also.

[sanjoy]

Actually, can you please also add a test case?

[lissyx]

Actually, can you please also add a test case?

I can't find any existing test file supporting what is in tensorflow/core/kernels/cudnn_rnn_ops.cc, and the coverage of CONTRIBUTING.md for tests is quite scarce, it just covers how to run them. Would you have some pointers on where I should add test case?

[lissyx]

Actually, can you please also add a test case?

I can't find any existing test file supporting what is in tensorflow/core/kernels/cudnn_rnn_ops.cc, and the coverage of CONTRIBUTING.md for tests is quite scarce, it just covers how to run them. Would you have some pointers on where I should add test case?

Managed to add something as a new tf_cuda_cc_test, so it required some new headers and moving the two structs into that. It gets more invasise of a patch than I hoped for at the begining, but locally tests runs well and catches the regression.

[kaixih]

I am wondering if it is possible to limit the test into the python test, like in the lstm_v2_test.py. All we need to do is to design a simple training as in the issue 41630 that uses [1, 2048, 2048, 1, 74, 2, 2048] in the first step and store it to the cache but [1, 2048, 2048, 1, 75, 2, 2048] in the second step. If this can be successfully trained on GPU, we can view it as passed. How do you think?

[lissyx]

I am wondering if it is possible to limit the test into the python test, like in the lstm_v2_test.py. All we need to do is to design a simple training as in the issue 41630 that uses [1, 2048, 2048, 1, 74, 2, 2048] in the first step and store it to the cache but [1, 2048, 2048, 1, 75, 2, 2048] in the second step. If this can be successfully trained on GPU, we can view it as passed. How do you think?

I can look into that, it would sounds like a better option than the current test I wrote.

[lissyx]

lstm_v2_test.py.

Do you know if there's any magic to run python tests required? I've followed CONTRIBUTING.md section about running unit tests, but

bazel test -s --verbose_failures --config=noaws --config=nogcp --config=nohdfs --config=nonccl -c opt --config=cuda -k //tensorflow/python/keras/layers
[...]
INFO: Analyzed target //tensorflow/python/keras/layers:layers (0 packages loaded, 0 targets configured).
INFO: Found 1 target and 0 test targets...
Target //tensorflow/python/keras/layers:layers up-to-date (nothing to build)
INFO: Elapsed time: 0.336s, Critical Path: 0.00s
INFO: 0 processes.
INFO: Build completed successfully, 1 total action
INFO: Build completed successfully, 1 total action

[kaixih]

Have you tried to simply run python lstm_v2_test.py?

[lissyx]

Have you tried to simply run python lstm_v2_test.py?

Even simpler: //tensorflow/python/keras/layers:lstm_v2_test bazel target, I assumed passing the directory would run all tests there.

[lissyx]

Making progresses:

  @test_util.run_gpu_only
  def test_lstm_variable_sequence_length(self):
    if test.is_built_with_rocm():
      self.skipTest('Skipping the test as ROCm MIOpen does not '
                    'support padded input yet.')
    input_shape = 2048
    rnn_state_size = 2048
    batch = 2

    timestep = 1

    (x_train, y_train), _ = testing_utils.get_test_data(
        train_samples=batch,
        test_samples=0,
        input_shape=(timestep, input_shape),
        num_classes=rnn_state_size,
        random_seed=random_seed.DEFAULT_GRAPH_SEED)

    timestep = 1

    (x2_train, y2_train), _ = testing_utils.get_test_data(
        train_samples=batch,
        test_samples=0,
        input_shape=(timestep, input_shape),
        num_classes=rnn_state_size,
        random_seed=random_seed.DEFAULT_GRAPH_SEED)

    inputs = keras.layers.Input(
        shape=[timestep, input_shape], dtype=dtypes.float32)
    masked_input = keras.layers.Masking()(inputs)

    with test_util.device(use_gpu=True):
      cudnn_layer = rnn.LSTM(rnn_state_size,
                             activation='tanh',
                             recurrent_activation='sigmoid',
                             recurrent_dropout=0,
                             unroll=False,
                             use_bias=True)
      cudnn_model = keras.models.Model(inputs, cudnn_layer(masked_input))

    cudnn_model.compile('rmsprop', 'mse')
    cudnn_model.fit(x_train, y_train)
    cudnn_model.fit(x2_train, y2_train)

With that, I can assert that it is going through the layers and reaches IsCompatibleWith:

IsCompatibleWith -> [num_layers, input_size, num_units, dir_count, max_seq_length, batch_size, cell_num_units]: [1, 2048, 2048, 1, 1, 2, 2048] 
IsCompatibleWith -> [num_layers, input_size, num_units, dir_count, max_seq_length, batch_size, cell_num_units]: [1, 2048, 2048, 1, 1, 2, 2048] 
IsCompatibleWith -> [num_layers, input_size, num_units, dir_count, max_seq_length, batch_size, cell_num_units]: [1, 2048, 2048, 1, 1, 2, 2048] 
IsCompatibleWith -> [num_layers, input_size, num_units, dir_count, max_seq_length, batch_size, cell_num_units]: [1, 2048, 2048, 1, 1, 2, 2048] 
IsCompatibleWith -> [num_layers, input_size, num_units, dir_count, max_seq_length, batch_size, cell_num_units]: [1, 2048, 2048, 1, 1, 2, 2048] 
IsCompatibleWith -> [num_layers, input_size, num_units, dir_count, max_seq_length, batch_size, cell_num_units]: [1, 2048, 2048, 1, 1, 2, 2048] 

However, changing timestep to be 74 on x_train and 75 on x2_train seems to bypass the GPU implementation, because I can't see any IsCompatibleWith output in the log but I see:

WARNING:tensorflow:Model was constructed with shape (None, 75, 2048) for input Tensor("input_1:0", shape=(None, 75, 2048), dtype=float32), but it was called on an input with incompatible shape (None, 74, 2048).
W0731 11:00:01.229415 139681924948352 functional.py:583] Model was constructed with shape (None, 75, 2048) for input Tensor("input_1:0", shape=(None, 75, 2048), dtype=float32), but it was called on an input with incompatible shape (None, 74, 2048).

Even directly calling CuDNNLSTM from tensorflow/python/keras/layers/cudnn_recurrent.py I'm blocked at Keras level when adding the 74/75 dimensional difference.

Forcing eager execution using @test_util.run_in_graph_and_eager_modes,

ValueError: Error when checking input: expected input_1 to have shape (75, 64) but got array with shape (74, 64)

And to the best of my knowledge, with current TensorFlow, since there is no more tensorflow.contrib Keras is the only way to build LSTM layers, right? So I'm unsure how to solve this :/

[lissyx]

Can we hope for a 1.15.4 to include that fix ?

[sanjoy]

Can we hope for a 1.15.4 to include that fix ?

CC @goldiegadde

[lissyx]

Can we hope for a 1.15.4 to include that fix ?

CC @goldiegadde

Gentle ping, can we hope for 1.15.4 or should we direct people hitting this issue to use the magic reset state env variable?

[lissyx]

@sanjoy @kaixih @goldiegadde Gentle ping? There has been several commit pushed to r1.15, so it would really be nice to include this one if you are preparing a 1.15.4. I can send the PR if you'd like.

[goldiegadde]

@lissyx Sorry for the delayed response, can you please open a PR against the 1.15 branch. cc @mihaimaruseac as well.

[mihaimaruseac]

Apologies, I missed this myself too. Yes, please, let's open a PR against the branch

[lissyx]

@lissyx Sorry for the delayed response, can you please open a PR against the 1.15 branch. cc @mihaimaruseac as well.

#42634