You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Copy file name to clipboardExpand all lines: content/learning-paths/servers-and-cloud-computing/tune-pytorch-cpu-perf-with-threads/background.md
+25-31Lines changed: 25 additions & 31 deletions
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -1,38 +1,30 @@
1
1
---
2
-
title: Background Information
2
+
title: Understand PyTorch threading for CPU inference
3
3
weight: 3
4
4
5
5
### FIXED, DO NOT MODIFY
6
6
layout: learningpathall
7
7
---
8
8
9
-
## Background Information
9
+
A well-known challenge in parallel programming is choosing the right number of threads for a given amount of work. When multiple threads are created to perform a task, the actual computation must be large enough to justify the overhead of coordinating those threads.
10
10
11
-
A well-known challenge in parallel programming is choosing the right number of threads for a given amount of work. When multiple threads are created to perform a task, the actual computation must be large enough to justify the overhead of coordinating those threads.
11
+
If a computation is split across many threads, the costs of creating the threads and synchronizing their results through shared memory can easily outweigh any performance gains from parallel execution. The same principle applies to generative AI workloads running on CPU.
12
12
13
-
For example, if a computation is split across many threads, the costs of:
14
-
- creating the threads, and
15
-
- synchronizing their results through shared memory
13
+
When work is distributed across multiple threads, communication and synchronization overhead increases the total amount of work the system must perform. This creates a trade-off between latency (the time to process a single request) and throughput (the number of requests processed per unit time).
16
14
17
-
can easily outweigh any performance gains from parallel execution. The same principle applies to generative AI workloads running on CPU.
15
+
PyTorch attempts to automatically choose an appropriate number of threads. However, as you will learn, in some cases you can manually tune this configuration to improve performance.
18
16
19
-
When work is distributed across multiple threads, communication and synchronization overhead increases the total amount of work the system must perform. This creates a trade-off between:
17
+
## Multithreading with PyTorch on CPU
20
18
21
-
-**Latency** – the time to process a single request, and
22
-
-**Throughput** – the number of requests processed per unit time.
19
+
When running inference, PyTorch uses an Application Thread Pool. PyTorch supports two types of parallelism: inter-op parallelism spawns threads to run separate operations in a graph in parallel (for example, one thread for a matmul and another thread for a softmax), while intra-op parallelism spawns multiple threads to work on the same operation.
23
20
24
-
PyTorch attempts to automatically choose an appropriate number of threads. However, as we will show, in some cases you may want to manually fine-tune this configuration to improve performance.
21
+
The diagram below is taken from the [PyTorch documentation](https://docs.pytorch.org/docs/stable/index.html).
25
22
26
-
## Multi-threading with PyTorch on CPU
23
+
![Diagram showing PyTorch's threading model with application thread pool, inter-op thread pool, and intra-op thread pool for CPU inference#center](./pytorch-threading.jpg"PyTorch threading model")
27
24
25
+
The `torch.set_num_threads()`[API](https://docs.pytorch.org/docs/stable/generated/torch.set_num_threads.html) sets the maximum number of threads to spawn in the Application Thread Pool.
28
26
29
-
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 an **Application Thread Pool**. In PyTorch, there are 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.
30
-
31
-
32
-
33
-
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.
34
-
35
-
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) for more detail). 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).
27
+
As of PyTorch 2.8.0, the default number of threads equals the number of CPU cores (see [PyTorch CPU Threading Documentation](https://docs.pytorch.org/docs/2.8/notes/cpu_threading_torchscript_inference.html) for more detail). PyTorch determines the ideal number of threads based on the workload size, as shown in this code snippet from [ParallelOpenMP.h](https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/ParallelOpenMP.h):
In PyTorch builds that use OpenMP, the maximum size of the application’s thread pool can be configured once at runtime using the `OMP_NUM_THREADS` environment variable. The actual number of threads used will scale up to this limit depending on the workload and the `grain_size`.
50
-
51
-
The short example below illustrates that the default settings on many-core systems may not provide optimal performance for all workloads.
41
+
In PyTorch builds that use OpenMP, the maximum size of the application's thread pool can be configured at runtime using the `OMP_NUM_THREADS` environment variable. The actual number of threads used scales up to this limit depending on the workload and the `grain_size`.
52
42
53
-
## Basic PyTorch Example
43
+
The example below demonstrates that the default settings on many-core systems might not provide optimal performance for all workloads.
54
44
55
-
Create a new file named `pytorch_omp_example.py` and paste in the Python script below. The script performs a matrix multiplication in eager mode on two 256×256 random matrices.
56
-
For this relatively small computation, we will:
45
+
## Basic PyTorch example
57
46
58
-
- Observe the default performance of PyTorch’s parallelism
59
-
- Print the parallel configuration using `torch.__config__.parallel_info()`.
47
+
Create a new file named `pytorch_omp_example.py` with the following Python script. The script performs a matrix multiplication in eager mode on two 256×256 random matrices, showing the default performance of PyTorch's parallelism and printing the parallel configuration:
60
48
61
49
```python
62
50
import os
@@ -95,13 +83,13 @@ if __name__ == "__main__":
95
83
main()
96
84
```
97
85
98
-
Run the python script
86
+
Run the Python script:
99
87
100
88
```bash
101
89
python pytorch_omp_example.py
102
90
```
103
91
104
-
You will observe the an output similar to the follow. As you can see the number of threads is set to core count of 96 and the time to execute is 2.24 ms.
92
+
The output is similar to:
105
93
106
94
107
95
```output
@@ -124,12 +112,16 @@ Environment variables:
124
112
ATen parallel backend: OpenMP
125
113
```
126
114
127
-
Now reduce the number of OpenMP threads using the `OMP_NUM_THREADS` value and observe a reduction is matrix multiply time to 0.64 ms.
115
+
The number of threads is set to the core count of 96, and the execution time is 2.24 ms.
116
+
117
+
Reduce the number of OpenMP threads using the `OMP_NUM_THREADS` environment variable:
128
118
129
119
```bash
130
120
OMP_NUM_THREADS=16 python pytorch_omp_example.py
131
121
```
132
122
123
+
The output shows a different `Matrix multiply time` of of 0.64 ms.
124
+
133
125
```output
134
126
PyTorch version: 2.10.0.dev20251124
135
127
OMP_NUM_THREADS in environment: 16
@@ -150,4 +142,6 @@ Environment variables:
150
142
ATen parallel backend: OpenMP
151
143
```
152
144
153
-
We will now move on from this trivial example to a much larger workload of a large language model (LLM).
145
+
The time varies with the number of threads and type of processor in your system.
146
+
147
+
In the next section, you'll apply these concepts to a much larger workload using a large language model (LLM).
This Learning Path uses 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's intended for experimentation rather than production use.
11
+
{{% /notice %}}
10
12
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.
13
+
## Create a Hugging Face account
12
14
13
-
### 1. Create HuggingFace Account
15
+
Create a [Hugging Face account](https://huggingface.co/) if you don't already have one. After creating your account, 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. Access is typically granted within 15 minutes.
14
16
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.
17
+
## Connect to an Arm system and install Docker
16
18
17
-
### 2. Connect to an Arm system and install Docker
19
+
If this is your first time using Arm-based cloud instances, see the [getting started guide](/learning-paths/servers-and-cloud-computing/csp/).
18
20
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.
21
+
The example code uses an AWS Graviton 4 (`m8g.24xlarge`) instance running Ubuntu 24.04 LTS, based on the Neoverse V2 architecture. You can use any Arm server with at least 16 cores. Keep note of your CPU count so you can adjust the example code as needed.
20
22
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.
23
+
Install Docker using the [Docker install guide](/install-guides/docker/) or the [official documentation](https://docs.docker.com/engine/install/ubuntu/). Follow the post-installation steps to configure Docker for non-root usage.
22
24
25
+
## Run the PyTorch-aarch64 Docker container image
23
26
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.
27
+
You have two options for the Docker container. You can use a container image from Docker Hub or you can build the container image from source. Using a ready-made container makes it easier to get started, and building from source provides the latest software. The container image on Docker Hub is updated about once a month.
25
28
29
+
Open a terminal or use SSH to connect to your Arm Linux system.
26
30
27
-
### 3. Build the PyTorch-AArch64 Docker Container
31
+
### Use container image from Docker Hub
28
32
29
-
Connect to the Arm instance and run the following to clone the repository.
33
+
Download the ready-made container image from Docker Hub:
You should see the following output in your terminal, confirming that you are in the correct directory inside the Docker container.
74
+
The shell prompt will appear, and you are ready to start.
50
75
51
-
```outpu
76
+
```output
52
77
aarch64_pytorch ~>
53
78
```
54
79
55
-
### 5. Login to HuggingFace
56
-
80
+
## Log in to Hugging Face
57
81
82
+
Create a new Read token on Hugging Face by navigating to [Create new Access Token](https://huggingface.co/settings/tokens/new?tokenType=read).
58
83
59
-
Create a new `read` token on HuggingFace by clicking [this link](https://huggingface.co/settings/tokens/new?tokenType=read).
84
+
![Screenshot of Hugging Face token creation page showing the 'Create new token' dialog with token type set to 'Read'#center](./hf-access-token.jpg"Hugging Face token creation")
60
85
61
-
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.
86
+
Provide a token name, create the token, and copy the generated value. From within the Docker container, run the following command and paste the token when prompted:
64
87
65
88
```bash
66
89
huggingface-cli login
67
90
```
68
91
69
-
> **Note**: The login will not persist once the docker session has ended
92
+
Messages indicating the token is valid and login is successful are printed.
93
+
94
+
Be aware that the login doesn't persist after the Docker container exits. You'll need to log in again if you restart the container.
0 commit comments