from_pretrained's torch_dtype "auto" mode doesn't handle nested models
Author: fpgaminer
Pull request: no
Labels: name: Core: Modeling

name: Feature request

### Feature request

It would be nice if when setting `torch_dtype` to `auto` when calling `from_pretrained`, it properly respects _nested_ `torch_dtype`s specified in the model's config.  Right now it only pulls the dtype from the root config.

### Motivation

`from_pretrained` has a `torch_dtype` argument which can be set to `auto`.  The documentation says:  `"auto" - A torch_dtype entry in the config.json file of the model will be attempted to be used.`  However it looks like the code doesn't handle the case where `config.json` contains nested models, for example in the case of `LlavaForConditionalGeneration` which contains two children models under the keys `text_config` and `vision_config` which each can have `torch_dtype` set separately and differently from the parent config.

In my use case, I'd like to have a `LlavaForConditionalGeneration` which is loaded in as `float32` except for `text_config` which I want loaded in as `bfloat16`.  Here is the code I'm debugging the issue with:

```
m = LlavaForConditionalGeneration.from_pretrained("./saved_models/llama-joycaption-alpha-two-hf-llava", torch_dtype="auto")
print("LlavaConfig")
print("-" * 80)
print(m.config)
print("-" * 80)
print()

params_by_dtype = defaultdict(list)
for k, v in m.state_dict().items():
	params_by_dtype[v.dtype].append(k)

print("Params by dtype")
print("-" * 80)
for dtype, params in params_by_dtype.items():
	print(f"{dtype}:")
	for p in params:
		print(f"  {p}")
	print()
	print("-" * 80)
	print()
```

The config that gets printed:

```
LlavaConfig {
  "_name_or_path": "./saved_models/llama-joycaption-alpha-two-hf-llava",
  "architectures": [
    "LlavaForConditionalGeneration"
  ],
  "ignore_index": -100,
  "image_seq_length": 729,
  "image_token_index": 128077,
  "model_type": "llava",
  "projector_hidden_act": "gelu",
  "text_config": {
    "_name_or_path": "unsloth/Meta-Llama-3.1-8B-Instruct",
    "architectures": [
      "LlamaForCausalLM"
    ],
    "bos_token_id": 128000,
    "eos_token_id": [
      128001,
      128008,
      128009
    ],
    "intermediate_size": 14336,
    "max_position_embeddings": 131072,
    "model_type": "llama",
    "num_key_value_heads": 8,
    "pad_token_id": 128004,
    "rms_norm_eps": 1e-05,
    "rope_scaling": {
      "factor": 8.0,
      "high_freq_factor": 4.0,
      "low_freq_factor": 1.0,
      "original_max_position_embeddings": 8192,
      "rope_type": "llama3"
    },
    "rope_theta": 500000.0,
    "torch_dtype": "bfloat16",
    "unsloth_version": "2024.9",
    "vocab_size": 128256
  },
  "torch_dtype": "float32",
  "transformers_version": "4.44.2",
  "vision_config": {
    "hidden_act": "gelu_pytorch_tanh",
    "hidden_size": 1152,
    "image_size": 384,
    "intermediate_size": 4304,
    "layer_norm_eps": 1e-06,
    "model_type": "siglip_vision_model",
    "num_attention_heads": 16,
    "num_hidden_layers": 27,
    "patch_size": 14
  },
  "vision_feature_layer": -2,
  "vision_feature_select_strategy": "full"
}
```

Based on the rest of what gets printed (which is too long to include here) all parameters get loaded as `float32`.

### Your contribution

As far as I can tell, calling `from_pretrained` loads the entire model at once; it doesn't call any logic in a nested way.  The auto detection logic seems to start here:

https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/modeling_utils.py#L3988

Which just sets a single dtype; it doesn't analyze to check for any nested dtypes.  Then the model weights get loaded:

https://github.com/huggingface/transformers/blob/1bd604d11c405dfb8b78bda4062d88fc75c17de0/src/transformers/modeling_utils.py#L4181

`_load_pretrained_model` doesn't seem to have any logic to flexibly load different dtypes.

So I'm not sure how to fix/improve this logic short of major surgery to `_load_pretrained_model`?

Comments:
Should be fixed by https://github.com/huggingface/transformers/pull/32238 when it's merged :)

The PR merged but I see now that you requested to load vision and text models in different dtypes. With the auto it actually works correctly and loads vision in fp32 and lm in bf16 which are the weights from how state dict was saved.

But I also checked if setting `model.vision_config.torch_dtype = "float32"` and `model.vision_config.torch_dtype = "float16"` will work, unfortunately it didn't. Oke, I'll take this issue and enable assigning different dtypes for each backbone model