diff --git a/.github/workflows/ci-gpu.yml b/.github/workflows/ci-gpu.yml index ec98fd3..2641b9e 100644 --- a/.github/workflows/ci-gpu.yml +++ b/.github/workflows/ci-gpu.yml @@ -94,14 +94,14 @@ jobs: source .venv_train_${{ matrix.python-version }}_${{ matrix.torch-cuda }}/bin/activate bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_train.log r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r - # 0.35: short run (16k samples, 2 epochs for stable LER across py versions) + # 0.35: short run (16k samples, 4 epochs for stable LER across py versions) python code/scripts/check_ler_from_log.py /tmp/ci_train.log --max-ler 0.35 env: EXPERIMENT_NAME: ci_short PREDECODER_TRAIN_SAMPLES: "16384" PREDECODER_VAL_SAMPLES: "2048" PREDECODER_TEST_SAMPLES: "2048" - PREDECODER_TRAIN_EPOCHS: "2" + PREDECODER_TRAIN_EPOCHS: "4" - name: Training + inference with multi-worker DataLoader (num_workers=2) shell: bash @@ -115,7 +115,7 @@ jobs: PREDECODER_TRAIN_SAMPLES: "16384" PREDECODER_VAL_SAMPLES: "2048" PREDECODER_TEST_SAMPLES: "2048" - PREDECODER_TRAIN_EPOCHS: "2" + PREDECODER_TRAIN_EPOCHS: "4" PREDECODER_INFERENCE_NUM_WORKERS: "2" # --------------------------------------------------------------------------- @@ -157,20 +157,20 @@ jobs: # TODO: matrix by CUDA major version [cu12, cu13] pip install -r code/requirements_public_train-cu12.txt - - name: Mid-tier training + inference with LER check (32k train, 2 epochs) + - name: Mid-tier training + inference with LER check (32k train, 4 epochs) shell: bash run: | . .venv_mid/bin/activate bash code/scripts/smoke_run.sh 2>&1 | tee /tmp/ci_mid.log r=${PIPESTATUS[0]}; [ $r -ne 0 ] && exit $r - # 0.2: mid-tier (32k/2 epochs); loosen if flaky + # 0.2: mid-tier (32k/4 epochs); loosen if flaky python code/scripts/check_ler_from_log.py /tmp/ci_mid.log --max-ler 0.2 env: EXPERIMENT_NAME: ci_mid PREDECODER_TRAIN_SAMPLES: "32768" PREDECODER_VAL_SAMPLES: "4096" PREDECODER_TEST_SAMPLES: "4096" - PREDECODER_TRAIN_EPOCHS: "2" + PREDECODER_TRAIN_EPOCHS: "4" - name: HE compile tests (torch.compile + autotune on GPU) run: | @@ -256,7 +256,7 @@ jobs: PREDECODER_TRAIN_SAMPLES: "16384" PREDECODER_VAL_SAMPLES: "2048" PREDECODER_TEST_SAMPLES: "2048" - PREDECODER_TRAIN_EPOCHS: "2" + PREDECODER_TRAIN_EPOCHS: "4" - name: Multi-GPU smoke training with parallel spacelike HE (2 GPUs, DDP) # Additive coverage on top of the default-config multi-GPU smoke above. @@ -291,7 +291,7 @@ jobs: PREDECODER_TRAIN_SAMPLES: "16384" PREDECODER_VAL_SAMPLES: "2048" PREDECODER_TEST_SAMPLES: "2048" - PREDECODER_TRAIN_EPOCHS: "2" + PREDECODER_TRAIN_EPOCHS: "4" # --------------------------------------------------------------------------- # GPU coverage: captures GPU-specific code paths missed by the CPU coverage job diff --git a/README.md b/README.md index 171468e..54af744 100644 --- a/README.md +++ b/README.md @@ -176,6 +176,24 @@ Inference note: - Some environments crash during `torch.compile`. - Disable compile: `TORCH_COMPILE=0 bash code/scripts/local_run.sh`. - Or try a safer mode: `TORCH_COMPILE=1 TORCH_COMPILE_MODE=reduce-overhead bash code/scripts/local_run.sh`. +- **Multi-GPU SIGSEGV inside `libnccl.so.2` on the CUDA 12.8 stack (known issue)**: + - The `torch 2.11.0+cu128` wheel pins `nvidia-nccl-cu12==2.28.9`, which has + been observed to segfault on any multi-rank NCCL operation (DDP gradient + sync, all-reduce) on H100x8 nodes running newer CUDA-13.x-family drivers. + The CUDA 13.0 stack (NCCL 2.29.7) is unaffected on the same hardware. + - Preferred: use the cu130 requirements/wheels on such nodes + (`TORCH_CUDA=cu130`, `code/requirements_public_train-cu13.txt`). + - If you must stay on cu128, upgrade NCCL in place. The torch wheel declares + an exact `nvidia-nccl-cu12==2.28.9` dependency, so pip will report a + dependency conflict for the upgrade — it is safe to proceed because + `libnccl.so.2` is ABI-compatible (`--no-deps` keeps the rest of the CUDA + package set untouched): + ```bash + pip install --no-deps --upgrade "nvidia-nccl-cu12>=2.29" + ``` + - `export NCCL_NVLS_ENABLE=0` has also been reported to avoid the crash + family on NVLink-Switch topologies; `NCCL_DEBUG=INFO` helps capture the + crash site for bug reports. - **Blackwell GPUs (RTX 5080/5090, GB200/GB300)**: - Stable PyTorch wheels (`cu124`) do not ship SM 12.0 kernels. Install the nightly build with the `cu128` index: @@ -281,7 +299,11 @@ PREDECODER_SAFETENSORS_CHECKPOINT=outputs//models/_ WORKFLOW=inference bash code/scripts/local_run.sh ``` -`MODEL_ID` is the public model identifier (1–5); see `model/registry.py` for the mapping. +`MODEL_ID` is the public model identifier — 1–5 for surface-code checkpoints, or +B for the color-code cascade model; see `model/registry.py` for the mapping. +Color-code *convolutional* checkpoints (model_id 1/2/4/5 trained with +`code: color`) are not supported by this converter — color training widens the +final conv layer, so the rebuilt surface architecture cannot load them. The pre-trained public models use `--model-id 1` (R=9) and `--model-id 4` (R=13). ### ONNX export and quantization (optional, post-training) @@ -335,7 +357,7 @@ ONNX_WORKFLOW=3 WORKFLOW=inference bash code/scripts/local_run.sh |---|---|---| | `CONFIG_NAME` | `config_public` | Use the defaults from the `conf/$CONFIG_NAME.yaml` file | | `DISTANCE` | Use the distance specified in the `conf/$CONFIG_NAME.yaml` file | code distance (surface or color) | -| `N_ROUNDS` | Calibration samples for INT8/FP8 post-training quantization. | number of rounds in memory experiment | +| `N_ROUNDS` | Use the number of rounds specified in the `conf/$CONFIG_NAME.yaml` file | number of rounds in memory experiment | Notes: @@ -351,8 +373,12 @@ CUDA-Q Realtime, you will need a test harness with valid inputs—both the exported neural network model and the corresponding syndrome data. The utility script `code/export/generate_test_data.py` is provided to generate -this exact data (both an `.onnx` file and several `.bin` files) so you can -easily consume it in the CUDA-Q QEC realtime AI decoder. +the syndrome-data side of this (several `.bin` files: detector samples, +observables, check matrices, priors, and PyMatching baseline predictions) so +you can easily consume it in the CUDA-Q QEC realtime AI decoder. The `.onnx` +model itself comes from the ONNX export step above; optionally pass it to this +script via `--onnx-model` to also record the pre-decoder outputs +(`predecoder_outputs.bin`) for the same samples. > **Important:** The `--distance` and `--n-rounds` arguments provided to this script **must match** the values used in the preceding section when running the @@ -512,12 +538,13 @@ The generator reads from `conf/config_public.yaml`: | `data.code_rotation` | code orientation (`XV`/`XH`/`ZV`/`ZH` or `O1`..`O4`) | | `data.noise_model` | 25-parameter noise model dict (optional) | | `test.meas_basis_test` | `X`, `Z`, or `both` (default `both`) | -| `test.p_error` | scalar noise level (default `0.003`) | +| `test.p_error` | scalar noise level (default `0.006` for surface, `0.003` for color; fixed by the public config) | | `test.num_samples` | shots per basis (default `262144`, ~20 MB per file) | The default sample count is large because the smoke run targets LER stable to -~3 significant digits; override `++test.num_samples=N` (or set the field in a -local override config) to shrink it for a faster iteration. Output goes to: +~3 significant digits; set `PREDECODER_INFERENCE_NUM_SAMPLES=N` in the +environment to shrink it for a faster iteration (the public config rejects +direct `test.num_samples` overrides). Output goes to: ```text outputs/offline_stim_run/stim_samples/samples_X.dets @@ -689,9 +716,10 @@ If you change any config settings, also change the experiment name so outputs ar #### Model selection - `model_id`: for `code: surface`, one of **{1,2,3,4,5}**; for `code: color`, - one of **{1,2,4,5}**. Color model 3 is intentionally unavailable in the + one of **{1,2,4,5,B}**. Color model 3 is intentionally unavailable in the public config because its receptive field is larger than the currently - supported color-code training window. + supported color-code training window. Model B (cascade/bottleneck, R=13) + is color-code only. #### Code family @@ -704,6 +732,7 @@ Each `model_id` has a fixed receptive field \(R\): - **model 3**: \(R=17\) - **model 4**: \(R=13\) - **model 5**: \(R=13\) +- **model B**: \(R=13\) (color code only) Optimizer learning rate is managed internally: surface code uses the model-specific public LR table, while color code always uses **1 × 10⁻⁵** for @@ -891,12 +920,24 @@ Training/validation data generation can load precomputed frames from: - `frames_data/` -If frames are missing, the code can fall back to on-the-fly generation, but it is slower. To precompute frames: +If frames are missing, the code can fall back to on-the-fly generation, but it +is slower. To precompute surface-code frames, run `qec.precompute_dem` once per +basis (one basis per invocation; `--rotation` accepts the internal orientations +`XV`/`XH`/`ZV`/`ZH`). Public training always runs at the chosen model's +receptive field, so match `--distance`/`--n_rounds` to it — 9 for models 1/2, +17 for model 3, 13 for models 4/5 (the default `model_id: 1` needs the d=9 +files): ```bash -python3 code/data/precompute_frames.py --distance 13 --n_rounds 13 --basis X Z --rotation O1 +PYTHONPATH=code python code/qec/precompute_dem.py \ + --distance 9 --n_rounds 9 --basis X --rotation XV --dem_output_dir frames_data +PYTHONPATH=code python code/qec/precompute_dem.py \ + --distance 9 --n_rounds 9 --basis Z --rotation XV --dem_output_dir frames_data ``` +For color codes, use `--code color` without `--rotation` — see +[Precompute the augmented DEM bundle](#precompute-the-augmented-dem-bundle). + Precomputed DEM/frame artifacts are structural: they encode which detector responses each possible error column can produce for a given distance, number of rounds, basis, and rotation. The active scalar or 25-parameter noise model @@ -931,7 +972,7 @@ internal Hydra schema, so they bypass the public validator. |-------------|---------| | `conf/config_color_model_1_s_LR3e-4.yaml` | Train a model-1-shaped color-code pre-decoder at `d=9, r=9` (superdense schedule). | | `conf/config_color_threshold_model_1_d13.yaml` | Threshold sweep against a trained color-code checkpoint at `d=13` (set `model_checkpoint_dir` to a training run's `models/` directory). | -| `conf/config_inference_color_model_5.yaml` | Run inference with a trained model-5-shaped color-code checkpoint via the public runner (`workflow.task=inference`; set `model_checkpoint_file` to your `.pt`). Defaults to `d=9, R=9, p=1e-3`; override `test.num_samples` / `test.p_error` / `test.meas_basis_test` for sweeps. | +| `conf/config_inference_color_model_5.yaml` | Run inference with a trained model-5-shaped color-code checkpoint via the public runner (`workflow.task=inference`; set `model_checkpoint_file` to your `.pt`). Model 5 has receptive field `R=13`; the test window defaults to `distance=9, n_rounds=9, p=1e-3`. Override `test.num_samples` / `test.p_error` / `test.meas_basis_test` for sweeps. | #### Precompute the augmented DEM bundle @@ -1001,9 +1042,12 @@ launcher routes them around the public-config validator, forwards `--config-name=` and `workflow.task=` to `code/workflows/run.py`, and `run_color` then dispatches to the training entry point. Pass training overrides (epochs, batch size, etc.) via `EXTRA_PARAMS`. (Color training can -also be driven from `conf/config_public.yaml` with `code: color`, -`workflow.task: train`, and `data.precomputed_frames_dir` pointing at your -precomputed bundle.) +also be driven from `conf/config_public.yaml` with `code: color` and +`workflow.task: train`; there the precomputed-frames path is not +user-configurable — the validator auto-sets it to the repo-relative +`frames_data/` directory, so generate your bundle there. Custom +`data.precomputed_frames_dir` paths require the standalone color configs +above.) #### Color-code limitations in this release diff --git a/code/export/checkpoint_to_safetensors.py b/code/export/checkpoint_to_safetensors.py index a64d0be..f8972cf 100644 --- a/code/export/checkpoint_to_safetensors.py +++ b/code/export/checkpoint_to_safetensors.py @@ -38,6 +38,7 @@ from export.safetensors_utils import _build_minimal_cfg, save_safetensors from model.factory import ModelFactory +from model.registry import get_model_spec def _load_checkpoint_state_dict(checkpoint_path: str, device: str) -> dict: @@ -82,9 +83,15 @@ def main(): ) parser.add_argument( "--model-id", - type=int, + type=str, required=True, - help="Public model ID (1..5).", + help=( + "Public model ID: 1..5 for surface-code checkpoints, or B for the " + "color-code cascade model. Color-code convolutional checkpoints " + "(model_id 1/2/4/5 trained with code=color) are not supported by " + "this converter: color training widens the final conv layer, so " + "the rebuilt surface architecture cannot load them." + ), ) parser.add_argument( "--fp16", @@ -108,6 +115,13 @@ def main(): ) args = parser.parse_args() + try: + spec = get_model_spec(args.model_id) + except ValueError as e: + parser.error(f"argument --model-id: {e}") + # Canonical registry key: int for 1..5, "B" for the cascade model. + model_id = spec.model_id + dtype = "fp16" if args.fp16 else "fp32" checkpoint_path = Path(args.checkpoint) @@ -123,8 +137,8 @@ def main(): print(f"Loading checkpoint: {checkpoint_path}") state_dict = _load_checkpoint_state_dict(str(checkpoint_path), args.device) - print(f"Building model architecture for model_id={args.model_id} ...") - cfg = _build_minimal_cfg(args.model_id) + print(f"Building model architecture for model_id={model_id} ...") + cfg = _build_minimal_cfg(model_id) model = ModelFactory.create_model(cfg).to(args.device) model.load_state_dict(state_dict, strict=True) @@ -135,7 +149,7 @@ def main(): print(" Converted to float16") print(f"Saving to: {output_path}") - save_safetensors(model, str(output_path), model_id=args.model_id, dtype=dtype) + save_safetensors(model, str(output_path), model_id=model_id, dtype=dtype) size_mb = output_path.stat().st_size / (1024**2) print(f"Done. File size: {size_mb:.2f} MB") diff --git a/code/export/safetensors_utils.py b/code/export/safetensors_utils.py index d51c18c..76868bd 100644 --- a/code/export/safetensors_utils.py +++ b/code/export/safetensors_utils.py @@ -22,7 +22,7 @@ import json from pathlib import Path -from typing import Optional +from typing import Optional, Union import torch from safetensors import safe_open @@ -34,12 +34,12 @@ from omegaconf import OmegaConf -def _build_minimal_cfg(model_id: int): +def _build_minimal_cfg(model_id: Union[int, str]): """Build a minimal inference config for model_id without a full Hydra setup.""" spec = get_model_spec(model_id) cfg = OmegaConf.create( { - "model_id": model_id, + "model_id": spec.model_id, "distance": spec.receptive_field, "n_rounds": spec.receptive_field, "data": { @@ -47,13 +47,18 @@ def _build_minimal_cfg(model_id: int): }, } ) + # The cascade/bottleneck model "B" is color-code only. Conv ids (1..5) + # build the surface architecture here; color-trained conv checkpoints have + # a widened final layer (filters[-1] >= 16) and cannot be loaded into it. + if spec.model_version == "predecoder_memory_cascade": + cfg.code = "color" return apply_public_defaults_and_model(cfg, spec) def save_safetensors( model: torch.nn.Module, path: str, - model_id: int, + model_id: Union[int, str], dtype: str = "fp32", extra_metadata: Optional[dict] = None, ) -> None: @@ -63,7 +68,7 @@ def save_safetensors( Args: model: The model to save (should already be on cpu or target device). path: Output file path (e.g. "model_fp32.safetensors"). - model_id: Public model ID (1..5). + model_id: Public model ID (1..5 or "B"). dtype: "fp32" or "fp16". extra_metadata: Optional dict of additional string metadata to embed. """ @@ -73,7 +78,7 @@ def save_safetensors( spec = get_model_spec(model_id) metadata = { - "model_id": str(model_id), + "model_id": str(spec.model_id), "quant_format": dtype, "model_version": spec.model_version, "receptive_field": str(spec.receptive_field), @@ -89,7 +94,7 @@ def save_safetensors( def load_safetensors( safetensors_path: str, - model_id: Optional[int] = None, + model_id: Optional[Union[int, str]] = None, device: str = "cuda", ): """ @@ -107,16 +112,15 @@ def load_safetensors( with safe_open(safetensors_path, framework="pt", device="cpu") as f: metadata = dict(f.metadata()) - # Resolve model_id + # Resolve model_id. Keep it as-is (int or string): get_model_spec() + # normalizes numeric strings to ints and "b" to "B" for the cascade model. if model_id is None: if "model_id" not in metadata: raise ValueError( f"SafeTensors file has no 'model_id' in metadata and model_id was not provided: " f"{safetensors_path}" ) - model_id = int(metadata["model_id"]) - else: - model_id = int(model_id) + model_id = metadata["model_id"] cfg = _build_minimal_cfg(model_id) model = ModelFactory.create_model(cfg).to(device) diff --git a/code/qec/surface_code/memory_circuit.py b/code/qec/surface_code/memory_circuit.py index 86f7187..7f63216 100644 --- a/code/qec/surface_code/memory_circuit.py +++ b/code/qec/surface_code/memory_circuit.py @@ -1233,6 +1233,249 @@ def set_error_rates(self): } +def triangular_color_code_circuit(d): + + # Function which creates the Cmat version of the triangular color code circuit + + # Validate integrality of counts (optional but helpful) + if (3 * d * d + 1) % 4 != 0 or (3 * (d * d - 1)) % 8 != 0 or (3 * d - 1) % 2 != 0: + raise ValueError("d does not yield integer counts for this construction.") + + num_data = (3 * d * d + 1) // 4 + num_plaquettes = (3 * (d * d - 1)) // 8 + num_ancilla = 2 * num_plaquettes + num_rows_data_qubits = (3 * d - 1) // 2 + + num_time_steps = 10 + + # Use int dtype to keep gate codes as integers (optional) + Cmat = np.ones((num_data + num_ancilla, num_time_steps), dtype=int) + + cnot_target = 10000 # Number representing the target of a CNOT gate + + # First num_data rows of Cmat are for data qubits, and the remaining rows are the ancillas. + # Follow the same convention as the Gidney paper. Inverted triangle, data qubits are labelled left to right, top to bottom + + for tt in range(num_time_steps): + + if tt == 0: + # Initialize ancillas in each plaquette + index = num_data + for _ in range(num_plaquettes): + Cmat[index, tt] = 3 # Prepare in |+> + index += 1 + Cmat[index, tt] = 4 # Prepare in |0> + index += 1 + + elif tt == num_time_steps - 1: + # Measure the ancillas in each plaquette + index = num_data + for _ in range(num_plaquettes): + Cmat[index, tt] = 5 # Measure in X + index += 1 + Cmat[index, tt] = 6 # Measure in Z + index += 1 + + elif tt == 1 or tt == num_time_steps - 2: + # CNOT between the two ancilla qubits: control |+>, target |0> + index = num_data + for _ in range(num_plaquettes): + Cmat[index, tt] = cnot_target + (index + 2) # control points to target row + Cmat[index + 1, tt] = cnot_target # mark as "target" + index += 2 + + elif tt == 2: + # First sequence (data = control) + data_qubit_index = d + index_ancilla = num_data + cols = d - 1 + for rr in range(num_rows_data_qubits - 1): + if rr % 3 == 0: + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 1 + elif rr % 3 == 1: + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + data_qubit_index += 1 + if cc == cols - 1: + index_ancilla += 3 + else: + index_ancilla += 1 + else: # rr % 3 == 2 + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 1 + + elif tt == 3: + # Second sequence (data = control) + data_qubit_index = 0 # <-- fixed indentation (no leading space) + index_ancilla = num_data + cols = d - 1 + for rr in range(num_rows_data_qubits - 1): + if rr % 3 == 0: + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + index_ancilla += 1 + if cc == cols - 1: + data_qubit_index += 3 + else: + data_qubit_index += 1 + cols -= 1 + elif rr % 3 == 1: + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + data_qubit_index += 1 + if cc == cols - 1: + index_ancilla += 3 + else: + index_ancilla += 1 + else: # rr % 3 == 2 + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 1 + + elif tt == 4: + # Third sequence (data = control) + data_qubit_index = 1 + index_ancilla = num_data + d - 1 + cols = d - 1 + for rr in range(num_rows_data_qubits - 2): + if rr % 3 == 0: + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + elif rr % 3 == 1: + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 2 + else: # rr % 3 == 2 + for cc in range(cols): + Cmat[data_qubit_index, tt] = cnot_target + (index_ancilla + 1) + Cmat[index_ancilla, tt] = cnot_target + index_ancilla += 1 + if cc == cols - 1: + data_qubit_index += 3 + else: + data_qubit_index += 1 + + elif tt == 5: + # First sequence (data = target) + data_qubit_index = d + index_ancilla = num_data + cols = d - 1 + for rr in range(num_rows_data_qubits - 1): + if rr % 3 == 0: + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 1 + elif rr % 3 == 1: + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + data_qubit_index += 1 + if cc == cols - 1: + index_ancilla += 3 + else: + index_ancilla += 1 + else: # rr % 3 == 2 + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 1 + + elif tt == 6: + # Second sequence (data = target) + data_qubit_index = 0 + index_ancilla = num_data + cols = d - 1 + for rr in range(num_rows_data_qubits - 1): + if rr % 3 == 0: + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + index_ancilla += 1 + if cc == cols - 1: + data_qubit_index += 3 + else: + data_qubit_index += 1 + cols -= 1 + elif rr % 3 == 1: + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + data_qubit_index += 1 + if cc == cols - 1: + index_ancilla += 3 + else: + index_ancilla += 1 + else: # rr % 3 == 2 + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 1 + + elif tt == 7: + # Third sequence (data = target) + data_qubit_index = 1 + index_ancilla = num_data + d - 1 + cols = d - 1 + for rr in range(num_rows_data_qubits - 2): + if rr % 3 == 0: + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + elif rr % 3 == 1: + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + data_qubit_index += 1 + index_ancilla += 1 + cols -= 2 + else: # rr % 3 == 2 + for cc in range(cols): + Cmat[index_ancilla, tt] = cnot_target + (data_qubit_index + 1) + Cmat[data_qubit_index, tt] = cnot_target + index_ancilla += 1 + if cc == cols - 1: + data_qubit_index += 3 + else: + data_qubit_index += 1 + + # Add idle locations (2 during preparation, 7 during measurement) + for jj in range(num_data): + Cmat[jj, 0] = 2 + Cmat[jj, num_time_steps - 1] = 7 + + return Cmat + + if __name__ == "__main__": d = 5 diff --git a/code/tests/test_safetensors_export.py b/code/tests/test_safetensors_export.py index 4b447d3..eaeecf5 100644 --- a/code/tests/test_safetensors_export.py +++ b/code/tests/test_safetensors_export.py @@ -113,6 +113,26 @@ def test_invalid_dtype_raises(self): with self.assertRaises(ValueError): save_safetensors(model, path, model_id=self.MODEL_ID, dtype="int8") + def test_round_trip_model_b(self): + """String model ids (color-code cascade model "B") round-trip end to end.""" + cfg = _build_minimal_cfg("B") + self.assertEqual(cfg.code, "color") + model = ModelFactory.create_model(cfg) + with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as f: + path = f.name + self.addCleanup(os.unlink, path) + + save_safetensors(model, path, model_id="B", dtype="fp32") + loaded, metadata = load_safetensors(path, model_id=None, device="cpu") + + self.assertEqual(metadata["model_id"], "B") + self.assertEqual(metadata["model_version"], "predecoder_memory_cascade") + self._assert_state_dicts_close(model.state_dict(), loaded.state_dict(), atol=0.0) + + # Lower-case alias normalizes through the registry + loaded_alias, _ = load_safetensors(path, model_id="b", device="cpu") + self.assertIsNotNone(loaded_alias) + class TestSafeTensorsRunPyIntegration(unittest.TestCase): """Test PREDECODER_SAFETENSORS_CHECKPOINT code path in workflows/run.py.""" diff --git a/code/workflows/run.py b/code/workflows/run.py index 5b686bf..2d5bd2e 100644 --- a/code/workflows/run.py +++ b/code/workflows/run.py @@ -233,6 +233,21 @@ def run_surface(cfg: DictConfig): raise ValueError(f"Invalid test.meas_basis_test={basis_cfg!r}; expected X, Z, or both.") num_samples = int(getattr(cfg.test, "num_samples", 1000)) + # The public config rejects test.* overrides, so the shot count is + # otherwise fixed at the merged default; honor the same env var the + # inference workflows use. Unlike those (which silently ignore bad + # values), fail fast here — this workflow's only purpose is sampling. + env_samples = os.environ.get("PREDECODER_INFERENCE_NUM_SAMPLES") + if env_samples: + try: + num_samples = int(env_samples) + except ValueError: + num_samples = -1 + if num_samples <= 0: + raise ValueError( + "PREDECODER_INFERENCE_NUM_SAMPLES must be a positive integer, " + f"got {env_samples!r}" + ) output_dir = os.path.join(HydraConfig.get().runtime.output_dir, "stim_samples") noise_model_cfg = getattr(cfg.data, "noise_model", None) noise_model_params = None @@ -718,6 +733,25 @@ def _load_model(cfg, dist): model_id=None, device=str(dist.device), ) + + # The cascade model "B" is color-code only: loading it into a surface + # runtime config would silently run the surface workflow with a color + # cascade model, so fail fast instead of warning. + loaded_spec_model_id = metadata.get("model_id") + if loaded_spec_model_id is not None: + from model.registry import get_model_spec + code_name = str(getattr(cfg, "code", "surface")).lower() + spec = get_model_spec(loaded_spec_model_id) + if ( + spec.model_version == "predecoder_memory_cascade" and + not code_name.startswith("color") + ): + raise ValueError( + f"SafeTensors checkpoint has model_id={loaded_spec_model_id} " + f"(color-code cascade model), but the runtime config has " + f"code={code_name!r}. Re-run with code=color to use this model." + ) + if dist.rank == 0: loaded_model_id = metadata.get("model_id", "unknown") dtype = metadata.get("quant_format", "fp32") diff --git a/conf/config_pre_decoder_memory_surface_model_1_d9.yaml b/conf/config_pre_decoder_memory_surface_model_1_d9.yaml index 405569e..5ec39c4 100644 --- a/conf/config_pre_decoder_memory_surface_model_1_d9.yaml +++ b/conf/config_pre_decoder_memory_surface_model_1_d9.yaml @@ -56,7 +56,7 @@ data: p_min: 0.001 # Minimum physical error rate (when using random p per batch) p_max: 0.006 # Maximum physical error rate (when using random p per batch) error_mode: circuit_level_surface_custom # Error model (Stim inference) - precomputed_frames_dir: /root/pre-decoder/dem_data # Path to DEM bundle dir (H.X/H.Z/p/A). Use precompute_frames.py --dem_output_dir to generate. + precomputed_frames_dir: /root/pre-decoder/dem_data # Path to DEM bundle dir (H.X/H.Z/p/A). Use qec/precompute_dem.py --dem_output_dir to generate. enable_correlated_pymatching: False # If true, use two-pass correlated MWPM matching (more accurate but ~2.7x slower) code_rotation: XV # Surface code orientation: XV, XH, ZV, ZH (default: XV)