Fail to deserialize model: Expected all tensors to be on the same device about tensorizer

[Eta0]

@lizzzcai as it turns out, it can be! I re-did the code around this in #96 to be more effective, initializing the weights on the intended device from the start rather than copying them there. This should be the most efficient method possible.

[Eta0]

Thank you for the bug report! I put up a change that should fix this issue here: #91.

This was caused by non-persistent buffers appearing on the CPU device instead of CUDA. Non-persistent buffers on a module are not touched during deserialization via deserializer.load_into_module(model), so they were remaining on the CPU while everything else was loaded to CUDA. The change adds an extra model = model.to(device) so that any remaining non-persistent buffers are moved to a device matching the deserializer's device.

[lizzzcai]

Hi @Eta0 thanks for your reply and fix, it is working from my testing.

Just curious and want to check if this model = model.to(device) can be further avoided as I try to measure the time spent and it takes around 0.14-0.2 sec for a 7b model ("mistralai/Mistral-7B-Instruct-v0.1").