Flax T5 model - code typo during AutoRegressive decoding?

[giganttheo]

System Info

• transformers version: 4.33.3
• Platform: Linux-5.15.120+-x86_64-with-glibc2.35
• Python version: 3.10.12
• Huggingface_hub version: 0.17.3
• Safetensors version: 0.3.3
• Accelerate version: not installed
• Accelerate config: not found
• PyTorch version (GPU?): 2.0.1+cu118 (False)
• Tensorflow version (GPU?): 2.13.0 (False)
• Flax version (CPU?/GPU?/TPU?): 0.7.4 (cpu)
• Jax version: 0.4.10
• JaxLib version: 0.4.10
• Using GPU in script?: No
• Using distributed or parallel set-up in script?: No

Who can help?

@sanchit-gandhi

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

The discussed file is transformers/src/transformers/models/t5/modeling_flax_t5.py at line 408

https://github.com/huggingface/transformers/blob/2aef9a96011133f6b399b598fd69cfeca936eb37/src/transformers/models/t5/modeling_flax_t5.py#L408C1-L409C1

Expected behavior

Hi,

During autoregressive decoding, keys and values are computed one token at a time and cache is used to recover the keys and values from previous calls.

In the _concatenate_to_cache method of the Attention module, an attention mask is computed in order for the new query to only attend to the previous key positions and not the remaining zero elements. This is what is explained in the comments in this function.

However, the new attention mask is not used afterwards, because its name is attention_attention_mask and not attention_mask which is the one being used in every other line.

From my understanding, this is likely a typo and I am not sure how it changes the behavior of the model, if at all.

[sanchit-gandhi]

Hey @giganttheo - this indeed looks like a bug! The attention mask should be overridden, as in Flax GPT Neo:

transformers/src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py

Line 219 in 2aef9a9

key, value, attention_mask = self._concatenate_to_cache(key, value, query, attention_mask)

Would you like to open a PR to correct this? We can check whether we get different outputs after the fix by running the slow tests:

transformers/tests/models/t5/test_modeling_flax_t5.py

Line 761 in 57f44dc

class FlaxT5ModelIntegrationTests(unittest.TestCase):

Which you can do so with:

RUN_SLOW=1 pytest -sv tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests

[giganttheo]

Thank you for your response, I ran the slow tests on google colab with a T4 GPU and the outputs give:

=============================== warnings summary ===============================
../../usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py:1373
  /usr/local/lib/python3.10/dist-packages/_pytest/config/__init__.py:1373: PytestConfigWarning: Unknown config option: doctest_glob
  
    self._warn_or_fail_if_strict(f"Unknown config option: {key}\n")

tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests::test_small_generation
tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests::test_small_generation_bfloat16
  /content/transformers/src/transformers/generation/flax_utils.py:322: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use and modify the model generation configuration (see https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )
    warnings.warn(

tests/models/t5/test_modeling_flax_t5.py::FlaxT5ModelIntegrationTests::test_summarization
  /content/transformers/src/transformers/models/t5/tokenization_t5.py:238: FutureWarning: This tokenizer was incorrectly instantiated with a model max length of 512 which will be corrected in Transformers v5.
  For now, this behavior is kept to avoid breaking backwards compatibility when padding/encoding with `truncation is True`.
  - Be aware that you SHOULD NOT rely on t5-base automatically truncating your input to 512 when padding/encoding.
  - If you want to encode/pad to sequences longer than 512 you can either instantiate this tokenizer with `model_max_length` or pass `max_length` when encoding/padding.
  - To avoid this warning, please instantiate this tokenizer with `model_max_length` set to your preferred value.
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================== 6 passed, 4 warnings in 242.20s (0:04:02) ===================

I am not very familiar with the contribution process so I am not sure what the next steps are for this, I tried following the huggingface contribution guide (https://huggingface.co/transformers/v4.2.2/contributing.html ), do you have any other resources to learn the basics of contribution guides for open source projects and transformers? I might have other contributions to do in the near future :)

[sanchit-gandhi]

That looks good - all the slow tests have passed 👍 We can see this from the last line on the print out:

================== 6 passed, 4 warnings in 242.20s (0:04:02) ===================

That means we get correct results with the changes. The next steps would be opening a pull request (PR) with the correction. Feel free to tag me on the PR and I'll get you a review asap. Once reviewed, we can merge into main. I would mention on the PR that you made the fix and then subsequently ran the slow tests and they passed - this will inform the reviewer that the correct T5 behaviour is maintained