|
| 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 | + |
| 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. |
0 commit comments