Skip to content

Key Switching return NaN #28

Description

@Xuan-Wang-Summer

Multi-GPU: peer copies into ncclMemAlloc memory silently write nothing, so every key-switching op returns NaN

This is generated in assistance of Claude.

Summary

On FIDESlib 2.1.2, every CKKS operation that key-switches or rescales returns NaN
when more than one GPU is used. Mult, Square, MultRescale, Conjugate,
Rotate and HoistedRotate all decrypt to -nan at N=2, N=3 and N=4, while the
same tests pass at N=1 with 45–47 bits of precision.

The cause is that the multi-GPU exchange buffers are allocated with ncclMemAlloc
but transferred with cudaMemcpyPeerAsync. ncclMemAlloc returns a CUDA VMM
allocation that is only mapped for NCCL's own collectives; the peer device is never
granted access via cuMemSetAccess. cudaMemcpyPeerAsync into such memory returns
cudaSuccess and writes nothing.
The return value is also discarded
(src/PeerUtils.cu:809), so nothing surfaces. Each GPU proceeds to key-switch
against a gather buffer holding only its own limbs, producing garbage.

A second, environment-dependent problem compounds this: on this machine
(RTX PRO 6000 Blackwell, no NVLink), calling cudaDeviceEnablePeerAccess makes
all subsequent peer copies silently no-op, even between plain cudaMalloc
buffers. See "Broken P2P" below — that part may be specific to RTX-class boards.

Two-line proof, straight out of the ModUp gather. After the peer exchange, each GPU's
bufferGATHER holds only the limbs it computed itself ('X' = data, '.' = untouched;
24 global limb slots, GPU0 owns 0–11 and GPU1 owns 12–23):

[POSTCOMM] gpu=0 d=2 slots=XXXXXXXXXXXX............
[POSTCOMM] gpu=1 d=2 slots=............XXXXXXXXXXXX

Nothing crossed. The copies were issued, with correct source/destination pointers and
devices, and returned success:

[PEERCPY] cur_dev=0 src_dev=0 dst_dev=1 bytes=2097152 rc=0 sync_rc=0
          src=0x...(dev0) dst=0x...(dev1) src[0]=334084425805265616 dst[0]=0  <-- MISMATCH

Environment

FIDESlib 2.1.2, commit 0ec4055 (tip of upstream/main, clean tree)
OpenFHE 1.5.1 (bundled patched build)
CUDA runtime 12.9 (nvcc V12.9.86), driver reports 13.2
NCCL 2.29.7
Driver 595.58.03
GPUs 4x NVIDIA RTX PRO 6000 Blackwell (sm_120), PCIe, no NVLink
OS / GCC / CMake Ubuntu 24.04.4, GCC 13.3.0, CMake 3.28.3

Reproduction

# Passes: 45-bit precision
FIDESLIB_USE_NUM_GPUS=1 ./build/fideslib-test \
  --gtest_filter='OpenFHEInterfaceTests/OpenFHEInterfaceTest.Mult/5'

# Fails: result is all -nan
FIDESLIB_USE_NUM_GPUS=2 ./build/fideslib-test \
  --gtest_filter='OpenFHEInterfaceTests/OpenFHEInterfaceTest.Mult/5'

All 8 parameter instantiations fail, covering dnum = 1..4 and the FLEXIBLEAUTOEXT,
FIXEDAUTO and FIXEDMANUAL scaling techniques, so the bug is independent of both.
The shared parameter set (TTALL64BOOT, test/ParametrizedTest.cuh:449) is
logN = 16, L = multDepth = 23, firstModSize = 60, scaleModSize = 59.

Bootstrap fails too (decode error at N=4, hang at N=2), but only because it is built
from these operations. Mult/5 runs in ~12 s with no bootstrap precomputation and is
the cheaper reproducer.

Root cause

bufferGATHER (the ModUp gather buffer) and top_limb_buffer / top_limb_buffer2
(the rescale exchange buffers) are allocated with ncclMemAlloc:

  • src/CKKS/LimbPartition.cu:2643bufferGATHER
  • src/CKKS/Context.cu:733,735top_limb_buffer, top_limb_buffer2

But when MEMCPY_PEER is on (the default, LimbPartitionMGPU.cu:47), the data is
moved with cudaMemcpyPeerAsync (src/PeerUtils.cu:809) or a graph memcpy node
(:798), not with NCCL collectives. ncclMemAlloc uses the CUDA VMM API
(cuMemCreate/cuMemMap) and only maps the allocation for NCCL's own transports.
Because the peer device was never granted access, the copy is dropped. It returns
cudaSuccess; cudaStreamSynchronize on the copy's stream also returns success.

The unchecked return at PeerUtils.cu:809 means this fails completely silently:

cudaMemcpyPeerAsync(dst, dst_dev, src, src_dev, elems * sizeof(float), s);

The failure is invisible in the test suite as well, because ASSERT_ERROR_OK cannot
detect NaN — see "Bonus bug" below.

Consistent with this, the two ops that need no cross-GPU exchange (ScalarAdd,
MultScalar without rescale) are the only ones that pass multi-GPU.

Fix

Allocate the buffers that the peer-copy transport writes into with plain cudaMalloc
(which cudaDeviceEnablePeerAccess covers) rather than ncclMemAlloc, whenever
MEMCPY_PEER is the active transport. Keep ncclMemAlloc for the NCCL path, which
needs registered memory.

With that change (plus disabling peer access on this box, below), the previously
failing operations pass at N=2 and N=4 with 46–47 bits of precision, matching N=1.

Independently, cudaMemcpyPeerAsync's return value should be checked. Had it been,
this would have been a loud failure rather than silent NaN — though on this hardware
it returns cudaSuccess anyway, so error checking alone would not have caught it.

Broken P2P: cudaDeviceCanAccessPeer cannot be trusted

Even with plain cudaMalloc buffers, peer copies on this machine silently do nothing
once cudaDeviceEnablePeerAccess has been called. Standalone reproducer, no
FIDESlib involved:

canAccessPeer 0->1=1  1->0=1
A: no-peer-access, stream on src dev: dst[0]=0xabababababababab LANDED
enablePeerAccess(0->1) = no error
enablePeerAccess(1->0) = no error
B: peer-access enabled:          dst[0]=0x0 NOT LANDED
C: synchronous cudaMemcpyPeer:   dst[0]=0x0 NOT LANDED
D: cudaMemcpyAsync(Default):     dst[0]=0x0 NOT LANDED
driver/runtime: driver=13020 runtime=12090

cudaDeviceCanAccessPeer reports 1 and nvidia-smi topo -p2p r says OK for every
pair, but direct P2P transfers do not work. Before peer access is enabled the driver
stages the copy through the host and it succeeds; enabling peer access selects the
direct path, which silently drops the write. Disabling peer access again restores
working copies, so the condition is recoverable at runtime.

Underlying cause on this host: the IOMMU

The machine is an AMD host with AMD-Vi active (/sys/class/iommuivhd0..3, and
each GPU sits in its own IOMMU group: 38, 62, 11, 7). The kernel command line carries
no iommu parameter, so it is at the BIOS default (enabled). NVIDIA GPU-to-GPU PCIe
P2P does not work through an active IOMMU on AMD platforms — which is precisely the
"capability query says yes, transfers silently drop" signature.

Remediation is a boot-time setting, not a code change:

# /etc/default/grub
GRUB_CMDLINE_LINUX_DEFAULT="quiet splash iommu=pt"     # try passthrough first
# or, if that is not enough:
GRUB_CMDLINE_LINUX_DEFAULT="quiet splash amd_iommu=off"
sudo update-grub && sudo reboot

Note also from nvidia-smi topo -m that GPU0/GPU1 are NODE (same root complex)
while GPU2 and GPU3 are SYS (different root complexes). Even with the IOMMU
disabled, P2P across SYS links is typically unsupported, so a runtime probe is
needed regardless.

Fix: probe P2P instead of trusting the query

Context.cu enabled peer access whenever cudaDeviceCanAccessPeer said yes. Replaced
with: enable peer access, then verify each enabled pair with a real 8 KB
cudaMemcpyPeer and check the bytes arrive. If any pair drops its copy, tear peer
access back down for every pair and set canP2P = false, which falls back to the
host-staged path that works everywhere. FIDESLIB_DISABLE_P2P=1 forces the fallback
without probing.

On this host the probe fires and multi-GPU then works with no environment variables:

FIDESlib: cudaDeviceCanAccessPeer reports P2P support, but peer copies are silently
dropped. Disabling peer access and falling back to host-staged copies. (On AMD hosts
this is usually the IOMMU: boot with amd_iommu=off or iommu=pt.)
GPU P2P? 0
[       OK ] OpenFHEInterfaceTests/OpenFHEInterfaceTest.Mult/5 (3182 ms)

(Pre-existing bug fixed along the way: the old loop called
cudaDeviceEnablePeerAccess(j, 0) in both branches. The second, guarded by
cudaDeviceCanAccessPeer(&canAccessPeer, j, i), is meant to enable access from j
toward i and needs cudaSetDevice(j) first — so access was only ever enabled in one
direction.)

Bonus bug: ASSERT_ERROR_OK cannot detect NaN

test/ParametrizedTest.cuh:559:

double diff = abs(resultGPU->GetRealPackedValue().at(i) - result->GetRealPackedValue().at(i));
acc += diff * diff;
Max = std::max(Max, diff);
...
ASSERT_LE(Max, pow(2.0, -result->GetLogPrecision() + 4));

std::max(a, b) returns (a < b) ? b : a. With a = 0.0 and b = NaN, 0.0 < NaN
is false, so it returns 0.0. Max stays 0.0 no matter how many slots are NaN, and
the assertion passes. acc does propagate the NaN but is only printed, never asserted.
The result is a test that prints its own contradiction and still reports [ OK ]:

Max error: 0 (Expected: 1.42109e-14), dev: -nan
Result GPU (-nan, -nan, -nan, ... ); Estimated precision: -nan bits
[       OK ] OpenFHEInterfaceTests/OpenFHEInterfaceTest.Rotate/5 (3167 ms)

Because of this, Rotate reported success on all-NaN output at N=2 and N=4. The other
key-switching ops were caught only incidentally: Mult/Square/MultRescale fail on
ASSERT_EQ_CIPHERTEXT (an exact integer-limb comparison), and Conjugate happens to
throw from OpenFHE's Decode() first. Rotate's ASSERT_EQ_CIPHERTEXT is commented
out (test/OpenFheInterfaceTests.cu:1156), leaving only the NaN-blind check.

Suggested fix:

     double diff = abs(resultGPU->GetRealPackedValue().at(i) - result->GetRealPackedValue().at(i));
+    ASSERT_TRUE(std::isfinite(diff)) << "Non-finite result at slot " << i;
     acc += diff * diff;
     Max = std::max(Max, diff);

This is arguably the more urgent of the two bugs: the suite cannot currently be trusted
to catch numerical breakage. Happy to split it into its own issue.

Ruled out

Each of these was tested and eliminated:

  • CUDA graph capture. FIDESLIB_USE_GRAPH_CAPTURE=0 still produces -nan.
  • A race / missing synchronization. Forcing cudaDeviceSynchronize() plus an
    OpenMP barrier around the gather does not help. The data is never produced, not late.
  • Gather addressing / off-by-one. GATHERmeta is global (L+1), and
    generateGatherLimb fills h_gatherptr[i] = bufferGATHER + N*i. The INTT's write
    offset and the peer copy's read offset provably coincide, and instrumentation confirms
    the owner's region is fully written (nonzero=262144/262144).
  • Transfer direction. transferKernel(src, dst, ...) — the call site is correct,
    despite reading backwards next to the commented-out cudaMemcpyPeerAsync above it.
  • Out-of-bounds access. compute-sanitizer --tool memcheck reports zero invalid
    reads/writes.
  • NCCL lacking Blackwell support. cuobjdump confirms sm_120 cubins are present.
    The cudaErrorNoKernelImageForDevice errors memcheck reports come from
    ncclInitKernelsForDevice probing with cudaFuncGetAttributes, and are benign.
  • Other tenants on the shared GPU box. N=1 passes on the most heavily loaded GPU;
    every 2-GPU pairing fails identically; the NaN is bit-reproducible.

Bootstrap N=2 hang — fixed (third ncclMemAlloc/collective site)

The N=2 bootstrap hang was a third instance of the same root cause. broadcastLimb0_mgpu
(LimbPartitionMGPU.cu:2755), called from ModRaise, broadcasts the raised
ciphertext's limb-0 with raw ncclSend/ncclRecv — the one remaining collective on the
bootstrap path. On this host NCCL's transport deadlocks (broken P2P / IOMMU), so the
non-root GPU's ncclRecv never completes and the device wedges. Bisected by probing each
GPU with cudaDeviceSynchronize after every step of ModRaise: all pass until
broadcastLimb0(), after which device 1's sync never returns.

Fixed by routing the broadcast through the same host-staged cudaMemcpyPeer path used for
the gather when MEMCPY_PEER is active (root pushes limb-0 to each peer's top-limb buffer,
then an OpenMP barrier). With this, no bootstrap configuration hangs at any GPU count.

The debugging also corrected an earlier mis-statement in this document: the hang was not
a hostpin_polling_kernel spin. pollingKernel at LimbPartitionMGPU.cu:1420 is behind
if (0 && GRAPH_CAPTURE) (dead code). The "100% GPU util" reading was other tenants on
this shared box; per-process nvidia-smi pmon showed our process idle-waiting, consistent
with a stalled NCCL recv, not a spin.

Bootstrap N=4 corruption — fixed (a moddown data race)

What first looked like "precision degrades with GPU count" (N=4 mostly failing decode) was
not a numeric bug — it was a cross-GPU data race in RNSPoly::moddown. The
multi-GPU mod-down issues peer copies whose completion is not ordered against the streams the
next operation runs on, so at higher GPU counts the next op reads a partially-updated limb
and corrupts the result nondeterministically. More GPUs ⇒ more peer copies ⇒ the race almost
always fires at N=4, rarely at N=2.

Proven a race (not precision) by: (1) two N=4 runs with the same fixed PRNG seed, keys and
input produce different ciphertexts; (2) inserting cudaDeviceSynchronize makes them
bit-identical; (3) a sync after moddown specifically took N=4 from ~1/8 to 8/8 passing.

Fixed by syncing all GPUs before moddown returns (multi-GPU branch, RNSPoly.cpp). See
experiments/multigpu-bootstrap-precision/FINDING.md for the full investigation.

N=1 N=2 N=4
OpenFHEBootstrap/* (random keys, after fix) 16/16 16/16 16/16

With the fix, N=4 is deterministic and bit-exact to single-GPU. Bootstrap now passes at
every GPU count; multi-GPU is numerically correct.

(The barrier is correctness-first — a full device sync in mod-down. A lighter version would
wait only on the specific peer/digit streams mod-down touches; see FINDING.md.)

Note on the NCCL transport

FIDESLIB_USE_MEMCPY_PEER=0 selects the ncclBroadcast path instead. It segfaults
with graph capture on, and hangs with it off. Not investigated further.

Separately, modupMGPU uses a single static std::atomic_uint64_t skip shared across
all GPU threads while the graph cache map_c_to_map_graph_exec[8] is per-id. If two
GPUs ever disagree about whether a cached graph exists, they disagree about whether to
capture. Not the cause of this NaN, but it looks unsound.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions