Best practices for autograd in PyTorch using cuSV

[mtjrider]

@wcqc writes:

When wrapping a cuQuantum API call by defining a custom torch.autograd.Function, say we are wrapping apply_matrix_batched(...) for example, a backward(ctx, *grad_outputs) definition is required by PyTorch (although None could be returned in some cases). In this case especially for parametrized gates such CRX, the gate matrices would be involved in the backward pass for computing/passing the gradients; however, it seems the autodiff rule for each gate needs to be manually defined separately? What would be the right way to wrap a cuquantum Python API call with torch.autograd.Function such that all (parametrized) gates are handled automatically without the need for manual gradient computations? Any pointers are welcome (including being referred to PyTorch discussions and asking there). Thanks.

My understanding: you want to know how best to define forward, backward w.r.t. custatevec.apply_matrix_batched. Is that correct?

[wcqc]

Yes. I've defined the forward pass and it seems to be working fine but it seems to in order to make everything work the backward() def would contain the derived versions of the gates used in forward(). I'm subclassing torch.autograd.Function.

[wcqc]

Apologies! The additional question was not meant to push — just an actual additional question came up while we were thinking hard about the backward pass in cuquantum. The above answer is helpful as always, we really appreciate your very detailed code and answers so far! Will try to be more specific when we have more concrete thoughts/code for this. Thanks again.

[mtjrider]

Of course! Keep us posted on any questions you have. The more you can share with us, the better! 😄

[wcqc]

Both will require that you manage intermediate state data separate from the actual state data, and this is difficult to discuss without some kind of toy example.

Just to confirm: Does this mean the backward pass won't be much more efficient than without cuQuantum as it currently stands, as all the gradients are essentially managed outside of cuquantum, e.g., in standard PyTorch (assuming one doesn't go ahead and implement gradient calculations in cuquantum)?

[DmitryLyakh]

Just to be precise, cuQuantum has two modules, cuStateVec (state-vector library) and cuTensorNet (tensor-network library). My understanding is that you are specifically asking about cuStateVec capabilities in terms of back-propagation, which has been answered above. In contrast, the cuTensorNet module provides gradients via back-propagation (experimental support), although that's probably not what you are looking for since it seems like you are only interested in state-vector algorithms. Just FYI.

[mtjrider]

@DmitryLyakh is correct. Thanks for the quick summary.

@wcqc the answer to your question as it pertains to custatevec.apply_matrix is nuanced.
apply_matrix operates on the data in-place. It's provided a pointer to the state, gate, etc., and it checks the status code provided by the interface binding. In other words, it is opaque w.r.t. PyTorch.

One implication for the aforementioned is that PyTorch cannot maintain an accurate DAG to compute chain-rule contributions to the gradient calculation. There is some loose documentation on this subject here.

If you implement backward on something like apply_matrix directly, the intermediate state vector data required to maintain the gradient will be proportional to the number of times a parametrized apply_matrix API is called. This may be optimal for certain circuits; however, if there are many such API calls, the performance will be limited by:

1. The available memory bandwidth and the latency associated with maintaining intermediate state vector data on a leaf-node basis (see PyTorch documentation linked above) ... and
2. The overhead of a single API call which cannot be interleaved with (1).

This can be mitigated by deferring and accumulating gradient contributions so that they span multiple calls to apply_matrix. At this moment, I'll reiterate that the choice to accumulate + defer is highly dependent on the circuit in question as well as the type of gate subject to parametrization.

[wcqc]

Deleted the question about CUDA quantum, just saw that it was doing parameter shift in the backward() impl. and backprop wasn't automatic either.

[mtjrider]

Without understanding your purpose, it's difficult to comment.
As an example, I feel tempted to mention adjoint differentiation.

PennyLane has a fun tutorial/demo on this subject, here.