Improvements of KVCache and refactoring of subclasses of classes in model.py

[mseeger]

Studying your codebase (trying to learn about transformers in depth), I noted a few things that can be improved:

• KVCache: Second dimension of buffers should always be n_query_groups. If 1 < n_query_groups < n_head, you are wasting memory. Easy to fix.
• KVCache: forward returns tensors with final dimension max_seq_length. This is wasteful for the subsequence dot production attention computation. Can shorten this to a length that just covers all positions in input_pos. Relatively easy to fix.
• Code in adapter.py, adapter_v2.py, lora.py does a lot of copy and paste, which makes changing anything in model.py hard. I'd refactor that, so that as much common code only lives in model.py.

Let me know if this makes sense.

Thanks for doing this project. I really understand the details of transformer models better now.

[mseeger]

Check #1870

[mseeger]

@Andrei-Aksionov I have one question about #1891. You changed input_pos_maxp1 from type int to type torch.Tensor. This is always a scalar int. Why is this type change needed?

In fact, the way that this is used in the code would fail if this was not a scalar int.

[Andrei-Aksionov]

There was a problem with the test for sequentially.py when it tried to move an integer to a device, which is not possible.

In fact, the way that this is used in the code would fail if this was not a scalar int.

Could you point out in which place it will fail?

[mseeger]

It is used to make ranges, such as a[:input_pos_maxp1]. But since the code seems to work, maybe this is okay.

it is odd that an int member of a class needs to be moved to a device

[Andrei-Aksionov]

It is used to make ranges, such as a[:input_pos_maxp1]. But since the code seems to work, maybe this is okay.

You can easily use torch.Tensor for these kinds of operations.

import torch

device = "mps"

idx = torch.tensor(2, device=device)
python_list = [1, 2, 3]
tensor = torch.tensor(python_list, device=device)

print(python_list[:idx])
>> [1, 2]

print(tensor[:idx])
>> tensor([1, 2], device='mps:0')

it is odd that an int member of a class needs to be moved to a device

It's just the existing code moves all the inputs to the target device via pre-hook.
Either change the code, or just use input_pos_maxp1 as a tensor.

Besides, when input_pos_maxp1 is not on the device it will be moved to there implicitly, and that can cause cudaStreamSynchronize, which can hurt the performance. Although a proper code profiling is required to determine how severe the impact is, it's better to stay on a safe side and just put everything on the target device beforehand.

[mseeger]

OK, this has been done.