Setting up my custom device map for a LLM

🤗Transformers
4

This post has some ideas: Error: The model weights are not tied. Please use the `tie_weights` method before using the `infer_auto_device` function, even after adding model.tie_weights()

Maybe this could work:

import transformers
from accelerate import init_empty_weights
from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("chaoyi-wu/PMC_LLAMA_7B")
with init_empty_weights():
    model = AutoModelForCausalLM.from_config(config)

from accelerate import infer_auto_device_map, init_empty_weights
device_map = infer_auto_device_map(model)
1 Like