🎓 A1Learner
🔍

PyTorch Threads and Interop

Thread settings control how many CPU cores PyTorch uses at once. Intra-op threads split one operation across cores, and inter-op threads run separate operations side by side.

Ask the machine

Two short calls report the current settings. On this machine both answers are 2.

import torch

print(torch.get_num_threads(), torch.get_num_interop_threads())

Set intra-op threads

The intra-op count can change at any point in a program. The next operations use the new count.

torch.set_num_threads(1)
print(torch.get_num_threads())

Interop is set once, early

The inter-op setter must run before other PyTorch work in the program. A late call raises an error.

import torch

torch.set_num_interop_threads(1)
print(torch.get_num_interop_threads())

One matmul, split work

With the count back at 2, one large matmul can split its rows across both cores. The result is the same either way.

torch.set_num_threads(2)
a = torch.randn(768, 768)
b = torch.randn(768, 768)
c = a @ b

Example

import time

import torch

torch.manual_seed(0)
a = torch.randn(768, 768)
b = torch.randn(768, 768)


def timed(reps=15):
    for _ in range(3):
        a @ b
    start = time.perf_counter()
    result = None
    for _ in range(reps):
        result = a @ b
    ms = (time.perf_counter() - start) * 1000 / reps
    return result, ms


print("Threads at start:", torch.get_num_threads())
print("Interop threads:", torch.get_num_interop_threads())
torch.set_num_threads(1)
one, one_ms = timed()
torch.set_num_threads(2)
two, two_ms = timed()
print("Results equal:", torch.equal(one, two))
print(f"One thread per matmul: {one_ms:.0f} ms")
print(f"Two threads per matmul: {two_ms:.0f} ms")

Key points

  • Intra-op threads: threads that split one operation, like one matmul
  • Inter-op threads: threads that run independent operations in parallel
  • Use torch.get_num_threads() to read the intra-op count
  • Use torch.set_num_threads(n) to change the intra-op count
  • The inter-op count can be set only once, before parallel work starts
  • This machine reports 2 intra-op threads and 2 inter-op threads
  • In short, thread counts cap how much of the CPU one program can use at once

In this example, two 768 by 768 matrices are multiplied 15 times. The runs use one thread, then two threads. A short warmup runs before each timing. The program first prints the counts this machine starts with. Both counts are 2, matching its 2 cores. The two results compare equal, because thread count changes scheduling, not maths. The two times vary on every run. On the recorded run, one thread was ahead, and a quieter run can flip that order. Two cores is the whole story on this machine, so no setting unlocks a large win here.

Note

More threads cannot beat the core count. This machine has 2 cores, so 2 is the honest ceiling.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does torch.set_num_threads(1) change?