L-BFGS optimizer doesn't work properly on CUDA

follower: <p>Hi, I have some problems when I use L-BFGS optimizer on pytorch.</p>
<p>My problems are below.</p>
<ul>
<li>L-BFGS optimizer with CUDA doesn’t converge or converge too early (converge on high loss value)</li>
<li>L-BFGS with CPU work perfectly.</li>
<li>If I set data types of all tensor to float64, It is work.</li>
</ul>
<p>I think that these problem depends on my environment because my code works good on Google Colab.<br>
But even in the new virtual conda env, same problems happen.</p>
<p>My env are below.</p>
<ul>
<li>Windows 10 64bit</li>
<li>Python 3.8.10</li>
<li>Pytorch 1.10.0</li>
<li>Nvidia RTX 3070</li>
<li>CUDA 11.1 &amp; cuDNN v8.2.1</li>
<li>Newest Nvidia GPU Driver</li>
</ul>
<p>This is my example code.</p>
<pre><code class="lang-auto">import torch
from torch.autograd import Variable
import numpy as np
import matplotlib.pyplot as plt

torch.manual_seed(1234)
np.random.seed(1234)

device = torch.device("cuda")  # Or torch.device("cpu")
dtype = torch.float32  # Or torch.float64

torch.set_default_dtype(dtype)

x_data = np.linspace(0, 1, 100)[:, None]
y_data = np.sin(x_data * 2.0 * np.pi) + np.sin(x_data * 20.0 * np.pi) * 0.1

x_input = Variable(torch.tensor(x_data, dtype=dtype)).to(device)
y_target = Variable(torch.tensor(y_data, dtype=dtype)).to(device)

model = torch.nn.Sequential(
    torch.nn.Linear(1, 20),
    torch.nn.Tanh(),
    torch.nn.Linear(20, 1),
).to(device)

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.LBFGS(model.parameters(), lr=1.0, max_iter=5000)


y_out_init = model(x_input)

print(f"Device : {device}, dtype : {dtype}")
iteration = 0


def closure():
    global iteration
    optimizer.zero_grad()
    y_out = model(x_input)
    loss = loss_fn(y_out, y_target)
    loss.backward()
    iteration += 1
    print(f"\rIteration : {iteration}, Loss : {loss.item():.5e}", end="")
    if iteration % 500 == 0:
        print("")
    return loss


optimizer.step(closure)
y_out_final = model(x_input).detach().cpu().numpy()

# visualization
plt.title("Curve fitting")
plt.xlabel("x")
plt.ylabel("y")

plt.plot(x_data, y_data, label="ground truth")
plt.plot(x_data, y_out_init.detach().cpu().numpy(), label="init curve")
plt.plot(x_data, y_out_final, label="optimized curve")

plt.legend()
plt.show()
</code></pre>
<p>Here are my outputs,</p>
<blockquote>
<p>Device : cpu, dtype : torch.float32<br>
Iteration : 421, Loss : 4.79544e-03</p>
</blockquote>
<blockquote>
<p>Device : cpu, dtype : torch.float64<br>
Iteration : 500, Loss : 4.79593e-03<br>
Iteration : 625, Loss : 4.79562e-03</p>
</blockquote>
<blockquote>
<p>Device : cuda, dtype : torch.float32<br>
Iteration : 500, Loss : 4.91581e-03<br>
Iteration : 1000, Loss : 4.93076e-03<br>
Iteration : 1500, Loss : 4.92084e-03<br>
Iteration : 2000, Loss : 4.91559e-03<br>
Iteration : 2465, Loss : 2.16753e+17</p>
</blockquote>
<blockquote>
<p>Device : cuda, dtype : torch.float64<br>
Iteration : 400, Loss : 4.79128e-03</p>
</blockquote>
<p>Does anyone have an idea what might be the problem?</p>

ptrblck: <p>Could you disable TF32 via <code>torch.backends.cuda.matmul.allow_tf32 = False</code> and rerun the code, please?</p>

follower: <p>Oh, Now it seems to work fine. Would you mind explaining about this?</p>
<p>And, should I always disable TF32 to use the L-BFGS optimizer?</p>
<p>Thanks for your comments anyway.</p>