Some improvements for KV caching

[mseeger]

• Ensure that KVCache buffers are only as large as config.n_query_groups
• Shrink buffers returned by KVCache to just cover input_pos entries
• Clean up children of classes in model.py, in particular remove forward copies

[mseeger]

Can somebody help with failing tests? I don't understand why tests for Windows are failing, but pass for all other systems. And I also don't understand why the GPU tests are failing.

[Andrei-Aksionov]

Hello @mseeger

Thank you for another PR.

Can somebody help with failing tests? I don't understand why tests for Windows are failing

yeah, there is always something with Windows.
Perhaps we are lucky and a simple torch bump from #1893 might help 🤷

And I also don't understand why the GPU tests are failing.

I'll check it tomorrow.

[Andrei-Aksionov]

Hello @mseeger

It's quite a PR 🫠 🙂.
I left a couple of comments.
Overall it looks really good. I like how PEFT variants now look like, allows focusing on differences easily now 😊.

(I'll take a look why GPU tests are failing later.)

[mseeger]

OK, I reacted to comments. I also did a small change in lora.py, where the CausalSelfAttention.__init__ was still copy and paste, now it calls the superclass init.

[Andrei-Aksionov]

OK, I reacted to comments. I also did a small change in lora.py, where the CausalSelfAttention.init was still copy and paste, now it calls the superclass init.

Cool, we are almost there 🙂.
There are a couple of unresolved comments left.

On my side I'll try to find and fix issues with failing GPU tests, hopefully this year 😃.

[Andrei-Aksionov]

Overall, the issue with GPU+Thunder is something specific to the latter.
I'll merge the PR as is and later discuss it with Thunder team.

Thanks again for the PR (and for the patience 😊).

Happy New Year! 🚀