Torchlayer error when running on GPU - Tensor is on CPU, but expected it to be on GPU

[zzh237]

Hi, the below code throws errors if I run on GPU,

@staticmethod
def backward(ctx, dy):  # pragma: no cover
        """Implements the backwards pass QNode vector-Jacobian product"""
        ctx.dy = dy
        vjp = dy.view(1, -1) @ ctx.jacobian.apply(ctx, *ctx.saved_tensors) ##error throws 
        vjp = torch.unbind(vjp.view(-1))
        return (None,) + tuple(vjp)

Error message:

36
  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py", line 225, in grad
37
    inputs, allow_unused, accumulate_grad=False)
38
  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/function.py", line 89, in apply
39
    return self._forward_cls.backward(self, *args)  # type: ignore
40
  File "/usr/local/lib/python3.6/dist-packages/pennylane/interfaces/torch.py", line 175, in backward
41
    vjp = dy.view(1, -1) @ ctx.jacobian.apply(ctx, *ctx.saved_tensors)
42
RuntimeError: Tensor for argument #3 'mat2' is on CPU, but expected it to be on GPU (while checking arguments for addmm)

If I change the code to the code like below, then it works! So could you change the code? Thanks.

 @staticmethod
    def backward(ctx, dy):  # pragma: no cover
        """Implements the backwards pass QNode vector-Jacobian product"""
        ctx.dy = dy
        if dy.is_cuda:
            cuda_device = dy.get_device()
        vjp = dy.view(1, -1) @ ctx.jacobian.apply(ctx, *ctx.saved_tensors).to(dy)
        vjp = torch.unbind(vjp.view(-1))
        return (None,) + tuple(vjp)

[mariaschuld]

Hey @zzh237. I am not sure I understand this issue - you seem to refer to some previous thread (i.e., what do you mean by "it still throws the error"?)...

Could you please edit and clarify your message, explaining what minimum working example you run, what you expect to see, and what the unexpected behaviour is? Information on your system and PL version will also help speeding up support. I tried to fix your formatting a bit already, but the idea is that issues are self-contained reports of a problem. Thanks! :)

[zzh237]

Hey @zzh237. I am not sure I understand this issue - you seem to refer to some previous thread (i.e., what do you mean by "it still throws the error"?)...

Could you please edit and clarify your message, explaining what minimum working example you run, what you expect to see, and what the unexpected behaviour is? Information on your system and PL version will also help speeding up support. I tried to fix your formatting a bit already, but the idea is that issues are self-contained reports of a problem. Thanks! :)

Hi @mariaschuld , thank you so much for the format! I have updated that.

[mariaschuld]

Let me try to understand...

You run some PennyLane code (please post this code too as a minimum working example - I can only guess that you are using PennyLane in combination with PyTorch?) and the code is supposed to be executed on your GPU. But when you run the code you get an error, which you can fix by changing those lines you show, i.e. by sending an object in the vjp calculation to the GPU? (I am not an expert on this, so I am slightly confused why creating a variable cuda_device = dy.get_device() which is unused in the remainder of the function helps?) Can you explain how you found this solution?

We are really keen to improve PennyLane's GPU capabilities, so happy to try and consider any changes. But we definitely need more context to help here!

[zzh237]

Here is the code:

 n_qubits = 4                # Number of qubits
q_depth = 2                 # Depth of the quantum circuit (number of variational layers)
q_delta = 0.01              # Initial spread of random quantum weights
dev = qml.device("default.qubit", wires=n_qubits)

@qml.qnode(dev, interface="torch")
def quantum_net(q_input_features, q_weights_flat):
	"""
	The variational quantum circuit.
	"""

	# Reshape weights
	q_weights = q_weights_flat.reshape(q_depth, n_qubits)

	# Start from state |+> , unbiased w.r.t. |0> and |1>
	H_layer(n_qubits)

	# Embed features in the quantum node
	RY_layer(q_input_features)

	# Sequence of trainable variational layers
	for k in range(q_depth):
		entangling_layer(n_qubits)
		RY_layer(q_weights[k])

	# Expectation values in the Z basis
	exp_vals = [qml.expval(qml.PauliZ(position)) for position in range(n_qubits)]
	return tuple(exp_vals)

class Net(nn.Module):
    """
    Torch module implementing the *dressed* quantum net.
    """

    def __init__(self):
        """
        Definition of the *dressed* layout.
        """

        super().__init__()

        self.q_params = nn.Parameter(q_delta * torch.randn(q_depth * n_qubits))


    def forward(self, input_features):
        """
        Defining how tensors are supposed to move through the *dressed* quantum
        net.
        """

        # obtain the input features for the quantum circuit
        # by reducing the feature dimension from 512 to 4
        q_in = torch.tanh(input_features) * np.pi / 2.0

        # Apply the quantum circuit to each element of the batch and append to q_out
        q_out = torch.Tensor(0, n_qubits)
        q_out = q_out.to(device)
        for elem in q_in:
            q_out_elem = quantum_net(elem, self.q_params).float().unsqueeze(0)
            q_out = torch.cat((q_out, q_out_elem))

        # return the two-dimensional prediction from the postprocessing layer
        return q_out
device = "cuda:0"
net = Net(layer_sizes).to(device)

for i in range(epoch, self.args.epochs):
            for x, y in self.trainloader:
                     net.train()
       
                      self.optimizer.zero_grad() 
       
                       print("### model is on GPU", next(self.model.parameters()).is_cuda) 
        
                       out = net(x)
        
                       loss = F.cross_entropy(out, y, reduction='mean') 
 
                       loss.backward()

Error message:

24
    loss.backward()
25
  File "/usr/local/lib/python3.6/dist-packages/torch/tensor.py", line 245, in backward
26
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
27
  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py", line 147, in backward
28
    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
29
  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/function.py", line 89, in apply
30
    return self._forward_cls.backward(self, *args)  # type: ignore
31
  File "/usr/local/lib/python3.6/dist-packages/pennylane/interfaces/torch.py", line 175, in backward
32
    vjp = dy.view(1, -1) @ ctx.jacobian.apply(ctx, *ctx.saved_tensors)
33
RuntimeError: Tensor for argument #3 'mat2' is on CPU, but expected it to be on GPU (while checking arguments for addmm)

I found that solution by looking at this:

Torchlayer error when running on GPU #709

[mariaschuld]

Ah, perfect, thanks! Let me get back to you on this.

[ADITYA964]

I am also facing this same problem when using penny lane and pytorch.

[mariaschuld]

Hey @zzh237 and @ADITYA964. It looks like while ironing out how PennyLane can fully run on GPUs is on our near-term to-do list, this will be a bigger effort. If you are keen to contribute, feel free to discuss solutions here and make a PR once we decided on a way forward!

I wonder if in the meantime the fixes in the PR you mentioned could help? Sorry that I cannot do more at this stage!

[ADITYA964]

@mariaschuld understood. For those people that have problem like me and @zzh237 , downgrade the version of pennylane .

pip install pennylane==0.14.1

It works perfect using this version.

[mariaschuld]

That is an important piece of information, thanks @ADITYA964!

@josh also tagging you here to keep this in mind going forward.