After running the following code:
import torch
from accelerate import infer_auto_device_map, init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM
config = AutoConfig.from_pretrained("facebook/opt-13b")
with init_empty_weights():
model = AutoModelForCausalLM.from_config(config)
device_map = infer_auto_device_map(model, no_split_module_classes=["OPTDecoderLayer"], dtype="float16")
model = AutoModelForCausalLM.from_pretrained(
"facebook/opt-13b",
device_map=device_map,
offload_folder="offload",
offload_state_dict=True,
torch_dtype=torch.float16)
I end up with the following error:
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
/tmp/ipykernel_23/3733710460.py in <module>
4 offload_folder="offload",
5 offload_state_dict=True,
----> 6 torch_dtype=torch.float16)
/opt/conda/lib/python3.7/site-packages/transformers/models/auto/auto_factory.py in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
463 model_class = _get_model_class(config, cls._model_mapping)
464 return model_class.from_pretrained(
--> 465 pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
466 )
467 raise ValueError(
/opt/conda/lib/python3.7/site-packages/transformers/modeling_utils.py in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
2527 # Dispatch model with hooks on all devices if necessary
2528 if device_map is not None:
-> 2529 dispatch_model(model, device_map=device_map, offload_dir=offload_folder, offload_index=offload_index)
2530
2531 if output_loading_info:
TypeError: dispatch_model() got an unexpected keyword argument 'offload_index'
I'm trying to load a large model from hugging face, and to do that, I split the weights into the GPUs from Kaggle. I don't understand exactly what is going on here, and how to deal with this issue.
please try updating your accelerate package, for example
pip install accelerate==0.18.0