Memory leak in qiskit/circuit/parametertable.py ?

floriankittelmann commented 10 months ago

Environment

Qiskit Terra version: 0.25.3
Qiskit version: 0.44.3
Qiskit Aer version: 0.12.2
Python version: 3.10.13
Operating system: MacOS 14.1 and Ubuntu

What is happening?

I am trying to learn some Quantum Machine Learning models with Qiskit and after a training of >2h my process is killed from the operating system, because the RAM usage got too high. The tracemalloc library hinted me to the class ParameterReferences inside qiskit/circuit/parametertable.py, which seems to accumulate memory space after every iteration:

  4%|██████▌                                                                                                                                                               | 5/126 [00:14<06:19,  3.14s/it]Top 10 lines
#1: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/dill/_dill.py:442: 547.1 KiB
    obj = StockUnpickler.load(self)
#2: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/linecache.py:137: 493.4 KiB
    lines = fp.readlines()
#3: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copy.py:280: 414.9 KiB
    y.__dict__.update(state)
#4: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:30: 262.2 KiB
    self._instance_ids = {}
#5: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copyreg.py:101: 186.4 KiB
    return cls.__new__(cls, *args)
#6: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:45: 183.9 KiB
    self._instance_ids = {self._instance_key(ref): ref[0] for ref in refs}
#7: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/scipy/optimize/_cobyla_py.py:293: 162.9 KiB
    xopt, info = cobyla.minimize(calcfc, m=m, x=np.copy(x0), rhobeg=rhobeg,
#8: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/abc.py:123: 144.4 KiB
    return _abc_subclasscheck(cls, subclass)
#9: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copy.py:76: 137.4 KiB
    return copier(x)
#10: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:27: 122.6 KiB
    return (id(ref[0]), ref[1])
972 other: 1755.4 KiB
Total allocated size: 4410.6 KiB
477331456
  5%|███████▉                                                                                                                                                              | 6/126 [00:18<06:38,  3.32s/it]Top 10 lines
#1: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/dill/_dill.py:442: 636.2 KiB
    obj = StockUnpickler.load(self)
#2: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copy.py:280: 499.5 KiB
    y.__dict__.update(state)
#3: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/linecache.py:137: 493.4 KiB
    lines = fp.readlines()
#4: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:30: 315.3 KiB
    self._instance_ids = {}
#5: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copyreg.py:101: 224.2 KiB
    return cls.__new__(cls, *args)
#6: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:45: 220.2 KiB
    self._instance_ids = {self._instance_key(ref): ref[0] for ref in refs}
#7: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copy.py:76: 165.0 KiB
    return copier(x)
#8: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/scipy/optimize/_cobyla_py.py:293: 162.8 KiB
    xopt, info = cobyla.minimize(calcfc, m=m, x=np.copy(x0), rhobeg=rhobeg,
#9: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:27: 144.3 KiB
    return (id(ref[0]), ref[1])
#10: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/abc.py:123: 143.3 KiB
    return _abc_subclasscheck(cls, subclass)
980 other: 2002.4 KiB
Total allocated size: 5006.6 KiB
480002048
  6%|█████████▏                                                                                                                                                            | 7/126 [00:22<06:50,  3.45s/it]Top 10 lines
#1: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/dill/_dill.py:442: 663.6 KiB
    obj = StockUnpickler.load(self)
#2: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copy.py:280: 584.0 KiB
    y.__dict__.update(state)
#3: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/linecache.py:137: 493.4 KiB
    lines = fp.readlines()
#4: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:30: 368.5 KiB
    self._instance_ids = {}
#5: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copyreg.py:101: 261.9 KiB
    return cls.__new__(cls, *args)
#6: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:45: 256.4 KiB
    self._instance_ids = {self._instance_key(ref): ref[0] for ref in refs}
#7: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/copy.py:76: 192.7 KiB
    return copier(x)
#8: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/qiskit/circuit/parametertable.py:27: 165.7 KiB
    return (id(ref[0]), ref[1])
#9: /Users/floriankittelmann/Library/Caches/pypoetry/virtualenvs/qiskit-q-gan-qdc4MjGM-py3.10/lib/python3.10/site-packages/scipy/optimize/_cobyla_py.py:293: 162.7 KiB
    xopt, info = cobyla.minimize(calcfc, m=m, x=np.copy(x0), rhobeg=rhobeg,
#10: /opt/homebrew/Cellar/python@3.10/3.10.13_1/Frameworks/Python.framework/Versions/3.10/lib/python3.10/abc.py:123: 142.3 KiB
    return _abc_subclasscheck(cls, subclass)
984 other: 2209.5 KiB
Total allocated size: 5500.8 KiB
485015552

How can we reproduce the issue?

Unfortunately every attempt to create a minimal working example failed on my side, because I don't find the exact circumstances why and where this happens. Nevertheless, I would like to provide a small example how I use the qiskit library and how my iterations look like

from qiskit.result import QuasiDistribution
from qiskit.primitives import BaseSampler, SamplerResult, BackendSampler
from qiskit.primitives.primitive_job import PrimitiveJob
import numpy as np
from qiskit.circuit.library import TwoLocal
from qiskit import QuantumCircuit, Aer
import linecache
import tracemalloc
from qiskit.algorithms.optimizers import COBYLA

def get_cpu_sampler() -> BackendSampler:
    simulator = Aer.get_backend("qasm_simulator")
    device_caption = 'CPU'
    simulator.set_options(device=device_caption)
    return BackendSampler(backend=simulator)

def get_np_from_quasi_dists(list_quasi_dists: list[QuasiDistribution]) -> np.ndarray:
    nof_states = 2 ** 3
    list_probabilities = []
    for quasi_dist in list_quasi_dists:
        probs_one_sample: list[float] = [
            quasi_dist.get(i)
            if quasi_dist.get(i) is not None else 1.0e-18
            for i in range(nof_states)
        ]
        list_probabilities.append(probs_one_sample)
    return np.array(list_probabilities)

class MyModel:

    def __init__(self, _sampler: BaseSampler):
        self.__qnn = QuantumCircuit(3)
        self.__qnn.compose(TwoLocal(
            num_qubits=3,
            rotation_blocks="ry",
            entanglement_blocks="cz",
            entanglement="full",
            reps=1
        ), inplace=True)
        self.__qnn.measure_all()
        self.__sampler = _sampler

    def num_params(self) -> int:
        return self.__qnn.num_parameters

    def __get_model_with_amplitude_encoding(self, features: np.ndarray) -> QuantumCircuit:
        nof_features = features.shape[0]
        if nof_features != 2**3:
            raise Exception("wrong input dimension")

        features = features.reshape((2**3,))
        init = QuantumCircuit(3)
        init.initialize(features / np.linalg.norm(features))

        qc = QuantumCircuit(3)
        qc = qc.compose(init, inplace=False)
        qc.barrier()
        qc = qc.compose(self.__qnn, inplace=False)
        qc.measure_all()
        return qc

    def forward(self, x: np.ndarray, params: np.ndarray) -> np.ndarray:
        list_qcs = [self.__get_model_with_amplitude_encoding(feature) for feature in x]

        job: PrimitiveJob = self.__sampler.run(
            circuits=list_qcs,
            parameter_values=[params for i in range(len(list_qcs))]
        )
        result: SamplerResult = job.result()
        return get_np_from_quasi_dists(result.quasi_dists)

def display_top(_snapshot, key_type='lineno', limit=10):
    # see https://docs.python.org/3/library/tracemalloc.html
    _snapshot = _snapshot.filter_traces((
        tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
        tracemalloc.Filter(False, "<unknown>"),
    ))
    top_stats = _snapshot.statistics(key_type)

    print("Top %s lines" % limit)
    for index, stat in enumerate(top_stats[:limit], 1):
        frame = stat.traceback[0]
        print("#%s: %s:%s: %.1f KiB"
              % (index, frame.filename, frame.lineno, stat.size / 1024))
        line = linecache.getline(frame.filename, frame.lineno).strip()
        if line:
            print('    %s' % line)

    other = top_stats[limit:]
    if other:
        size = sum(stat.size for stat in other)
        print("%s other: %.1f KiB" % (len(other), size / 1024))
    total = sum(stat.size for stat in top_stats)
    print("Total allocated size: %.1f KiB" % (total / 1024))

sampler = get_cpu_sampler()
model = MyModel(sampler)

input_rand = np.random.rand(3, 2**3)
labels_rand = np.random.rand(3, 2**3)

optimizer = COBYLA()

tracemalloc.start()

def min_func(params: np.ndarray) -> float:
    out = model.forward(input_rand, params)
    loss = float((np.square(labels_rand - out)).mean())
    display_top(tracemalloc.take_snapshot())
    return loss

optimizer.minimize(min_func, np.random.rand(model.num_params()))

What should happen?

No accumulation of memory space

Any suggestions?

Can you think of any circumstances, in which an object of the class ParameterReferences might accumulate memory space? Can you give me some context about this class, such that it helps me to find the bug in my code? Any hint could be helpful, thank you in advance

jakelishman commented 10 months ago

A few things that may or may not help:

I think you've got a duplicated measure_all in your example codeblock at the moment, which is causing invalid circuits to be passed to a Sampler (multiple classical registers), so nothing can actually run from it. I just removed one of these to test.
On my Mac with the same version configurations as you've specified here, I don't see any memory leakage with the example you're giving.
Seeing a lot of allocations in ParameterReferences may not necessarily mean that class is actually the culprit - it might just be that it's the largest "owner" of memory, and something else is causing a bunch of references to it to remain.
It would be a bit surprising if there's an actual leak in ParameterTable or ParameterReferences because these objects are purely touched by Python code, and as long as there's no extant references to them being held by your code, they should get happily freed when their refcounts hit zero, or garbage collected if there are any cycles (which there largely shouldn't be). It is possible that Qiskit is somewhere causing extra references to these to be stored, however, but I can't think off the top of my head where this might be.

As a suggestion for you to look: maybe try and audit your code base and make sure that you're not keeping references to all temporary QuantumCircuit objects during your optimisation pass. If you're storing references to them in a "history" list or anything like that, that could be the reason that the references aren't getting freed correctly.

floriankittelmann commented 10 months ago

@jakelishman Thank you very much for your answer. I also think, that the ParameterTable class is not the actual root cause of the problem. But unfortunately, I do not save my QuantumCircuit in any kind of history list. I only create a list of QuantumCircuit in the beginning, such that I don't need to generate the circuits in each forward pass. Like this:

class MyModel:
    def __init__(self, _sampler: BaseSampler, _dataset: list[np.ndarray]):
        self.__list_circuits = [self.__get_model_with_amplitude_encoding(feature) for feature in _dataset] 
        ...

    def forward(self, list_idx: list[int], params: np.ndarray) -> np.ndarray:
        list_qcs = [self.__list_circuits[i] for i in list_idx]
        ...

I just tested that as well and it seems not to accumulate memory space in my minimal working example either. Anyway, I needed to change my trainings to TorchConnector from Qiskit Machine Learning, because the Adam Optimizer using Parameter Shift Rule for gradient calculation resulted in much better results. So my whole training procedure changed and I don't experience similar problems anymore. I guess I wont do any further debugging

Qiskit / qiskit