Fix infer_auto_device_map when tied weights share the same prefix name

[fxmarty]

As per title, thanks to @Giuseppe5 & @nickfraser notice.

On main currently, the detection of tied_param_goups and tied_params is wrong as e.g. if we have a group ["compute.weight", "compute.weight_submodule.parameter"] and we are currently treating the parameter compute.weight, tied_param_goups will wrongfully be empty as all(name in k for k in tied_group) is True.

This result in an error in this example:

import torch
import accelerate

class SubModule(torch.nn.Module):
  def __init__(self, ref_to_parameter):
    super().__init__()
    self.parameter = ref_to_parameter
  def forward(self, x):
    return self.x + torch.max(self.parameter)

class LinearModuleAndSubModule(torch.nn.Linear):
  def __init__(self, in_features, out_features):
    super().__init__(in_features, out_features)
    self.weight_submodule = SubModule(self.weight)
  def forward(self, x):
    return torch.nn.functional.linear(self.weight_submodule(x), self.weight)

class Model(torch.nn.Module):
  def __init__(self):
    super().__init__()
    self.compute = LinearModuleAndSubModule(3, 8)
  def forward(self, x):
    return self.compute(x)

model = Model()
print(model)

from accelerate.utils import infer_auto_device_map
# Low memory device, just to force splitting and trigger the error
device_memory = {0: 4, 'cpu': 96000}
device_map = infer_auto_device_map(model, device_memory, verbose=True)

The error does not exist if we use the name self.brrrweight_submodule instead of self.weight_submodule.

[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.

[muellerzr]

Thanks for the fix and the test!

[SunMarc]

Thx for fixing and adding a test !