Use torch.bool instead of torch.int64 for non-persistant causal mask buffer

[fxmarty]

Adding self.register_buffer("causal_mask", torch.triu(causal_mask, diagonal=1), persistent=False) in @ArthurZucker's rewrite of llama & gemma adds a 500 MB overhead when serializing to ONNX/TorchScript IR/PyTorch ExportedProgram (from https://pytorch.org/docs/stable/export.html), for max_position_embeddings=8182.

Essentially, these IRs do not support non-persistent buffers. One quick fix is to use torch.bool instead of torch.int64, but bool is still 8-bits in pytorch (pytorch/pytorch#41571) & the overhead is still ~70 MB.

The lowered overhead is acceptable to me, but this won't scale to 10M context length.

[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.

[xenova]

Can confirm this shrunk my tiny-random-GemmeForCausalLM ONNX export from ~500MB to ~70MB (PR). Ideally, there would be no overhead, but I think this helps a ton for now!

[amyeroberts]

LGTM - thanks for digging into this and fixing!

Happy to merge once slow model tests for gemma and llama are confirmed to be passing.

[fxmarty]

@amyeroberts Running on A100, I can confirm that no additional tests are failing with RUN_SLOW=1 CUDA_VISIBLE_DEVICES=2 pytest tests/ -k "llama or gemma" -s -vvvvv compared to running on main.