Feat: Add landmark attention

[NanoCode012]

Closes #109

NOTE: Only Llama!

• Test it runs
• Run tests from landmark

[NanoCode012]

Testing via https://github.com/epfml/landmark-attention/blob/main/llama/run_test.py

n=1k

Running test 50/50 for n = 1000...
Number of tokens in this prompt:  338
Expected number in the prompt: 16835, base output: 16835
Success!
Number of tokens in this prompt:  338
Expected number in the prompt: 16835, mem output: 16835
Success!
Accuracy base for n = 1000: 100.0%
Accuracy mem for n = 1000: 100.0%

n=5k

Running test 50/50 for n = 5000...
Number of tokens in this prompt:  1406
Expected number in the prompt: 18704, base output: 18704
Success!
Number of tokens in this prompt:  1406
Expected number in the prompt: 18704, mem output: 18704
Success!
Accuracy base for n = 5000: 100.0%
Accuracy mem for n = 5000: 96.0%

n=8k

Running test 50/50 for n = 8000...
Number of tokens in this prompt:  2206
Expected number in the prompt: 38907, base output: 38907
Success!
Number of tokens in this prompt:  2206
Expected number in the prompt: 38907, mem output: 38907
Success!
Accuracy base for n = 8000: 98.0%
Accuracy mem for n = 8000: 96.0%

n=10k

Running test 50/50 for n = 10000...
Number of tokens in this prompt:  2738
Expected number in the prompt: 1333, base output: .
.
.
.
.

Fail.
Number of tokens in this prompt:  2738
Expected number in the prompt: 1333, mem output: 1333
Success!
Accuracy base for n = 10000: 0.0%
Accuracy mem for n = 10000: 100.0%

n=12k

Running test 50/50 for n = 12000...
Number of tokens in this prompt:  3272
Expected number in the prompt: 10145, base output:  and and and and and
Fail.
Number of tokens in this prompt:  3272
Expected number in the prompt: 10145, mem output: 10145
Success!
Accuracy base for n = 12000: 0.0%
Accuracy mem for n = 12000: 98.0%

n=14k

Running test 10/10 for n = 14000...
Number of tokens in this prompt:  3805
Expected number in the prompt: 15978, mem output: 15978
Success!
Accuracy mem for n = 14000: 90.0%

n=18k

Running test 10/10 for n = 18000...
Number of tokens in this prompt:  4872
Expected number in the prompt: 48396, mem output: 48396
Success!
Accuracy mem for n = 18000: 100.0%

n=20k

Running test 10/10 for n = 20000...
Number of tokens in this prompt:  5406
Expected number in the prompt: 41403, mem output: 41403
Success!
Accuracy mem for n = 20000: 100.0%

n=25k

Running test 10/10 for n = 25000...
Number of tokens in this prompt:  6738
Expected number in the prompt: 30430, mem output: 30430
Success!
Accuracy mem for n = 25000: 100.0%

n=38k

Running test 10/10 for n = 38000...
Number of tokens in this prompt:  10207
Expected number in the prompt: 27396, mem output: 27396
Success!
Accuracy mem for n = 38000: 100.0%

n=45k

Running test 10/10 for n = 45000...
Number of tokens in this prompt:  12072
Expected number in the prompt: 44988, mem output: 44988
Success!
Accuracy mem for n = 45000: 100.0%

n=50k

Running test 10/10 for n = 50000...
Number of tokens in this prompt:  13406
Expected number in the prompt: 31773, mem output: 1234
Fail.
Accuracy mem for n = 50000: 80.0%

n=60k

Running test 10/10 for n = 60000...
Number of tokens in this prompt:  16072
Expected number in the prompt: 31528, mem output: 31528
Success!
Accuracy mem for n = 60000: 100.0%

n=75k

Running test 10/10 for n = 75000...
Number of tokens in this prompt:  20071
Expected number in the prompt: 22901, mem output: 22901
Success!
Accuracy mem for n = 75000: 100.0%

Thanks to AutoMeta for compute for testing.

[winglian]

🔥

[NanoCode012]

This feature should allow longer context.

The test above uses the author's delta weights to confirm feature parity.

[winglian]

reassigning this here inside the function causes this error:

ERROR:root:local variable 'LlamaForCausalLM' referenced before assignment
Traceback (most recent call last):
  File "/workspace/axolotl/src/axolotl/utils/models.py", line 214, in load_model
    model = LlamaForCausalLM.from_pretrained(
UnboundLocalError: local variable 'LlamaForCausalLM' referenced before assignment
^CTraceback (most recent call last):
  File "/workspace/axolotl/src/axolotl/utils/models.py", line 214, in load_model
    model = LlamaForCausalLM.from_pretrained(
UnboundLocalError: local variable 'LlamaForCausalLM' referenced before assignment

[NanoCode012]

Hm, weirdly I didn't get this error.

Second, the new PR should fix this import issue. But, it hasn't been tested.

[NanoCode012]

I need to go afk for now, so I'll test the PR when I get back since it seem like a big bug.