remove repetitive entries from device lists

[xloem]

I found when loading that each tensor for a device was reloaded for every layer in an n**2 manner.

It seemed the device list in load_state_dict was loaded as device map keys, but then treated as a unique set despite having an element for every entry in the map, i.e. each device was in the list many times over.

I searched the file for devices, found two uses of lists as if they were sets, and added code to remove duplicate entries from both instances. I used a dict rather than a set so that order of entries is preserved, keeping the first device as the primary one.

EDIT: the other list I found already had unique entries and that change has been removed.

With this change, for me, the loading code no longer reloads every tensor for each map entry, instead loading them once for each device.

[HuggingFaceDocBuilderDev]

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

[sgugger]

Could you please post a small reproducer of the code that failed before this PR?

[xloem]

For me, the below code runs both timed tests in about 2 seconds with this PR, whereas before it the second test hangs and ends up measuring 1200 seconds:

import timeit
import accelerate, safetensors.torch, torch

import warnings; warnings.filterwarnings('ignore')

DEV = 3

keys = [f'{letter}{number}' for letter in bytes(list(range(ord('A'), ord('z')))).decode() for number in range(10)]

state_dict = {
        key : torch.rand(1024) for key in keys
}
device_map = {
        key : DEV for key in keys
}
safetensors.torch.save_file(state_dict, 'test.safetensors', metadata=dict(format='pt'))
print('Loading state_dict with safetensors.torch.load_file:')
def timed_reference_call():
    safetensors.torch.load_file('test.safetensors', device=DEV)
result = timeit.timeit(timed_reference_call, setup = timed_reference_call, number=10)
print(result)
print('Loading state_dict with accelerate.utils.modeling.load_state_dict:')
def timed_test_call():
    accelerate.utils.modeling.load_state_dict('test.safetensors', device_map)
result = timeit.timeit(timed_test_call, setup = timed_test_call, number=10)
print(result)

[sgugger]

Thanks for explainig. I didn't understand your change at first, let's use a set instead of going through the keys of an artificial dict.

[sgugger]

Mmm I don't know what's happening at GitHub right nw, but I can't authorize the tests to run. Will wait for a bit and hopefully the button will come back in a couple of hours. Stay tuned!

[sgugger]

All green, thanks again for your contribution!