[BC 4.37 -> 4.38] for Llama family, memory and speed

[ArthurZucker]

What does this PR do?

Fixes the BC issues between the two versions in term of memory consumption.
This fix is made a lot easier by all the tests, so thanks a lot @gante!

fixes #29412, fixes #29484 , fixes #29644, fixes #29651

[gante]

In general looks good to me, although I'm not 100% sure on the causal_mask *= torch.arange(target_length, device=device) > cache_position[0] line + assisted generation -- going to have a look

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[ArthurZucker]

torch script can be fixed by @fxmarty in a follow up PR + patch IMO!

[fxmarty]

@ArthurZucker that would be great if torchscript/fx tests pass &

optimum-cli export onnx --model fxmarty/tiny-llama-fast-tokenizer llama_onnx

does not break

[younesbelkada]

Thanks for the offline explanation ! this shouldn't affect FA2 as we always return attention_mask without processing it for FA2 modules in _update_causal_mask, as you explained offline causal_mask *= torch.arange(target_length, device=device) > cache_position[0] is used to mask out the cached hidden states

[gante]

@ArthurZucker the line causal_mask *= torch.arange(target_length, device=device) > cache_position[0] should become causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)

In the previous version, in a situation with 17 cached tokens and 4 new assistant tokens, the causal mask would become

(Pdb) causal_mask
tensor([[[[     0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,      0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,      0.,      0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,      0.,      0.,      0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.]]]],
       device='cuda:0', dtype=torch.float16)

i.e. not upper triangular. After the suggested change, it becomes

(Pdb) causal_mask
tensor([[[[     0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0., -65504., -65504., -65504., -65504.],
          [     0.,      0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0., -65504., -65504., -65504.],
          [     0.,      0.,      0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0., -65504., -65504.],
          [     0.,      0.,      0.,      0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0., -65504.],
          [     0.,      0.,      0.,      0.,      0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.,
               -0.,     -0.,     -0.,     -0.,     -0.,     -0.,     -0.]]]],
       device='cuda:0', dtype=torch.float16)

[ArthurZucker]

Any idea why our tests don't complain?

[gante]

Any idea why our tests don't complain?

I think we don't have hard correctness checks for assisted generation 🙈 only API checks

[poedator]

congratulations with this PR fixing many important things!
Unfortunately, it broke my heart custom 4D masks support.
try RUN_SLOW=1 python -m pytest -v ./tests/test_modeling_utils.py::Mask4DTestHard::test_partial_stacked_causal_mask

@gante and myself introduced this test recently in #29731 - please make it a part of the test suite for all things related to attention_masks and StaticCache

here is what happens (numbers based on test_partial_stacked_causal_mask after this line 2170:
in modeling_llama.py::_update_causal_mask()
the causal_mask has shape (1, 1, sequence_length, target_length) (1,1,9,12) (3 tokens in cache)
if the custom 4D mask enters with same shape, it triggers offset = 3 and then causes error when copied over causal mask.
if the custom 4D mask enters with shape (1, 1, 12, 12), offset stays at zero but then again this causes error when copied over causal mask, for mask_slice and causal_mask shapes don't match.
There are also changes here which may affect this test.

I hesitate to offer a PR because it may break other things that you try to do with this part of code. But, please, make the test work. BTW, it may be OK to change the test, for instance passing the whole bigger mask including cached items by editing this line to mask_1b = mask_1

Special note on StaticCache: I like this feature and I want to use custom 4D masks with it. So far this is not tested. I'd be glad to contribute such test once this issue is fixed. It will look like test_partial_stacked_causal_mask, only with StaticCache.

cc @ArthurZucker @gante

[ArthurZucker]

I can try to fix it, and yes I thought it would be tested automatically. It should be part of the tokenization_common or at least LlamaIntegrationTests or something. cc @gante if you can take a look I'll gladly review a PR

[gante]

^ this PR should fix it 🤗