GPs and Apple GPU

[ghost]

Hello everyone!

I'm coming to you because I'm having a bit of a problem implementing variational GPs.

I have a MacBook Pro M1 Pro with a built-in GPU and I want to do computations on it. So I replace the .cuda() with .to(device = 'mps') after first passing the data in float32.

I get the following error:

File ~/miniforge3/envs/StageM2/lib/python3.10/site-packages/gpytorch/variational/variational_strategy.py:104, in VariationalStrategy._cholesky_factor(self, induc_induc_covar)
    102 @cached(name="cholesky_factor", ignore_args=True)
    103 def _cholesky_factor(self, induc_induc_covar: LinearOperator) -> TriangularLinearOperator:
--> 104     L = psd_safe_cholesky(to_dense(induc_induc_covar).type(_linalg_dtype_cholesky.value()))
    105     return TriangularLinearOperator(L)

TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64. Please use float32 instead.

which I think is related to the implementation of variational_strategy.py, where you have to force float64 for reasons of numerical stability?

Although I'm doing exactly the example on the SVGP site in the docs, here's a snippet of the code:

data = torch.Tensor(loadmat('elevators.mat')['data'])
X = data[:, :-1]
X = X - X.min(0)[0]
X = 2 * (X / X.max(0)[0]) - 1
y = data[:, -1]
train_n = int(floor(0.8 * len(X)))
train_x = X[:train_n, :].contiguous().to(device = 'mps') #already in float32
train_y = y[:train_n].contiguous().to(device ='mps')
test_x = X[train_n:, :].contiguous().to(device='mps')
test_y = y[train_n:].contiguous().to(device='mps')

train_dataset = TensorDataset(train_x, train_y)
train_loader = DataLoader(train_dataset, batch_size=1024, shuffle=True)

test_dataset = TensorDataset(test_x, test_y)
test_loader = DataLoader(test_dataset, batch_size=1024, shuffle=False)

class GPModel(ApproximateGP):
    def __init__(self, inducing_points):
        variational_distribution = CholeskyVariationalDistribution(inducing_points.size(0))
        variational_strategy = VariationalStrategy(self, inducing_points, variational_distribution, learn_inducing_locations=True)
        super(GPModel, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

inducing_points = train_x[:500, :]
model = GPModel(inducing_points=inducing_points).float().to(device='mps')
likelihood = gpytorch.likelihoods.GaussianLikelihood().to(device='mps')

num_epochs = 10
model.train()
likelihood.train()

optimizer = torch.optim.Adam([
    {'params': model.parameters()},
    {'params': likelihood.parameters()},
], lr=0.01)
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))
epochs_iter = tqdm(range(num_epochs), desc="Epoch")
for i in epochs_iter:
    minibatch_iter = tqdm(train_loader, desc="Minibatch", leave=False)
    for x_batch, y_batch in minibatch_iter:
        optimizer.zero_grad()
        output = model(x_batch)
        loss = -mll(output, y_batch)
        minibatch_iter.set_postfix(loss=loss.item())
        loss.backward()
        optimizer.step()

Is my method the wrong one, or is it not possible in its current state to do the training on the Apple Silicon GPU?

To the best of my knowledge, I haven't found a topic dealing with a similar problem on the github. If, by mistake, you've already replied to this, I apologise!

Details :

• Python 3.10
• GPyTorch 1.11
• Torch 2.0.1
• MacOS 13.1

Thanks and have a nice week!

[Balandat]

Hmm I didn't know that but it looks like Apple GPUs just don't support 64 bit floating point computation. Which is a bit of a bummer b/c for GPs where kernel matrices are often poorly conditioned this can be important. I haven't looked at your particular problem, but as a general workaround one could try to offload the computations where this precision is important to the CPU and leave most other computations on the GPU. I assume this might actually not be too bad for Apple silicon since CPU and GPU share the same "unified" memory so there wouldn't really be overhead from moving things between memory.

[ghost]

Thank you for your reply.

Would making these changes mean modifying the source code of variational_strategy and kernels to make CPU / GPU switches during training? Wouldn't there be other ways?

Thank you very much!

[Balandat]

Would making these changes mean modifying the source code of variational_strategy and kernels to make CPU / GPU switches during training?

Yes, likely.

Wouldn't there be other ways?

I haven't looked at the details, but unless you do some pretty hacky things with mocks I am not sure how to achieve that.