Skip to content

Commit b0b2400

Browse files
author
Your Name
committed
add learning path about tuning thread count to improve LLM performance on pytorch-cpu
1 parent 9c75a4b commit b0b2400

9 files changed

Lines changed: 404 additions & 0 deletions

File tree

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
---
2+
title: Background Information
3+
weight: 3
4+
5+
### FIXED, DO NOT MODIFY
6+
layout: learningpathall
7+
---
8+
9+
## Background Information
10+
11+
A well-known tradeoff in parallel programming is choosing the right granularity of work. When multiple threads are created to perform a task, the amount of actual computation must justify the cost of coordinating those threads. For example, if a computation is split across many threads, the overhead of creating those threads and synchronizing their results through shared memory can easily outweigh any performance gains if each thread does too little work. The same principle applies to generative AI systems.
12+
13+
When distribution work across multiple threads, the communication and synchronisation overhead results in more overall work required. This introduces a tradeoff between latency, time to execute a single request, and throughput, number of requests processed in a given time interval. PyTorch looks to automatically find the correct number of threads, however as we will demonstrate, in some cases you may wish to manually fine tune.
14+
15+
## Multi-threading with PyTorch on CPU
16+
17+
18+
The diagram below is taken from the [PyTorch documentation](https://docs.pytorch.org/docs/2.8/notes/cpu_threading_torchscript_inference.html). When running inference, PyTorch uses a thread pool. In PyTorch, we have inter-op parallelism, which is spawning threads to run separate operations in a graph in parallel (e.g., 1 thread for a matmul and another thread for a softmax). Additionally there's intra-op parallelism can be used to spawn multiple threads to work on the same operation.
19+
20+
![threading-in-pytorch](./pytorch-threading.jpg)
21+
22+
In PyTorch, the `torch.set_num_threads()` [API](https://docs.pytorch.org/docs/stable/generated/torch.set_num_threads.html) is used to set the maximum number of threads to spawn in the Application Thread Pool. As of PyTorch 2.8.0, the default number of threads is equal to the number of CPU cores (see [PyTorch CPU Threading Documentation](https://docs.pytorch.org/docs/2.8/notes/cpu_threading_torchscript_inference.html)). PyTorch looks to find the ideal number of threads as described with the following code snippet taken from the PyTorch source code, [ParallemOpenMP.h](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/ParallelOpenMP.h).
23+
24+
```cpp
25+
int64_t num_threads = omp_get_num_threads();
26+
if (grain_size > 0) {
27+
num_threads = std::min(num_threads, divup((end - begin), grain_size));
28+
}
29+
30+
...
31+
inline int64_t divup(int64_t x, int64_t y) {
32+
return (x + y - 1) / y;
33+
}
34+
```
35+
36+
In PyTorch build with OpenMP, we can set the application thread pool once at runtime using the `OMP_NUM_THREADS` environment variable. The basic demonstration below will show that the default argument on many core systems may not be ideal for your workload. In the use case of running a language model with a low parameter count on a many-core CPU system, this default configuration may not be the most performant configuration.
37+
38+
## Basic Example
39+
40+
Create a new file called `pytorch_omp_example.py` and paste in the Python script below. We are performing a matrix multiplication in eager mode on two 256-by-256 random matrices. This is clearly a very small operation but we will observe the default behaviour and print the parallel configuration using the `torch.__config__.parallel_info()` function.
41+
42+
```python
43+
import os
44+
import time
45+
import torch
46+
47+
48+
def main():
49+
print(f"PyTorch version: {torch.__version__}")
50+
51+
# Read OMP_NUM_THREADS from the environment
52+
omp_threads = os.environ.get("OMP_NUM_THREADS")
53+
print(f"OMP_NUM_THREADS in environment: {omp_threads}")
54+
55+
# If it's set and looks like a number, use it to set PyTorch's intra-op threads
56+
if omp_threads and omp_threads.isdigit():
57+
torch.set_num_threads(int(omp_threads))
58+
59+
# Show how many threads PyTorch will actually use for intra-op parallelism
60+
print(f"torch.get_num_threads(): {torch.get_num_threads()}\n")
61+
62+
# A simple operation to illustrate parallelism:
63+
size = 256
64+
a = torch.randn(size, size)
65+
b = torch.randn(size, size)
66+
67+
start = time.time()
68+
c = a @ b # matrix multiplication (runs in a parallel region on CPU)
69+
end = time.time()
70+
71+
print(f"Result shape: {c.shape}")
72+
print(f"Matrix multiply time: {end - start:.5f} seconds")
73+
print(f"\nThreading Information = {torch.__config__.parallel_info()}")
74+
75+
if __name__ == "__main__":
76+
main()
77+
```
78+
79+
80+
Running the python script above you will observe the following output. As you can see the number of threads is set to core count of 96 and the time to execute is 2.24 ms.
81+
82+
```bash
83+
python pytorch_omp_example.py
84+
```
85+
86+
```output
87+
PyTorch version: 2.10.0.dev20251124
88+
OMP_NUM_THREADS in environment: None
89+
torch.get_num_threads(): 96
90+
91+
Result shape: torch.Size([256, 256])
92+
Matrix multiply time: 0.00224 seconds
93+
94+
Threading Information = ATen/Parallel:
95+
at::get_num_threads() : 96
96+
at::get_num_interop_threads() : 96
97+
OpenMP 201511 (a.k.a. OpenMP 4.5)
98+
omp_get_max_threads() : 96
99+
Intel(R) MKL-DNN v3.11.0 (Git Hash 0b8a866c009b03f322e6526d7c33cfec84a4a97a)
100+
std::thread::hardware_concurrency() : 96
101+
Environment variables:
102+
OMP_NUM_THREADS : [not set]
103+
ATen parallel backend: OpenMP
104+
```
105+
106+
Now when reduce the number of OpenMP threads using the `OMP_NUM_THREADS` value and observe a reduction is matrix multiply time to 0.64 ms.
107+
108+
```bash
109+
OMP_NUM_THREADS=16 python pytorch_omp_example.py
110+
```
111+
112+
```output
113+
PyTorch version: 2.10.0.dev20251124
114+
OMP_NUM_THREADS in environment: 16
115+
torch.get_num_threads(): 16
116+
117+
Result shape: torch.Size([256, 256])
118+
Matrix multiply time: 0.00064 seconds
119+
120+
Threading Information = ATen/Parallel:
121+
at::get_num_threads() : 16
122+
at::get_num_interop_threads() : 96
123+
OpenMP 201511 (a.k.a. OpenMP 4.5)
124+
omp_get_max_threads() : 16
125+
Intel(R) MKL-DNN v3.11.0 (Git Hash 0b8a866c009b03f322e6526d7c33cfec84a4a97a)
126+
std::thread::hardware_concurrency() : 96
127+
Environment variables:
128+
OMP_NUM_THREADS : 16
129+
ATen parallel backend: OpenMP
130+
```
131+
132+
We will now move on from this trivial example to a LLM.
115 KB
Loading
Loading
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
---
2+
title: Fine Tune LLM performance on CPU with multithreading
3+
4+
minutes_to_complete: 20
5+
6+
who_is_this_for: ML Engineers looking to fine tune the inference performance of PyTorch running on CPU systems
7+
8+
learning_objectives:
9+
- Understand how PyTorch uses multiple threads for CPU inference and the various tradeoffs involved
10+
- Tune the thread count to improve performance for specific models and systems
11+
12+
prerequisites:
13+
- Intermediate understanding of Python and PyTorch
14+
- Access to an Arm-based system
15+
16+
author: Kieran Hejmadi
17+
18+
### Tags
19+
skilllevels: Introductory
20+
subjects: ML
21+
armips:
22+
- Neoverse
23+
tools_software_languages:
24+
- Python
25+
- PyTorch
26+
- Bash
27+
operatingsystems:
28+
- Linux
29+
30+
31+
further_reading:
32+
- resource:
33+
title: Arm Tool Solutions
34+
link: https://github.com/ARM-software/Tool-Solutions/tree/main
35+
type: website
36+
37+
38+
### FIXED, DO NOT MODIFY
39+
# ================================================================================
40+
weight: 1 # _index.md always has weight of 1 to order correctly
41+
layout: "learningpathall" # All files under learning paths have this same wrapper
42+
learning_path_main_page: "yes" # This should be surfaced when looking for related content. Only set for _index.md of learning path content.
43+
---
Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
---
2+
# ================================================================================
3+
# FIXED, DO NOT MODIFY THIS FILE
4+
# ================================================================================
5+
weight: 21 # The weight controls the order of the pages. _index.md always has weight 1.
6+
title: "Next Steps" # Always the same, html page title.
7+
layout: "learningpathall" # All files under learning paths have this same wrapper for Hugo processing.
8+
---
9+
10+
If you would like to see other dials available to improve the performance of a LLM on Arm or other types of AI models. Please look at the examples page within the [Tools Solutions repository](https://github.com/ARM-software/Tool-Solutions/tree/main/ML-Frameworks/pytorch-aarch64/examples.)
Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
---
2+
title: Setup Environment
3+
weight: 2
4+
5+
### FIXED, DO NOT MODIFY
6+
layout: learningpathall
7+
---
8+
9+
## Build
10+
11+
In this learning path, we will use Arm’s downstream canary release of PyTorch, which includes ready-to-use examples and scripts. While this release offers access to the latest downstream features, it is intended for experimentation rather than production use.
12+
13+
### 1. Create HuggingFace Account
14+
15+
Create up a [huggingface account](https://huggingface.co/) if you do not already have one. Once created, request access to the [1B](https://huggingface.co/google/gemma-3-1b-it) and [270M](https://huggingface.co/google/gemma-3-270m-it) variants of Google's Gemma-3 model. It will take around 15 minutes to be granted access.
16+
17+
### 2. Connect to an Arm system and install Docker
18+
19+
lease see our [getting started guide](https://learn.arm.com/learning-paths/servers-and-cloud-computing/csp/) if it is your first time using Arm-based cloud instances.
20+
21+
In this example I will be using an AWS Graviton 4 (`m8g.24xlarge`)instance running Ubuntu 24.04 LTS. This is based on the Neoverse V2 architecture.
22+
23+
24+
Install docker through the [official documentation](https://docs.docker.com/engine/install/ubuntu/) or our [Arm install guide](https://learn.arm.com/install-guides/docker/docker-desktop-arm-linux/). Make sure to follow the post-installation steps.
25+
26+
27+
### 3. Build the PyTorch-AArch64 Docker Container
28+
29+
Connect to the Arm instance and run the following to clone the repository.
30+
31+
```bash
32+
git clone https://github.com/ARM-software/Tool-Solutions.git
33+
cd Tool-Solutions/ML-Frameworks/pytorch-aarch64/
34+
```
35+
Run the following bash script to build the container
36+
37+
```bash
38+
./build.sh -n $(nproc - 1)
39+
```
40+
41+
> **Note**: On a 96-core instance `m8g.24` AWS this will take approximately 20 minutes to build.
42+
43+
Once the build has finished, run the following command replacing `<version>` with the version of torch and torchao just built.
44+
45+
```bash
46+
./dockerize.sh ./results/torch-<version>linux_aarch64.whl ./results/torchao-<version>-py3-none-any.whl
47+
```
48+
49+
You should see the following output in your terminal, confirming that you are in the correct directory inside the Docker container.
50+
51+
```outpu
52+
aarch64_pytorch ~>
53+
```
54+
55+
### 5. Login to HuggingFace
56+
57+
58+
59+
Create a new `read` token on HuggingFace by clicking [this link](https://huggingface.co/settings/tokens/new?tokenType=read).
60+
61+
![hf-token](./HF-access-token.jpg)
62+
63+
Provide a suitable token name, press create token and copy the generated token value. From within docker container, enter the following command and paste the token to login.
64+
65+
```bash
66+
huggingface-cli login
67+
```
68+
69+
> **Note**: The login will not persist once the docker session has ended
61.7 KB
Loading
111 KB
Loading

0 commit comments

Comments
 (0)