Put ancestors on the same device as the KV cache

[raphaelchinchilla]

Addressed the issue reported in #656 with the solution proposed by @rlouf .

I am really not sure whether this is the best solution.

[lapp0]

Thanks for contributing!

I think we should synchronize the device to ancestors since I believe (correct me if I'm wrong) it's possible that the kv_cache layers are on different devices.

Also could you please smoke test this with all 3 samplers on a 2 GPU system with device="auto" set?

[raphaelchinchilla]

Thanks for contributing!

I am not sure whether this is really a contribution, I just did what @rlouf suggested :)

I think we should synchronize the device to ancestors since I believe (correct me if I'm wrong) it's possible that the kv_cache layers are on different devices.

Honestly, I understand very little about the package (I started using it yesterday), and I do not feel I can actually give my opinion on whether you are wrong about it or not. It does feel like too simple of a contribution.

Also could you please smoke test this with all 3 samplers on a 2 GPU system with device="auto" set?

Sorry, which 3 samplers are you talking about?

[rlouf]

I just checked transformers and that's how they do it. ancestors is orders of magnitude smaller than layer so even with several layers that's very likely the most efficient. Anyway, if it solves @raphaelchinchilla 's problem let's merge it.

@raphaelchinchilla thank you for pushing the change, much appreciated! 🙏