伊辛解码(Ising Decoding)
本仓库提供用于构建、定制和部署可扩展量子纠错解码器的 AI 训练方案:
- 神经网络在空间和时间维度上消费检测器综合征
- 它预测能够降低综合征密度 / 改善解码的修正值
- 标准解码器(PyMatching)产生最终的逻辑判定
公开发布版暴露了一个单一的用户面向配置文件 和一个单一的运行器脚本。
目录
- 论文发表
- 高层工作流程
- [快速入门(训练 + 推理)](#快速入门(训练 + 推理))
- 依赖项
- 故障排除
- 推理(预训练模型)
- 模型导出和下游工具
- [将 .pt 检查点转换为 SafeTensors](#将 .pt 检查点转换为 SafeTensors)
- [ONNX 导出和量化](#ONNX 导出和量化)
- [为 CUDA-Q QEC 生成数据](#为 CUDA-Q QEC 生成数据)
- [从 Stim 检测器样本进行离线解码](#从 Stim 检测器样本进行离线解码)
- [使用 cudaq-qec 进行解码器消融研究](#使用 cudaq-qec 进行解码器消融研究)
- 配置和高级用法
- 日志和输出
- [测试和 CI](#测试和 CI)
- [测试(CPU + GPU)](#测试(CPU + GPU))
- [CI(GitHub Actions)](#CI(GitHub Actions))
- 结果
- 许可证
论文发表
本实现配套论文:
Christopher Chamberland, Jan Olle, Muyuan Li, Scott Thornton, 和 Igor Baratta,
"表面码的快速且精确的基于 AI 的预解码器"
arXiv:2604.12841,2026年。
doi:10.48550/arXiv.2604.12841
如果您在研究或已发表的工作中使用本仓库,请引用该论文。
高层工作流程
+----------------------------------------+ 使用:
| 1. 训练或下载模型 | - Ising-Decoding 仓库(训练)
| | - Hugging Face(下载)
+------------------+---------------------+
|
v
+----------------------------------------+ 使用:
| 2. 评估性能 | - Ising-Decoding 仓库
| (运行推理测试) |
+------------------+---------------------+
|
+--v---------------v---------------------+ 使用:
| 3. 研究实时性能 | - Ising-Decoding 仓库(3a, 3b)
| | - CUDA-Q QEC(3c)
| +-------------------------------+ |
| | 3a. 启用 ONNX_WORKFLOW & | |
| | 选择量化格式 | |
| +--------------+--------------+ |
| | |
| +--------------v--------------+ |
| | 3b. 运行 generate_test_data | |
| +--------------+--------------+ |
| | |
| +--------------v--------------+ |
| | 3c. 将 .onnx 和 .bin 文件 | |
| | 导入 CUDA-Q QEC | |
| +-------------------------------+ |
+----------------------------------------+
快速入门(训练 + 推理)
从仓库根目录:
bash
code/scripts/local_run.sh
该脚本在本地运行 Hydra 工作流(无需 SLURM),并读取一个用户面向的配置文件:
conf/config_public.yaml
依赖项
目标 Python 版本:3.11, 3.12, 3.13
提供两个最小化的 requirements 文件:
code/requirements_public_inference.txt(Stim + PyTorch 路径)code/requirements_public_train-cuXY.txt(训练路径,其中 XY = 12 或 13)
安装示例(推荐但可选:使用虚拟环境):
bash
# 可选:创建并激活虚拟环境
python -m venv .venv
source .venv/bin/activate
# 可选:安装支持 CUDA 的 PyTorch(示例:选择任意可用的 cuXXX)
# 选择一个匹配你的 CUDA 运行时的版本;cu130 已知可用。
export TORCH_CUDA=cu130
# 仅推理(训练安装包含推理依赖的超集)
pip install -r code/requirements_public_inference.txt
# 训练(包含推理依赖,根据需要调整为 cu13)
pip install -r code/requirements_public_train-cu12.txt
bash code/scripts/check_python_compat.sh
提示:要强制使用支持 CUDA 的 PyTorch,设置 TORCH_CUDA=cuXXX(推荐 cu13x)或 TORCH_WHL_INDEX=https://download.pytorch.org/whl/cuXXX 后再运行安装。
快速入门:
bash
# 训练(读取 conf/config_public.yaml)
bash code/scripts/local_run.sh
# 推理(从 outputs/<exp>/models/* 加载已保存的模型)
WORKFLOW=inference bash code/scripts/local_run.sh
推理注意事项:
- 在裸机上,保持默认的 DataLoader workers。
- 在容器中,设置更大的共享内存大小(例如
docker run --shm-size=1g ...)。 - 如果无法更改
--shm-size,设置PREDECODER_INFERENCE_NUM_WORKERS=0以避免共享内存 worker 崩溃。 - 默认评估较重(每个 basis
cfg.test.num_samples=262144次射击);预计推理会花费一定时间。
故障排除
避免短运行时出现 steps_per_epoch=0:
保持 PREDECODER_TRAIN_SAMPLES >= per_device_batch_size * accumulate_steps * world_size。
注意:batch 调度在 epoch 0 后跳至 2048,因此 epoch 1 使用 2048 * 2 * world_size 的有效 batch 大小。
对于快速短运行,使用 GPUS=1 和 PREDECODER_TRAIN_SAMPLES >= 4096。
训练启动期间出现段错误(torch.compile):
某些环境在 torch.compile 期间崩溃。
禁用 compile:TORCH_COMPILE=0 bash code/scripts/local_run.sh
或尝试更安全的模式:TORCH_COMPILE=1 TORCH_COMPILE_MODE=reduce-overhead bash code/scripts/local_run.sh
Blackwell GPU(RTX 5080/5090, GB200/GB300):
稳定的 PyTorch wheels(cu124)不附带 SM 12.0 内核。
安装带有 cu128 索引的 nightly 版本:
bash
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128
Windows(Git Bash / WSL):
Triton 不支持原生 Windows,这会导致 torch.compile 失败。运行前禁用它:
bash
export TORCH_COMPILE_DISABLE=1 # PyTorch 级别标志
# 或者,对于仓库脚本等效地:
export PREDECODER_TORCH_COMPILE=0
当直接在脚本外运行(在 notebook 或 local_run.sh 之外)时,设置 Python 路径以便仓库模块可导入:
bash
export PYTHONPATH="code"
推理期间找不到预训练模型:
find_best_model 首先在 {output}/models/best_model/ 内搜索,然后回退到 {output}/models/。如果您将下载的 .pt 文件放在别处,请将其移动到上述目录之一,或直接指向它:
bash
PREDECODER_MODEL_CHECKPOINT_FILE=path/to/Ising-Decoder-SurfaceCode-1-Accurate.pt \
WORKFLOW=inference bash code/scripts/local_run.sh
推理(预训练模型)
如果您不在本地训练,可以使用预训练模型运行推理。
(可选)创建 venv 并安装推理依赖:
bash
python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
pip install -r code/requirements_public_inference.txt
获取预训练模型
本仓库附带两个预训练模型文件(使用 Git LFS 跟踪):
models/Ising-Decoder-SurfaceCode-1-Fast.pt(感受野 R=9)models/Ising-Decoder-SurfaceCode-1-Accurate.pt(感受野 R=13)
这些检查点面向由公共配置编码的统一电路级退极化 设置。
自定义、非均匀的 25 参数噪声模型由下面的流水线支持训练;它们是训练时的定制选项,而非附带检查点的属性。
克隆通过 git lfs pull 获取文件。可选地,设置 PREDECODER_MODEL_URL 为 LFS/raw URL 以在工作树外获取文件(例如在最小检出或 CI 中)。
在 conf/config_public.yaml 中设置:
yaml
EXPERIMENT_NAME=predecoder_model_1
model_id: 1
运行推理:
bash
WORKFLOW=inference EXPERIMENT_NAME=predecoder_model_1 bash code/scripts/local_run.sh
推理输出写入 outputs/<EXPERIMENT_NAME>/,完整日志在 outputs/<EXPERIMENT_NAME>/run.log。
模型导出和下游工具
将 .pt 检查点转换为 SafeTensors(可选,训练后)
默认情况下,训练在 outputs/<EXPERIMENT_NAME>/models/ 下生成 .pt 检查点,推理直接加载它们。SafeTensors 导出是可选的 ------ 在下游工具需要 SafeTensors 格式时使用。
步骤 1 --- 转换最佳训练检查点:
bash
PYTHONPATH=code python code/export/checkpoint_to_safetensors.py \
--checkpoint outputs/<EXPERIMENT_NAME>/models/<checkpoint>.pt \
--model-id <MODEL_ID> [--fp16]
输出写在检查点旁边(例如 <checkpoint>_fp16.safetensors)。
步骤 2 --- 从 SafeTensors 文件运行推理:
bash
PREDECODER_SAFETENSORS_CHECKPOINT=outputs/<EXPERIMENT_NAME>/models/<checkpoint>_fp16.safetensors \
WORKFLOW=inference bash code/scripts/local_run.sh
MODEL_ID 是公共模型标识符(1-5);映射见 model/registry.py。
预训练的公共模型使用 --model-id 1(R=9)和 --model-id 4(R=13)。
ONNX 导出和量化(可选,训练后)
在训练后(或从附带的 .safetensors 文件开始),您可以将模型导出到 ONNX,并选择性地应用 INT8 或 FP8 训练后量化以进行部署。
您还可以在推理时更改表面码距离和轮数。也就是说 ------ 更改这些参数中的任何一个时,您不需要重新训练新模型;由于模型是 3D 卷积神经网络,模型将简单地在一个新的解码卷积体上运行。
要使用新的距离运行,只需将 DISTANCE=<your distance> 添加到下面的命令中。
要使用新的轮数运行,只需将 N_ROUNDS=<your number of rounds> 添加到下面的命令中。
在运行 local_run.sh 推理之前,设置 ONNX_WORKFLOW 和(可选)(QUANT_FORMAT, DISTANCE, N_ROUNDS) 环境变量:
| ONNX_WORKFLOW | 行为 |
|---|---|
| 0(默认) | 仅 PyTorch 推理,不导出 ONNX |
| 1 | 导出 ONNX 模型并用 PyTorch 运行推理 |
| 2 | 导出 ONNX 模型并通过 TensorRT 运行推理 |
| 3 | 加载预存在的 TensorRT 引擎文件并运行推理 |
bash
# 仅导出 ONNX(无 TensorRT)
ONNX_WORKFLOW=1 WORKFLOW=inference bash code/scripts/local_run.sh
# 导出 ONNX + 应用 INT8 量化 + 运行 TensorRT 推理
ONNX_WORKFLOW=2 QUANT_FORMAT=int8 WORKFLOW=inference bash code/scripts/local_run.sh
# 导出 ONNX + 应用 FP8 量化 + 运行 TensorRT 推理
ONNX_WORKFLOW=2 QUANT_FORMAT=fp8 WORKFLOW=inference bash code/scripts/local_run.sh
# 使用预构建的 TensorRT 引擎(跳过导出)
ONNX_WORKFLOW=3 WORKFLOW=inference bash code/scripts/local_run.sh
量化变量:
| 变量 | 默认值 | 描述 |
|---|---|---|
| QUANT_FORMAT | unset | int8 或 fp8。未设置表示无量化(FP32 ONNX)。 |
| QUANT_CALIB_SAMPLES | 256 | INT8/FP8 训练后量化的校准样本数。 |
电路变量:
| 变量 | 默认值 | 描述 |
|---|---|---|
| CONFIG_NAME | config_public | 使用 conf/$CONFIG_NAME.yaml 文件中的默认值 |
| DISTANCE | 使用 conf/$CONFIG_NAME.yaml 中指定的距离 |
表面码距离 |
| N_ROUNDS | 使用 conf/$CONFIG_NAME.yaml 中指定的轮数 |
内存实验中的轮数 |
注意:
- TensorRT 工作流(
ONNX_WORKFLOW=2或3)需要tensorrt和modelopt。 - FP8 量化失败是致命的。INT8 失败会静默回退到 FP32 ONNX 模型。
- ONNX 和引擎文件写入当前工作目录。
ONNX_WORKFLOW也被decoder_ablation工作流遵守 ------ 见下文。
为 CUDA-Q QEC 实时预解码器测试应用生成数据
在 CUDA-Q Realtime 等端到端下游系统中评估神经预解码器时,您需要一个测试框架,其中包含有效输入 ------ 既包括导出的神经网络模型,也包括相应的综合征数据。
提供了实用脚本 code/export/generate_test_data.py 来生成这些精确的数据(包括一个 .onnx 文件和几个 .bin 文件),以便您可以轻松地在 CUDA-Q QEC 实时 AI 解码器中消费它。
重要 :此脚本提供的 --distance 和 --n-rounds 参数必须 与前一部分运行 ONNX 导出时使用的值匹配(例如 ONNX_WORKFLOW=2)。
有关如何将这些文件导入 CUDA-Q Realtime C++ 流水线的详细演练,请参阅下游文档:Realtime AI Predecoder Pipeline。
bash
python3 code/export/generate_test_data.py --distance 13 --n-rounds 104 --num-samples 10000 --basis X --p-error=0.003 --simple-noise
示例输出:
Building circuit: D=13, T=104, basis=X, rotation=XV, p=0.003
Circuit built in 0.007s
Building detector error model and PyMatching matcher...
DEM + matcher built in 0.083s
Detectors: 17472, Observables: 1
Extracting check matrices (beliefmatching)...
H shape: (17472, 93864), O shape: (1, 93864), priors shape: (93864,)
Sampling 10000 shots...
Sampled in 1.006s
Decoding with PyMatching (baseline)...
Errors: 30/10000, LER: 0.0030
Decode time: 5.439s (543.9 us/shot)
Writing outputs to test_data/d13_T104_X/
Done.
H_csr.bin 808,944 bytes
O_csr.bin 2,932 bytes
detectors.bin 698,880,008 bytes
metadata.txt 162 bytes
observables.bin 40,008 bytes
priors.bin 750,916 bytes
pymatching_predictions.bin 40,008 bytes
从 Stim 检测器样本进行离线解码
这是用于解码在内存模拟器外生成的检测器样本的基于文件的路径。它为两个不同的受众而存在:
- 您已有检测器样本 (来自 QPU、第三方模拟器或先前缓存的运行),并希望将其输入我们提供的相同解码器。跳到"自带检测器样本"部分。
- 您需要一个可复现的端到端冒烟测试。使用下面的本地生成器,然后运行相同的解码命令。
文件契约
每个 basis 恰好两个文件:
<root>/
samples_X.dets # Stim 稀疏检测器样本格式
metadata_X.json # 电路 + 噪声指纹
samples_Z.dets
metadata_Z.json
samples_*.dets 使用 Stim 的稀疏格式并附加逻辑可观测量,因此一行 shot D3 D8 L0 表示检测器 3 和 8 在该次射击中触发,逻辑可观测量 0 翻转。Stim 不在样本格式中编码内存 basis,因此 X 和 Z 始终存在于单独的文件中;当 cfg.test.meas_basis_test=both 时,LER 循环遍历两者。解析器(resolve_stim_sample_paths)也接受替代布局 <root>/<basis>/samples.dets + metadata.json 以及扁平布局 <root>/samples.dets + metadata.json。
metadata JSON 的形式为 qec.surface_code.stim_sample_io.build_stim_sample_metadata 写入的内容:
json
{
"schema_version": 2,
"artifact": "stim_detector_samples",
"format": "dets",
"append_observables": true,
"distance": 7,
"n_rounds": 7,
"basis": "X",
"code_rotation": "XV",
"num_detectors": 168,
"num_observables": 1,
"num_shots": 262144,
"p_error": 0.003,
"noise_model": "25-param",
"noise_model_sha256": "abcd...",
"noise_model_params": {
"p_prep_X": 0.002,
"...": 0.0
}
}
p_error、noise_model、noise_model_sha256 和 noise_model_params 是可选但推荐的;当存在时,解码器会对其活动噪声模型与记录指纹进行交叉检查,如果两者不一致,默认会引发错误。在此模式之前写入的文件(无噪声字段)保持原样加载。
code_rotation 接受规范名称(XV、XH、ZV、ZH)和公共别名(O1...O4)。
自带检测器样本
如果您有来自其他地方(QPU、外部模拟器)的 .dets 数据,契约正是上面所述的三件事:
- 以 Stim 稀疏格式写入
samples_{basis}.dets并附加可观测量。 - 写入匹配上述模式的
metadata_{basis}.json。最简单的方式是调用qec.surface_code.stim_sample_io中的build_stim_sample_metadata(...)和write_metadata_json(...);您也可以手动编写。 - 确保
conf/config_public.yaml反映您的样本来源的实验:distance、n_rounds、data.code_rotation和data.noise_model必须完全匹配。解码器从这些重建 Stim 内存电路,并在解码前根据它验证文件。
然后指向启动器的目录:
bash
PREDECODER_STIM_SAMPLES_DIR=/path/to/your/dets \
PREDECODER_DECODE_MODE=pymatching_only \
WORKFLOW=inference bash code/scripts/local_run.sh
默认情况下验证是严格的:距离、轮数、basis、方向、检测器计数、可观测量存在、p_error 或 noise_model_sha256 不匹配时,在解码前会为每个不匹配引发一个显式错误。要仅将噪声不匹配降级为警告(例如在标定研究中扫描 p_error 时),设置 PREDECODER_STIM_STRICT_NOISE=0。结构不匹配始终是致命的。
生成本地参考文件
bash
WORKFLOW=generate_stim_data \
EXPERIMENT_NAME=offline_stim_run \
bash code/scripts/local_run.sh
生成器从 conf/config_public.yaml 读取:
| 配置字段 | 作用 |
|---|---|
| distance | 表面码距离 |
| n_rounds | 测量轮数 |
| data.code_rotation | 码方向(XV/XH/ZV/ZH 或 O1...O4) |
| data.noise_model | 25 参数噪声模型字典(可选) |
| test.meas_basis_test | X、Z 或 both(默认 both) |
| test.p_error | 标量噪声水平(默认 0.003) |
| test.num_samples | 每个 basis 的射击数(默认 262144,每个文件约 20 MB) |
默认样本数量较大,因为冒烟运行目标为 LER 稳定到约 3 位有效数字;覆盖 ++test.num_samples=N(或在本地覆盖配置中设置该字段)以缩小以进行更快的迭代。输出到:
outputs/offline_stim_run/stim_samples/samples_X.dets
outputs/offline_stim_run/stim_samples/metadata_X.json
outputs/offline_stim_run/stim_samples/samples_Z.dets
outputs/offline_stim_run/stim_samples/metadata_Z.json
generate_stim_data 工作流仅写入 Stim 样本构件。CUDA-Q .bin 构件(detectors.bin、H_csr.bin 等)存在于单独的输出目录中,由 python code/export/generate_test_data.py 直接生成;参见"[CUDA-Q 部分](#CUDA-Q 部分)"部分。
解码文件
仅 PyMatching ------ 可作为与伊辛预解码器进行 apples-to-apples 比较的基线。在此模式下,启动器将神经模型替换为 torch.nn.Identity(),不需要检查点:
bash
PREDECODER_STIM_SAMPLES_DIR=outputs/offline_stim_run/stim_samples \
PREDECODER_DECODE_MODE=pymatching_only \
WORKFLOW=inference bash code/scripts/local_run.sh
伊辛预解码器后接 PyMatching ------ 需要模型检查点。将 PREDECODER_MODEL_CHECKPOINT_FILE(或配置中的 model_checkpoint_file)指向已发布的模型之一,或先在相同的 EXPERIMENT_NAME 下运行训练:
bash
PREDECODER_STIM_SAMPLES_DIR=outputs/offline_stim_run/stim_samples \
PREDECODER_DECODE_MODE=ising_decoding_pymatching \
EXTRA_PARAMS="++model_checkpoint_file=models/Ising-Decoder-SurfaceCode-1-Fast.pt" \
WORKFLOW=inference bash code/scripts/local_run.sh
两个命令都无需更改 conf/config_public.yaml;现有配置控制模型、距离、轮数、方向和噪声模型,Stim 文件元数据在解码前根据重建的电路进行检查。
要持久化每次射击的比较数组,还需设置:
bash
PREDECODER_DECODE_OUTPUT_DIR=offline_decode_outputs
设置后,pymatching_only 写入:
{basis}_observables.npy{basis}_pymatching_predictions.npy
ising_decoding_pymatching 写入这些加上:
{basis}_predecoder_residual_detectors.npy{basis}_ising_decoding_pymatching_predictions.npy
该目录在第一次写入时延迟创建,因此指向尚不存在的路径是安全的。
冒烟脚本
code/scripts/offline_smoketest.sh
脚本默认 EXPERIMENT_NAME=offline_stim_run(匹配上面的示例路径),生成 Stim 文件,用 pymatching_only 解码,并且(如果 models/Ising-Decoder-SurfaceCode-1-Fast.pt 在磁盘上)用 ising_decoding_pymatching 再次解码。然后它解析推理循环在其摘要块最后一行发出的结构化 [Inference Summary] JSON 标记。该标记默认关闭以保持交互式和 notebook 运行整洁;冒烟测试通过在每个推理调用前导出 PREDECODER_EMIT_INFERENCE_SUMMARY=1 来选择启用。如果您想将这些结果导入其他工具,也请设置相同的环境变量。
一个 d=7、n_rounds=7、O1、262,144 次射击/次运行的示例如下所示。将时序/加速视为冒烟信号,而非基准测试:
[offline_smoketest.sh] Avg LER 0.002678 (no pre-decoder) -> 0.002285 (after); PyMatching speedup 1.815x
使用 cudaq-qec 进行解码器消融研究(可选)
decoder_ablation 工作流比较神经预解码器留下的残余综合征上的多个全局解码器。它支持预解码器的 PyTorch 和 TensorRT 后端以及来自 cudaq-qec 包(cudaq_qec)的 GPU 加速全局解码器。
PyTorch 预解码器 + cudaq-qec 全局解码器:
bash
# 需要:cudaq-qec (cudaq_qec), ldpc, beliefmatching, scipy
WORKFLOW=decoder_ablation bash code/scripts/local_run.sh
TRT 预解码器 + cudaq-qec 全局解码器(完整 GPU 流水线):
用于 inference 的相同 ONNX_WORKFLOW 变量也适用于此处。当 TRT 引擎激活时,神经预解码器通过 TensorRT 运行(快速、量化推理),而 cudaq-qec 解码器在 GPU 上处理残余综合征 ------ 将快速 TRT 推理与 GPU 加速全局解码端到端结合。
bash
# 导出 ONNX,构建 TRT 引擎,运行消融(TRT 预解码器 + cudaq-qec)
ONNX_WORKFLOW=2 WORKFLOW=decoder_ablation bash code/scripts/local_run.sh
# INT8 量化 TRT 预解码器 + cudaq-qec
ONNX_WORKFLOW=2 QUANT_FORMAT=int8 WORKFLOW=decoder_ablation bash code/scripts/local_run.sh
# 加载先前构建的引擎,然后运行消融
ONNX_WORKFLOW=3 WORKFLOW=decoder_ablation bash code/scripts/local_run.sh
消融研究报告每个解码器的逻辑错误率、cudaq-qec BP 变体的收敛统计、残余综合征权重分布和时序分解。结果写入 outputs/<EXPERIMENT_NAME>/plots/。
基准测试的解码器变体
| 解码器 | 来源 | 说明 |
|---|---|---|
| No-op | --- | 仅预解码器输出,无全局校正 |
| Union-Find | ldpc | 快速,次优 LER(逻辑错误率) |
| BP-only | ldpc | 置信传播,无 OSD |
| BP+LSD-0 | ldpc | 带有局部统计解码的 BP |
| Uncorr-PM | PyMatching | 不相关的最小权重完美匹配 |
| Corr-PM | PyMatching | 相关 MWPM(最佳经典基线) |
| cudaq-BP | cudaq-qec | GPU 上和积 BP |
| cudaq-MinSum | cudaq-qec | GPU 上最小和 BP |
| cudaq-BP+OSD-0/7 | cudaq-qec | BP + 有序统计解码 |
| cudaq-MemBP | cudaq-qec | 基于内存的最小和 BP |
| cudaq-MemBP+OSD | cudaq-qec | 内存 BP + OSD |
| cudaq-RelayBP | cudaq-qec | 顺序中继组合 |
当 cudaq_qec 可导入时,cudaq-qec 解码器会自动加载;如果包不存在,研究将优雅地降级为非 cudaq 解码器。
配置和高级用法
GPU 选择
默认 :如果您不设置 CUDA_VISIBLE_DEVICES 或 GPUS,将使用所有 GPU。
使用一个特定的 GPU(推荐用于精确选择):
bash
CUDA_VISIBLE_DEVICES=1 GPUS=1 bash code/scripts/local_run.sh
使用多个 GPU(前 N 个可见设备):
bash
GPUS=4 bash code/scripts/local_run.sh
显式多 GPU 选择 (比 GPUS 更细粒度):
bash
CUDA_VISIBLE_DEVICES=4,5,6,7 GPUS=4 bash code/scripts/local_run.sh
公共配置(conf/config_public.yaml)
外部用户应仅编辑 conf/config_public.yaml 。
如果您更改任何配置设置,也请更改实验名称,以免输出混合。
模型选择
model_id:{1,2,3,4,5} 之一
每个 model_id 有一个固定的感受野(R):
| 模型 | 感受野 |
|---|---|
| model 1 | R=9 |
| model 2 | R=9 |
| model 3 | R=17 |
| model 4 | R=13 |
| model 5 | R=13 |
训练建议
Models 1, 4, 5(不相关匹配):
- 至少训练 100 个 epoch。更少的 epoch 会导致欠训练模型。
- 每个 epoch 的射击数:使用 8 个 GPU 训练时,每个 epoch 使用 6700 万 次射击(
PREDECODER_TRAIN_SAMPLES=67108864)。每个 epoch 使用更少的射击数会产生更差的结果。
距离 / 轮数语义
顶层 distance / n_rounds 是评估目标 (您在推理中关心的)。
训练在模型感受野上运行:distance = n_rounds = R。
码方向
data.code_rotation:O1, O2, O3, O4
为了具体说明,以下是 distance-3 布局及相应的逻辑算子支撑(● = 在逻辑中,· = 不在逻辑中)。
O1
CODE LAYOUT:
(z)
D D D
[X] [Z] (x)
D D D
(x) [Z] [X]
D D D
(z)
LOGICAL X (lx):
● ● ●
· · ·
· · ·
LOGICAL Z (lz):
● · ·
● · ·
● · ·
O2
CODE LAYOUT:
(x)
D D D
(z) [X] [Z]
D D D
[Z] [X] (z)
D D D
(x)
LOGICAL X (lx):
● · ·
● · ·
● · ·
LOGICAL Z (lz):
● ● ●
· · ·
· · ·
O3
CODE LAYOUT:
(x)
D D D
[Z] [X] (z)
D D D
(z) [X] [Z]
D D D
(x)
LOGICAL X (lx):
● · ·
● · ·
● · ·
LOGICAL Z (lz):
● ● ●
· · ·
· · ·
O4
CODE LAYOUT:
(z)
D D D
(x) [Z] [X]
D D D
[X] [Z] (x)
D D D
(z)
LOGICAL X (lx):
● ● ●
· · ·
· · ·
LOGICAL Z (lz):
● · ·
● · ·
● · ·
噪声模型(公共默认)
data.noise_model:一个25 参数电路级 噪声模型(SPAM、空闲和 CNOT Pauli 信道)。
附带的配置使用统一电路级退极化 映射,其中所有 25 个值都来自单个物理错误率 p(例如 p_prep_{X,Z}=2*p/3、p_idle_cnot_{X,Y,Z}=p/3 和 p_cnot_*=p/15)。
您可以编辑 data.noise_model 以在非均匀/自定义 25 参数模型上训练。在这种情况下,Torch 训练生成器从活动 25p 模型刷新采样概率向量,而不是折叠回标量统一退极化路径。
训练噪声上缩放(表面码)
当训练表面码预解码器时,您指定的噪声参数可能非常小(例如 p = 1e-4),这会产生极其稀疏的综合征和缓慢的收敛。为解决此问题,训练流水线自动上缩放 所有 25 个噪声模型参数,使得最大的有效故障信道概率等于 6 x 10^-3 的固定目标(刚好低于约 7.5 x 10^-3 的表面码阈值)。
考虑的七个信道("大 P"):
| 信道 | 值 |
|---|---|
| P_prep_X | p_prep_X |
| P_prep_Z | p_prep_Z |
| P_meas_X | p_meas_X |
| P_meas_Z | p_meas_Z |
| P_idle_cnot | p_idle_cnot_X + p_idle_cnot_Y + p_idle_cnot_Z |
| P_idle_spam(有效) | 0.5 x (p_idle_spam_X + p_idle_spam_Y + p_idle_spam_Z) |
| P_cnot | 所有 15 个 p_cnot_* 的总和 |
max_group = max(P_prep_X, P_prep_Z, P_meas_X, P_meas_Z, P_idle_cnot, P_idle_spam_effective, P_cnot)
两个设计说明:
- X / Z 制备和测量保持分开 。它们是独立的单 Pauli 故障信道 ------ 求和
p_prep_X + p_prep_Z(或p_meas_X + p_meas_Z)会重复计算有效信道概率,并且会为一个原本处于目标的退极化噪声模型膨胀max_group。 p_idle_spam_*在比较前减半 。SPAM 窗口空闲由一个两步模型构建(每步一个状态制备和一个辅助复位半),因此原始配置总和代表两个退极化步骤。缩放决策使用单步有效值0.5 x p_idle_spam_raw;原始值仍以idle_spam_raw在日志中报告。
上缩放规则:
- 如果
max_group < 6e-3:所有 25 个 p 乘以6e-3 / max_group仅用于训练数据生成。评估始终按原样使用原始的用户指定噪声模型。 - 如果
max_group >= 6e-3:参数不修改(训练日志发出警告,以防这表示配置错误)。 - 非表面码类型(
code_type != "surface_code")从不进行上缩放。
算法简述:
流水线计算上面的七个信道,取 p_max = max(...),并以 0.006 / p_max 重新缩放整个 25 参数向量,使得 p_max 提升到 0.6%(6 x 10^-3)。原始噪声模型保留用于评估不变。
我们发现,在较密集的综合征上训练,然后在较稀疏的数据上评估,比直接在稀疏数据上训练产生更好的结果。
跳过上缩放噪声
如果您需要使用精确的噪声参数进行训练(例如用于基准测试或对照实验),您可以通过配置或环境变量禁用上缩放:
配置 (conf/config_public.yaml):
yaml
data:
skip_noise_upscaling: true
noise_model:
p_prep_X: 0.002
# ... 其余 25 个参数
环境变量:
bash
PREDECODER_SKIP_NOISE_UPSCALING=1 bash code/scripts/local_run.sh
任一方法都会导致训练流水线逐字使用用户指定的噪声模型 ------ 不应用任何缩放。训练日志将确认:
[Train] noise_model upscaling SKIPPED (skip_noise_upscaling=true or PREDECODER_SKIP_NOISE_UPSCALING=1).
预计算帧(推荐)
训练/验证数据生成可以从以下位置加载预计算帧:
frames_data/
如果帧缺失,代码可以回退到即时生成,但速度较慢。要预计算帧:
bash
python3 code/data/precompute_frames.py --distance 13 --n_rounds 13 --basis X Z --rotation O1
预计算的 DEM/帧构件是结构性的:它们编码每个可能的错误列可以为给定距离、轮数、basis 和旋转产生哪些检测器响应。活动的标量或 25 参数噪声模型控制每列的采样概率。因此,当仅概率变化时,可以重用缓存的结构构件;训练生成器在加载时从活动噪声模型刷新概率向量。
恢复训练和运行推理
推理使用来自 outputs/<experiment_name>/models/ 的训练模型,因此在从训练切换到推理时保持相同的 EXPERIMENT_NAME。
训练自动恢复 :如果运行中断,再次启动相同的训练命令(相同的 EXPERIMENT_NAME)将自动加载找到的最新检查点并继续训练(直到固定的 100 个 epoch)。要强制干净重启,设置 FRESH_START=1,但我们建议改更改 EXPERIMENT_NAME。
HE 加速(高级):并行类空
类空同调等价(HE)通道独立地对每个 (batch, round) diff 帧进行规范化。默认情况下,规范化过程按顺序处理稳定子。使用 data.use_parallel_spacelike: True 时,缓存构建计算稳定子-重叠图的 2-划分,使得两种颜色类别在 torch.compile 友好的内循环中并行减少。这减少了每个 HE 通道的 Python <-> 编译图交叉,并向 GPU 暴露更多并行性。
如何启用
在任何配置中:
yaml
data:
use_compile: True # 需要才能看到加速
use_parallel_spacelike: True
或在 CLI 上:
bash
EXTRA_PARAMS="data.use_compile=True data.use_parallel_spacelike=True" \
bash code/scripts/local_run.sh
优点(何时启用)
- GPU 上更快的类空 HE (对于旋转单 basis 表面码),通过摊销每次迭代的 Python 开销并让两种颜色类别一起通过
torch.compile。 - 在支持的码上与顺序路径综合征等价 :并行路径保持 HE 不变量并产生有效的非递增代表,同时避免顺序稳定子顺序。输出不保证与顺序路径按位相同;两者都是同一陪集的有效代表。
- 在
code/tests/mid/test_homological_equivalence.py下添加了覆盖率。 - 与
data.use_weight2组合 ------ 权重-2 固定等价通道按颜色应用。
缺点 / 注意事项(何时关闭)
- 仅限旋转单 basis 表面码。2-着色假设稳定子-重叠图是二分图,这对于此处针对的旋转表面码按构造成立。彩色码、非旋转布局、子系统码和混合 basis 矩阵可以产生奇数环;在这种情况下,缓存构建会拒绝并给出诊断命名违规稳定子对,而不是静默回退。
use_compile=True是加速所必需的;没有它,划分已构建但优化的编译内循环不会进入。torch.compile有冷启动成本。第一次编译调用可能会在 Inductor/CUDA 图捕获运行时暂停,而形状变化(如不同的 batch 大小或轮数)可能触发重新编译。- 缓存构建成本和内存略有增长 。打包的
parallel_partition_packed视图在缓存构建时物化一次,因此热路径只进行 dtype 转换。 - 面向 GPU。并行路径为 CUDA 设计;在 CPU 上您可能不会看到相对于顺序路径的加速。
日志和输出
内容写入位置
运行组织在:
outputs/<experiment_name>/
models/ (检查点 + 模型文件)
tensorboard/
config/ (每次运行使用的配置快照)
run.log (最新运行的日志副本)
logs/<experiment_name>_<timestamp>/
<workflow>.log (完整 stdout/stderr)
code/scripts/local_run.sh 自动将配置快照到:
outputs/<experiment_name>/config/<config_name>_<timestamp>.yaml
outputs/<experiment_name>/config/<config_name>_<timestamp>.overrides.txt
TensorBoard(训练指标)
TensorBoard 日志位于 outputs/<experiment_name>/tensorboard/ 下。
关键标量(如 TensorBoard 所示):
| 指标 | 描述 |
|---|---|
| Loss/train_step | 训练损失(BCEWithLogits)每优化步骤记录。越低越好。 |
| LearningRate/train | 每训练步骤的当前学习率(预热/调度后)。 |
| BatchSize | 每个 epoch 的有效 batch 大小:per_device_batch_size * accumulate_steps * world_size。我们累积 2 步:一步用于 X basis 电路,另一步用于 Z basis。 |
| Metrics/LER | 评估目标上的逻辑错误率(在训练时评估期间计算)。越低越好。 平均:在 cfg.test.num_samples 蒙特卡罗射击上计算,每个 basis(X 和 Z)。 默认:cfg.test.num_samples = 262144(当前公共版本硬编码)。 分布式:每个 rank 使用 cfg.test.num_samples // world_size 个射击/每个 basis(任何余数被丢弃)。 |
| Metrics/LER_Reduction_Factor | 预解码器后 LER 与基线 LER 的比率("相对改进"因子)。>1 表示改进。如果两者均为 0,我们记录 1.0。 平均:来自相同的 LER 评估运行(与 Metrics/LER 的射击计数相同)。 |
| Metrics/PyMatching_Speedup | 预解码器的平均 PyMatching 加速:latency_baseline / latency_after。>1 表示预解码后 PyMatching 解码更快。 平均:延迟在小子集(cfg.test.latency_num_samples,默认 10000)上使用单射击 PyMatching(batch_size=1,matcher.decode)测量,以微秒/轮报告。 |
| Metrics/SDR | 综合征密度降低因子:syndrome_density_before / syndrome_density_after。>1 表示预解码器降低了综合征密度。 |
| EarlyStopping/epochs_since_best | 自最佳验证指标以来的 epoch 数(我们使用 LER 作为验证指标)。 |
| EarlyStopping/best_metric | 迄今为止观察到的最佳(最低)验证损失。 |
评估默认值(公共版本)
- 训练期间的验证损失使用即时生成器。
- 测试/推理指标(LER / SDR / 延迟)默认为 Stim 路径。
测试和 CI
测试(CPU + GPU)
仅 CPU 的测试速度快,推荐用于快速验证:
bash
PYTHONPATH=code python -m unittest discover -s code/tests -p "test_*.py"
当没有 GPU 可用时,GPU 测试自动跳过。在 GPU 机器上,所有测试都会运行,包括由 torch.cuda.is_available() 保护的测试:
bash
PYTHONPATH=code python -m unittest discover -s code/tests -p "test_*.py"
噪声模型测试的有用环境变量:
RUN_SLOW=1启用 >=100k 射击统计测试NOISEMODEL_FAST_SHOTS控制快速层射击(默认 10000)NOISEMODEL_SLOW_SHOTS控制慢速层射击(默认 100000)
快速 GPU 运行示例:
bash
NOISEMODEL_FAST_SHOTS=2000 PYTHONPATH=code python -m unittest code/tests/test_noise_model.py
本地测试覆盖率:
要查看哪些代码被测试执行并获取报告:
bash
pip install -r code/requirements_public_inference.txt -r code/requirements_ci.txt
PYTHONPATH=code coverage run -m unittest discover -s code/tests -p "test_*.py"
coverage report
coverage html -d htmlcov
# 在浏览器中打开 htmlcov/index.html
CI 运行相同的套件并生成覆盖率报告,将 htmlcov/ 和 coverage.xml 作为作业构件发布。
CI(GitHub Actions)
CI 在 .github/workflows/ci.yml 中定义,并在推送到 main、pull-request/* 分支(通过 copy-pr-bot)、合并组检查和手动调度时运行:
| 作业 | 运行器 | 检查内容 |
|---|---|---|
| spdx-header-check | CPU | 所有源文件的 SPDX 许可证头 |
| unit-tests | CPU | 完整的 unittest discover 套件(GPU 测试自动跳过) |
| unit-tests-coverage | CPU | 相同的套件,带 coverage 报告 |
| python-compat | CPU | 跨 Python 3.11 / 3.12 / 3.13 的导入/安装检查 |
| gpu-tests | GPU | 自托管 GPU 运行器上的完整测试套件 |
| gpu-tests (train+inference) | GPU | 带 LER 检查的简短训练 + 推理 |
结果
在物理错误率 p = 0.003 和 0.006 下,X basis 解码的逻辑错误率(LER)与时间的关系:
许可证
本项目在 Apache License 2.0 下发布。
本仓库中的每个源文件都带有以下形式的 SPDX 版权和许可证头:
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
这些头的存在由 spdx-header-check CI 作业自动执行(参见 .github/workflows/ci.yml)。
与本项目捆绑或由其需要的第三方开源组件及其各自的版权声明和许可证文本列在 NOTICE 中。