基于 TensorCircuit-NG 的千比特万参数全梯度量子算法训练

17 Jul 2026

引言
三种物理视角下的性能对比
收缩路径指标：FLOPs、最大张量与 Write
视角一：全局收缩树——把整图硬算出来是什么体验
omeco 与 cotengra：两个路径搜索框架
切片（Slicing）的效率限制
视角二：局部因果滑窗——利用局域物理性
视角三：空间转移矩阵——无损边界的精确扫描
反直觉的加速：用计算换带宽
一些避坑细节
完整千比特 VQE：从单步梯度到端到端优化
结语

引言

如果要在经典计算机上模拟 1000 个量子比特，并对 40,000 个参数求取全梯度，你需要多大的算力？对于一个 1000 比特、10 层的线路，如果仅仅是计算单个态振幅，其计算开销也许可控。然而，当试图去求解这个系统的变分量子本征求解器（VQE）及其全部参数梯度时，所面临的挑战难度急剧上升，乍看起来至少需要超算集群。

这背后的原因在于：

哈密顿量的引入：计算 VQE 期望值需要计算 \(E(\boldsymbol\theta)=\langle\psi(\boldsymbol\theta)\vert H\vert\psi(\boldsymbol\theta)\rangle\)。这意味着电路的张量网络从原本单侧状态收缩，变成了 bra-MPO-ket 的三层网络夹心结构，图的规模直接翻倍不止。
反向传播的内存墙：求取上万个参数的梯度需要使用反向传播。在收缩图的反向计算中，编译器为了计算梯度，必须在显存中保留整个收缩路径上的大量中间结果。这比单纯的前向收缩或振幅计算要困难得多，对内存压力也更大。

在这篇技术博客中，我们展示了如何在单张 NVIDIA RTX 6000D GPU 上，通过三种不同的物理与工程路线，在保证精确梯度计算的前提下，将 1000 比特 VQE 的每步执行时间缩短至单卡 1 秒。

三种物理视角下的性能对比

我们以一维开边界的横场 Ising 模型（TFIM）为基准。电路的每一层在比特链上依次作用于相邻对 \((0,1),(1,2),\ldots,(n-2,n-1)\) 的对角纠缠门：

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

再作用单比特旋转。TFIM 的哈密顿量为：

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

我们也可将哈密顿量表示为精确且紧凑的 Matrix Product Operator（MPO，Bond Dimension 为 3）。

下表展示了三种不同收缩模式的实测对比（基于 NVIDIA RTX 6000D，使用 complex64 精度）：

视角（算法模式）	首次编译	首次完整梯度	稳态完整梯度	峰值显存	适用条件与价值
全局收缩树	5 h 20 min	5.88 s	1.09 s	51.97 GB	直接收缩完整 bra–MPO–ket 图；用于整图计算与交叉基准
局部因果滑窗	10.79 s	148.68 s	148.36 s	8.86 GB	局部项存在固定宽度的严格因果锥；局部项可独立分发
局部因果滑窗 (8卡并行)	10.50–11.28 s/卡	19.72–20.50 s	18.71 s	8.86 GB/卡	每卡分发 125 个局部项，实测获得 7.93× 加速
空间转移矩阵 (B3)	6.367 s	1.328 s	1.054 s	8.62 GB	一维空间重复结构
空间转移矩阵 (B5)	7.375 s	1.320 s	1.133 s	5.44 GB	同一严格算法的低显存模式

从表中可以看出不同方案的互补性：

全局树的稳态执行时间较短（1.09 秒），但需要 5 个多小时的冷编译以及 52 GB 级别的峰值显存。
局部滑窗虽然能将编译时间缩短到十秒级别，却因重复计算了相邻局部项共享的环境，导致单卡执行变慢。
空间转移矩阵通过将一维空间上的重复性结构显式地交给 jax.lax.scan，成功把编译成本从系统长度中解耦，兼顾了编译与执行效率。

contraction_tradeoff_bubbles

收缩路径指标：FLOPs、最大张量与 Write

在张量网络收缩中，路径质量由多项资源指标刻画：

FLOPs（浮点运算数）：完成这条路径预计需要的总标量乘加次数。例如 log10 FLOPs = 12.0 代表约 \(10^{12}\) 次标量运算。
最大张量大小（ Width \(w\)）：收缩路径中产生的最大中间张量的元素数，它对应前向传播时的瞬时显存压力。如果所有指标维度为 2，宽度 \(w\) 对应 \(2^w\) 个元素。
Total Write（总写入量）：代表整条路径中所有中间张量“写出”的元素总数。在反向传播中，这些写出的张量都会作为残差保留下来，用于梯度的反向计算。

对于同时计算期望和梯度的任务，必须同时关注这三个指标，而最终仍应以实测的编译时间、执行时间和 GPU 峰值显存为准。

视角一：全局收缩树——把整图硬算出来是什么体验

对完整 bra–MPO–ket 网络直接进行 value_and_grad，可以量化直接收缩的真实代价。

这一步需要先改变网络表示，再进行路径搜索。每个 \(R_{ZZ}\) 门被写成精确的 Rank-2 分解形式是很必要的：

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

保持原有的左到右 Ladder 门序是关键。如果把可交换的 \(R_{ZZ}\) 重排为偶键和奇键交替的 Brick-wall 结构，由于破坏了消元路径，张量网络的宽度会被瞬间推高到 39.585；而保持 Ladder 门序，宽度仅为 23.585。

在这张大图上，omeco 路径搜索找到了 log10 FLOPs = 12.0466、log2 write = 32.5912 的全局收缩路径。在单张 RTX 6000D 上，编译完成后的梯度执行仅需 1.09 s，峰值显存为 51 GB。这说明整图 VQE 梯度可以在单卡上完成，但资源瓶颈集中在全局 HLO 的编译时间与反向残差的存储空间上。

omeco 与 cotengra：两个路径搜索框架

既然全局图的资源瓶颈集中在缩并路径搜索与编译上，我们不禁要问：现有的顶尖工具表现如何？在这个大张量网络上，我们对比了 omeco 和 cotengra 两个框架在搜索收缩树时的表现，他们各有其技术侧重：

omeco 核心由 Rust 实现，其 TreeSA 算法在树空间中做高强度的模拟退火，适合在大型图上快速给出质量较好的无切片种子树。
cotengra 则更则有对收缩树操作和进一步微调的更丰富接口。例如基于现有路径重配子树，或者在固定树上精细搜索切片指标。

为了对比二者的表现，我们在一个 100 比特、10 层的 TFIM 图（包含 6460 个张量，8441 条指标）上进行了受控搜索对比：

搜索方法与预算	搜索时间	`log10 FLOPs`	`log2` 最大张量	`log2 total write`	读法说明
`omeco` TreeSA (16 试验 × 64 步)	15.68 s	11.068	24.585	29.517	16 次独立退火试验
`cotengra` CMA-ES (8个超参候选)	37.59 s	12.901	31.585	35.646	8 组路径搜索超参数候选
`cotengra` SA (12 步 × 12 次重排)	9.22 s	11.827	29.585	34.014	低预算 SA
`cotengra` SA (24 步 × 24 次重排)	34.09 s	10.873	23.585	30.367	较高预算 SA

这组结果表明，omeco TreeSA 在更短时间跑了明显更多的退火步，并给出了复杂度更低的树，性能明显更好。这是单树场景，对于将树切片的话，两个搜索框架在不同场景各有优劣，但对于需要深切片的场景，现有方案表现都不理想。

切片（Slicing）的效率限制

在显存不足时，将收缩树进行切片是控制显存的必要手段。但当显存足够放置无切片路径时，强行进行切片到多卡并不能线性地提升计算效率。

从张量网络收缩的本质来看，切片相当于把被选中的收缩指标强行拖延到了收缩路径的最后。在无切片的全局最优路径中，这些指标原本可以在中间阶段就被尽早消去，从而避免中间张量维度的膨胀。一旦强制对它们进行切片，不仅破坏了原本最优的消元顺序，还迫使每个子任务在最后求和前重复计算那些原本可以共享的中间结果。

在 1000 比特全局路径上，强制切为 8、32、128 个任务后，每个 slice 的 log2 write 仅仅从 32.5912 微弱地降到了 32.5632、32.5491、32.5365。由于没能有效拆开反向传播的主要内存结构，且收缩顺序偏离了最优路径，导致计算大量重复，总 log10 FLOPs 从 12.0466 显著上升至 12.9420、13.5391、14.1355。这意味着，在单卡显存允许的范围内，使用无切片路径通常能获得更高的硬件利用效率。

three_methods_schematic

视角二：局部因果滑窗——利用局域物理性

对于深度 \(L=10\) 的线路，考虑到双比特门对易，可以从 ladder 结构等价成 brickwall 结构，每个局部 TFIM 项在反向传播时都存在一个精确的反向因果锥。由于每层对角纠缠门彼此对易，该因果窗的宽度为 \(2L+2\)，在 \(L=10\) 时仅为 22 个格点。在全局树搜索是劣势的 brickwall 结构，在光锥方案里变成了优势。

利用 lax.scan 遍历哈密顿量的局部 Pauli 项，并使用 jax.checkpoint 包裹单项，编译规模便由窗口大小与线路深度决定，与系统总格点数解耦：

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

该方案的优势在于编译时间短（10 秒级），且单卡显存占用仅为 8.86 GB。若将其分发至 8 张 GPU 并行计算（每张卡负责 125 个局部项），稳态梯度时间可缩短至 18.71 秒。但该方案由于因果窗重叠，带来了大量的重复计算开销。

此外，该方法严重依赖于对角门的可对易性。如果换成非对易双比特门，因果锥宽度在 ladder 双比特门结构下将迅速扩展至 \(O(n)\)，该方法的优势将不复存在。

视角三：空间转移矩阵——无损边界的精确扫描

空间转移矩阵方案在空间比特方向上将系统分为左端块、重复的中间块和右端块。中间块接收左侧完整的边界张量，收缩块内 bra–MPO–ket 子网，再输出给下一个块。

转移矩阵的递推关系为：

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

在深度 \(L=10\) 时，空间切口穿过 ket 的 10 条 bond、bra 的 10 条 bond 和 1 条维数为 3 的 MPO bond，精确边界的大小为：

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

这大约相当于 24 MiB 的内存流量。这解释了该方法为什么能在 1000 比特尺度上轻松进行线性扫描。

在 3-site block 划分下，该模式的冷编译时间仅为 6.367 秒，执行时间为 1.054 秒每步，显存峰值为 8.62 GB。

每步扫描覆盖的空间宽度，即 block size，本身也是一个可调的时间—空间 tradeoff。较小的 block 通常能缩小块内收缩图，降低编译压力以及反向传播需要保存或重算的块内残差，但会增加传递完整边界张量的 scan 步数；较大的 block 可以减少 scan 步数，却可能提高块内收缩的最大张量、total write、编译时间和梯度峰值显存。因此，最优 \(b\) 不能仅由边界传递次数决定，而应结合路径搜索结果以及实际测得的编译时间、单步吞吐和峰值显存共同选择。在本文的测试中，3-site block 是吞吐更优的配置；5-site block 则把峰值显存进一步降至 5.44 GB，代价是执行时间从 1.054 秒小幅增加至 1.133 秒，直观展示了这一可调性。

反直觉的加速：用计算换带宽

值得注意的是，在转移矩阵中间块循环需要配合 jax.checkpoint 重算机制，这将显存峰值从 42.24 GB 降至 8.62 GB，且稳态时间也从 1.118 秒略微缩短至 1.054 秒。这种“通过重算减少中间残差，进而降低内存流量”的做法能够提升速度，可能与硬件特性有关。显卡尤其是消费级显卡的显存带宽相对于单精度计算能力可能更为瓶颈，因此减轻读写负担增加计算负担反而能带来更显著的收益。对于不同的显卡推荐调节优化缩并树路径的时间和空间复杂度的惩罚比例，也是同样的道理。

更值得思考的是，这样一个简单的空间方向切片与分块扫描方案，竟然在实际运行时间、编译时间和内存占用上都远好于经过高强度模拟退火搜索得到的整棵全局收缩树。这说明，对于特别大的张量网络，尤其是在需要反向传播、必须同时考虑中间残差驻留和梯度显存占用时，现有缩并路径搜索框架的考虑仍不充分，未来还有很大的发展空间。

一些避坑细节

TF32 精度的规避 对于 NVIDIA GPU 上的复数矩阵乘法，仅在 JAX 中设置最高精度并不足以避开 TF32 的自动降级。需要在程序运行前显式指定：
```
export NVIDIA_TF32_OVERRIDE=0
```
否则状态范数会降至 0.98–0.996，导致能量计算出现明显偏差，这也是机器学习和量子模拟对精度容忍度差别的一个具体例子。
Rayon 线程池栈溢出 当 omeco 在搜索超深树（深度 > 1000）并进行切片代码生成时，由于 deep tree 的递归遍历，可能在 Rayon 线程中触发 Segmentation fault。需要在 Python 进程启动或 Rayon 全局线程池初始化前显式设置：
```
export RUST_MIN_STACK=67108864
```

完整千比特 VQE：从单步梯度到端到端优化

为了验证这套技术方案不仅能计算一次能量与梯度，而且能够真正支撑完整的变分优化，我们进一步运行了一个更复杂的 1000 比特 VQE。线路包含 10 层最近邻 \(R_{ZZ}\) Ladder，每层的单比特部分由连续的 \(R_X\)、\(R_Z\)、\(R_X\) 三个旋转组成：

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

在构造完整 bra–MPO–ket 网络前，我们先将每层连续的 \(R_X\)–\(R_Z\)–\(R_X\) 精确合并为一个 \(2\times2\) 单比特张量，再使用同一套无近似截断的空间转移矩阵方案计算能量和全部 40,000 个参数的梯度。整个优化在单张 NVIDIA RTX 6000D GPU 上完成，共进行了 2000 次 Adam 更新。

对于本文采用的 1000-site 临界开边界 TFIM，严格基态能量可以由自由费米子谱直接得到。2000 次 Adam 更新后的 VQE 能量为 \(E_{\rm VQE}=-1268.4365\)，总能量误差为 \(E_{\rm VQE}-E_0=4.440\)，对应相对误差约为 \(3\times10^{-3}.\)

在没有超参调试的情况下，一个包含复杂门序列、完整 1000 比特参数和精确全梯度的千比特 VQE 已经达到千分之几的相对能量精度。这一结果表明，本文的张量网络表示、路径搜索、空间扫描、自动微分和 GPU 执行并非孤立的性能技巧，已经能够组成一条从线路定义、精确能量与梯度计算，到实际优化收敛和严格结果验证的端到端技术链路。

vqe_energy_density_error_2000

结语

1000 比特 VQE 的全梯度模拟与完整优化实践表明，通过结合低秩门分解、收缩路径搜索、无损空间转移矩阵、JAX 自动微分与编译器机制，我们不仅可以在单张 GPU 上快速计算千比特 VQE 的完整梯度，也可以真正执行数千步算法训练，并在单卡小时量级、未经系统精调时达到约千分之三的相对能量误差。

以上结果展示了整套方案从单次模拟走向端到端大规模变分计算的潜力，也为进一步优化线路结构、学习率调度和更高精度探索留下了清晰空间。千比特 VQE 的跑通证明了一点：在量子模拟的深水区，算法层面的数学美感必须与底层的编译器机制（XLA）、自动微分和 GPU 硬件特性等深度耦合。TensorCircuit-NG 提供的不仅仅是一个模拟器，而是一条从量子算法定义到端到端超大规模变分计算的完整技术链路和软硬件打通的关键基础设施。