Enable torch-only NMOptimizer execution modes

[vedika-saravanan]

Description

This PR updates NMOptimizer to run through torch-backed contractions without relying on cuTensorNet. The optimizer now uses an internal reduced noise-contraction path so all supported execution modes share the same differentiable contraction flow.

Runtime / performance impact

This changes NMOptimizer execution to use torch-backed contractions and an internal reduced noise-contraction path. The reduced path avoids the full unreduced tensor-network contraction that can hit torch einsum label/memory limits at larger problem sizes.

Self-review checklist

Please confirm each item before requesting review. Check [x] or strike
through and explain.

Before requesting review

• I reviewed my own full diff in GitHub or my editor.
• PR is in Draft if it is not yet ready for review.
• Temporary / debugging changes have been removed.
• Local test logs reviewed; no unexplained warnings or errors.
• CI logs reviewed; no unexplained warnings or errors.
• Full CI has been run.

Scope and size

• PR is under ~1000 lines, or an exception is justified in the description.
• Refactoring-only changes are isolated in their own PR(s).
• No existing tests were disabled or modified just to make this PR pass
  (if so, an issue has been raised).

Tests

• New functionality has new tests.
• Tests fail if the new functionality is broken (including crashes), not
  just when it is missing.
• Negative tests added where exceptions are expected.
• Truth data added where simple EXPECT_* / assert checks are
  insufficient for algorithmic correctness.
• CI runtime impact considered; team notified if significant.

Documentation

• Public-facing APIs have Doxygen docs.
• User-visible behavior changes have public docs, or a follow-up is
  tracked.

Code style

• Naming follows the existing convention (snake_case vs camelCase) for
  the area being modified.

Dependencies

• No new third-party dependencies, or the team has been notified and
  OSRB tickets filed.

[npancotti]

Review — NVIDIA/cudaqx PR #621: "Enable torch-only NMOptimizer execution modes"

Focus: API and user experience (per request), backed by validation + a noise-learning
experiment on Google's surface-code DEMs, all run on a single H100 (Docker, GPU 0) on
fi-Kermit.nvidia.com.

Overall: the direction is good and the core is well-tested. My comments are mostly about
docs/UX polish plus one important scalability caveat: did not manage to fit exact d=5 contraction on a H100. Do we assume contengrust? or some specific optimizer?

API / UX findings

1. Docstring formatting bug in NMOptimizer.__init__ (public API docs)

In the Args: block, the note beginning "Per-error noise tensors are contracted with their
adjacent code tensors..." is indented like a parameter but has no parameter name. It will
render as a malformed/──orphaned argument in Sphinx/Doxygen. Move it into the class docstring
or a Notes: section.

2. execute="codegen" docstring drift

The docstring still advertises codegen as partial-evaluating "the path into a flat Python
function." That's still true for predict, but _build_loss_codegen (the fused
cross-entropy loss) was removed and all modes now go through _build_loss_wrapped. Users
who chose codegen for the fused-loss perf profile won't get it anymore. Either restore the
fused loss or update the docstring to say the loss is no longer fused.

3. Evaluating LER requires constructing a full optimizer

The rewritten example computes held-out LER via NMOptimizer(...).logical_error_rate(),
i.e. it builds the whole reduced TN + codegen + autograd graph just to run inference. This is
a real UX/cost smell — and it bites hard at scale: at d5/r05 each such construct costs
~12 minutes and can OOM (see §8). Previously evaluation used
qec.get_decoder("tensor_network_decoder", ...). Recommend a lightweight, gradient-free
inference/eval entry point (or documenting the get_decoder path for evaluation).

4. The example lost its self-checking gates

The old example raised RuntimeError if training didn't reduce loss / didn't beat the uniform
baseline. The rewrite only prints. Because this file also runs as an example test in CI
(scripts/ci/test_examples.sh qec), it no longer fails on a real regression. If the gate was
dropped because online resampling made it flaky, add a loose statistical gate (e.g. assert
learned_LER <= uniform_LER + tol) rather than removing the check entirely.

5. Contractor naming is opaque

torch vs oe_torch vs oe_torch_compiled (all with backend="torch") are hard to tell
apart, and NMOptimizer now hard-pins oe_torch_compiled while the plain torch config is
effectively unused by the optimizer. Add a short table/docstring explaining what each does and
when it's selected, or mark the internal ones as such.

6. CUDA fallback UX — good change

device="cuda" with no CUDA now warns and falls back to CPU (clearer than the old
backend-swap warning). Confirmed working (see below). Minor: the device docstring says
"cuda" or "cpu" but the example passes "cuda:0"; document the "cuda:N" form.

7. optimize_path return-value change

It now returns self._reduced_info and just stores path_batch = optimize or "auto", instead
of computing full_tn.contraction_info(...). Anyone inspecting the returned info (e.g.
.path) gets different semantics than the parent. Worth a note in the docstring / changelog.

8. (Most important) Be explicit about the scalability envelope

The PR says the reduced path "avoids the full unreduced tensor-network contraction that can hit
torch einsum label/memory limits at larger problem sizes." True, but it does not lower the
treewidth of the code+syndrome network, which is the real wall. Empirically (§Experiment), a
d=5, r=5 surface-code DEM does not fit on an 80 GB H100 in float64 at any batch — even
batch=1 OOMs — and setup alone is ~12 min. The description should state the practical envelope
(≈ d3 / small r on one GPU) so users don't expect d5 to "just work."

---

Validation (single H100, container ghcr.io/nvidia/cudaqx:latest, torch 2.9.0+cu126)

Unit tests

• test_nm_optimizer.py: 77 / 77 pass (CPU + CUDA), including the PR's own
  test_reduced_path_matches_full_network_reference, CPU/GPU parity, and
  test_forward_parity_with_tn_decoder (NMOptimizer forward matches the reference cuTensorNet
  decoder).
• test_tensor_network_decoder.py: 34 / 35 pass; the PR-modified
  test_valid_contractor_configs passes 9/9.
• The 1 failure (test_decoder_decode_single) is not a PR defect — the released 0.14.0
  core returns DecoderResult.result as a Python list, while the test (on main) expects an
  np.ndarray. Pure core-version skew from testing the PR's Python layer against a released
  wheel.

Docs example (tn_noise_learning.py)

Runs end-to-end and learns:

• Prior MAE: 1.39e-2 → 4.20e-3
• Held-out LER: uniform 0.0369 → learned 0.0230, essentially matching true-DEM 0.0231.

Toolchain notes (environment, not the PR)

• Tests were run by overlaying the PR's Python decoder layer onto the released cudaq_qec
  (the release predates NMOptimizer), with a small BatchDecoderResult compat shim for the
  parity tests (substitutes only the result container; the real contractions run).
• torch pin matters: with the latest torch (2.13) the backward hit a Triton
  bmm_outer_product kernel that crashed (LLVM x87 assertion) at nontrivial batch sizes.
  Pinning torch 2.9.0 (as the PR CI does) fixed it. Worth keeping the CI torch pin
  documented for downstream users.

---

Experiment — learning Google's DEM from an uninformed prior

Setup: parse Google's circuit_noisy.stim → DEM (decompose_errors=True) →
(H, L, true_priors) via beliefmatching; generate a well-specified factorized dataset from
those priors; start NMOptimizer from a uniform, uninformed prior p0 = 0.01; fit with
logit-space Adam via make_compiled_step; ~1000-shot batches (online resampling).
Data: Zenodo 6804040 (Google, "Suppressing quantum errors by scaling a surface code logical
qubit"), Z-basis.

d3 / r03 (surface_code_bZ_d3_r03_center_5_3) — main run — SUCCESS

• DEM: 24 detectors, 219 error mechanisms.
• build 29.8 s; train 200 steps in 5.3 s (26.5 ms/step); peak GPU 0.22 GB @ batch 1000.
• Loss 197.9 → 157.1; prior MAE vs Google DEM 1.03e-2 → 6.50e-3.
• Held-out LER (20k shots): uniform(0.01) 0.0935 → learned 0.0683, vs true-DEM
  0.0664 (gap +0.0018). i.e. from an uninformed start the learner recovers ~all of the
  achievable decoding gain and lands within 0.2% LER of Google's own DEM.

d5 / r05 (surface_code_bZ_d5_r05_center_5_5) — scalability — DOES NOT FIT (let's try to resolve this)

• DEM: 120 detectors, 1677 error mechanisms (vs 24/219 at d3).

• Reduced-TN build/path-opt alone ≈ 12 min per construction (1797 pairwise contractions;
  1677 sequential contract_ind calls + 5 path-optimizer trials).

• Differentiable contraction memory (float64), batch reduced per instruction:

  batch	largest intermediate	result
  1000	2^31 × 1000 (≈2.1e12)	OOM (needed ~17 TB)
  4	2^32 (34 GB)	OOM (alloc 32 GB)
  2	2^32 (34 GB)	OOM
  1	2^30 (8.6 GB single)	OOM — 74 GB accumulated across the path before failing

• Conclusion: at d5/r05 the exact differentiable contraction has an intrinsic ~2^32-element
  intermediate (treewidth-driven, ~batch-independent at small batch) and the 1797-step path
  keeps many multi-GB intermediates live for autograd → OOM even at batch=1 on 80 GB
  float64. Reducing the batch does not rescue d5/r05 — the code+syndrome network is simply
  too large. Remaining levers (not exercised, since you asked to hold dtype/rounds/path fixed
  and reduce batch): float32 (~halves footprint), fewer rounds, a lower-treewidth path
  (cotengra + kahypar), or slicing/multi-GPU.

Net: the torch reduced path is correct and efficient at d3 (recovers Google's DEM to within
0.2% LER, 27 ms/step, 0.22 GB), but exact contraction hits a hard wall by d5/r05 on one 80 GB
GPU. That envelope should be stated in the PR.

---

Housekeeping

• Left running on fi-Kermit: container pr621 (GPU 0) with the extracted dataset, scripts,
  and logs (~/d3.log, ~/d5sweep.log). Remove with docker rm -f pr621 when done.