Skip to content

Reduce LLaMA memory usage - #18181

Merged
kunal-vaishnavi merged 7 commits into
microsoft:mainfrom
kunal-vaishnavi:kvaishnavi/llama-reduce-memory-usage
Nov 1, 2023
Merged

Reduce LLaMA memory usage#18181
kunal-vaishnavi merged 7 commits into
microsoft:mainfrom
kunal-vaishnavi:kvaishnavi/llama-reduce-memory-usage

Conversation

@kunal-vaishnavi

Copy link
Copy Markdown
Contributor

Description

This PR reduces the memory usage when exporting and benchmarking LLaMA.

Motivation and Context

  • Exporting: The PyTorch model is deleted from memory after a successful export instead of deleting it from memory after exporting + converting the ONNX model to the desired precision.
  • Benchmarking: In the ONNX model with GroupQueryAttention, the KV cache inputs use the same GPU memory for both the prompt and token generation benchmarks.

Comment thread onnxruntime/python/tools/transformers/models/llama/llama_inputs.py Fixed
Comment thread onnxruntime/python/tools/transformers/models/llama/llama_inputs.py Fixed
Comment thread onnxruntime/python/tools/transformers/models/llama/llama_inputs.py Fixed
Comment thread onnxruntime/python/tools/transformers/models/llama/llama_inputs.py Fixed
Comment thread onnxruntime/python/tools/transformers/models/llama/llama_inputs.py Fixed
Comment thread onnxruntime/python/tools/transformers/models/llama/llama_inputs.py Fixed
Comment thread onnxruntime/python/tools/transformers/models/llama/llama_inputs.py
Comment thread onnxruntime/python/tools/transformers/models/llama/benchmark.py
Comment thread onnxruntime/python/tools/transformers/models/llama/benchmark.py Outdated
frank-dong-ms
frank-dong-ms previously approved these changes Oct 31, 2023

@tianleiwu tianleiwu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM.

BTW, there is no need to convert torch.tensor->numpy for io_binding. You can directly use torch tensor in io binding. See example in

@kunal-vaishnavi
kunal-vaishnavi merged commit d1b85f5 into microsoft:main Nov 1, 2023
tianleiwu pushed a commit that referenced this pull request Nov 1, 2023
### Description
This PR reduces the memory usage when exporting and benchmarking LLaMA.



### Motivation and Context
- Exporting: The PyTorch model is deleted from memory after a successful
export instead of deleting it from memory after exporting + converting
the ONNX model to the desired precision.
- Benchmarking: In the ONNX model with GroupQueryAttention, the KV cache
inputs use the same GPU memory for both the prompt and token generation
benchmarks.
kleiti pushed a commit to kleiti/onnxruntime that referenced this pull request Mar 22, 2024
### Description
This PR reduces the memory usage when exporting and benchmarking LLaMA.



### Motivation and Context
- Exporting: The PyTorch model is deleted from memory after a successful
export instead of deleting it from memory after exporting + converting
the ONNX model to the desired precision.
- Benchmarking: In the ONNX model with GroupQueryAttention, the KV cache
inputs use the same GPU memory for both the prompt and token generation
benchmarks.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants