fix: propagate precision correctly to enable non-bf16 inference

[Icedgarr]

This PR fixes some incompatibilities that I encountered when instantiating TSS from fam/llm/fast_inference.py with older and less powerful GPUs (e.g. Google Colab T4 GPU).

fam/llm/fast_inference_utils.py was putting the model to the device (cuda) with dtype.bfloat16 instead of using the precision parameter that contains the selected dtype (by default float16 or bfloat16 depending on the GPU architecture).

The linear layer of the Attention class in fam/llm/fast_model.py was also missing the dtype definition using the one provided in the config.

[vatsalaggarwal]

thanks a lot for this! one minor comment, ready to merge o/w!

[vatsalaggarwal]

why're these required? I think with the fix here, this shouldn't be needed?

[Icedgarr]

The following line was throwing an error due to the use of mixed types, q was float16 but k and v are bfloat16.
y = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)

Error: TorchRuntimeError: Failed running call_function <built-in function scaled_dot_product_attention>(*(FakeTensor(..., device='cuda:0', size=(2, 16, s0, 128)), FakeTensor(..., device='cuda:0', size=(2, 16, 2048, 128), dtype=torch.float16), FakeTensor(..., device='cuda:0', size=(2, 16, 2048, 128), dtype=torch.float16)), **{'attn_mask': FakeTensor(..., device='cuda:0', size=(1, 1, s0, 2048), dtype=torch.bool), 'dropout_p': 0.0}): Expected query, key, and value to have the same dtype, but got query.dtype: float key.dtype: c10::Half and value.dtype: c10::Half instead.

However, I have just checked and this alone did not solve the issue, it worked after I run the code with the torch dynamo disabled as well (doing export TORCHDYNAMO_DISABLE=1). It may be some of the operations done to the key and value tensors (I suspect this one: k = k.repeat_interleave(self.n_head // self.n_local_heads, dim=1) because it is the only one performed on k and v, but not q).

If you prefer, I can remove this change from the PR and if I or someone else find the root cause and a way to solve it we can create another PR.

[vatsalaggarwal]

So queries are the right dtype but keys and values are not? That sounds like it might be related to kv-cache not being the right dtype ... but we seem to be setting it correctly here... did you already have a look there?

[Icedgarr]

That's right, I see that the kv-cache runs when I execute the code, so it is likely to be what changes the dtypes, which according to the code you reference should not happen. If you agree I'll remove this part of the PR and investigate a bit further tomorrow to try to fix this other issue.

[Icedgarr]

I have reverted the last commit since it was not required for this fix.

[RahulVadisetty91]

need modification to add support for other devices or more robust error handling.