Loading BloomForCausalLM from sharded checkpoints
Hey everyone,I am currently working on my master thesis and have used the Transformers library succesfully for most of the experiments I wanted to conduct. The only thing I am stuck with is loading a sharded version of Bloom-7b1, which I am loading using the accelerate load_checkpoint_and_dispatch following the guideBig model inference, which results in the following error.File “load_bloom_test.py”, line 22, inmodel = load_checkpoint_and_dispatch(File “/opt/conda/lib/python3.8/site-packages/accelerate/big_modeling.py”, line 375, in load_checkpoint_and_dispatchload_checkpoint_in_model(File “/opt/conda/lib/python3.8/site-packages/accelerate/utils/modeling.py”, line 699, in load_checkpoint_in_modelset_module_tensor_to_device(model, param_name, param_device, value=param)File “/opt/conda/lib/python3.8/site-packages/accelerate/utils/modeling.py”, line 105, in set_module_tensor_to_devicenew_module = getattr(module, split)File “/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py”, line 1260, ingetattrraise AttributeError(“‘{}’ object has no attribute ‘{}’”.format(AttributeError: ‘BloomForCausalLM’ object has no attribute ‘word_embeddings’The code I am using works fine for other models, such as Galactica (for all variants, including galactica-120b) and looks like this:import torch

from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch

# Loading model from config on meta device and using sharded checkpoints
model_name = "models/bloom-7b1"
config = AutoConfig.from_pretrained(model_name)
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

# Calculating device map
device_map = infer_auto_device_map(
    model, 
    no_split_module_classes=["BloomBlock"],
    dtype=torch.float16
)

# Loading the checkpoint according to auto device_map
model = load_checkpoint_and_dispatch(
    model,
    model_name,
    device_map=device_map
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
inputs = tokenizer("Hello, my name is", return_tensors="pt")
inputs = inputs.to(0)
output = model.generate(inputs["input_ids"])
print(tokenizer.decode(output[0]))I have redownloaded the model a couple of times and also tried usingbigscience/bloom-7b1to download it directly from the hub during execution of the script.I would like this approach to work, as it has worked fine for all other models and I have not managed to get a distributed setup working using a different approach when using a model that doesn’t fit on a single GPU anymore.Any help with either using a different loading approach or a solution to my problem would be greatly appreciated
2023-03-07T15:14:34Z
I also get and error, but slightly different.Traceback (most recent call last):
  File "run_inference.py", line 12, in <module>
    model = load_checkpoint_and_dispatch(
  File "/home/anaconda3/envs/research/lib/python3.8/site-packages/accelerate/big_modeling.py", line 427, in load_checkpoint_and_dispatch
    load_checkpoint_in_model(
  File "/home/anaconda3/envs/research/lib/python3.8/site-packages/accelerate/utils/modeling.py", line 748, in load_checkpoint_in_model
    raise ValueError(f"{param_name} doesn't have any device set.")
ValueError: word_embeddings.weight doesn't have any device set.
2023-03-07T15:19:20Z
The same code for gpt-j-6B works.Reference:Handling big models for inference@sguggerCould you check it for us?Also, OPT raises an error.Traceback (most recent call last):
  File "run_inference.py", line 13, in <module>
    model = load_checkpoint_and_dispatch(
  File "/home/anaconda3/envs/research/lib/python3.8/site-packages/accelerate/big_modeling.py", line 427, in load_checkpoint_and_dispatch
    load_checkpoint_in_model(
  File "/home/anaconda3/envs/research/lib/python3.8/site-packages/accelerate/utils/modeling.py", line 748, in load_checkpoint_in_model
    raise ValueError(f"{param_name} doesn't have any device set.")
ValueError: decoder.embed_tokens.weight doesn't have any device set.The code I used it for testing.from accelerate import init_empty_weights, load_checkpoint_and_dispatch
from transformers import AutoConfig, AutoModelForCausalLM

# checkpoint = "bigscience/bloom-7b1"
# checkpoint = "EleutherAI/gpt-j-6B"
checkpoint = "facebook/opt-6.7b"
config = AutoConfig.from_pretrained(checkpoint)

with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)


model = load_checkpoint_and_dispatch(
    model,
    # "bloom-7b1",
    # "sharded-gpt-j-6B",
    "opt-6.7b",
    device_map="sequential",
    # no_split_module_classes=["BloomBlock"],
    # no_split_module_classes=["GPTJBlock"],
    no_split_module_classes=["OPTDecoderLayer"],
)

print(model.REDACTED_TOKEN_map)I tried using differentdevice_map, but it didn’t work…
2023-03-07T17:15:05Z
You won’t be able to use the checkpoint on the Hub directly in Accelerate, you need to pass via Transformer, since the model on the Hub is the checkpoint of the base model, so not the model with the LM head.
2023-03-08T01:50:04Z
Could I ask what is the “Transformer” you mentioned? Does it indicate the “nn.Transformer” in PyTorch native or is there any other “Transformer” class in the accelerate? I am new to this library.
2023-03-08T02:05:03Z
Transformers the library you are using to load the model.AutoModelForCausalLM.from_pretrainedacceptsdevice_map="sequential".
2023-03-08T16:29:13Z
It works usingAutoModelForCausalLM.from_pretrained("facebook/opt-6.7b", device_map="auto").Thank youIt really helps.Also, I think that the inferred device map doesn’t need no module split.It already handles depending on the architecture. Is it right?
2023-03-08T16:39:32Z
Yes Transformers puts that info for you