Use the correct GPU device internally with `default.qubit.torch` #1982

antalszava · 2021-12-03T23:22:43Z

Context:

The way how the default.qubit.torch device handles CUDA compatibility has been updated recently and made it into a v0.19.1 release.

One extra case to cover is related to the use case of having multiple GPUs, e.g.,:

import pennylane as qml
import torch

dev = qml.device("default.qubit.torch", wires=1)

@qml.qnode(dev, interface='torch')
def circuit(weights):
    qml.RX(weights, wires=0)
    return qml.expval(qml.PauliZ(0))

weights = torch.tensor(0.0, device='cuda:1')

circuit(weights)

Description of the Change:

Refactors the logic that stores the device used by the gate parameters to store the exact CUDA device being used including its index.

Benefits:

Use cases, where not the default GPU device (with index 0) is being used, should work.

Possible Drawbacks:
N/A

Related GitHub Issues:
N/A

github-actions · 2021-12-03T23:22:58Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

antalszava · 2021-12-03T23:23:03Z

[sc-11965]

…I/pennylane into multi_gpu_correct_device

antalszava · 2021-12-03T23:27:37Z

pennylane/devices/default_qubit_torch.py

@@ -126,7 +126,8 @@ def circuit(x):
            If ``shots > 0`` is used, the ``diff_method="backprop"``
            QNode differentiation method is not supported and it is recommended to consider
            switching device to ``default.qubit`` and using ``diff_method="parameter-shift"``.
-        torch_device='cpu' (str): the device on which the computation will be run, ``'cpu'`` or ``'cuda'``
+        torch_device='cpu' (str): the device on which the computation will be
+        run, e.g., ``'cpu'`` or ``'cuda'``


This could now even be cuda:1, etc. set internally.

antalszava · 2021-12-03T23:29:25Z

tests/devices/test_default_qubit_torch.py

-            qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=2), requires_grad=True
+            qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=2),
+            requires_grad=True,
+            device=torch_device,


This test was raising a warning for the mismatch between Torch devices; specifying the device argument here helped.

antalszava · 2021-12-03T23:30:06Z

tests/gpu/test_torch.py

@@ -36,7 +36,7 @@ def test_device_to_cuda(self):
        res = dev.execute(tape)

        assert res.is_cuda
-        assert dev._torch_device == "cuda"
+        assert "cuda" in dev._torch_device


Note: we're storing the name of the CUDA device including its index.

antalszava · 2021-12-03T23:30:44Z

tests/interfaces/test_batch_torch.py

@@ -399,7 +399,7 @@ def cost(a):
            dev = qml.device("default.qubit.autograd", wires=2)
            return dev.batch_execute([tape])[0]

-        expected = qml.grad(cost)(0.1)
+        expected = qml.grad(cost, argnum=0)(0.1)


This test was raising a warning wrt. to not marking trainable parameters correctly as per the ongoing deprecation we have.

codecov · 2021-12-03T23:43:44Z

Codecov Report

Merging #1982 (0cda50d) into v0.20.0-rc0 (53b3b2b) will increase coverage by 0.28%.
The diff coverage is 100.00%.

@@               Coverage Diff               @@
##           v0.20.0-rc0    #1982      +/-   ##
===============================================
+ Coverage        98.81%   99.10%   +0.28%     
===============================================
  Files              227      221       -6     
  Lines            17291    16853     -438     
===============================================
- Hits             17086    16702     -384     
+ Misses             205      151      -54

Impacted Files	Coverage Δ
pennylane/devices/default_qubit_torch.py	`92.00% <100.00%> (+5.26%)`	⬆️
pennylane/beta/__init__.py	`0.00% <0.00%> (-100.00%)`	⬇️
pennylane/devices/default_qubit_jax.py	`86.20% <0.00%> (-8.80%)`	⬇️
pennylane/qnode_old.py	`95.52% <0.00%> (-0.09%)`	⬇️
pennylane/transforms/metric_tensor.py	`99.48% <0.00%> (-0.01%)`	⬇️
pennylane/measure.py	`99.03% <0.00%> (ø)`
pennylane/optimize/qng.py	`100.00% <0.00%> (ø)`
pennylane/fourier/__init__.py	`100.00% <0.00%> (ø)`
pennylane/templates/__init__.py	`100.00% <0.00%> (ø)`
pennylane/transforms/control.py	`100.00% <0.00%> (ø)`
... and 13 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 53b3b2b...0cda50d. Read the comment docs.

mlxd · 2021-12-06T11:42:36Z

Hi @antalszava Happy to report the tests pass for me locally.
Still no access to a multi-GPU box to test GPU choice, but it should be safe on single-GPU instances.

albi3ro

Also add # pragma: no cover everywhere that code isn't going to be getting tested without gpu's.

This might be out of scope for this PR, but instead of doing string operations to separate type versus index,

we accept either a torch device or a string that will create a torch device
we wrap the input with torch.device. So then self._torch_device is a torch.device instance
We have a clean separation between type and index and don't have to search in a string.

But at the end of the day, this fixes the problem in front of us, and is good enough. If it runs and solves the problem, feel free to merge in.

albi3ro · 2021-12-06T15:21:48Z

pennylane/devices/default_qubit_torch.py

+                # Using hasattr in case we don't have a Torch tensor as input
+                if hasattr(data, "is_cuda"):
+                    if data.is_cuda:
+                        return ":".join([data.device.type, str(data.device.index)])


Suggested change

# Using hasattr in case we don't have a Torch tensor as input

if hasattr(data, "is_cuda"):

if data.is_cuda:

return ":".join([data.device.type, str(data.device.index)])

if getattr(data, "is_cuda, False):

return str(data.device)

There seems to be a subtlety here, likely because of

In [4]: dev = torch.device("cuda") In [5]: str(dev) Out[5]: 'cuda' In [6]: ":".join([dev.type, str(dev.index)]) Out[6]: 'cuda:None'

This makes 2 tests fail, I'd rather keep it as is. We could think more about the string vs. Torch device storage in the device a bit more in the future.

pennylane/devices/default_qubit_torch.py

Co-authored-by: Christina Lee <christina@xanadu.ai>

…I/pennylane into multi_gpu_correct_device

* Fix metric tensor for circuits with inverted gates (#1987) * fix metric tensor for inverse gates * changelog * modify test for block diagonal * undo tmp change * AmplitudeEmbedding: Fix warning from implicit cast of complex to real (#1990) * Fix warning from implicit cast of complex to real * Changelog * Fix warning without failing other tests * V0.20.0: remove deprecated features (#1981) * remove the decorator * remove the decorator tests * remove the remaining template decorator use * Remove the default.tensor devices * no mention of default.tensor.tf * remove fourier.spectrum * no default.tensor entry points * no beta devices ref in docs * remove the diag_approx keyword argument of qml.metric_tensor and qml.QNGOptimizer * changelog * changelog formatting * Remove the metric_tensor method of the old QNode; remove tests and adjust an example for QNG * remove unused variable * Fix error mitigation tests * Resolve error from texted-based circuit drawer for tapes (#1994) * Add support for drawing tapes to textdrawer Tapes appearing in the operations list of circuits (or other tapes) are now drawn as an opaque tape object, and its contents are appended to the circuit drawing as a subcircuit. * Fix charset of tape subcircuits * Add persistent indices to ops in nested circuits When a circuit to be printed contains tapes nested within, the index printed to identify each tape, matrix argument, etc. is now passed to lower levels of nesting (and passed back up) to retain unique names for circuit elements. * Add tests for drawing nested tapes * Changelog * Review: docstrings, arg name, generator * Make multi_dispatch treat scipy same as numpy (#2001) * Make multi_dispatch treat scipy same as numpy * Changelog * Add test to assert warning is no longer raised * Use pytest.mark instead of warnings module for filter Adding a warning filter using the warnings module seems to persist across unit tests. * Remove empty circuits and device tests to check supported operations (#1979) * check supported operations in more places * add PauliX to measurment qnodes * Incrementing the version number to `0.21.0-dev` (#1988) * changelogs * changelog * version number bumps * clean changelog file * RC to master (#1995) * Fix metric tensor for circuits with inverted gates (#1987) * fix metric tensor for inverse gates * changelog * modify test for block diagonal * undo tmp change * AmplitudeEmbedding: Fix warning from implicit cast of complex to real (#1990) * Fix warning from implicit cast of complex to real * Changelog * Fix warning without failing other tests * V0.20.0: remove deprecated features (#1981) * remove the decorator * remove the decorator tests * remove the remaining template decorator use * Remove the default.tensor devices * no mention of default.tensor.tf * remove fourier.spectrum * no default.tensor entry points * no beta devices ref in docs * remove the diag_approx keyword argument of qml.metric_tensor and qml.QNGOptimizer * changelog * changelog formatting * Remove the metric_tensor method of the old QNode; remove tests and adjust an example for QNG * remove unused variable * Fix error mitigation tests Co-authored-by: David Wierichs <davidwierichs@gmail.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> * eliminate empty circuits, minor rewriting * black, changelog, fix probs * fix merging with rc branch * state non empty circuit Co-authored-by: antalszava <antalszava@gmail.com> Co-authored-by: David Wierichs <davidwierichs@gmail.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> * Update doc/releases/changelog-dev.md * Update pennylane/drawer/representation_resolver.py * Update pennylane/drawer/representation_resolver.py * have a single self.label_offsets dictionary for linting purposes * organize tape drawing into its own method * no too many return statements check * format * [WIP] Jit probs bug (#1998) * trying to use jax unique function to allow for compilation * testing * Overrided the estimate_probability method in the default_qubit_jax device to allow for jit-ing when working with qml.probs * Apply suggestions from code review * missed comment * lint * Added tests * fixed bug with fill value * cleaning up * Fix missing function call in metric tensor example (#2008) * Update docs of `qml.fourier.reconstruct` and `qml.QubitDensityMatrix` (#2005) * reconstruct fixes * QubitDensityMatrix fix * revert merge master commit * minor doc fixes (#2009) * fix typo in CVNeuralNetLayers UsageDetails * Fix some minor doc example for drawer (#2010) * Merge release candidate into master (#2003) * Fix metric tensor for circuits with inverted gates (#1987) * fix metric tensor for inverse gates * changelog * modify test for block diagonal * undo tmp change * AmplitudeEmbedding: Fix warning from implicit cast of complex to real (#1990) * Fix warning from implicit cast of complex to real * Changelog * Fix warning without failing other tests * V0.20.0: remove deprecated features (#1981) * remove the decorator * remove the decorator tests * remove the remaining template decorator use * Remove the default.tensor devices * no mention of default.tensor.tf * remove fourier.spectrum * no default.tensor entry points * no beta devices ref in docs * remove the diag_approx keyword argument of qml.metric_tensor and qml.QNGOptimizer * changelog * changelog formatting * Remove the metric_tensor method of the old QNode; remove tests and adjust an example for QNG * remove unused variable * Fix error mitigation tests * Resolve error from texted-based circuit drawer for tapes (#1994) * Add support for drawing tapes to textdrawer Tapes appearing in the operations list of circuits (or other tapes) are now drawn as an opaque tape object, and its contents are appended to the circuit drawing as a subcircuit. * Fix charset of tape subcircuits * Add persistent indices to ops in nested circuits When a circuit to be printed contains tapes nested within, the index printed to identify each tape, matrix argument, etc. is now passed to lower levels of nesting (and passed back up) to retain unique names for circuit elements. * Add tests for drawing nested tapes * Changelog * Review: docstrings, arg name, generator * Make multi_dispatch treat scipy same as numpy (#2001) * Make multi_dispatch treat scipy same as numpy * Changelog * Add test to assert warning is no longer raised * Use pytest.mark instead of warnings module for filter Adding a warning filter using the warnings module seems to persist across unit tests. * Remove empty circuits and device tests to check supported operations (#1979) * check supported operations in more places * add PauliX to measurment qnodes * Incrementing the version number to `0.21.0-dev` (#1988) * changelogs * changelog * version number bumps * clean changelog file * RC to master (#1995) * Fix metric tensor for circuits with inverted gates (#1987) * fix metric tensor for inverse gates * changelog * modify test for block diagonal * undo tmp change * AmplitudeEmbedding: Fix warning from implicit cast of complex to real (#1990) * Fix warning from implicit cast of complex to real * Changelog * Fix warning without failing other tests * V0.20.0: remove deprecated features (#1981) * remove the decorator * remove the decorator tests * remove the remaining template decorator use * Remove the default.tensor devices * no mention of default.tensor.tf * remove fourier.spectrum * no default.tensor entry points * no beta devices ref in docs * remove the diag_approx keyword argument of qml.metric_tensor and qml.QNGOptimizer * changelog * changelog formatting * Remove the metric_tensor method of the old QNode; remove tests and adjust an example for QNG * remove unused variable * Fix error mitigation tests Co-authored-by: David Wierichs <davidwierichs@gmail.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> * eliminate empty circuits, minor rewriting * black, changelog, fix probs * fix merging with rc branch * state non empty circuit Co-authored-by: antalszava <antalszava@gmail.com> Co-authored-by: David Wierichs <davidwierichs@gmail.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> * Update doc/releases/changelog-dev.md * Update pennylane/drawer/representation_resolver.py * Update pennylane/drawer/representation_resolver.py * have a single self.label_offsets dictionary for linting purposes * organize tape drawing into its own method * no too many return statements check * format Co-authored-by: David Wierichs <davidwierichs@gmail.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> Co-authored-by: antalszava <antalszava@gmail.com> * Style2 Co-authored-by: Christina Lee <christina@xanadu.ai> Co-authored-by: David Wierichs <davidwierichs@gmail.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> Co-authored-by: antalszava <antalszava@gmail.com> * Tape validation for the JAX interface (batch execute) (#2011) * add tape validation to the JAX interface to provide users with a descriptive error message * formatting * changelog * Fix behaviour of `qml.gradients.gradient_transform` for multiple array arguments (#1989) * Fix gradient_transform for complex QNode shapes Previously, the gradient transform did not account for the shape of classical jacobians produced by `qml.jacobian` for complex QNode argument scenarious, leading to a shape mismatch error during tensor contraction of the classical and quantum jacobian. This is resolved by distinguishing between the case of a single classical Jacobian for a single QNode argument, and multiple classical Jacobians stacked together when all QNode arguments have the same shape. The two cases require different tensor manipulations as `qml.jacobian` behaves differently in each. Two test cases are added that include multiple vector arguments to a QNode and non-trainable QNode arguments mixed with trainable ones. * Format * Changelog * Fix coverage by removing some "fallback" branches and adding dummy test. * Update doc/releases/changelog-dev.md Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> * Review: small corrections * Proposal: leave old behaviour as fallback branch * Review: remove tf & torch logic, comment * Review: raise warning for unknown cjac shape Co-authored-by: Antal Szava <antalszava@gmail.com> Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> * Fix docstrings for `CommutingEvolution` (#2012) * Example and docstrings * Undo. * Use the correct GPU device internally with `default.qubit.torch` (#1982) * changes * format * docstring * lint * Update pennylane/devices/default_qubit_torch.py Co-authored-by: Christina Lee <christina@xanadu.ai> * Update pennylane/devices/default_qubit_torch.py Co-authored-by: Christina Lee <christina@xanadu.ai> * no cover; extend docstring * format; correct typo * revert * no cover for GPU part * format Co-authored-by: Christina Lee <christina@xanadu.ai> * updated docstring example (#2019) * Fix Barrier example * `v0.20.0` release notes (#1977) * current release * rename latest changelog * reorder the entries, create sections * version, release notes file * Update doc/releases/changelog-0.20.0.md * Update doc/releases/changelog-0.20.0.md Co-authored-by: Josh Izaac <josh146@gmail.com> * Update doc/releases/changelog-0.20.0.md Co-authored-by: Josh Izaac <josh146@gmail.com> * init module: breaking * Apply suggested changes * Apply suggested changes * rendering issues * mutable kwarg * apply suggestions * move to breaking changes * Lie opt example * Update doc/releases/changelog-0.20.0.md Co-authored-by: Josh Izaac <josh146@gmail.com> * transformation example in the notes * updates * Fix examples * fix refs * contributors * More PRs listed * PauliError * PauliError * custom ops as attribute * Update changelog-0.20.0.md * Update changelog-0.20.0.md Co-authored-by: Josh Izaac <josh146@gmail.com> * readd changelog item * correct version numbers * nl * disallow qml.probs with default.qubit.jax when using a shot vector * no else Co-authored-by: David Wierichs <davidwierichs@gmail.com> Co-authored-by: David Ittah <dime10@users.noreply.github.com> Co-authored-by: Christina Lee <christina@xanadu.ai> Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca> Co-authored-by: Romain <rmoyard@gmail.com> Co-authored-by: Josh Izaac <josh146@gmail.com>

antalszava added 4 commits December 3, 2021 18:15

changes

5dfff62

format

9c63064

docstring

50a1d72

Merge branch 'master' into multi_gpu_correct_device

ca0e287

antalszava added 2 commits December 3, 2021 18:23

lint

634e019

Merge branch 'multi_gpu_correct_device' of https://github.com/XanaduA…

4f780e5

…I/pennylane into multi_gpu_correct_device

antalszava commented Dec 3, 2021

View reviewed changes

antalszava requested review from albi3ro and mlxd December 3, 2021 23:31

Merge branch 'master' into multi_gpu_correct_device

a178292

mlxd approved these changes Dec 6, 2021

View reviewed changes

Merge branch 'master' into multi_gpu_correct_device

1acc9ed

antalszava changed the base branch from master to v0.20.0-rc0 December 6, 2021 17:16

albi3ro reviewed Dec 7, 2021

View reviewed changes

Base automatically changed from v0.20.0-rc0 to master December 9, 2021 19:18

antalszava changed the base branch from master to v0.20.0-rc0 December 10, 2021 22:25

antalszava and others added 7 commits December 10, 2021 17:28

Update pennylane/devices/default_qubit_torch.py

329983a

Co-authored-by: Christina Lee <christina@xanadu.ai>

Update pennylane/devices/default_qubit_torch.py

0f2e7c7

Co-authored-by: Christina Lee <christina@xanadu.ai>

no cover; extend docstring

7ead23c

format; correct typo

ce387a8

Merge branch 'multi_gpu_correct_device' of https://github.com/XanaduA…

b27d07d

…I/pennylane into multi_gpu_correct_device

revert

70ef3fe

no cover for GPU part

2d3e316

format

0cda50d

antalszava merged commit 021be57 into v0.20.0-rc0 Dec 11, 2021

antalszava deleted the multi_gpu_correct_device branch December 11, 2021 00:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use the correct GPU device internally with `default.qubit.torch` #1982

Use the correct GPU device internally with `default.qubit.torch` #1982

antalszava commented Dec 3, 2021

github-actions bot commented Dec 3, 2021

antalszava commented Dec 3, 2021

antalszava Dec 3, 2021

antalszava Dec 3, 2021

antalszava Dec 3, 2021

antalszava Dec 3, 2021

codecov bot commented Dec 3, 2021 •

edited

Loading

mlxd commented Dec 6, 2021

albi3ro left a comment

albi3ro Dec 6, 2021

antalszava Dec 10, 2021

Use the correct GPU device internally with default.qubit.torch #1982

Use the correct GPU device internally with default.qubit.torch #1982

Conversation

antalszava commented Dec 3, 2021

github-actions bot commented Dec 3, 2021

antalszava commented Dec 3, 2021

antalszava Dec 3, 2021

Choose a reason for hiding this comment

antalszava Dec 3, 2021

Choose a reason for hiding this comment

antalszava Dec 3, 2021

Choose a reason for hiding this comment

antalszava Dec 3, 2021

Choose a reason for hiding this comment

codecov bot commented Dec 3, 2021 • edited Loading

Codecov Report

mlxd commented Dec 6, 2021

albi3ro left a comment

Choose a reason for hiding this comment

albi3ro Dec 6, 2021

Choose a reason for hiding this comment

antalszava Dec 10, 2021

Choose a reason for hiding this comment

Use the correct GPU device internally with `default.qubit.torch` #1982

Use the correct GPU device internally with `default.qubit.torch` #1982

codecov bot commented Dec 3, 2021 •

edited

Loading