4.35.0 requires troublesome pytorch install for JAX users
Author: colehaus
Pull request: no

### System Info

- `transformers` version: 4.35.0
- Platform: Linux-5.19.0-45-generic-x86_64-with-glibc2.35
- Python version: 3.11.6
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.2
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.7.4 (gpu)
- Jax version: 0.4.20
- JaxLib version: 0.4.20
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@sanchit-gandhi

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

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

### Reproduction

`FlaxT5ForConditionalGeneration.from_pretrained("t5-base")`

### Expected behavior

I was using transformers 4.31.0 and loading pre-trained models worked pretty smoothly. I recently tried to upgrade to 4.35.0 ad encountered a cluster of issues:

- The same `from_pretrained` call fails at runtime now and says that pytorch must be installed (I assume this is due to the `safetensors` changes mentioned here: https://github.com/huggingface/transformers/releases/tag/v4.35.0)
- I don't see any way to disable the new loading behavior and revert to the old loading behavior
- I install pytorch (sort of unfortunate that a third-party lib is mandating the install of an entire separate DL stack)
- However, now `jax.devices()` warns about a failure to initialize CUDA and only returns the CPU as a device. This is because `pytorch` pins specific versions of CUDA support libs (https://github.com/pytorch/pytorch/blob/ceb07656c234b78bd045312f871b9cdca8c3c3ba/.github/scripts/generate_binary_build_matrix.py#L32) which are incompatible with those JAX wants.
- I also tried some other permutations of pytorch and JAX versions (between 0.4.14 and 0.4.20 IIRC). None of them are compatible. They all fail at runtime with either JAX's failure to initialize CUDA or throw an error about JAX being compiled against newer versions of the libs.

TL;DR: 4.35.0 requires `torch` alongside `jax` to load models and it's hard to do this in a way that satisfies both sets of constraints for CUDA support libs.

This was ultimately resolvable for me by installing the CPU version of pytorch, but I thought I'd write out this issue for anyone else that encounters it and so that perhaps this interaction can be handled and/or documented more fluently.

Comments:
(Also, after the upgrade I get this warning:

Some weights of the model checkpoint at t5-base were not used when initializing FlaxT5ForConditionalGeneration: {('decoder', 'block', '0', 'layer', '1', 'EncDecAttention', 'relative_attention_bias', 'kernel')}
- This IS expected if you are initializing FlaxT5ForConditionalGeneration from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing FlaxT5ForConditionalGeneration from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).

Based on https://github.com/huggingface/transformers/issues/10221, I think the model is as expected and the issue is just that `model.safetensors` has an extra/inappropriate tensor there? If that's not the case, I can make a proper issue for this.)

cc @LysandreJik @sanchit-gandhi 

Thanks for reporting! This shouldn't be the case. We'll take a look and release a patch.

@sanchit-gandhi, @amyeroberts, curious to hear your thoughts:

Right now the implementation defaults to `safetensors` to load weights (it has been the case for PyTorch since the initial implementation and for TF since https://github.com/huggingface/transformers/pull/19900.

This is problematic as the overwhelming majority of `safetensors` weights are saved from PyTorch and will need to be loaded in PyTorch. As @colehaus has shown above, installing PyTorch solves the problem, but is a requirement we really do not want for users wanting to run Flax/TF models.

AFAIK we have no way of knowing the metadata of the checkpoint (and hence to know from which framework it was saved) without first fetching the entire checkpoint. I don't think we'll want to download safetensors checkpoints before reading the metadata (which will be PT in 99% of cases), and then fetching a .msgpack.

Given this, I think the solution is to default to `.msgpack` checkpoints (and their sharded counterparts) even if there are safetensors weights in the repository.

EDIT: Alternative solutions that might be cleaner
- After talking to @Narsil I understand that we could eventually just download the first few MBs of the checkpoints and have access to the metadata, to check there
- The necessity to have `torch` installed for the conversion doesn't seem to have a hard necessity. I can look at relaxing that requirement, which would make everything much simpler.

As an immediate workaround @colehaus, please downgrade to v4.34; we'll aim to resolve this cleanly as this is clearly unintended behavior.

@LysandreJik Aligned with what you suggest here - as a first step defaulting to `.msgpack` seems like a pragmatic solution to fix quickly. 

I think we can do both of the alternative solutions. Grabbing the metadata and checking before downloading would have been my suggestion too if it was possible. 

@colehaus, could you please let me know if checking out this PR solves your issue? https://github.com/huggingface/transformers/pull/27460

It's a bandaid while we implement the metadata fetching/relaxing of the `torch` requirement for the conversion.

Installing from source should now solve the issue @colehaus. I'll release a patch probably tomorrow or Wednesday.

how to convert a .msgpack (flax/jax format) model weights to .bin (huggingface format) or .pth (pytorch format)? thanks a lot

See my comment here: https://github.com/huggingface/transformers/issues/26813#issuecomment-1835703292