Mixed-precision reconstruction, fused-kernel dispatch, and CUDA graph capture for tomography INR#22
Draft
cedriclim1 wants to merge 92 commits into
Draft
Mixed-precision reconstruction, fused-kernel dispatch, and CUDA graph capture for tomography INR#22cedriclim1 wants to merge 92 commits into
cedriclim1 wants to merge 92 commits into
Conversation
Add a has_quantem_cuda capability flag (mirroring has_torch/has_cupy) and a use_cuda_kernels config option (default true) as a kill-switch. ObjectPixelated.get_tv_loss now goes through tv_loss_vol_sq in tomography/utils.py, which dispatches to quantem-cuda's fused squared-anisotropic TV kernel (identical math, analytic backward, 10-14x faster fwd+bwd at 256^3-512^3) when the package is installed, the object is a CUDA fp32 tensor, and the kill-switch is on; otherwise it falls back to the existing pure-torch expression. quantem-cuda is not yet on PyPI, so the [cuda] extra is deferred to a one-line follow-up after the 0.1.0 publish (an unpublished package in optional-dependencies breaks uv lock for everyone). Until then, installing quantem-cuda from source enables the dispatch. Tests cover the fallback paths (CPU, no quantem-cuda, kill-switch) and kernel parity for values and gradients when it is installed.
interpolate_ms_features_tilted routes each multiscale level through the fused kernel (rotate + grid_sample + Hadamard product, analytic backward) when quantem-cuda is installed, tensors are CUDA fp32, and the use_cuda_kernels config flag is on; the torch path is unchanged and remains the fallback. Covered by GPU/CPU parity, gradient parity, per-scale dispatch-spy, and kill-switch tests.
The per-pixel DataLoader path costs ~43 ms/batch on CPU (8192 Python __getitem__ calls + collate + H2D copy) while the GPU step itself takes ~20 ms, so single-GPU reconstructions idle the GPU half the time. DeviceBatchSampler keeps the tilt stack resident on the device and builds each batch with index arithmetic and two tensor lookups, yielding the same batch dicts (and drop_last/val-split semantics) as the DataLoader path. DDP runs keep the DataLoader + DistributedSampler path. Also disable bf16 autocast in the validation loop to match the training pass: the so3 pose solve (lu_factor) has no BFloat16 kernel, and a bf16 val loss is not comparable to the fp32 train loss.
Every rank derives the same epoch permutation from seed + epoch (CPU generator, identical across ranks and reproducible) and takes an equal-size contiguous shard, with the ragged tail dropped so per-rank batch counts always match and gradient sync cannot hang. The training loop's existing sampler.set_epoch call drives reshuffling, exactly like DistributedSampler; single-process runs auto-advance the epoch instead. The train/val split now uses a fixed-seed generator: identical across ranks (no leakage between a rank's train shard and another's val shard) and stable across save/reload, so resumed runs keep validating on the same held-out pixels. Verified with a 2-GPU torchrun run: equal batch counts on both ranks, finite identical reduced losses, no deadlock.
…raint crash - TomographyINRDataset.__getitem__ decomposed pixel indices by shape[1] (height) instead of shape[2] (width), and __len__ used max(shape)^2 instead of H*W; both corrupted the pixel mapping and overcounted the dataset for rectangular tilt images. - ObjectPixelated.apply_soft_constraints built soft_loss as a leaf tensor with requires_grad=True and then added the TV loss in-place, which raises RuntimeError whenever tv_vol > 0. Build it without requires_grad and add out-of-place so the loss backprops through get_tv_loss.
Evaluate the four TV taps (base, +x, +y, +z) in one batched 4N-point forward instead of four serial 10k-point calls, which were kernel-launch bound (volume TV 1.62 -> 0.52 ms, full training step 7.16 -> 5.23 ms at the 200^3 benchmark config). Also gate the plane and volume TV terms on their own weights: tv_plane > 0 with tv_vol == 0 previously applied no plane TV at all, and tv_vol > 0 computed a zero-weighted plane TV. The batched call uses the unwrapped model, consistent with the base tap and _get_plane_tv_loss.
- interpolate_ms_features_tilted: build the per-plane coordinate pairs with unbind/stack views instead of allocating an index tensor on device every call (a host-to-device copy) and running an (T, 3, B, 3) expand+gather. Bitwise-identical output; ~1.1x on the isolated forward (B=4096, T=8, 3 scales) plus one fewer H2D transfer per model call. - ObjectTensorDecomp.get_volume_tv_loss: evaluate the base points and the three axis-shifted copies in one batched forward (4N points) instead of four sequential model calls. Bitwise-identical loss; 3.1x faster forward and 3.3x with backward (10k samples, KPlanesTILTED backbone).
Each training step previously made two model calls on a tensor-decomp object: the main batch forward and a second 40k-point call for the volume-TV finite-difference taps. Both backwards accumulate gradients into the same plane parameters, so autograd runs the full gradient-accumulation traffic twice; profiling showed that traffic (copy_/add_/fill_, ~1.8 ms/step) exceeds the interpolation backward itself. The object model now exposes sample_tv_tap_coords(), the training loop concatenates the returned tap points onto the main batch for a single model call via forward_with_tv_taps() (mask and hard constraints apply to the main chunk only; tap densities stay raw, matching the previous TV semantics), and the tap densities reach get_volume_tv_loss through ReconstructionContext. The two-call path remains as a fallback for direct callers and non-tensor-decomp models. Full step 5.67 -> 4.79 ms (-15.5%) at the 200^3 benchmark config; loss and gradient parity verified against the two-call path, including taps crossing the [-1, 1] boundary.
…e crash - ObjectINR.forward masked densities outside [-1, 1] in x and y only, so tilted rays whose sample points leave the volume along z still picked up extrapolated density and biased the integrated projections at high tilt. The mask now bounds all three axes. - ObjectINR.apply_soft_constraints read ctx.coords.device before the coords-is-None assert, raising AttributeError instead of returning a zero loss when no soft constraints are active and no coords are passed. Fall back to the model device when coords are absent.
- TomographyINRDataset.__getitem__ wrapped each of the three index fields in torch.tensor(), allocating three scalar tensors per item on the dataloader hot path; default_collate builds the same int64 batch tensors from plain ints. 2.9x faster batch loading (12.1 -> 4.2 ms/batch, batch_size=1024, 60x256x256 stack), which dominated the per-batch GPU compute (~0.4 ms). - transform_batch_rays applied the three Euler rotations as nine elementwise passes over the full (B, S) ray tensors. Compose them into one (B, 3, 3) matrix and apply with a single batched matmul: same result to 4e-7 (float32 op reordering), up to 1.15x at large batches and far fewer kernel launches.
- TomographyConventional._reconstruction_epoch wrote the aligned measurement into proj_forward, which radon_torch overwrites immediately after, while the error term keeps reading the original tilt stack -- inline_alignment was a no-op. The aligned image is now persisted into the tilt stack (the tensor the error term reads), and the measurement is transposed to match the slice/detector orientation of the forward projection it is correlated with. - differentiable_rotz/rotx_vectorized raised for more than one angle: the per-slice vmap built affine_grid with a (T, 2, 3) matrix against a slice batch of 1. A rotation about an axis applies the same 2-D transform to every slice along it, so that axis can ride along as grid_sample channels in a single call -- which both fixes multi-angle/per-volume batching and removes the vmap. Scalar-angle outputs are unchanged (verified to 1e-6).
…g in rot_ZXZ - Every SchedulerParams dataclass (Plateau, Exponential, Cyclic, Linear, CosineAnnealing) mutated its own fields inside params(): the first call permanently baked derived values (min_lr, gamma, base/max_lr, total_iters, T_max) into the instance, so a config shared between the object and pose optimizers -- or reused across reconstruct() calls with a different num_iter -- silently kept the first call's values. Derived values are now computed locally; explicitly-set fields still take precedence. - rot_ZXZ re-wrapped all three Euler angles with torch.tensor() whenever any one of them was a non-tensor, copying and detaching tensor angles and silently cutting gradient flow (plus a UserWarning). Each angle is now converted independently, leaving tensors untouched.
…able_tilts setter - Siren/HSiren created and seeded a torch.Generator for the winner initialization but drew the perturbation with torch.randn_like, which ignores generators -- so winner_initialization=72 (used by TomographyLiteINR) silently had no effect on reproducibility and every seed produced the same global-RNG noise. Draw from torch.randn with the seeded generator instead. - TomographyDatasetBase.learnable_tilts had a setter that wrote a private attribute the getter never read, so assignments appeared to succeed while silently doing nothing. The value is derived from the tilt series; the setter is removed so assignment now raises instead of lying.
…s resolutions - TomographyPixDataset.to() and TomographyINRDataset.to() rebuilt the z1/z3/ shift parameters from the initial-value buffers (_z1_angles etc.), which are never updated during training -- so any device move after training, including Tomography.from_file(...).to(device), silently reset every learned pose to zero. Both now go through a shared helper that moves the current parameter values once they exist and only uses the buffers on first materialization. - KPlanes/KPlanesTILTED allocate all three plane grids from one (3[, *T], C, res[1], res[0]) tensor, ignoring res[2]; an anisotropic resolution silently gave the XZ/YZ planes the wrong grid along z. Constructor now raises a clear ValueError instead (isotropic behavior unchanged; supporting true anisotropy needs per-plane parameters, which would change the checkpoint layout).
…-softloss Fix INR dataset pixel indexing for non-square tilt images and TV soft-constraint crash
Speed up KPlanesTILTED feature interpolation and batched volume TV loss
Fold volume-TV tap evaluation into the main forward pass
…vice Fix ObjectINR out-of-volume masking along z and soft-constraint device crash
Reduce INR dataloader and ray-transform overhead
Fix no-op SIRT inline alignment and multi-angle rotation operators
…-grads Make scheduler params() pure and fix gradient-detaching angle wrapping in rot_ZXZ
Fix ignored winner-initialization seed in Siren and remove dead learnable_tilts setter
quantem-cuda moved its kernels into per-module submodules mirroring quantem: the K-Planes interpolation now lives in quantem.cuda.core.ml (KPlanesTILTED is a quantem.core.ml model shared across applications) and the TV kernels in quantem.cuda.core. Update the dispatch imports and the test monkeypatch targets accordingly.
docs/notebooks/quantem_cuda_kernels.ipynb spells out exactly what the quantem-cuda TV kernels implement (tv_loss_sq_3d: unnormalized squared anisotropic sum, tv_vol parity; tv_loss_iso_3d: corner-restricted isotropic mean, eps inside the sqrt) and what they do not (no anisotropic-L1 kernel, no per-axis weights, iso needs D >= 2), with parity checks, a ptychography section (complex multislice objects, per-axis weighting recipe, honest op-level timings), the fused TILTED K-Planes interpolation contract, and the transparent-dispatch config surface. Executed on an RTX PRO 6000 so the outputs are visible in review. .gitignore gains an exemption for curated docs notebooks.
Fix learned pose parameters being reset by to() and refuse anisotropic KPlanes resolutions
Device-resident batch sampling for INR tomography reconstruction
…kernel ObjectConstraints._calc_tv_loss now routes fp32 CUDA 3-D arrays through quantem.cuda.core.tv_loss_l1_3d (per-axis |diff| sums), composing the same anisotropic-L1 functional — per-axis (tv_weight_z, tv_weight_xy) weights and active-axis normalization included — with identical gradients (sign(0) = 0). Weighted size-1 axes fall through to torch so degenerate inputs behave exactly as before; the kill switch and capability flag are the same as the other dispatch points. ~4x fwd+bwd at multislice (16, 1024, 1024) sizes, neutral at single-slice sizes.
…ity gate docs/notebooks/quantem_cuda_kernels.ipynb now documents tv_loss_l1_3d (the ptychography functional) with a verified parity composition and updated timings/dispatch tables. docs/notebooks/ptycho_tv_kernel_quality.ipynb reruns the tutorial ducky reconstruction (simulated dataset, single-slice xy TV and multislice z+xy TV) with and without the kernel and gates on SSIM of the reconstructed phase against an fp-perturbation chaos floor: the pipeline is bit-exact deterministic, so a 1-ulp TV rounding difference diverges trajectories exactly like a 2e-7 relative defocus perturbation does. Both configs pass (SSIM 0.997 vs floor 0.996 single-slice; 0.992 vs 0.988 multislice), with final data-fidelity losses agreeing within the floor's scatter.
…n' into bench/cuda-dispatch
Single benchmark+profiling branch for the fasttomo/quantem benchmarking project so nsys reports and benchmark runs share one code lineage (bench/cuda-dispatch @ 20f767b + NVTX instrumentation @ fcb8a2a). PROFILING IS PROJECT-SCOPED: this merge commit is the designated removal point — revert -m 1 (or exclude the NVTX commits) before any upstream contribution; PR #21 stays profiling-free. # Conflicts: # src/quantem/tomography/tomography.py
…URE) torch.cuda.profiler.start/stop around n steady-state grad steps, paired with nsys --capture-range=cudaProfilerApi in the harness nsys launcher. Inert when the env var is unset (single early return; profiler import stays lazy). Opus-reviewed: APPROVE (production-inert, DDP lockstep-safe, capture-once).
…er + tomography changes + tests) as used for BIM-82 benchmark runs; base=46ef632
Port the continuous-rotation dataset extensions from the auto-et-continuous working tree onto the conference-snapshot lineage: per-line angle arrays on the INR dataset, atom_trace module, and the accompanying dataset/sampler/ holdout tests (63 tests).
Add keyword-only autocast_dtype (None/"bf16"/"fp16") replacing the two hardcoded disabled-autocast blocks; fp16 routes through GradScaler with the pose-gradient DDP sync ordered before unscale so found_inf stays consistent across ranks, and optimizer stepping gated identically to step_optimizers. Default None path is operation-identical to before. Adds CPU tests plus a CUDA-gated fp16 reconstruction test.
r9svd's U @ Vh matmul was autocast-eligible, so bf16/fp16 autocast fed a half tensor into torch.det, which has no half LU kernel. Wrap r9_to_rotmat in a disabled-autocast region: rotation matrices are precision-critical and staying fp32 also keeps the fused CUDA kernel's fp32 dispatch gate active under mixed precision. Quaternion path is elementwise-only and unaffected. Adds a CPU bf16-autocast regression test.
The fused CUDA kernels consume (3T, H, W, C); keeping the parameters physically NHWC makes the wrapper's per-call permute a metadata-only view instead of a full-grid copy each forward/backward (top level is 184 MB at T=8 C=48 200x200x3 scales). Logical NCHW shape, the grid_sample fallback, state_dict round-trips, and Adam state layout are unchanged; legacy row-major whole-module checkpoints are upgraded in __setstate__.
Add keyword-only grad_scaler override to reconstruct(): None keeps the fp16-implies-scaler policy, True/False force it (True without autocast is rejected). Enables bf16+scaler experiments; scaler mechanics unchanged.
…vel op Prefer quantem_cuda::kplanes_tilted_fuse_ms when exactly three fp32 grids are present and the op is exported; scale gates (or unit defaults) pass as scalars. Older wheels, other scale counts, non-fp32 inputs, and instrumented single-level entry points keep the per-level path. Dispatch tests cover the gate values and every fallback branch.
reconstruct(cuda_graphs=True) captures the forward/loss/backward region into a replayable graph with static input buffers (per-step copy_ + replay), and captures a capturable fused-Adam step as a second graph sharing the pool — the first Adam update runs eagerly to materialize optimizer state so the captured step tensor advances correctly on replay. box_fixed_ds gains a graph-only fixed-padded-ray branch (masked, sparsity denominator preserved); KPlanesTILTED with frozen poses hoists rotation matrices into a temporary non-persistent buffer override; capture-unsafe host-scalar constructions on the captured path were replaced with device-native allocations. DDP raises NotImplementedError; pose learning, partial batches, and trainable SO(3) fall back to eager. fp32 -16% per-epoch vs eager on the static benchmark with matching end metrics; the equivalence tests assert nonzero parameter movement, exact Adam step-count agreement, and tiered loss tolerances.
__init__ assigns the ParameterList via nn.Module.__setattr__, which registers it as the public child 'grids' and never calls the same-named property setter, so the property's backing _grids attribute never existed. Eager lookup happened to find the registered child, but dynamo resolves the data descriptor directly and raised AttributeError inside any traced forward. Drop the property; the registered module child is the documented contract and state-dict naming is unchanged. Adds a fullgraph trace test for KPlanesTILTED.
torch.cuda.nvtx.range_push raises RuntimeError on CPU-only wheels, which broke every reconstruct call (and the CI test matrix) without CUDA. Bind nvtx to a no-op shim when torch.cuda.is_available() is false; GPU runs keep the real nvtx ranges unchanged.
Capability-gated on three fp32 CUDA grid levels and quantem-cuda exposing plane_tv_loss; QUANTEM_PLANE_TV_FUSED=0 restores the eager chain, which remains the fallback for all other cases.
Under bf16 autocast the multiscale features must enter sigma_net as bf16 with no separate cast; autocast-off stays fp32.
…n capability get_tv_loss uses the fused op's raw scalar times the constraint weight exactly once and excludes the standalone term for that call; TruncExpActivation advertises nonnegative output so the redundant positivity clamp is skipped only for capability-marked activations.
FusedHiddenMLP keeps nn.Sequential parameter identity and falls back to the eager path for any non Linear/ReLU/Linear/ReLU/Linear topology, autocast-off, CPU, or missing extension.
…d-fork schedule grad_clip_max_norm=None/0 skips the validated-inactive clip pass. The combined plane-TV op becomes opt-out and is requested only inside the reconstruction forward scope, so pretraining and generic forwards never pay it. An opt-in side stream overlaps the consistency loss with the soft-constraint window; S3IM and all RNG stay on the main stream.
…ed-fork opt-in QUANTEM_FUSED_MLP defaults on with only "0" disabling, matching the TV knob. reconstruction_forward_context clears the stored TV scalar when the scoped forward raises. The prediction-fork side stream becomes opt-in and documents its RNG-free loss assumption; its equivalence test uses rtol=1e-4 with the atomic accumulation-order rationale. The CUDA-graph parity test gains a three-level arm covering the default-on combined plane-TV path, and model gradient parity uses measured error-vs-truth bounds.
Atomic accumulation-order drift compounds with step count; the 100-step arm uses rtol=1e-3 while the short arms keep 1e-4.
Two identical fork-disabled runs of the fixture diverge by 1.5e-2 at 100 steps from atomic accumulation-order chaos alone (zero at 10 steps), so exact-trajectory equivalence is only assertable on the short arms. The intrinsic-drift control test documents the measurement.
torch.bfloat16, torch.float16, and torch.float32 normalize to the existing bf16/fp16/off forms; strings and None keep working, anything else raises with both accepted spellings listed.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem addressed
The tomography INR reconstruction loop ran fp32-only with per-level kernel dispatch and no capture support. This PR carries the quantem side of the kernel/floating-point optimization campaign: production step time 42.0 → 26.5 ms (−37%) on RTX 6000 Pro with foreground-Pearson parity inside the fp32 seed band on both benchmark arms, and ≥0.95 GT correlation demonstrated at convergence.
Note for reviewers: this branch is based on
bench/conference-snapshot-20260722, so the diff vsmainincludes that lineage. The campaign's own commits are6d102cd..e61206f(9 commits) — review those.Changes
TomographyINRDataset, +63 tests) plusatom_trace.autocast_dtypeknob onreconstruct()(None/"bf16"/"fp16") with GradScaler for fp16 (DDP-safe unscale ordering), and an independentgrad_scaleroverride.torch.det; rotations stay fp32, which also keeps the fused-kernel dispatch gate active under mixed precision).interpolate_ms_features_tiltedprefers the new three-level op when available, with per-level fallback for other scale counts, older wheels, and instrumented paths.reconstruct(cuda_graphs=True)): forward/loss/backward captured with static buffers; capturable fused-Adam step as a second graph; frozen-pose KPlanesTILTED supported via a pre-capture rotation hoist. −16% per-epoch on fp32 eager (−2–4% on the optimized stack). DDP and pose learning fall back safely.KPlanes.gridsproperty that made any dynamo tracing of the model fail.Recommended configuration
bf16 autocast + fused Adam +
QUANTEM_KPLANES_BWD_VARIANT=5+QUANTEM_KPLANES_BWD_ZERO_TAU=6e-8(with the companion quantem-cuda PR).recon_continuous-style scripts also gain ~3–4% fromfused=TrueAdam independently.Negative results (documented, not landed)
Scoped torch.compile of the step region (+12% — inductor cannot fuse across custom-op boundaries) and manual get_coords op-count reduction (neutral; the loop is 88.6% GPU-busy so host launches are already hidden). Archived for a future A100 re-evaluation, where the 0.77x compile behavior seen on Perlmutter also needs re-testing.
Portability
The dispatch pattern is capability-based (
config.get("has_quantem_cuda")+ op-presence checks) at the model layer, so INR ptychography and other workflows using the same planar-decomposition models inherit the fused kernels and mixed-precision machinery without integration work.Update (round 2): adds the fused plane-TV loss dispatch and bf16 feature-boundary coverage. Requires the matching quantem-cuda branch (cedriclim1/quantem-cuda#2). Measured on RTX PRO 6000: 26.5 → 17.5 ms/step (−34%) on both benchmark arms, foreground Pearson inside the fp32 seed band at the 2830-step gate on each change individually and combined.
Update (rounds 3–4): combined multiscale+plane-TV op (TV gradient accumulated directly into the grid-gradient buffers), fused cuBLASLt MLP dispatch for the hybrid sigma head (default-on, eager fallback for other topologies), optional gradient-clip skip, reconstruction-scoped TV fusion, and an opt-in experimental side-stream schedule. Requires the matching quantem-cuda branch. Measured at the validated production configuration: 26.6 → 23.3 ms/step (−12%); a full 36,100-update reconstruction reproduces the reference quality (global Pearson 0.9916) in ~20 min.