Training a 1,000-Qubit, 40,000-Parameter Quantum Algorithm with Full Gradients Using TensorCircuit-NG

Introduction

If you wanted to simulate 1,000 quantum qubits and compute exact full gradients for 40,000 parameters on a classical computer, how much compute power would you need?

For a 1,000-qubit, 10-layer circuit, calculating a single state amplitude might have a manageable computational overhead. However, when you attempt to solve the system's Variational Quantum Eigensolver (VQE) and compute all of its parameter gradients, the difficulty skyrockets. At first glance, this sounds like a job that strictly requires a supercomputing cluster.

The underlying reasons for this explosion in complexity are twofold:

1. The Introduction of the Hamiltonian: Computing the VQE expectation value requires evaluating $E(\boldsymbol\theta)=\langle\psi(\boldsymbol\theta)\vert H\vert\psi(\boldsymbol\theta)\rangle$. This transforms the tensor network from a single-sided state contraction into a "bra-MPO-ket" three-layer sandwich structure, immediately more than doubling the graph size.
2. The Memory Wall in Backpropagation: Deriving gradients for tens of thousands of parameters requires backpropagation. To compute these gradients in a contraction graph, the compiler has to keep massive amounts of intermediate results from the forward pass in VRAM. This is vastly more difficult—and puts far more pressure on memory—than a simple forward contraction or amplitude calculation.

In this technical blog, we'll demonstrate how we cut the per-step execution time of a 1,000-qubit VQE down to just 1 second on a single NVIDIA RTX 6000D GPU. We achieve this by exploring three different physics- and engineering-driven contraction strategies, all without sacrificing exact gradient precision.

Performance Comparison: Three Physical Perspectives

We use a 1D Transverse Field Ising Model (TFIM) with open boundary conditions as our benchmark. Each layer of the circuit applies diagonal entangling gates sequentially to adjacent pairs $(0,1),(1,2),\ldots,(n-2,n-1)$ along the qubit chain:

$$R_{ZZ}(\theta)=\exp(-i\theta Z\otimes Z/2)$$

Followed by single-qubit rotations. The TFIM Hamiltonian is:

$$H_{\rm TFIM}=\sum_{i=0}^{n-2}X_iX_{i+1}+\sum_{i=0}^{n-1}Z_i$$

We can also represent this Hamiltonian as an exact and compact Matrix Product Operator (MPO, with a Bond Dimension of 3).

The table below shows the benchmark results for three different contraction modes (using an NVIDIA RTX 6000D at complex64 precision):

Perspective (Algorithm Mode)	Initial Compilation	First Full Gradient	Steady-State Full Gradient	Peak VRAM	Applicability & Value
Global Contraction Tree	5 h 20 min	5.88 s	1.09 s	51.97 GB	Directly contracts the full bra-MPO-ket graph; useful for full-graph compute and cross-benchmarking.
Local Causal Sliding Window	10.79 s	148.68 s	148.36 s	8.86 GB	Relies on strict causal cones of fixed width for local terms; terms can be dispatched independently.
Local Window (8-GPU Parallel)	10.50–11.28 s / GPU	19.72–20.50 s	18.71 s	8.86 GB / GPU	Dispatches 125 local terms per GPU; yielded a measured 7.93× speedup.
Spatial Transfer Matrix (B3)	6.367 s	1.328 s	1.054 s	8.62 GB	Leverages repeating 1D spatial structures.
Spatial Transfer Matrix (B5)	7.375 s	1.320 s	1.133 s	5.44 GB	A lower-VRAM variant of the same exact STM algorithm.

These approaches complement one another:

• The Global Tree delivers excellent steady-state execution time (1.09s) but suffers from a brutal 5+ hour cold compilation time and massive 52 GB peak VRAM usage.
• The Local Window approach crushes compile time down to ~10 seconds, but redundantly computes the shared environments of adjacent local terms, making single-GPU execution sluggish.
• The Spatial Transfer Matrix (STM) successfully decouples compilation costs from the system length by explicitly passing the repetitive 1D spatial structure to jax.lax.scan, perfectly balancing compilation and execution efficiency.

Contraction Path Metrics: FLOPs, Max Tensor, and Write

In tensor network contraction, path quality is defined by several resource metrics:

• FLOPs (Floating Point Operations): The total number of scalar multiply-add operations required for the path. For example, log10 FLOPs = 12.0 means roughly $10^{12}$ scalar ops.
• Max Tensor Size (Width $w$): The number of elements in the largest intermediate tensor generated during contraction. It directly dictates the instantaneous VRAM pressure during the forward pass. (If all indices have a dimension of 2, width $w$ corresponds to $2^w$ elements).
• Total Write: The total number of elements written out across all intermediate tensors. During backpropagation, all of these written tensors must be retained in memory as residuals to compute gradients.

When calculating expectations and gradients simultaneously, you must keep an eye on all three metrics. Ultimately, however, empirical compile time, execution time, and peak GPU VRAM are the ground truth.

Perspective 1: Global Contraction Tree — Brute-Forcing the Entire Graph

By calling value_and_grad directly on the complete bra-MPO-ket network, we can quantify the true cost of naive global contraction.

To make this feasible, we first alter the network representation. It is critical to write each $R_{ZZ}$ gate in an exact Rank-2 decomposed form:

$$R_{ZZ}(\theta)=\cos(\theta/2)I\otimes I-i\sin(\theta/2)Z\otimes Z$$

Maintaining the original left-to-right "ladder" gate ordering is also crucial. If we reorder the commuting $R_{ZZ}$ gates into an alternating even/odd bond "brick-wall" structure, we destroy the elimination path. The tensor network width instantly spikes to 39.585. Keeping the ladder ordering keeps the width at a manageable 23.585.

Using the omeco path search on this massive graph, we found a global contraction path with log10 FLOPs = 12.0466 and log2 write = 32.5912.

On a single RTX 6000D, post-compilation gradient execution takes only 1.09s, peaking at 51 GB of VRAM. This proves that full-graph VQE gradients can technically be run on a single card, but the bottlenecks are purely the global HLO compilation time and the memory needed to store backward residuals.

omeco vs. cotengra: Two Path-Search Frameworks

Since the global graph's resource bottleneck lies heavily in contraction path search and compilation, how do top-tier tools fare? We compared two frameworks—omeco and cotengra—on this massive tensor network.

• omeco (Rust-based) uses a TreeSA algorithm to perform intense simulated annealing in tree space. It is excellent for quickly finding high-quality un-sliced seed trees on large graphs.
• cotengra offers richer interfaces for manipulating and fine-tuning contraction trees, such as reconfiguring subtrees based on existing paths or performing granular slice-index searches on a fixed tree.

We ran a controlled search comparison on a 100-qubit, 10-layer TFIM graph (6,460 tensors, 8,441 indices):

Search Method & Budget	Search Time	log10 FLOPs	log2 Max Tensor	log2 Total Write	Notes
omeco TreeSA (16 trials × 64 steps)	15.68 s	11.068	24.585	29.517	16 independent SA trials
cotengra CMA-ES (8 hyperparam candidates)	37.59 s	12.901	31.585	35.646	8 sets of search params
cotengra SA (12 steps × 12 reconfigs)	9.22 s	11.827	29.585	34.014	Low-budget SA
cotengra SA (24 steps × 24 reconfigs)	34.09 s	10.873	23.585	30.367	Higher-budget SA

The data shows that omeco TreeSA ran significantly more annealing steps in less time and returned lower-complexity trees, outperforming the alternatives for this un-sliced scenario. However, for scenarios requiring deep slicing, neither framework is currently ideal.

The Efficiency Limits of Slicing

When VRAM is tight, slicing the contraction tree is a necessary evil to keep memory footprints in check. But if you have enough VRAM for an un-sliced path, forcing a slice to distribute across multiple GPUs does not scale efficiency linearly.

Fundamentally, slicing a tensor network means forcibly delaying the contraction of selected indices to the very end of the path. In an optimal, un-sliced path, these indices are eliminated early on, preventing intermediate tensors from blowing up in size. Slicing disrupts this optimal elimination order. It forces every sub-task to redundantly compute intermediate results that could have been shared before the final summation.

When we force-sliced the 1,000-qubit global path into 8, 32, and 128 tasks, the log2 write per slice barely dropped—from 32.5912 to 32.5632, 32.5491, and 32.5365 respectively. Because it failed to meaningfully dismantle the primary memory structure of the backward pass, while heavily deviating from the optimal contraction order, redundant computation skyrocketed. Total log10 FLOPs jumped from 12.0466 to 12.9420, 13.5391, and 14.1355.

Takeaway: If it fits in VRAM, using an un-sliced path usually maximizes hardware utilization.

Perspective 2: Local Causal Sliding Window — Exploiting Local Physics

For a circuit depth of $L=10$, we can leverage the commutativity of 2-qubit gates to convert the ladder structure into a brick-wall equivalent. Consequently, every local TFIM term possesses an exact backward causal cone during backpropagation. Since the diagonal entangling gates commute, the width of this causal window is strictly $2L+2$, which is only 22 sites for $L=10$. The brick-wall structure, which was a liability for global tree search, becomes a massive advantage here.

By using lax.scan to iterate over the local Pauli terms of the Hamiltonian and wrapping single terms in jax.checkpoint, the compilation scale becomes solely dependent on window size and circuit depth, totally decoupled from the total number of system sites:

$$T_{\rm window}=O(n C_{\rm local}(L)),\qquad M_{\rm window}=O(M_{\rm local}(L))$$

This approach boasts rapid compile times (~10 seconds) and a lean VRAM footprint of 8.86 GB. If dispatched across 8 GPUs (125 local terms per GPU), the steady-state gradient time drops to 18.71s. However, the overlapping causal windows introduce massive computational redundancy.

Furthermore, this method heavily relies on the commutativity of diagonal gates. If we swap them for non-commuting 2-qubit gates, the causal cone under a ladder structure expands to $O(n)$, rendering this method useless.

Perspective 3: Spatial Transfer Matrix — Lossless Exact Scanning

The spatial transfer matrix strategy slices the system along the spatial qubit dimension into a left-end block, repeating middle blocks, and a right-end block. The middle block receives the complete boundary tensor from the left, contracts its internal bra-MPO-ket subnetwork, and outputs to the next block.

The recurrence relation is:

$$B_{k+1}=\mathcal{T}_k(\boldsymbol\theta_k)B_k$$

At depth $L=10$, the spatial cut crosses 10 ket bonds, 10 bra bonds, and 1 MPO bond (dimension 3). The exact boundary size is:

$$D_{\rm boundary}=3\times2^{2L}=3\times4^L = 3\times4^{10}$$

This amounts to roughly 24 MiB of memory traffic, which perfectly explains why this method can effortlessly perform linear scanning at the 1,000-qubit scale.

Using a 3-site block division, this mode requires just 6.367s for cold compilation, executes at 1.054s per step, and peaks at 8.62 GB of VRAM.

The spatial width of each scan step (the block size $b$) acts as a tunable time-space tradeoff:

• Smaller blocks shrink the internal contraction graph, reducing compile pressure and the residuals saved for backprop, but increase the number of scan steps needed to pass the boundary tensor.
• Larger blocks reduce scan steps but inflate maximum tensor size, total write, compile time, and peak VRAM.

In our tests, a 3-site block offered optimal throughput. A 5-site block reduced peak VRAM further to 5.44 GB, at the slight cost of increasing execution time to 1.133s.

Counterintuitive Acceleration: Trading Compute for Bandwidth

Fascinatingly, pairing the transfer matrix's middle block loop with jax.checkpoint (recomputation) dropped the peak VRAM from 42.24 GB down to 8.62 GB, while also slightly reducing steady-state execution time from 1.118s to 1.054s.

"Reducing memory traffic by recomputing intermediate residuals" actually speeding up execution highlights a hardware reality: on modern consumer GPUs, memory bandwidth is often a much harder bottleneck than raw single-precision FLOPs. Taking on a heavier compute burden to alleviate read/write strain yields net performance gains. This same logic applies when tuning the penalty ratio between time and space complexity during contraction tree searches.

Even more thought-provoking is the fact that this relatively simple spatial slicing and block-scanning strategy heavily outperformed the global tree—derived from intense simulated annealing—in execution time, compile time, and memory footprint. This implies that for truly massive tensor networks requiring backpropagation (where intermediate residuals and VRAM are paramount), existing contraction path algorithms still have significant blind spots and ample room for innovation.

Pitfalls to Avoid

1. Bypassing TF32 Precision Limitations

When doing complex matrix multiplications on NVIDIA GPUs, simply setting the highest precision in JAX is not enough to avoid automatic degradation to TF32. You must explicitly set this environment variable before runtime:

Bash

   export NVIDIA_TF32_OVERRIDE=0

If you don't, the state norm can slip to 0.98–0.996, causing noticeable deviations in energy calculations. This is a stark reminder of the differing precision tolerances between standard Machine Learning and Quantum Simulation.

1. Rayon Thread Pool Stack Overflows

When omeco searches ultra-deep trees (depth > 1000) and generates slicing code, the recursive traversal of the deep tree can trigger a Segmentation Fault in Rayon threads. You must explicitly bump up the stack size before initiating Python or the Rayon global pool:

Bash

   export RUST_MIN_STACK=67108864

Full 1,000-Qubit VQE: From Single-Step Gradients to End-to-End Optimization

To prove that this isn't just a parlor trick for a single energy/gradient calculation, we ran a complete, complex 1,000-qubit VQE optimization. Our circuit features 10 layers of nearest-neighbor $R_{ZZ}$ ladders, with each layer's single-qubit component consisting of sequential $R_X$, $R_Z$, and $R_X$ rotations:

$$\vert{}0\rangle^{\otimes 1000}\xrightarrow{H^{\otimes 1000}}\left[R_{ZZ}\text{ Ladder}\rightarrow R_X\rightarrow R_Z\rightarrow R_X\right]^{10}.$$

Before constructing the full bra-MPO-ket network, we exactly compressed each consecutive $R_X$–$R_Z$–$R_X$ sequence into a single $2\times2$ single-qubit tensor. We then utilized the exact same Spatial Transfer Matrix scheme (with zero approximation truncations) to compute the energy and the gradients for all 40,000 parameters. The entire optimization ran on a single NVIDIA RTX 6000D, completing 2,000 Adam updates.

For this 1,000-site critical open-boundary TFIM, the exact ground state energy can be derived directly from the free-fermion spectrum. After 2,000 updates, our VQE energy was $E_{\rm VQE}=-1268.4365$, resulting in a total energy error of $E_{\rm VQE}-E_0=4.440$. This equates to a relative error of about $3\times10^{-3}$.

Without any hyperparameter tuning, a 1,000-qubit VQE featuring complex gate sequences, full parameters, and exact full gradients hit a relative energy precision in the thousandths. This confirms that our tensor network representations, path searches, spatial scanning, autodiff, and GPU execution aren't isolated micro-optimizations—they form a cohesive, end-to-end technical pipeline capable of translating circuit definitions into actual, strictly-verified optimization convergence.

Conclusion

Our work on full-gradient simulation and optimization for a 1,000-qubit VQE demonstrates that by hybridizing low-rank gate decomposition, contraction path search, lossless spatial transfer matrices, JAX automatic differentiation, and compiler mechanics, we can do the impossible. We can compute the full gradient of a 1,000-qubit VQE rapidly on a single GPU, execute thousands of training steps within hours, and achieve ~0.3% relative energy error without systemic fine-tuning.

These results illustrate the immense potential of moving from single-shot simulations to massive end-to-end variational computation, leaving plenty of headroom for exploring better circuit structures, learning rate schedules, and higher precision.

Getting a 1,000-qubit VQE to run natively proves one thing: in the deep waters of quantum simulation, the mathematical elegance of algorithmic design must be deeply coupled with the underlying mechanics of compilers (XLA), automatic differentiation, and GPU hardware architectures. TensorCircuit-NG is more than just a simulator; it is a critical piece of infrastructure that bridges software and hardware, providing a complete technical pipeline from quantum algorithm definition to ultra-large-scale variational optimization.