Add support for weights_only flag when loading state_dict

[jerryzh168]

Summary:
This is to enable loading a state_dict with wrapper tensor subclasses (used in torchao to for quantized weights)

Test Plan:
tested locally with torchao weights, also need #32306:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import TorchAoConfig
from torchao.utils import benchmark_model
import torchao

DEVICE_TYPE = "cuda"

def init_model_and_benchmark(model_id, torch_dtype=torch.bfloat16, quantization_config=None):
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    if quantization_config is not None:
        model = AutoModelForCausalLM.from_pretrained(model_id, device_map=DEVICE_TYPE, torch_dtype=torch.bfloat16, quantization_config=quantization_config)
    else:
        model = AutoModelForCausalLM.from_pretrained(model_id, device_map=DEVICE_TYPE, torch_dtype=torch.bfloat16, weights_only=False)

    # sanity check: run the model
    input_text = "What are we having for dinner?"
    input_ids = tokenizer(input_text, return_tensors="pt").to(DEVICE_TYPE)
    output = model.generate(**input_ids, max_new_tokens=1000)
    print(tokenizer.decode(output[0], skip_special_tokens=True))

    NUM_WARMUP = 1
    NUM_RUNS = 5

    if quantization_config is not None:
        torchao.quantization.utils.recommended_inductor_config_setter()

    model = torch.compile(model, mode="max-autotune")

    benchmark_model(model.generate, NUM_WARMUP, kwargs=input_ids, device_type=DEVICE_TYPE)
    print("running benchmark")
    results = benchmark_model(model.generate, NUM_RUNS, kwargs=input_ids, device_type=DEVICE_TYPE)
    return model, results

model_id = "jerryzh168/test-model"
torchao.quantization.utils.recommended_inductor_config_setter()
bf16_model, bf16_time = init_model_and_benchmark(model_id)
print(f"bf16: {bf16_time}")

Reviewers:

Subscribers:

Tasks:

Tags:

[amyeroberts]

Thanks for adding this support!

Overall, I think this looks good. Let's get a second 👍 from @ArthurZucker too as it's touching core code

[jerryzh168]

I learned from @mikaylagawarecki that weights_only is going to be set to True by default in the future. but this flag could still be helpful if people are using older versions of pytorch and want to use torchao I think

[jerryzh168]

Hi @ArthurZucker can you take a look at this PR?

[jerryzh168]

please let me know if we want to add a test, but load test is a bit harder to write because it needs to upload some models like:

transformers/tests/test_modeling_utils.py

Line 930 in abbffc4

model = BertModel.from_pretrained("hf-internal-testing/tiny-random-bert")

[ArthurZucker]

LGTM sorry for being late on the review here!

[ArthurZucker]

Can you rebase on main and resolve the conflicts?

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

[jerryzh168]

thanks for the review @ArthurZucker and @amyeroberts I just updated the PR, please feel free to merge when the CI is green

[amyeroberts]

@jerryzh168 Could you run make fix-copies and push the changes? This should resolve the failing quality checks

[jerryzh168]

thanks, updated, please take a look again @amyeroberts

[ArthurZucker]

let's revert unrelated changes and let's go

[ArthurZucker]

this unrelated change should not be included here ! 🤗

[jerryzh168]

@ArthurZucker this is fix from running make fix-copies 😅

[ArthurZucker]

ok no worries then let's merge!

[ArthurZucker]

Sorry for the delay would kindly ask to rebase to make sure we are up to date as other things were merged !

[ArthurZucker]

Thanks once more @jerryzh168