Autocast on cpu dramatically slow

Laurent1: <p>Hi,</p>
<p>On a toy regression model with pytorch 2.1.2 on cpu, <code>torch.autocast</code> is really slow.<br>
Without <code>with torch.autocast(device_type="cpu", dtype=torch.bfloat16, cache_enabled=True)</code> in the code below, I get</p>
<p>Epoch 0: Loss: 5.163879e-01 in 2.6s.<br>
Epoch 1: Loss: 5.981598e+01 in 2.5s.<br>
Epoch 2: Loss: 6.057042e-01 in 2.5s.<br>
Epoch 3: Loss: 3.923391e+00 in 2.5s.<br>
Epoch 4: Loss: 1.904698e+00 in 2.5s.<br>
training done in 12.7s.</p>
<p>With <code>torch.autocast</code>, I get</p>
<p>Epoch 0: Loss: 5.164276e-01 in 230.8s.<br>
Epoch 1: Loss: 5.980967e+01 in 231.1s.<br>
Epoch 2: Loss: 6.053763e-01 in 233.0s.<br>
Epoch 3: Loss: 3.924905e+00 in 232.8s.<br>
Epoch 4: Loss: 1.904750e+00 in 232.0s.<br>
training done in 1159.7s.</p>
<p>So, the code with autocast is 2 orders of magnitude slower.</p>
<p>Am I using autocast correctly in the code below ? thanks</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import numpy as np
from time import time
torch.set_num_threads(1)

input_size = 1
output_size = 1
hidden_size = 512
num_data = 100000

# seeds
seed = 1234
torch.manual_seed(seed)
torch.use_deterministic_algorithms(True)
np.random.seed(seed)

# hyper-parameters
num_epochs = 5
learning_rate = 0.01

# toy dataset
x_train = np.random.rand(num_data,input_size)
y_train = np.cos(2*np.pi*x_train) + 0.1*np.random.randn(num_data,input_size)

# regression model
model = nn.Sequential(nn.Linear(input_size, hidden_size),
                       nn.GELU(),
                       nn.Linear(hidden_size, hidden_size),
                       nn.GELU(),
                       nn.Linear(hidden_size, output_size))

# loss and optimizer
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay = 1e-2)  

# train the model
x_train = torch.from_numpy(x_train.astype(np.float32))
y_train = torch.from_numpy(y_train.astype(np.float32))

global_start_time = time()
for epoch in range(num_epochs):
    start_time = time()
    
    # forward pass
    with torch.autocast(device_type="cpu", dtype=torch.bfloat16, cache_enabled=True):
        outputs = model(x_train)
        loss = criterion(outputs, y_train)
    
    # backward and optimize
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    
    print(f'Epoch {epoch}: Loss: {loss.item():.6e} in {time()-start_time:.1f}s.')
print(f'training done in {time()-global_start_time:.1f}s.')
</code></pre>
<p>PyTorch built with:</p>
<ul>
<li>GCC 9.3</li>
<li>C++ Version: 201703</li>
<li>Intel(R) oneAPI Math Kernel Library Version 2022.1-Product Build 20220311 for Intel(R) 64 architecture applications</li>
<li>Intel(R) MKL-DNN v3.1.1 (Git Hash 64f6bcbcbab628e96f33a62c3e975f8535a7bde4)</li>
<li>OpenMP 201511 (a.k.a. OpenMP 4.5)</li>
<li>LAPACK is enabled (usually provided by MKL)</li>
<li>NNPACK is enabled</li>
<li>CPU capability usage: AVX2</li>
<li>Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CXX_COMPILER=/opt/rh/devtoolset-9/root/usr/bin/c++, CXX_FLAGS= -D_GLIBCXX_USE_CXX11_ABI=0 -fabi-version=11 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOCUPTI -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=old-style-cast -Wno-invalid-partial-specialization -Wno-unused-private-field -Wno-aligned-allocation-unavailable -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Werror=cast-function-type -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_DISABLE_GPU_ASSERTS=ON, TORCH_VERSION=2.1.2, USE_CUDA=0, USE_CUDNN=OFF, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=OFF, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF,</li>
</ul>

ptrblck: <p>I don’t see anything obviously wrong in your code so could you create an issue on GitHub, please?</p>

Laurent1: <p>Thanks for your reply.<br>
I will create an issue.</p>

Laurent1: <p><a href="https://github.com/pytorch/pytorch/issues/118499" rel="noopener nofollow ugc"> #118499 </a></p>

Aknw_Fen: <p>Sad that the issue is still unresolved, I was wondering whether your CPU has bf16 or a related flag (found in <code>/proc/cpuinfo</code> or <code>lscpu</code>)?</p>
<p>Mine does not, only f16c; I don’t get as bad results, but still pretty bad:</p>
<pre><code class="lang-auto">Epoch 0: Loss: 5.164276e-01 in 5.2s.

...similar in other epochs

training done in 25.4s.

----noautocast----
Epoch 0: Loss: 5.163879e-01 in 1.9s.

....ditto

training done in 9.6s.
</code></pre>
<p>and with <code>torch.float16</code> it takes forever.</p>