Consider top-level buffers when computing infer_auto_device_map

[younesbelkada]

What does this PR do?

This PR adds list(model._buffers) inside modules_to_treat when computing the auto_device_map. This scenario occured when I tried to add accelerate support for BART-like models when the final_logits_bias is registered as a buffer and is different than a uint type. It seems that we need to assign a device to this buffer.

The other solution is to "force-ignore" the buffer in check_device_map here since the tensors that are in model._buffers are stored in the state_dict.

cc @sgugger @muellerzr

slow tests from tests/test_bigmodeling.py pass!

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[sgugger]

Can you explain how some buffers don't end up in model._buffers because I don't fully understand that part.

[younesbelkada]

So if I understood it correctly, if you have some modules such as nn.BatchNorm in your model (as it is done in the accelerate CI test), the buffers running_mean and running_var will not be stored inside model._buffers but in model.named_buffers(). That is why I had to "filter out" the buffers by considering only the ones that are inside model._buffers and model.named_buffers()

Here is an example that I have quickly tried:

model = nn.Sequential(nn.Linear(1, 1), nn.BatchNorm1d(1), nn.Embedding(1, 1), nn.LayerNorm(1)) 
print(list(model.named_buffers()))
>>>[('1.running_mean', tensor([0.])), ('1.running_var', tensor([1.])), ('1.num_batches_tracked', tensor(0))]
print(list(model._buffers))
>>> []
model.register_buffer("position_bias", torch.ones(1))
print(list(model._buffers))
>>> ['position_bias']
print(list(model.named_buffers()))
>>> [('position_bias', tensor([1.])), ('1.running_mean', tensor([0.])), ('1.running_var', tensor([1.])), ('1.num_batches_tracked', tensor(0))]

[sgugger]

I think in this case, it's just the difference between named_buffers(recurse=True) and named_buffers(recurse=False). I'm not convinced this fix is the right fix, so would like to learn more what is failing.

[younesbelkada]

Ah yes I see, you're probably right here! Let me dig a bit more and get back to you here

[younesbelkada]

@sgugger I might have more clue on what is failing
I think that the problem comes from the fact that the infer_auto_device_map does take into account only modules and submodules, I have made a script below to better illustrate the problem

import torch.nn as nn
import torch
from accelerate.utils import infer_auto_device_map
from accelerate.big_modeling import dispatch_model

class SubModule(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_buffer("position_bias", torch.ones(1, 1000))

class Model(nn.Module):
    def __init__(self, wrap_module=True):
        super().__init__()
        self.l1 = nn.Linear(1000, 1000)
        self.l2 = nn.Linear(1000, 1000)
        self.l3 = nn.Linear(1000, 1000)

        self.bn1 = nn.BatchNorm1d(1000)
        self.bn2 = nn.BatchNorm1d(1000)

        if wrap_module:
            self.position_bias = SubModule()
        else:
            self.register_buffer("position_bias", torch.ones(1, 1000))

# Test 1: wrapping with a module - this will pass
model = Model()
device_map = infer_auto_device_map(model, {0:"10MB", "cpu":"100MB"})
model = dispatch_model(model, device_map)

# Test 2: below will fail
model = Model(wrap_module=False)
device_map = infer_auto_device_map(model, {0:"10MB", "cpu":"100MB"})
model = dispatch_model(model, device_map)

Let me know what do you think!

I guess this failed for BartPreTrainedModel since the position_bias buffer is on the parent module itself

[sgugger]

Ah, in this case it looks very much like the problem #747 fixed for top-level parameters, so the fix should be pretty similar here too!

[sgugger]

Perfect, thanks!

[younesbelkada]

The whole testing suite (including slow tests) is green! 🟢 Merging !