Closed johnPertoft closed 1 year ago
Thanks for reporting the issue. I remember having the same problem for the second issue (which comes from the buffers in the attention layer of GPT-NeoX) in Transformers but I don't recall right now how I fixed it. I'm on vacation until the beginning of Jan, so will look deeper then!
In the meantime, a workaround is to just load the model with Transformers:
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"EleutherAI/gpt-neox-20b",
device_map="auto",
offload_folder="/tmp/gpt-neox-20b-offload-accelerate",
)
You can then run the forward pass as long as you put it under a torch.no_grad()
(otherwise it OOMs because of the activations saved for the backward pass).
Thanks for your patience! I dug into this a bit more and the PR linked above should fix your initial code sample (as long as you add a torch.no_grad()
for the forward pass, otherwise you will probably get OOM).
System Info
Information
Tasks
no_trainer
script in theexamples
folder of thetransformers
repo (such asrun_no_trainer_glue.py
)Reproduction
Run the following script (slightly adapted from https://huggingface.co/docs/accelerate/usage_guides/big_modeling) but with the GPTNeoX model.
It crashes in
load_checkpoint_and_dispatch
with the following traceI then tried to update https://github.com/huggingface/accelerate/blob/cf22df91899dd57ead1df6c30983261596da7c99/src/accelerate/utils/modeling.py#L677 to always load the state_dict into cpu (
map_location="cpu"
) which got me past that error.But I subsequently ran into this crash instead
Expected behavior