From a284fef77c4f7417fa077f3a18e05d6c5ca9213f Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 13:45:54 +0100 Subject: [PATCH 1/9] ModelParallelism example --- examples/parallel.py | 55 ++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- 2 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 examples/parallel.py diff --git a/examples/parallel.py b/examples/parallel.py new file mode 100644 index 00000000..ce61f3c4 --- /dev/null +++ b/examples/parallel.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import torch + +import pyqtorch as pyq +from pyqtorch.parametric import Parametric + +N_DEVICES = 2 + +assert torch.cuda.device_count() == N_DEVICES + + +class ParallelCircuit(torch.nn.Module): + def __init__(self, c0: pyq.QuantumCircuit, c1: pyq.QuantumCircuit, params_c0, params_c1): + super().__init__() + self.c0 = c0.to("cuda:0") + self.c1 = c1.to("cuda:1") + self.params_c0 = torch.nn.ParameterDict({k: v.to("cuda:0") for k, v in params_c0.items()}) + self.params_c1 = torch.nn.ParameterDict({k: v.to("cuda:1") for k, v in params_c1.items()}) + + def forward(self, state: torch.Tensor, values: dict = dict()) -> torch.Tensor: + state = self.c0.forward(state.to("cuda:0"), self.params_c0) + return self.c1.forward(state.to("cuda:1"), self.params_c1) + + +def hea(n_qubits: int, n_layers: int, param_name: str) -> list: + ops = [] + for layer in range(n_layers): + ops += [pyq.RX(i, f"{param_name}_0_{layer}_{i}") for i in range(n_qubits)] + ops += [pyq.RY(i, f"{param_name}_1_{layer}_{i}") for i in range(n_qubits)] + ops += [pyq.RX(i, f"{param_name}_2_{layer}_{i}") for i in range(n_qubits)] + ops += [pyq.CNOT(i % n_qubits, (i + 1) % n_qubits) for i in range(n_qubits)] + return ops + + +n_qubits = 2 +c0 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "theta")) +c1 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "phi")) + + +def init_params(circ: pyq.QuantumCircuit) -> torch.nn.ParameterDict: + return torch.nn.ParameterDict( + { + op.param_name: torch.rand(1, requires_grad=True) + for op in circ.operations + if isinstance(op, Parametric) + } + ) + + +params_c0 = init_params(c0) +params_c1 = init_params(c1) + +circ = ParallelCircuit(c0, c1, params_c0, params_c1) +new_state = circ.forward(pyq.zero_state(n_qubits)) diff --git a/pyproject.toml b/pyproject.toml index 5f13d939..d54b983e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ authors = [ ] requires-python = ">=3.8,<3.13" license = {text = "Apache 2.0"} -version = "1.0.6" +version = "1.0.7" classifiers=[ "License :: OSI Approved :: Apache Software License", "Programming Language :: Python", From 746d7e83183ddf16302df66b490b3462e06cf22b Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 14:02:02 +0100 Subject: [PATCH 2/9] working example --- examples/parallel.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/examples/parallel.py b/examples/parallel.py index ce61f3c4..ed52ed4a 100644 --- a/examples/parallel.py +++ b/examples/parallel.py @@ -7,6 +7,7 @@ N_DEVICES = 2 +assert torch.cuda.is_available() assert torch.cuda.device_count() == N_DEVICES @@ -17,10 +18,13 @@ def __init__(self, c0: pyq.QuantumCircuit, c1: pyq.QuantumCircuit, params_c0, pa self.c1 = c1.to("cuda:1") self.params_c0 = torch.nn.ParameterDict({k: v.to("cuda:0") for k, v in params_c0.items()}) self.params_c1 = torch.nn.ParameterDict({k: v.to("cuda:1") for k, v in params_c1.items()}) + self.observable = pyq.Z(0).to("cuda:1") def forward(self, state: torch.Tensor, values: dict = dict()) -> torch.Tensor: state = self.c0.forward(state.to("cuda:0"), self.params_c0) - return self.c1.forward(state.to("cuda:1"), self.params_c1) + state = self.c1.forward(state.to("cuda:1"), self.params_c1) + projected = self.observable.forward(state) + return pyq.inner_prod(state, projected).real def hea(n_qubits: int, n_layers: int, param_name: str) -> list: @@ -53,3 +57,14 @@ def init_params(circ: pyq.QuantumCircuit) -> torch.nn.ParameterDict: circ = ParallelCircuit(c0, c1, params_c0, params_c1) new_state = circ.forward(pyq.zero_state(n_qubits)) + +optimizer = torch.optim.Adam({**circ.params_c0, **circ.params_c1}.values(), lr=0.01, foreach=False) +epochs = 10 + +for epoch in range(epochs): + optimizer.zero_grad() + y_pred = circ.forward(pyq.zero_state(n_qubits)) + loss = y_pred - torch.rand(1, device=y_pred.device) + loss.backward() + print(f"{epoch}:{loss.item()}") + optimizer.step() From ccc7677853be460bb47d67c2a7dde403bbdc856f Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 14:49:41 +0100 Subject: [PATCH 3/9] example with FM --- examples/parallel.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/examples/parallel.py b/examples/parallel.py index ed52ed4a..166c73f4 100644 --- a/examples/parallel.py +++ b/examples/parallel.py @@ -1,6 +1,10 @@ from __future__ import annotations +from functools import reduce +from operator import add + import torch +from torch.nn.functional import mse_loss import pyqtorch as pyq from pyqtorch.parametric import Parametric @@ -14,13 +18,17 @@ class ParallelCircuit(torch.nn.Module): def __init__(self, c0: pyq.QuantumCircuit, c1: pyq.QuantumCircuit, params_c0, params_c1): super().__init__() + self.feature_map = pyq.QuantumCircuit( + n_qubits, [pyq.RX(i, "x") for i in range(n_qubits)] + ).to("cuda:0") self.c0 = c0.to("cuda:0") self.c1 = c1.to("cuda:1") self.params_c0 = torch.nn.ParameterDict({k: v.to("cuda:0") for k, v in params_c0.items()}) self.params_c1 = torch.nn.ParameterDict({k: v.to("cuda:1") for k, v in params_c1.items()}) self.observable = pyq.Z(0).to("cuda:1") - def forward(self, state: torch.Tensor, values: dict = dict()) -> torch.Tensor: + def forward(self, state: torch.Tensor, inputs: dict = dict()) -> torch.Tensor: + state = self.feature_map.forward(state, {k: v.to("cuda:0") for k, v in inputs.items()}) state = self.c0.forward(state.to("cuda:0"), self.params_c0) state = self.c1.forward(state.to("cuda:1"), self.params_c1) projected = self.observable.forward(state) @@ -52,19 +60,35 @@ def init_params(circ: pyq.QuantumCircuit) -> torch.nn.ParameterDict: ) +# Target function and some training data +def fn(x, degree): + return 0.05 * reduce(add, (torch.cos(i * x) + torch.sin(i * x) for i in range(degree)), 0) + + +x = torch.linspace(0, 10, 100) +y = fn(x, 5) + + params_c0 = init_params(c0) params_c1 = init_params(c1) circ = ParallelCircuit(c0, c1, params_c0, params_c1) -new_state = circ.forward(pyq.zero_state(n_qubits)) +state = pyq.zero_state(n_qubits) + + +def exp_fn(inputs: dict[str, torch.Tensor]) -> torch.Tensor: + return pyq.expectation( + circ, state, {**circ.params_c0, **circ.params_c1, **inputs}, circ.observable, "ad" + ) + optimizer = torch.optim.Adam({**circ.params_c0, **circ.params_c1}.values(), lr=0.01, foreach=False) epochs = 10 for epoch in range(epochs): optimizer.zero_grad() - y_pred = circ.forward(pyq.zero_state(n_qubits)) - loss = y_pred - torch.rand(1, device=y_pred.device) + y_pred = circ.forward(pyq.zero_state(n_qubits), {"x": x}) + loss = mse_loss(y_pred, y.to("cuda:1")) loss.backward() print(f"{epoch}:{loss.item()}") optimizer.step() From a9bb159c05b5642397e887e232b5c6335d4689da Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 14:50:52 +0100 Subject: [PATCH 4/9] move state --- examples/parallel.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/parallel.py b/examples/parallel.py index 166c73f4..ac7bbcb4 100644 --- a/examples/parallel.py +++ b/examples/parallel.py @@ -28,8 +28,10 @@ def __init__(self, c0: pyq.QuantumCircuit, c1: pyq.QuantumCircuit, params_c0, pa self.observable = pyq.Z(0).to("cuda:1") def forward(self, state: torch.Tensor, inputs: dict = dict()) -> torch.Tensor: - state = self.feature_map.forward(state, {k: v.to("cuda:0") for k, v in inputs.items()}) - state = self.c0.forward(state.to("cuda:0"), self.params_c0) + state = self.feature_map.forward( + state.to("cuda:0"), {k: v.to("cuda:0") for k, v in inputs.items()} + ) + state = self.c0.forward(state, self.params_c0) state = self.c1.forward(state.to("cuda:1"), self.params_c1) projected = self.observable.forward(state) return pyq.inner_prod(state, projected).real From 562823186a7bec5f943054618733576a930fccff Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 15:45:13 +0100 Subject: [PATCH 5/9] pipeline parallel --- examples/parallel.py | 105 +++++++++++++++++++++++++------------------ 1 file changed, 62 insertions(+), 43 deletions(-) diff --git a/examples/parallel.py b/examples/parallel.py index ac7bbcb4..702722a7 100644 --- a/examples/parallel.py +++ b/examples/parallel.py @@ -1,8 +1,10 @@ from __future__ import annotations +import timeit from functools import reduce from operator import add +import numpy as np import torch from torch.nn.functional import mse_loss @@ -10,27 +12,38 @@ from pyqtorch.parametric import Parametric N_DEVICES = 2 +N_QUBITS = 2 +N_POINTS = 100 assert torch.cuda.is_available() assert torch.cuda.device_count() == N_DEVICES -class ParallelCircuit(torch.nn.Module): - def __init__(self, c0: pyq.QuantumCircuit, c1: pyq.QuantumCircuit, params_c0, params_c1): +def init_params(circ: pyq.QuantumCircuit, device: torch.device) -> torch.nn.ParameterDict: + return torch.nn.ParameterDict( + { + op.param_name: torch.rand(1, requires_grad=True, device=device) + for op in circ.operations + if isinstance(op, Parametric) + } + ) + + +class ModelParallelCircuit(torch.nn.Module): + def __init__(self, n_qubits: int = N_QUBITS): super().__init__() self.feature_map = pyq.QuantumCircuit( n_qubits, [pyq.RX(i, "x") for i in range(n_qubits)] ).to("cuda:0") - self.c0 = c0.to("cuda:0") - self.c1 = c1.to("cuda:1") - self.params_c0 = torch.nn.ParameterDict({k: v.to("cuda:0") for k, v in params_c0.items()}) - self.params_c1 = torch.nn.ParameterDict({k: v.to("cuda:1") for k, v in params_c1.items()}) + self.c0 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "theta")).to("cuda:0") + self.c1 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "phi")).to("cuda:1") + self.params_c0 = init_params(self.c0, device="cuda:0") + self.params_c1 = init_params(self.c1, device="cuda:1") self.observable = pyq.Z(0).to("cuda:1") - def forward(self, state: torch.Tensor, inputs: dict = dict()) -> torch.Tensor: - state = self.feature_map.forward( - state.to("cuda:0"), {k: v.to("cuda:0") for k, v in inputs.items()} - ) + def forward(self, x: torch.Tensor) -> torch.Tensor: + state = pyq.zero_state(N_QUBITS) + state = self.feature_map.forward(state.to("cuda:0"), {"x": x.to("cuda:0")}) state = self.c0.forward(state, self.params_c0) state = self.c1.forward(state.to("cuda:1"), self.params_c1) projected = self.observable.forward(state) @@ -47,50 +60,56 @@ def hea(n_qubits: int, n_layers: int, param_name: str) -> list: return ops -n_qubits = 2 -c0 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "theta")) -c1 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "phi")) +def fn(x: torch.Tensor, degree: int) -> torch.Tensor: + return 0.05 * reduce(add, (torch.cos(i * x) + torch.sin(i * x) for i in range(degree)), 0) -def init_params(circ: pyq.QuantumCircuit) -> torch.nn.ParameterDict: - return torch.nn.ParameterDict( - { - op.param_name: torch.rand(1, requires_grad=True) - for op in circ.operations - if isinstance(op, Parametric) - } - ) +x = torch.linspace(0, 10, N_POINTS) +y = fn(x, 5) +circ = ModelParallelCircuit(N_QUBITS) -# Target function and some training data -def fn(x, degree): - return 0.05 * reduce(add, (torch.cos(i * x) + torch.sin(i * x) for i in range(degree)), 0) + +optimizer = torch.optim.Adam({**circ.params_c0, **circ.params_c1}.values(), lr=0.01, foreach=False) +epochs = 10 -x = torch.linspace(0, 10, 100) -y = fn(x, 5) +def train(circ) -> None: + for epoch in range(epochs): + optimizer.zero_grad() + y_pred = circ.forward(x) + loss = mse_loss(y_pred, y.to("cuda:1")) + loss.backward() + print(f"{epoch}:{loss.item()}") + optimizer.step() -params_c0 = init_params(c0) -params_c1 = init_params(c1) +class PipelineParallelCircuit(ModelParallelCircuit): + def __init__(self, split_size=N_POINTS // 2, *args, **kwargs): + super(PipelineParallelCircuit, self).__init__(*args, **kwargs) + self.split_size = split_size -circ = ParallelCircuit(c0, c1, params_c0, params_c1) -state = pyq.zero_state(n_qubits) + def forward(self, x: torch.Tensor) -> torch.Tensor: + state = pyq.zero_state(N_QUBITS) + splits = iter(x.split(self.split_size, dim=0)) + s_next = next(splits) + s_prev = self.feature_map(state, {"x": s_next.to("cuda:0")}) + s_prev = self.c0.forward(state, self.params_c0).to("cuda:1") + ret = [] + for s_next in splits: + s_prev = self.c1.forward(s_prev, self.params_c1) + ret.append(pyq.inner_prod(s_prev, self.observable.forward(s_prev)).real) -def exp_fn(inputs: dict[str, torch.Tensor]) -> torch.Tensor: - return pyq.expectation( - circ, state, {**circ.params_c0, **circ.params_c1, **inputs}, circ.observable, "ad" - ) + s_prev = self.feature_map(state, {"x": s_next.to("cuda:0")}) + s_prev = self.c0.forward(state, self.params_c0).to("cuda:1") + s_prev = self.c1.forward(s_prev, self.params_c1) + ret.append(pyq.inner_prod(s_prev, self.observable.forward(s_prev)).real) + + return torch.cat(ret) -optimizer = torch.optim.Adam({**circ.params_c0, **circ.params_c1}.values(), lr=0.01, foreach=False) -epochs = 10 -for epoch in range(epochs): - optimizer.zero_grad() - y_pred = circ.forward(pyq.zero_state(n_qubits), {"x": x}) - loss = mse_loss(y_pred, y.to("cuda:1")) - loss.backward() - print(f"{epoch}:{loss.item()}") - optimizer.step() +setup = "model = PipelineParallelCircuit()" +pp_run_times = timeit.repeat("train(circ)", setup, number=1, repeat=10, globals=globals()) +pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times) From 291ecb0403d7af11d9eb7b0ef6cc98c6d12a41ce Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 15:47:13 +0100 Subject: [PATCH 6/9] state to dev --- examples/parallel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/parallel.py b/examples/parallel.py index 702722a7..3c2d60d7 100644 --- a/examples/parallel.py +++ b/examples/parallel.py @@ -42,7 +42,7 @@ def __init__(self, n_qubits: int = N_QUBITS): self.observable = pyq.Z(0).to("cuda:1") def forward(self, x: torch.Tensor) -> torch.Tensor: - state = pyq.zero_state(N_QUBITS) + state = pyq.zero_state(N_QUBITS).to("cuda:0") state = self.feature_map.forward(state.to("cuda:0"), {"x": x.to("cuda:0")}) state = self.c0.forward(state, self.params_c0) state = self.c1.forward(state.to("cuda:1"), self.params_c1) @@ -90,7 +90,7 @@ def __init__(self, split_size=N_POINTS // 2, *args, **kwargs): self.split_size = split_size def forward(self, x: torch.Tensor) -> torch.Tensor: - state = pyq.zero_state(N_QUBITS) + state = pyq.zero_state(N_QUBITS).to("cuda:0") splits = iter(x.split(self.split_size, dim=0)) s_next = next(splits) s_prev = self.feature_map(state, {"x": s_next.to("cuda:0")}) From 40ed91e457bee6fe674715d3da9c6ae3784a6248 Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 16:24:20 +0100 Subject: [PATCH 7/9] working benchmark --- examples/parallel.py | 92 +++++++++++++++++++++++++++----------------- 1 file changed, 57 insertions(+), 35 deletions(-) diff --git a/examples/parallel.py b/examples/parallel.py index 3c2d60d7..b99e9fd8 100644 --- a/examples/parallel.py +++ b/examples/parallel.py @@ -14,6 +14,15 @@ N_DEVICES = 2 N_QUBITS = 2 N_POINTS = 100 +N_EPOCHS = 1 + + +def fn(x: torch.Tensor, degree: int) -> torch.Tensor: + return 0.05 * reduce(add, (torch.cos(i * x) + torch.sin(i * x) for i in range(degree)), 0) + + +x = torch.linspace(0, 10, N_POINTS) +y = fn(x, 5) assert torch.cuda.is_available() assert torch.cuda.device_count() == N_DEVICES @@ -29,6 +38,37 @@ def init_params(circ: pyq.QuantumCircuit, device: torch.device) -> torch.nn.Para ) +def hea(n_qubits: int, n_layers: int, param_name: str) -> list: + ops = [] + for layer in range(n_layers): + ops += [pyq.RX(i, f"{param_name}_0_{layer}_{i}") for i in range(n_qubits)] + ops += [pyq.RY(i, f"{param_name}_1_{layer}_{i}") for i in range(n_qubits)] + ops += [pyq.RX(i, f"{param_name}_2_{layer}_{i}") for i in range(n_qubits)] + ops += [pyq.CNOT(i % n_qubits, (i + 1) % n_qubits) for i in range(n_qubits)] + return ops + + +class SingleDeviceCircuit(torch.nn.Module): + def __init__(self, n_qubits: int = N_QUBITS): + super().__init__() + self.feature_map = pyq.QuantumCircuit( + n_qubits, [pyq.RX(i, "x") for i in range(n_qubits)] + ).to("cuda:1") + self.c0 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "theta")).to("cuda:1") + self.c1 = pyq.QuantumCircuit(n_qubits, hea(n_qubits, 1, "phi")).to("cuda:1") + self.params_c0 = init_params(self.c0, device="cuda:1") + self.params_c1 = init_params(self.c1, device="cuda:1") + self.observable = pyq.Z(0).to("cuda:1") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + state = pyq.zero_state(N_QUBITS).to("cuda:1") + state = self.feature_map.forward(state.to("cuda:1"), {"x": x.to("cuda:1")}) + state = self.c0.forward(state, self.params_c0) + state = self.c1.forward(state.to("cuda:1"), self.params_c1) + projected = self.observable.forward(state) + return pyq.inner_prod(state, projected).real + + class ModelParallelCircuit(torch.nn.Module): def __init__(self, n_qubits: int = N_QUBITS): super().__init__() @@ -50,37 +90,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return pyq.inner_prod(state, projected).real -def hea(n_qubits: int, n_layers: int, param_name: str) -> list: - ops = [] - for layer in range(n_layers): - ops += [pyq.RX(i, f"{param_name}_0_{layer}_{i}") for i in range(n_qubits)] - ops += [pyq.RY(i, f"{param_name}_1_{layer}_{i}") for i in range(n_qubits)] - ops += [pyq.RX(i, f"{param_name}_2_{layer}_{i}") for i in range(n_qubits)] - ops += [pyq.CNOT(i % n_qubits, (i + 1) % n_qubits) for i in range(n_qubits)] - return ops - - -def fn(x: torch.Tensor, degree: int) -> torch.Tensor: - return 0.05 * reduce(add, (torch.cos(i * x) + torch.sin(i * x) for i in range(degree)), 0) - - -x = torch.linspace(0, 10, N_POINTS) -y = fn(x, 5) - -circ = ModelParallelCircuit(N_QUBITS) - - -optimizer = torch.optim.Adam({**circ.params_c0, **circ.params_c1}.values(), lr=0.01, foreach=False) -epochs = 10 - - def train(circ) -> None: - for epoch in range(epochs): + optimizer = torch.optim.Adam( + {**circ.params_c0, **circ.params_c1}.values(), lr=0.01, foreach=False + ) + for epoch in range(N_EPOCHS): optimizer.zero_grad() y_pred = circ.forward(x) loss = mse_loss(y_pred, y.to("cuda:1")) loss.backward() - print(f"{epoch}:{loss.item()}") optimizer.step() @@ -90,26 +108,30 @@ def __init__(self, split_size=N_POINTS // 2, *args, **kwargs): self.split_size = split_size def forward(self, x: torch.Tensor) -> torch.Tensor: - state = pyq.zero_state(N_QUBITS).to("cuda:0") + init_state = pyq.zero_state(N_QUBITS).to("cuda:0") splits = iter(x.split(self.split_size, dim=0)) s_next = next(splits) - s_prev = self.feature_map(state, {"x": s_next.to("cuda:0")}) - s_prev = self.c0.forward(state, self.params_c0).to("cuda:1") + s_prev = self.feature_map(init_state, {"x": s_next.to("cuda:0")}) + s_prev = self.c0.forward(s_prev, self.params_c0).to("cuda:1") ret = [] for s_next in splits: s_prev = self.c1.forward(s_prev, self.params_c1) ret.append(pyq.inner_prod(s_prev, self.observable.forward(s_prev)).real) - s_prev = self.feature_map(state, {"x": s_next.to("cuda:0")}) - s_prev = self.c0.forward(state, self.params_c0).to("cuda:1") + s_prev = self.feature_map(init_state, {"x": s_next.to("cuda:0")}) + s_prev = self.c0.forward(s_prev, self.params_c0).to("cuda:1") s_prev = self.c1.forward(s_prev, self.params_c1) ret.append(pyq.inner_prod(s_prev, self.observable.forward(s_prev)).real) - return torch.cat(ret) -setup = "model = PipelineParallelCircuit()" -pp_run_times = timeit.repeat("train(circ)", setup, number=1, repeat=10, globals=globals()) -pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times) +if __name__ == "__main__": + res = {} + for model_cls in [SingleDeviceCircuit, ModelParallelCircuit, PipelineParallelCircuit]: + setup = "circ = model_cls(N_QUBITS)" + pp_run_times = timeit.repeat("train(circ)", setup, number=1, repeat=10, globals=globals()) + pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times) + res[model_cls.__name__] = pp_mean + print(res) From 0f26c78a2cc3a713cbd348b1fcbfdf503f17937e Mon Sep 17 00:00:00 2001 From: seitzdom Date: Thu, 22 Feb 2024 16:50:46 +0100 Subject: [PATCH 8/9] better logging --- examples/parallel.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/examples/parallel.py b/examples/parallel.py index b99e9fd8..7637cf88 100644 --- a/examples/parallel.py +++ b/examples/parallel.py @@ -128,10 +128,15 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: if __name__ == "__main__": - res = {} + res = {"n_qubits": N_QUBITS} for model_cls in [SingleDeviceCircuit, ModelParallelCircuit, PipelineParallelCircuit]: - setup = "circ = model_cls(N_QUBITS)" - pp_run_times = timeit.repeat("train(circ)", setup, number=1, repeat=10, globals=globals()) - pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times) - res[model_cls.__name__] = pp_mean + try: + setup = "circ = model_cls(N_QUBITS)" + pp_run_times = timeit.repeat( + "train(circ)", setup, number=1, repeat=10, globals=globals() + ) + pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times) + res[model_cls.__name__] = f"mean_runtime: {pp_mean}, std_runtime: {pp_std}" + except Exception as e: + res[model_cls.__name__] = f"failed, reason: {e}" print(res) From 5df42a521b025c4527dda50793a176af372c2581 Mon Sep 17 00:00:00 2001 From: seitzdom Date: Fri, 23 Feb 2024 18:17:40 +0100 Subject: [PATCH 9/9] modelparallelcircuit constructor --- pyqtorch/circuit.py | 78 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) diff --git a/pyqtorch/circuit.py b/pyqtorch/circuit.py index 1835ecc6..00935da6 100644 --- a/pyqtorch/circuit.py +++ b/pyqtorch/circuit.py @@ -3,6 +3,7 @@ from logging import getLogger from typing import Any, Iterator +import torch from torch import Tensor from torch import device as torch_device from torch.nn import Module, ModuleList @@ -106,3 +107,80 @@ def expectation( return AdjointExpectation.apply(circuit, observable, state, values.keys(), *values.values()) else: raise ValueError(f"Requested diff_mode '{diff_mode}' not supported.") + + +class PipedCircuit(QuantumCircuit): + def __init__(self, n_qubits: int, operations: list[Module], dev_idx: int): + super().__init__(n_qubits, operations) + self = self.to(torch_device(f"cuda:{dev_idx}")) + + def run(self, state: State = None, values: dict[str, Tensor] = {}) -> State: + if state is None: + state = self.init_state() + else: + state = state.to(self.device) + values = {k: v.to(self.device) for k, v in values.items()} + for op in self.operations: + state = op(state, values) + return state + + +class ModelParallelCircuit(QuantumCircuit): + def __init__(self, circ: QuantumCircuit, n_devices: int): + if not all([isinstance(subc, QuantumCircuit) for subc in circ.operations]): + msg = "Make sure the passed QuantumCircuit only contains other QuantumCircuits." + logger.error(msg) + raise ValueError(msg) + if not torch.cuda.is_available(): + msg = f"{self.__class__.__name__} requires at least two GPU devices." + logger.error(msg) + raise ValueError(msg) + dev_count = torch.cuda.device_count() + if dev_count < n_devices: + msg = f"Requested {n_devices} GPU devices however only {dev_count} devices available." + logger.error(msg) + raise ValueError(msg) + n_circs = len(circ.operations) + dev_indices = [i for i in range(n_devices) for _ in range(n_circs // n_devices)] + operations = [ + PipedCircuit(c.n_qubits, c.operations, dev_idx) + for c, dev_idx in zip(circ.operations, dev_indices) + ] + super().__init__(circ.n_qubits, operations) + + +# class PipelineParallelCircuit(ModelParallelCircuit): +# def __init__(self, split_size=N_POINTS // 2, *args, **kwargs): +# super(PipelineParallelCircuit, self).__init__(*args, **kwargs) +# self.split_size = split_size + +# def forward(self, x: torch.Tensor) -> torch.Tensor: +# init_state = zero_state(N_QUBITS).to("cuda:0") +# splits = iter(x.split(self.split_size, dim=0)) +# s_next = next(splits) +# s_prev = self.feature_map(init_state, {"x": s_next.to("cuda:0")}) +# s_prev = self.c0.forward(s_prev, self.params_c0).to("cuda:1") +# ret = [] + +# for s_next in splits: +# s_prev = self.c1.forward(s_prev, self.params_c1) +# ret.append(inner_prod(s_prev, self.observable.forward(s_prev)).real) + +# s_prev = self.feature_map(init_state, {"x": s_next.to("cuda:0")}) +# s_prev = self.c0.forward(s_prev, self.params_c0).to("cuda:1") + +# s_prev = self.c1.forward(s_prev, self.params_c1) +# ret.append(inner_prod(s_prev, self.observable.forward(s_prev)).real) +# return torch.cat(ret) + + +# def train(circ) -> None: +# optimizer = torch.optim.Adam( +# {**circ.params_c0, **circ.params_c1}.values(), lr=0.01, foreach=False +# ) +# for epoch in range(N_EPOCHS): +# optimizer.zero_grad() +# y_pred = circ.forward(x) +# loss = mse_loss(y_pred, y.to("cuda:1")) +# loss.backward() +# optimizer.step()