

TITLE
[model loading] framework-agnostic dtype parameter

COMMENTS
3

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
This is a split off from one of the discussions at https://github.com/huggingface/transformers/pull/13209:

1. It all started with trying to load torch models under either the desired dtype or the the dtype of the pretrained model - and thus avoid 2x memory usage needs e.g. if the model needs to be just fp16. So we added `torch_dtype` to `from_pretrained` and `from_config`.
2. Then we started storing `torch_dtype` in the config file for future possibly automatic loading model in the optimal "regime".
3. This resulted in a discrepancy where the same symbol sometimes means `torch.dtype` at other times a string like "float32" as we can't store `torch.dtype` in json.
4. then in https://github.com/huggingface/transformers/pull/13209#discussion_r693292542 we started discussing how `dtype` is really the same across pt/tf/flux and perhaps we should just use `dtype` in the config and variables and have it consistently to be a string ("float32") and convert it to the right dtype object of the desired framework at the point of use, e.g. `getattr(torch, "float32")`

A possible solution is to deprecate `torch_dtype` and replace it with `dtype` string both in config and in the function argument.

Possible conflicts with the naming:

1. we already have the `dtype` attribute in modeling_utils, which returns `torch.dtype` based on the first param's dtype.

   https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_utils.py#L205

   The context is different, but still this is something to consider to avoid ambiguity.

I may have missed some other areas. So please share if something else needs to be added.

Additional notes:
- wrt flux: https://github.com/huggingface/transformers/pull/13209#discussion_r694511759
>  #13098 - the idea of the PR is exactly to disentangle parameter dtype from matmul/computation dtype. In Flax, it's common practice that the dtype parameter defines the matmul/computation dtype, see: https://flax.readthedocs.io/en/latest/_autosummary/flax.linen.Dense.html#flax.linen.Dense.dtype instead of the parameter dtype and not the parameter dtype.
> So for Flax, I don't really think it would make sense to use a config.dtype to define weights dtype as it would be quite confusing with Flax's computation dtype parameter.


@LysandreJik, @sgugger, @patrickvonplaten 