Reduce the memory usage of logits from O(context_length) to O(1) #4688

iseeyuan · 2024-08-13T14:22:14Z

The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference.

This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output .

Test command:

python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory

Before: 284 MB activation, with 262 MB on logits
After: 162 MB activation, with 0.128 MB on logits

Verified with llamma_runner, before and after it generates the same text with temperature=0.

Now the dominant memory usage would be KV cache.

TODO:

Improve KV cache memory usage using pf16 or quantization.
This PR only fixes logits. Further activation memory optimization with one token output.

Additional tests:
llava

python -m unittest examples.models.llava.test.test_llava -k test_prefill_logits
python -m unittest examples.models.llava.test.test_llava -k test_generated_output
python -m unittest examples.models.llava.test.test_llava -k test_llava_export

pytorch-bot · 2024-08-13T14:22:18Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/4688

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 Cancelled Job

As of commit be438eb with merge base d7c069f ():

CANCELLED JOB - The following job was cancelled. Please retry:

trunk / test-custom-ops-macos (cmake) / macos-job (gh)
##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2024-08-14T00:21:20Z

@iseeyuan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Differential Revision: D61246566

facebook-github-bot · 2024-08-14T00:39:00Z

This pull request was exported from Phabricator. Differential Revision: D61246566

larryliu0820 · 2024-08-14T16:28:00Z

Can you fix llava CI job?

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-15T00:55:10Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-16T00:29:20Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-17T15:07:32Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-18T22:14:06Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-19T21:07:19Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-19T21:48:49Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Reviewed By: larryliu0820 Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-21T22:51:39Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Reviewed By: larryliu0820 Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-21T23:34:45Z

This pull request was exported from Phabricator. Differential Revision: D61246566

facebook-github-bot · 2024-08-21T23:46:50Z

This pull request was exported from Phabricator. Differential Revision: D61246566

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Pull Request resolved: #4688 Reviewed By: larryliu0820 Differential Revision: D61246566 Pulled By: iseeyuan

Summary: The logits size is big, with size [context_length x vocab_size]. But we always use the last (new) logits, because the model generates one new token in each Transformer inference. This PR changes the transformer to return the logits of the last token only. In the runner code, we don't have to fetch the logits for the last token specifically, but directly use the output . Test command: ``` python -m examples.models.llama2.export_llama --checkpoint /Users/myuan/data/llama/story110m/checkpoint.pt --params /Users/myuan/data/llama/story110m/params.json -kv --use_sdpa_with_kv_cache -X -qmode 8da4w --group_size 128 -d fp32 --max_seq_length 1024 --profile_memory ``` Before: 284 MB activation, with 262 MB on logits After: 162 MB activation, with 0.128 MB on logits Verified with llamma_runner, before and after it generates the same text with temperature=0. Now the dominant memory usage would be KV cache. TODO: - Improve KV cache memory usage using pf16 or quantization. - This PR only fixes logits. Further activation memory optimization with one token output. Reviewed By: larryliu0820 Differential Revision: D61246566 Pulled By: iseeyuan

facebook-github-bot · 2024-08-22T23:34:32Z

This pull request was exported from Phabricator. Differential Revision: D61246566

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 13, 2024

iseeyuan force-pushed the logistic_size branch from 277310e to 1e0ff89 Compare August 13, 2024 14:53

iseeyuan changed the title ~~[testing] Reduce the memory usage of logistics from O(context_length) to O(1)~~ Reduce the memory usage of logistics from O(context_length) to O(1) Aug 13, 2024

iseeyuan changed the title ~~Reduce the memory usage of logistics from O(context_length) to O(1)~~ Reduce the memory usage of logits from O(context_length) to O(1) Aug 13, 2024

iseeyuan force-pushed the logistic_size branch from 1e0ff89 to b5bc307 Compare August 14, 2024 00:19

facebook-github-bot force-pushed the logistic_size branch from b5bc307 to 4e1f741 Compare August 14, 2024 00:38

facebook-github-bot added the fb-exported label Aug 14, 2024

facebook-github-bot force-pushed the logistic_size branch from 4e1f741 to 45d639e Compare August 15, 2024 00:54

cccclai added the ciflow/trunk label Aug 15, 2024

facebook-github-bot force-pushed the logistic_size branch from 45d639e to 6ef8bef Compare August 16, 2024 00:29

facebook-github-bot force-pushed the logistic_size branch from 6ef8bef to 9f105a2 Compare August 17, 2024 15:07

facebook-github-bot force-pushed the logistic_size branch from 9f105a2 to 296fd23 Compare August 18, 2024 22:13

facebook-github-bot force-pushed the logistic_size branch from 296fd23 to 4d7963f Compare August 19, 2024 21:07

facebook-github-bot force-pushed the logistic_size branch from 4d7963f to 3101b7a Compare August 19, 2024 21:48

larryliu0820 approved these changes Aug 19, 2024

View reviewed changes

facebook-github-bot force-pushed the logistic_size branch from 3101b7a to 2bca64c Compare August 21, 2024 22:51

facebook-github-bot force-pushed the logistic_size branch from 2bca64c to 5e15211 Compare August 21, 2024 23:34

iseeyuan force-pushed the logistic_size branch from 5e15211 to 7384903 Compare August 21, 2024 23:46

facebook-github-bot force-pushed the logistic_size branch from 7384903 to be438eb Compare August 22, 2024 23:34

facebook-github-bot merged commit 11e8ed3 into main Aug 23, 2024
89 of 92 checks passed

msaroufim mentioned this pull request Aug 27, 2024

Llama3.1 and kv_cache quantization pytorch/ao#738

Merged

felipemello1 mentioned this pull request Aug 29, 2024

Bring to torchtune memory wins for generate.py pytorch/torchtune#1448

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reduce the memory usage of logits from O(context_length) to O(1) #4688

Reduce the memory usage of logits from O(context_length) to O(1) #4688

iseeyuan commented Aug 13, 2024 •

edited

Loading

pytorch-bot bot commented Aug 13, 2024 •

edited

Loading

facebook-github-bot commented Aug 14, 2024

facebook-github-bot commented Aug 14, 2024

larryliu0820 commented Aug 14, 2024

facebook-github-bot commented Aug 15, 2024

facebook-github-bot commented Aug 16, 2024

facebook-github-bot commented Aug 17, 2024

facebook-github-bot commented Aug 18, 2024

facebook-github-bot commented Aug 19, 2024

facebook-github-bot commented Aug 19, 2024

facebook-github-bot commented Aug 21, 2024

facebook-github-bot commented Aug 21, 2024

facebook-github-bot commented Aug 21, 2024

facebook-github-bot commented Aug 22, 2024

Reduce the memory usage of logits from O(context_length) to O(1) #4688

Reduce the memory usage of logits from O(context_length) to O(1) #4688

Conversation

iseeyuan commented Aug 13, 2024 • edited Loading

pytorch-bot bot commented Aug 13, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/4688

❌ 1 Cancelled Job

facebook-github-bot commented Aug 14, 2024

facebook-github-bot commented Aug 14, 2024

larryliu0820 commented Aug 14, 2024

facebook-github-bot commented Aug 15, 2024

facebook-github-bot commented Aug 16, 2024

facebook-github-bot commented Aug 17, 2024

facebook-github-bot commented Aug 18, 2024

facebook-github-bot commented Aug 19, 2024

facebook-github-bot commented Aug 19, 2024

facebook-github-bot commented Aug 21, 2024

facebook-github-bot commented Aug 21, 2024

facebook-github-bot commented Aug 21, 2024

facebook-github-bot commented Aug 22, 2024

iseeyuan commented Aug 13, 2024 •

edited

Loading

pytorch-bot bot commented Aug 13, 2024 •

edited

Loading