Check if the buffers fit GPU memory after device map auto inferred

[notsyncing]

What does this PR do?

Hello, I'm trying using accelerate to offload a large model (https://huggingface.co/TheBloke/WizardCoder-33B-V1.1-GPTQ) to CPU, with following code (requires #2383 if using Intel GPU, and huggingface/transformers#28755):

import datetime

import torch
import intel_extension_for_pytorch

import accelerate
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_id = "/mnt/external2/LLMs/WizardCoder-33B-V1.1-GPTQ"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Without offload_buffers (OOM)
pipe = pipeline("text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "auto", "max_memory": {0: "16GB", "cpu": "128GB"}})

# With offload_buffers (won't OOM)
#pipe = pipeline("text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16, "device_map": "auto", "max_memory": {0: "16GB", "cpu": "128GB"}, "offload_buffers": True})
print(str(pipe.model.hf_device_map))

print(str(datetime.datetime.now()) + " Generating...")
results = pipe("public void helloWorld() {")

print(str(datetime.datetime.now()) + " Output:")
print(results)

I have one Intel Arc A770 16GB GPU in my machine, but the code above always OOM on the GPU.

After some digging, I found that this model TheBloke/WizardCoder-33B-V1.1-GPTQ contains a huge buffer (model is 17GB, but buffers are more than 16GB), so it will not fit into my GPU at all without offload_buffers.

So I created this PR, to check if the buffer can fit on the GPU, and raise an exception for such case.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Thanks for digging into the buffer issue @notsyncing ! Could you add some tests to see that the buffer calculation is right and that we raise the error when the buffers don't fit the remaining space on the gpu. Can you have a second look @muellerzr ?

[SunMarc]

Hi @notsyncing, thanks for your work ! I've replied to your questions. This is not pressing but are you planning to finish this PR ? Otherwise, I can finish it and add you as co-author =)

[notsyncing]

Hi @notsyncing, thanks for your work ! I've replied to your questions. This is not pressing but are you planning to finish this PR ? Otherwise, I can finish it and add you as co-author =)

Yes, I'm planning to finish it. But for the multi-gpu scenario, I cannot test it because I have only one gpu. Would you mind helping testing this? Thanks!

[muellerzr]

We can run the relevant tests for you when you are ready/the normal CI here is green and report back any failures (on a multi-gpu runner), yes :)

[notsyncing]

Basically, what I mean is if the user let a device_map that only most of the modules on the second gpus and leave of the space available on the first gpu, the buffer will most likely fit in the first gpu (if we suppose the first gpu is the execution device). For example, this is the case with device_map="balanced_low_0" strategy.

Now I added a check to see if the buffer can fit any GPU, and only print the warning when every GPU does not fit. Will this work in your scenario?

[notsyncing]

Could you add some tests to see that the buffer calculation is right and that we raise the error when the buffers don't fit the remaining space on the gpu.

As for the tests, it seems difficult to write a test for this if it only prints a warning. I propose that we could use a environment variable (like ACCELERATE_BUFFER_SIZE_CHECK = "raise") to make it raise an exception instead:

if os.getenv("ACCELERATE_BUFFER_SIZE_CHECK") == "raise":
    raise ValueError("xxx")
else:
    logger.warn("yyy")

Would this be better?

[muellerzr]

You can use warnings.warning() to do so, and unittest has a method to check for a warning: with self.assertWarns(Warning):

[muellerzr]

Also there's no need to force push, we squash the commit history when merging

[SunMarc]

Now I added a check to see if the buffer can fit any GPU, and only print the warning when every GPU does not fit. Will this work in your scenario?

Thanks ! Let's do that first and see the feedback. We can always change this behavior if needed. Most probably, the user will receive this warning if all his gpus are full and some of the layers are offloaded to cpu/disk. Hence, the buffer won't fit any of the gpus.

[SunMarc]

Thanks for the iteration. I left a few comments. Could you also add tests that would trigger the warning for a single and multi gpu setup ? We can test the multi gpu test if needed !

[notsyncing]

@SunMarc I have added two test cases in test_modeling_utils.py, and fixed a problem that my previous code did not consider the reserved memory for the largest layer. Looks like we don't need an actual multi-gpu setup to test this.

[SunMarc]

Thanks for working on this PR and adding the tests ! I ran them and they passed. I left a couple of minor comments/suggestions. We can merge this PR after fixing these =)

[notsyncing]

Thanks for working on this PR and adding the tests ! I ran them and they passed. I left a couple of minor comments/suggestions. We can merge this PR after fixing these =)

@SunMarc all fixed.

[SunMarc]

Thanks for iterating ! LGTM. I re-opened a comment I made. Could you have a quick look ? Also, can you have a second look @muellerzr ?

[notsyncing]

Thanks for iterating ! LGTM. I re-opened a comment I made. Could you have a quick look ? Also, can you have a second look @muellerzr ?

Added the missing assertions. I did not notice them before, sorry 😂

[muellerzr]

Thanks! Overall this looks sound to me :)