Skip to content

feat(libtorch): run the bridge on the host when the SDK is CPU-only - #36

Open
NicolasRouquette wants to merge 3 commits into
lean-dojo:mainfrom
NicolasRouquette:libtorch-cpu-device
Open

NicolasRouquette wants to merge 3 commits into
lean-dojo:mainfrom
NicolasRouquette:libtorch-cpu-device

Conversation

@NicolasRouquette

@NicolasRouquette NicolasRouquette commented Oct 4, 2026 •

Copy link
Copy Markdown
Contributor

What

The LibTorch bridge pins c10::kCUDA, so the backend has one device and the default build has
none: unavailable.c fails every buffer operation. Before the move to LibTorch, the default build
linked the _stub.c host implementations of the same ABI, which is what kept the device buffer
suite runnable on hosted CI ("CPU hosted CI does not validate GPU execution" in
csrc/libtorch/README.md) and what let a CPU-only deployment run code written against
Runtime.Autograd.LibTorch.Buffer.

This PR lets the same bridge run on the host through ATen's CPU kernels. Nothing is reimplemented:
the device is chosen once, and every operation already goes through torchlean::options().

  • Selection. TORCHLEAN_LIBTORCH_DEVICE=cpu selects the host, =cuda insists on a CUDA
    device. Unset, a visible CUDA device is selected when the SDK has CUDA support; otherwise the
    host is selected only when the SDK has no CUDA support at all. A CUDA-enabled SDK without a
    visible device stays RuntimeStatus.nativeUnavailable, as today, rather than silently running
    on the host.
  • CMake. An SDK without the torch_cuda target is no longer an error; it compiles the backend
    with TORCHLEAN_LIBTORCH_CUDA=0, with no CUDA headers and no toolkit, and links neither NVRTC
    nor the CUDA driver (the toolkit is looked for, and both are linked, only for a CUDA-enabled
    SDK). The CPU-only pip wheel
    (pip install --index-url https://download.pytorch.org/whl/cpu torch==2.11.0, 746 MB, no
    toolkit) is a complete SDK for it.
  • Lean. Runtime.Autograd.LibTorch.DeviceKind (cuda | host) and deviceKind.
    RuntimeStatus keeps its three states; .nativeAvailable now means "a device is selected",
    CUDA or host, so requireNativeRuntime and every existing match are unchanged. Because the
    host satisfies it, a check that means "a GPU ran this" needs more: Buffer.requireCudaDevice
    requires the bridge and deviceKind = .cuda, the suite's TORCHLEAN_REQUIRE_CUDA=1 calls it
    and fails on the host, and TORCHLEAN_REQUIRE_NATIVE=1 is the check for a host session (the
    bridge linked and a device selected), so a CPU-only SDK job cannot pass on the unlinked default
    backend. The suite prints the device and checks that the selection agrees with the visible
    count and with the environment variable.
  • On the host: the native allocator counters and driver memory read zero; synchronize and
    emptyCache return at once; setMemoryFraction and setDevice are rejected through
    require_cuda(feature), the one call an entry point without a host implementation makes
    first, so it fails naming the feature rather than failing inside a CUDA call on a CPU tensor;
    the suite's native memory probes (accounting, saved attention buffers, OOM recovery) are
    skipped. Code that needs the CUDA headers compiles under #if TORCHLEAN_LIBTORCH_CUDA; code
    that needs the selected device calls require_cuda.
  • Runtime-compiled kernels (the generated operations of NN.Kernel, compiled through NVRTC)
    follow both rules: the NVRTC compile and launch path compiles under the macro, and run_kernel
    calls require_cuda at its entry. Configured binary arithmetic reaches its kernels through a
    pure entry point, where an error would end the process, so its encoded upload calls
    require_cuda as well. On the host a generated operation therefore fails through IO, naming
    the requirement. The CUDA suite runs its custom-computation and custom-autograd checks on a CUDA
    device; on the host it checks both refusals and runs the custom-autograd checks that need no
    generated kernel.

-K cuda=true keeps its meaning, "link the LibTorch backend"; the SDK decides the device. The
extern names (torchlean_cuda_*) and the Config.device := .cuda session name also keep their
meaning of "the bridge"; if you would rather have a libtorch-cpu device name for sessions, that
is a small follow-up and I am happy to do it, but I kept the user-facing vocabulary unchanged here.

Verification

On 5555080f (three commits on main 72754452, Lean v4.35.0):

run SDK device env result
curated suite, A4500 pip torch 2.11.0+cu128 cuda TORCHLEAN_REQUIRE_CUDA=1 all curated tests passed
same binary pip torch 2.11.0+cu128 host TORCHLEAN_LIBTORCH_DEVICE=cpu TORCHLEAN_REQUIRE_NATIVE=1 all curated tests passed; the three memory probes skipped
same binary, strict pip torch 2.11.0+cu128 host TORCHLEAN_LIBTORCH_DEVICE=cpu TORCHLEAN_REQUIRE_CUDA=1 fails: a CUDA device was required, but the LibTorch bridge runs on the host: the SDK is CPU-only, or TORCHLEAN_LIBTORCH_DEVICE=cpu selected the host
curated suite, no toolkit (TORCHLEAN_LIBTORCH_CUDA=0, ldd shows libtorch_cpu and no CUDA library) pip torch 2.11.0+cpu host TORCHLEAN_REQUIRE_NATIVE=1 all curated tests passed; the three memory probes skipped
same binary, strict pip torch 2.11.0+cpu host TORCHLEAN_REQUIRE_CUDA=1 fails, same message
default build (-Kcuda=false), CPU suite none — — 10200 jobs; all curated tests passed, CUDA kernels skipped as before
default build, native required none — TORCHLEAN_REQUIRE_NATIVE=1 fails: the LibTorch bridge was requested, but this executable was built without LibTorch; rebuild and run with `-K cuda=true`

The host run covers the same sections as the CUDA run, including the CUDA kernel coverage suite,
the float32 contract checks, determinism, attention, convolution and pooling, and the stress
tests; the differences between the two logs are the skipped memory probes and the generated-kernel
checks, which the host replaces by its refusal checks. Results are
bit-identical across the two devices only for programs made of the IEEE basic operations, gathers
and lookups (a downstream table-lookup retrieval was checked bit for bit, host against GPU);
ATen's transcendental functions come from different libraries on the host and on CUDA and differ
at the ulp level (a downstream Newton inverse differed by at most 1e-6 in float32), and reductions
differ in order. The suite's tolerances were met on both.

Repo lint: python3 scripts/checks/repo_lint.py --fail-on-warn OK.

Rationale

The curated device suite can run on a hosted runner with the CPU wheel and no toolkit, so the real
buffer ABI is exercised on every push instead of only on a GPU box, and TORCHLEAN_REQUIRE_NATIVE=1
makes that job fail rather than skip if the bridge is not linked. It also
gives downstream code written against the buffer ABI a CPU deployment target again, which the old
_stub.c files provided and unavailable.c does not.

This PR enables executing the same binary on GPU or on CPU.

Files

csrc/libtorch/torchlean.cpp, torchlean_libtorch.h, CMakeLists.txt, unavailable.c,
README.md; scripts/libtorch_build.py, scripts/README.md;
NN/Runtime/Autograd/Engine/LibTorch/{Controls,Buffer,Trusted}.lean, NN/Kernel/Runtime.lean;
NN/Tests/Suite.lean, NN/Tests/Runtime/Cuda/{Stress,Elementwise}.lean; lakefile.lean,
README.md, home_page/installation/index.md, the execution-modes chapter of the blueprint.

@Robertboy18

Copy link
Copy Markdown
Member

Hey Nicolas, thanks! CPU-only LibTorch support is worth having. One issue before merging: NN/Tests/Suite.lean uses requireNativeRuntime for TORCHLEAN_REQUIRE_CUDA=1, but nativeAvailable now includes host execution. That lets a CUDA-required run pass on CPU. Could you make that strict check require an actual CUDA device, while retaining a separate native-runtime check for CPU sessions? Please also check the combination with #35, whose device-information tests currently assume CUDA. Our pending NVRTC/custom-kernel work will need conditional CUDA compilation and a clear unsupported error on the host; that is an integration requirement for our local changes, not a claim that this PR fails to build independently. Holding off until the strict CUDA check is fixed.

The LibTorch bridge pinned `c10::kCUDA`, so a build without a CUDA device had no backend at all:
`unavailable.c` fails every buffer operation, and the device buffer suite could only run on a
GPU host. The bridge now selects its device once, at its first call, and runs on the host through
ATen's CPU kernels when the SDK has no CUDA support or when `TORCHLEAN_LIBTORCH_DEVICE=cpu` asks
for it. Nothing is reimplemented: every ATen operation already goes through `torchlean::options()`.

Selection: `TORCHLEAN_LIBTORCH_DEVICE=cpu` selects the host and `=cuda` insists on a CUDA device.
Unset, a visible CUDA device is selected when the SDK has CUDA support; otherwise the host is
selected only when the SDK has no CUDA support at all, so a CUDA-enabled SDK without a visible
device stays `RuntimeStatus.nativeUnavailable` rather than silently running on the host.

CMake no longer requires the `torch_cuda` target. A CPU-only SDK compiles the backend with
`TORCHLEAN_LIBTORCH_CUDA=0`, without CUDA headers, the CUDA toolkit, NVRTC or the driver API; the
CPU-only pip wheel is a complete SDK for it. `sdk.txt` records the choice.

Runtime-compiled kernels (the generated operations of `NN.Kernel`) need NVRTC and the driver API,
so they compile only with a CUDA-enabled SDK, and `run_kernel` refuses at its entry unless a CUDA
device is selected: on the host a generated operation fails through `IO` with an error that names
the requirement, before any CUDA call. The CUDA suite runs its custom-computation checks on a CUDA
device and checks that refusal on the host.

Lean: `Runtime.Autograd.LibTorch.DeviceKind` and `deviceKind` report the selection.
`RuntimeStatus` keeps its three states, with `.nativeAvailable` meaning that a device is
selected, CUDA or host, so `requireNativeRuntime` and every existing match are unchanged. On the
host the native allocator counters and driver memory read zero, `synchronize` and `emptyCache`
return at once, `setMemoryFraction` and `setDevice` are rejected, and the suite skips its native
memory probes. The suite prints the device and checks that the selection agrees with the visible
device count and with the environment variable.
…nd CUDA

Programs made of the IEEE basic operations, gathers and lookups give bit-identical results on
the two devices; transcendental functions come from different libraries on the host and on
CUDA and differ at the ulp level, and reductions differ in order.
`nativeAvailable` covers the host now that the bridge can run there, so `requireNativeRuntime`,
which the suite called for `TORCHLEAN_REQUIRE_CUDA=1`, let a run on a CPU-only SDK or under
`TORCHLEAN_LIBTORCH_DEVICE=cpu` pass a check meant to say that a GPU ran it.
`Buffer.requireCudaDevice` asks for the device kind as well, and the suite uses it for that
variable. `TORCHLEAN_REQUIRE_NATIVE=1` keeps the bridge-only meaning for CPU-only SDK sessions,
so that such a job fails on the unlinked default backend instead of passing with the device suite
skipped.

In the bridge, `require_cuda("<feature>")` is the entry check for an operation with no host
implementation (device selection, the allocator memory fraction and runtime-compiled kernels use
it), so it fails with a message naming the feature rather than with a CUDA call on a CPU tensor. The
READMEs, the installation page and the execution-modes chapter say what each variable rejects, and
quote the new messages.

Configured binary arithmetic computes only in runtime-compiled kernels, which the tape reaches
through pure entry points, where an error ends the process. The encoded upload therefore calls
`require_cuda` too, so on the host such a computation fails through `IO` when its tensor is
uploaded. The host branch of the CUDA suite checks both refusals and runs the custom-autograd checks
that need no generated kernel: the native-precision recordings and the resident-buffer checks.
@NicolasRouquette

Copy link
Copy Markdown
Contributor Author

Rebased onto current main (72754452, Lean 4.35); the tip is 5555080f, three commits: the
change, the bit-identity docs, and the review commit for your points.

The strict check. nativeAvailable covers the host on purpose, so the strict check could not
stay on it:

  • Buffer.requireCudaDevice (new) requires the bridge and deviceKind = .cuda. The suite's
    TORCHLEAN_REQUIRE_CUDA=1 now calls it, and the device-switching and memory probes the suite
    spawns inherit the strict meaning through the same variable. On the host it fails with
    a CUDA device was required, but the LibTorch bridge runs on the host: the SDK is CPU-only, or TORCHLEAN_LIBTORCH_DEVICE=cpu selected the host.
  • requireNativeRuntime keeps its meaning (bridge linked, a device selected, CUDA or host); the
    suite exposes it as TORCHLEAN_REQUIRE_NATIVE=1, the check for a CPU-only SDK session, so that
    such a job fails on the unlinked default backend (the LibTorch bridge was requested, but this executable was built without LibTorch; rebuild and run with -K cuda=true) instead of passing
    with the device suite skipped. The wording keeps "built without LibTorch", which
    NN/Tests/Backend/Profile.lean on main checks for. The host-device recipe in
    csrc/libtorch/README.md uses it, and scripts/README.md, the installation page and the
    execution-modes chapter say what each variable rejects.
  • The CLI's Config.device := .cuda still goes through requireNativeRuntime, "the bridge", as
    the description says; I left that unchanged.

Your NVRTC work. It landed on main while this was waiting, so the conditional compilation
you asked for is now part of this PR rather than a suggestion:

  • CMake looks for the CUDA toolkit, and links CUDA::nvrtc and CUDA::cuda_driver, only when the
    SDK exports torch_cuda; a CPU-only SDK builds without either.
  • In torchlean.cpp the NVRTC compile and launch path (check_driver through run_kernel) is
    under #if TORCHLEAN_LIBTORCH_CUDA, and run_kernel calls require_cuda("a runtime-compiled kernel") at its entry, before cudaFree(nullptr) could create a context or a host pointer could
    reach cuLaunchKernel. A CPU-only SDK gets a run_kernel that fails the same way. On the host
    the bridge raises LibTorch: a runtime-compiled kernel needs a CUDA device, but the bridge runs on the host (a CPU-only SDK, or TORCHLEAN_LIBTORCH_DEVICE=cpu), which reaches Lean as an IO
    error (rendered with c10's backtrace, and with its first letter lowercased by
    IO.Error.toString, as every bridge error is).
  • Configured binary arithmetic needed one more step. The tape reaches its kernels through a pure
    entry point (torchlean_kernel_run_encoded, via torchlean::invoke), where an error is
    lean_internal_panic; on the host that ended the process. The encoded upload
    (torchlean_cuda_buffer_of_encoded_io) now calls require_cuda("configured binary arithmetic"), so such a computation fails through IO when its tensor is uploaded, before any
    pure call.
  • NN/Tests/Runtime/Cuda/Elementwise.lean runs checkCustomComputations and
    checkCustomAutograd on a CUDA device. On the host it checks both refusals (a generated
    operation, and a recorded configured-format one) and runs the custom-autograd checks that need
    no generated kernel: the native-precision recordings and the resident-buffer checks.
    NN/Kernel/Runtime.lean's docstring says the host fails like a build without LibTorch.
  • The same two hooks are named in csrc/libtorch/README.md for whatever comes next: code that
    needs the CUDA headers compiles under #if TORCHLEAN_LIBTORCH_CUDA (CMake sets it from the
    SDK), and an entry point with no host implementation calls require_cuda("<feature>") before
    anything else. setDevice and setMemoryFraction use it too.

Combination with #35. #35 is stacked on this branch (its first three commits are this PR's);
details there. Merging this PR first leaves #35 with its own two commits.

What else the rebase changed. main renamed allocatorStats to memory and deleted
NN/Tests/API/BufferUpdates.lean; the device-agreement check this PR adds to NN/Tests/Suite.lean
no longer calls into it. Buffers now carry a dtype, so the bridge's checks read "a real buffer on
the selected device".

Verification. RTX A4500, pip torch 2.11.0+cu128 with a CUDA 12.8 toolkit, and pip torch
2.11.0+cpu with no toolkit; Lean v4.35.0, on 5555080f:

run SDK device env result
curated suite cu128 cuda TORCHLEAN_REQUIRE_CUDA=1 all curated tests passed
same binary cu128 host TORCHLEAN_LIBTORCH_DEVICE=cpu TORCHLEAN_REQUIRE_NATIVE=1 all curated tests passed
same binary, strict cu128 host TORCHLEAN_LIBTORCH_DEVICE=cpu TORCHLEAN_REQUIRE_CUDA=1 fails as quoted above
CPU-only wheel, no toolkit cpu host TORCHLEAN_REQUIRE_NATIVE=1 all curated tests passed
same binary, strict cpu host TORCHLEAN_REQUIRE_CUDA=1 fails, same message
default build none — — all curated tests passed, CUDA kernels skipped
default build none — TORCHLEAN_REQUIRE_NATIVE=1 fails as quoted above

Builds: -K cuda=true 8531 jobs against either SDK, the default build 10200 jobs, the same counts
as main. python3 scripts/checks/repo_lint.py --fail-on-warn: OK. The PR description is updated
to match.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants