| Environment | Requirements |
|---|---|
| Development / CI | Python 3.9+, PyTorch ≥ 2.1 |
| AWS Trn1 / Inf2 | Above + torch-neuronx ≥ 2.1, neuronx-cc ≥ 2.0 |
# CPU-only (laptop, CI)
pip install lqcd-neuron
# With Neuron backend (Trn1 / Inf2 instance)
pip install lqcd-neuron[neuron]
# Editable install from source
git clone https://github.com/JGalego/lqcd-neuron
cd lqcd-neuron
pip install -e ".[dev]"from lqcd_neuron.core import LatticeGeometry
geom = LatticeGeometry(T=8, Z=4, Y=4, X=4)
print(geom.volume) # 512
print(geom.gauge_shape) # (8, 4, 4, 4, 4, 3, 3)
print(geom.spinor_shape) # (8, 4, 4, 4, 4, 3)T, Z, Y, X are the lattice extents in the temporal and three spatial directions.
Nc=3 (SU(3)) and Ns=4 (Dirac spinor) are the defaults.
import torch
from lqcd_neuron.core import GaugeField, ColorSpinorField
# Cold start — all links are identity matrices
U_cold = GaugeField.cold(geom)
# Hot start — random SU(3) gauge field sampled from Haar measure
U = GaugeField.random(geom, seed=42)
# Point source spinor δ(x) δ_{α0} δ_{c0}
b = ColorSpinorField.point_source(geom, t=0, z=0, y=0, x=0, spin=0, color=0)
# Gaussian random spinor
psi = ColorSpinorField.gaussian(geom, seed=7)Cast to Neuron's preferred bfloat16:
U_bf16 = U.to(dtype=torch.complex32) # lossless bfloat16 complexfrom lqcd_neuron.observables import plaquette, wilson_action, polyakov_loop
# Wilson plaquette P ∈ [0, 1]
P = plaquette(U)
print(f"Plaquette: {P:.6f}")
# Wilson action S_W = β Σ (1 - P_{μν}/Nc)
S = wilson_action(U, beta=6.0)
print(f"Wilson action: {S:.3f}")
# Polyakov loop ⟨L⟩ — order parameter for deconfinement
L = polyakov_loop(U)
print(f"Polyakov loop: {L:.4f}")from lqcd_neuron.dirac import WilsonDirac
# D_W with bare mass m = 0.1
D = WilsonDirac(mass=0.1)
# D ψ
Dpsi = D(psi.tensor, U.tensor)
# M†M ψ (for CG normal equations)
MtMpsi = D.normal(psi.tensor, U.tensor)from lqcd_neuron.dirac import CloverWilsonDirac
D_clv = CloverWilsonDirac(mass=0.1, csw=1.0)
D_clv.set_gauge(U.tensor) # pre-compute clover matrices
out = D_clv(psi.tensor, U.tensor)from lqcd_neuron.solvers import ConjugateGradient
# Build right-hand side: b̃ = D† b
b_rhs = D.dagger(b.tensor, U.tensor)
# Solve M†M x = b̃
solver = ConjugateGradient(tol=1e-8, maxiter=500, verbose=True)
x, info = solver.solve(lambda v: D.normal(v, U.tensor), b_rhs)
print(f"Converged: {info.converged}")
print(f"Iterations: {info.iterations}")
print(f"Final |r|/|b|: {info.final_residual:.2e}")For non-Hermitian systems use BiCGStab:
from lqcd_neuron.solvers import BiCGStab
solver = BiCGStab(tol=1e-8, maxiter=500)
x, info = solver.solve(lambda v: D(v, U.tensor), b.tensor)from lqcd_neuron.neuron import NeuronCompiler, is_neuron_available
if is_neuron_available():
compiler = NeuronCompiler(dtype="bfloat16")
# Recommended: pass `gauge_field=U` so the compiler can pre-fuse the
# spin/colour hopping kernels and bake U into the .neff as a buffer.
# The wrapper still accepts (psi, U) for API compatibility — U is
# already on the NeuronCore so the second argument is ignored.
D_neuron = compiler.compile_dslash(
D, lattice_shape=geom.shape, nc=geom.nc, gauge_field=U.tensor,
)
Dpsi_neuron = D_neuron(psi.tensor, U.tensor)The compiled module is a drop-in replacement for D — same signature, same output
shape. The solver loop stays on the host:
def matvec_neuron(v):
return D.dagger(D_neuron(v, U.tensor), U.tensor) # D†D_neuron
x_neuron, info = ConjugateGradient(tol=1e-8).solve(matvec_neuron, b_rhs)For propagator calculations that need multiple right-hand sides, compile a
batched operator instead. This amortises the fixed per-call NeuronCore
dispatch cost across all B spinors and is the highest-throughput path:
import torch
B = 12 # e.g. one RHS per spin × colour source
D_batched = compiler.compile_dslash_batched(
D, lattice_shape=geom.shape, batch_size=B,
gauge_field=U.tensor, nc=geom.nc,
)
# psi_batch shape: (B, T, Z, Y, X, Ns, Nc)
psi_batch = torch.stack([psi.tensor for _ in range(B)], dim=0)
out_batch = D_batched(psi_batch)On instances with more than one NeuronCore (inf2.xlarge has 2,
trn1.32xlarge has 32), shard the batch across cores via
torch_neuronx.DataParallel:
from lqcd_neuron.neuron import get_device
num_cores = get_device().num_cores # auto-detected from /dev/neuron*
per_core_batch_size = 8
B_global = num_cores * per_core_batch_size
D_mc = compiler.compile_dslash_multicore(
D, lattice_shape=geom.shape, gauge_field=U.tensor,
num_cores=num_cores, # default: all detected cores
per_core_batch_size=per_core_batch_size,
nc=geom.nc,
)
psi_global = torch.stack([psi.tensor for _ in range(B_global)], dim=0)
out_global = D_mc(psi_global) # shape: (B_global, T, Z, Y, X, Ns, Nc)The gauge field is baked into each core's .neff, so only the spinor
shard crosses PCIe per core per call. This stacks with multi-RHS
batching: each core still runs the fused 12×12 batched kernel on its
per_core_batch_size slice.
At neuronx-cc HLO
instruction budget ([NCC_EVRF007]). compile_dslash detects this and
auto-routes through compile_dslash_sharded, which splits the lattice
along the T axis into num_shards slabs and compiles one NEFF per slab.
Halos are gathered host-side under periodic BCs. No call-site change is
needed:
# 32^4 — auto-shards into 8 slabs of (T_local=4, 32, 32, 32)
D_big = compiler.compile_dslash(
D, lattice_shape=(32, 32, 32, 32),
nc=geom.nc, gauge_field=U_big.tensor,
)
out = D_big(psi_big.tensor, U_big.tensor)To control the partition explicitly:
D_big = compiler.compile_dslash_sharded(
D, lattice_shape=(32, 32, 32, 32),
gauge_field=U_big.tensor,
num_shards=8, # must divide T; default picks the smallest
# power-of-2 keeping V_local ≲ 150k sites
nc=geom.nc,
)The shards currently dispatch sequentially on a single NeuronCore, so
this path restores compilability at large compile_dslash_batched) and the multicore
data-parallel path do not yet support sharding and will still fail on
lattices that overflow the per-NEFF budget.
Note: the
.neffproduced by the gauge-baked path is specific to the exact gauge configuration passed in. Re-compile (cheap once warm) whenUchanges between solves.
pytest tests/ -vAll tests run on CPU without Neuron hardware.
The shipped examples/bench_dslash.py script
measures applications-per-second for WilsonDirac.forward() across a sweep
of lattice sizes and multi-RHS batch sizes, with derived GFLOP/s and GB/s
columns.
make bench # CPU baseline only, default lattice sweep
make bench NEURON=1 # CPU + Neuron + Batched + Multicore
make bench NEURON=1 LATTICE=16x8x8x8 # restrict to one lattice
make bench NEURON=1 BATCH=1,8,32,64 # custom batch-size sweepThe BATCH knob (or --batch-sizes 8,16,32 directly on the script) controls
the multi-RHS batch sizes used for the Batched and Multicore columns.
Default is 8,16,32, which spans the dispatch-overhead → bandwidth-saturation
transition for inf2.* instances. Per-core HBM on NeuronCore-v2 is 32 GiB,
so even B=64 at V=24⁴ fits with comfortable headroom — feel free to push
higher.
Note: the Multicore column compiles with per_core_batch_size=B, so its
effective per-call RHS count is num_cores * B (1536 RHS per call at
B=64 on inf2.24xlarge). Throughput is reported per-RHS, so columns
are directly comparable.
The compiled-once-per-(lattice, batch) results are reported as one row per combination, with a visual separator between lattices when sweeping multiple batches.
make connect-bench NEURON=1 BATCH=1,8,32 # SSH to the instance, run there
make bench-job NEURON=1 BATCH=1,8,32,64 # one-shot Inf2, results emailed
make bench-job WALLCLOCK=480 # long sweep: 8h kill switch (0 disables)make bench-job provisions an ephemeral Inf2 from the bench launch template
defined in infra/main.tf, runs the benchmark, archives
the full log to S3, emails a summary (with a 7-day presigned download URL)
via SNS, and self-terminates. A 120-minute wallclock kill switch fires at
boot as a safety net; raise it with WALLCLOCK=<minutes> (or set
WALLCLOCK=0 to disable) for long sweeps that would otherwise trigger
The system is going down for poweroff … wall messages. See the
"Fire-and-forget bench job"
section of the top-level README for the full setup (notably:
notification_email in terraform.tfvars plus a one-click confirmation of
the SNS subscription).
While the run is in flight, partial results land in S3 as soon as each lattice finishes — no need to wait for the full sweep:
make bench-runs # list runs uploaded so far
make bench-tail RUN=<run_id> | jq . # stream per-lattice JSONL
make bench-tail RUN=<run_id> LOG=1 # follow the live bench.log| Regime | Lattice | Suggested BATCH |
|---|---|---|
| Dispatch-overhead | ≤ 8⁴ | 1,8,32 |
| Sweet spot | 16⁴ – 24⁴ | 1,8,32,64 |
| HBM-tight | 32⁴ | 1,4,16 |
The Speedup column reports the best Neuron column / CPU; GFLOP/s and
GB/s are derived from the Multicore column when available, otherwise
Batched, otherwise Neuron.
The Neuron SDK ships two stock observability tools:
neuron-top— interactivehtop-style snapshot of NeuronCore usage.neuron-monitor— JSON-stream metrics emitter, one record per period.
# Real-time utilisation snapshot
neuron-top
# Continuous JSON metrics stream (one record per second)
neuron-monitor --period 1
# Kernel-level timelines (view in Perfetto UI)
neuron-profile ...For kernel-level timelines (closer to NVIDIA Nsight Systems), use
neuron-profile from the Neuron SDK and view the trace in Perfetto.
See src/lqcd_neuron/params.py for all parameter dataclasses:
| Class | Key fields |
|---|---|
GaugeParam |
lattice_size, nc, precision, t_boundary, anisotropy |
InvertParam |
dslash_type, inv_type, mass, kappa, tol, maxiter |
CloverParam |
csw, precision |
NeuronCompileParam |
dtype, optimize_level, num_neuroncores |