Chapter 5: running train_model_simple() on CPU is really slow, why?

[frx-wintermute]

Hello,
I am running train_model_simple() from chapter 5 on CPU. The results seem to correspond to what the book shows (final train loss 0.391, final validation loss 6.452, and the generated sample text seems to be identical to the one displayed in the book).

However, the book says that a MacBook Air completed the training in about 5 minutes.
I am seeing times of the order of 3 hours or something.

Running on Debian GNU/Linux forky, with PyTorch 2.12.0 .
CPU is an AMD Ryzen 5 5500GT with 6 physical core (2 hardware threads per core, hence a total of 12 hardware threads).
Memory capacity is 32 GiB.

What I am seeing is that PyTorch seems to run with 6 threads (equal to the number of physical cores, by default):

$ python3
Python 3.13.14 (main, Jun 10 2026, 18:10:12) [GCC 15.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.get_num_threads())
6

But, during the training, I can see (with top) that the process uses about 110 % of CPU, only occasionally reaching 230 % (for a very short time). If I hit [Shift+H], so that top shows threads on separate lines, I see that there are 6 threads, but 5 are almost idle most of the time:

    PID  PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND  
 325584  20   0 4453872   3.4g 124612 R  99.7  12.6   2:12.84 model_p+ 
 325587  20   0 4453872   3.4g 124612 S   2.0  12.6   0:04.41 model_p+ 
 325588  20   0 4453872   3.4g 124612 S   2.0  12.6   0:04.41 model_p+ 
 325589  20   0 4453872   3.4g 124612 S   2.3  12.6   0:04.42 model_p+ 
 325590  20   0 4453872   3.4g 124612 S   2.0  12.6   0:04.41 model_p+ 
 325591  20   0 4453872   3.4g 124612 S   2.0  12.6   0:04.42 model_p+ 

What am I missing?
Any ideas?

[darekwojciechowski]

Those five idle threads are the whole story — PyTorch is configured for six but only actually using one. get_num_threads() reports the setting, not that the parallel backend engaged, and 110% CPU on a 6-core Zen 3 says it didn't.

Which also means the M3 comparison is a red herring. Your CPU shouldn't be 36× slower than a MacBook Air here; a working build lands in the tens of minutes at worst. Three hours means something's broken, not that the hardware is weak.

Quickest way to confirm, with the training loop out of the picture entirely:

Those five idle threads are the whole story — PyTorch is configured for six but only actually using one. get_num_threads() reports the setting, not that the parallel backend engaged, and 110% CPU on a 6-core Zen 3 says it didn't.

Which also means the M3 comparison is a red herring. Your CPU shouldn't be 36× slower than a MacBook Air here; a working build lands in the tens of minutes at worst. Three hours means something's broken, not that the hardware is weak.

Quickest way to confirm, with the training loop out of the picture entirely:

import time, torch

a = torch.randn(4096, 4096)
b = torch.randn(4096, 4096)
torch.mm(a, b)

t0 = time.perf_counter()
for _ in range(5):
    torch.mm(a, b)
dt = (time.perf_counter() - t0) / 5
print(f"{dt*1000:.0f} ms   {2*4096**3/dt/1e9:.0f} GFLOP/s")

Watch top while it runs. Healthy, you should see close to 600% CPU and a few hundred GFLOP/s. If it sits at ~100% and gives you tens of GFLOP/s, your BLAS is running serial and that fully explains the training time.

If that's what you get, check where torch came from — python3 -c "import torch; print(torch.__file__)". Anything under /usr/lib/python3/dist-packages is Debian's packaged build, which on forky links against the system OpenBLAS and is a common cause of exactly this. The upstream wheel in a venv is usually the fix:

python3 -m venv .venv
.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu

Re-run the benchmark after and compare. Also worth a quick env | grep -i omp first — a stray OMP_NUM_THREADS=1 from a shell profile produces the identical symptom and is a much cheaper fix.

Your losses match the book, so the model code is fine. It's purely a threading problem.

[frx-wintermute]

Hello @darekwojciechowski, thanks a lot for your answer. It helped me to pinpoint the root cause of the performance issue.

TL;DR: Everything is fine with Debian packages, after switching to OpenMP-enabled OpenBLAS (see below for details).

I tried the small benchmark you suggested:

$ cat mm_test.py 
#!/usr/bin/env python3

import time, torch

print(f"Will attempt to run with {torch.get_num_threads()} threads",
      flush=True)

a = torch.randn(4096, 4096)
b = torch.randn(4096, 4096)
torch.mm(a, b)

t0 = time.perf_counter()
for _ in range(5):
    torch.mm(a, b)
dt = (time.perf_counter() - t0)/5
print(f"{dt*1000:.0f} ms   {2*4096**3/dt/1e9:.0f} GFLOP/s")

I ran this benchmark with the Debian-packaged PyTorch and I obtained the following results:

$ ./mm_test.py | tee debian_packaged_torch.out
Will attempt to run with 6 threads
12049 ms   11 GFLOP/s

I saw (with top) that the process sat at about 100 % CPU (only one thread seemed to be active).

I tried to install PyTorch in a virtual environment and to run the benchmark with the pip-installed PyTorch:

$ python3 -m venv pytorch
$ source pytorch/bin/activate
(pytorch) $ pip install torch --index-url https://download.pytorch.org/whl/cpu
(pytorch) $ pip install numpy
(pytorch) $ ./mm_test.py | tee pip_torch.out
Will attempt to run with 6 threads
298 ms   462 GFLOP/s
(pytorch) $ deactivate

I saw (with top) that the process went well beyond 400 % CPU, maybe it reached almost 600 % for the very short execution time (6 threads seemed to be active).

Then I asked myself "why does the Debian package run about 40 times slower than the pip-installed PyTorch? Is it really due to system OpenBLAS, which is running serial?".
And then I realized that I didn't have any OpenBLAS Debian package installed:

$ aptitude search '~i openblas'
$

On the other hand, I had libblas3 installed:

$ aptitude search '~i libblas'
i A libblas-dev                     - Basic Linear Algebra Subroutines 3, static
i A libblas3                        - Basic Linear Algebra Reference implementat
$

Hence, I switched to the OpenMP-enabled OpenBLAS:

# aptitude install libblas3_ libblas-dev_ libopenblas-openmp-dev+M libopenblas0-openmp+M

For anyone not familiar with Debian, I can explain the meaning of the above command: what this command does is to remove the two packages libblas3 and libblas-dev, while installing the two packages libopenblas-openmp-dev and libopenblas0-openmp at the same time (and marking them as automatically installed). The switch is possible, since all these packages provide the same ABI and are drop-in replacements of each other. This is explicitly documented in the package metadata and the package manager takes it into account, when satisfying dependencies. Why did I have libblas3 installed, when I installed python3-torch? Because other previously-installed packages pulled in libblas3 as a dependency, without preferring OpenBLAS packages.

And now, with OpenMP-enabled OpenBLAS, I can enjoy the Debian-packaged PyTorch without performance issues:

 $ ./mm_test.py | tee debian_packaged_torch_withopenblasopenmp.out
Will attempt to run with 6 threads
220 ms   625 GFLOP/s

Even with better performance than the pip-installed PyTorch!

I re-ran train_model_simple() from chapter 5 on CPU and now the training took about 4 minutes, while giving the same results!
Yay!

Thanks again for your insight, it really pointed me into the right direction!