index out of bounds for --compile_prefill with int4 and int8 about gpt-fast

[lopuhin]

I'm not yet sure if this might be related to having an extra pad token in the model -- but if I remove it the error is still there and all the weights have original shape. HF is misbehaving so can't download the unmodified model right now.

[lopuhin]

Reopening as this does not look related to the custom tokenizer -- I can reproduce the issue on the original llama-2-13b-chat-hf model.

[lopuhin]

A possible fix inspired by Lightning-AI/litgpt#774 is to add a clone() call (committed in 636cd76)

diff --git a/generate.py b/generate.py
index 7f30de0..cb4d7e6 100644
--- a/generate.py
+++ b/generate.py
@@ -161,7 +161,7 @@ def generate(
     seq = empty
     input_pos = torch.arange(0, T, device=device)
 
-    next_token = prefill(model, prompt.view(1, -1), input_pos, **sampling_kwargs)
+    next_token = prefill(model, prompt.view(1, -1), input_pos, **sampling_kwargs).clone()
     if is_speculative:
         prefill(draft_model, prompt.view(1, -1), input_pos, **sampling_kwargs)

This removes the error, but there is no speedup for prompt of 708 tokens when doing --compile_prefill, but it's clear that prefill slow-down is noticeable compared to a shorter prompt.