How to use trust_remote_code=True with load_checkpoint_and_dispatch?

🤗Accelerate
4

Try this and let me know:

from transformers import AutoTokenizer, AutoModelForCausalLM

import transformers

import torch

from accelerate import init_empty_weights

from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, file_utils

from accelerate import init_empty_weights, load_checkpoint_and_dispatch, infer_auto_device_map

config = AutoConfig.from_pretrained(‘./model_path’,trust_remote_code=True)

config.attn_config[“attn_impl”] = “triton”

config.init_device = “cuda:0”

model = AutoModelForCausalLM.from_pretrained(“tiiuae/falcon-7b-instruct”, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True)

tokenizer = AutoTokenizer.from_pretrained(“EleutherAI/gpt-neox-20b”, padding_side=“left”)

tokenizer.pad_token_id = tokenizer.eos_token_id

print(“downlaoded tokenizer”)

model.save_pretrained(“./model_path/”)

tokenizer.save_pretrained(“./model_path/”)

print(“saved tokenizer”)

checkpoint =“./model_path”

with init_empty_weights():

config = AutoConfig.from_pretrained(‘./model_path’,trust_remote_code=True)

model = AutoModelForCausalLM.from_config(config,trust_remote_code=True)

max_memory = {0: “8GiB”, “cpu”: “4GiB”}

model = load_checkpoint_and_dispatch(model, checkpoint, device_map=“auto”, dtype=torch.bfloat16, offload_folder=“offload”)

model = model.tie_weights()

tokenizer = AutoTokenizer.from_pretrained(checkpoint, padding_side=“left”)

print(type(model))