feat(nvsnap): multi-GPU checkpoint/restore on criu-v2 - #1184
Draft
balajinvda wants to merge 12 commits into
Draft
Conversation
A checkpoint is currently pinned to the physical GPU it was captured on: restore recreates its device state on the same device or not at all. That makes a restore un-schedulable anywhere except the slot it came from, which is the wrong constraint for a platform whose whole argument is that a restore can land wherever there is capacity. The driver has supported remapping since r580 through gpuPairs on the restore args, taking a source device UUID and a target device UUID per GPU. Neither our CLI nor upstream's exposed it. This adds --gpu-map, applied on restore, resume, and the restore half of toggle, and ignored with a diagnostic on lock and checkpoint so a map given on the capture half is not silently dropped. Each side of a pair is a device index or a UUID, accepted with or without the GPU- prefix and dashes so operators can paste what nvidia-smi or the device plugin prints. Indices are resolved to UUIDs during parsing: an index only means something relative to one process's CUDA_VISIBLE_DEVICES on one node, and the entire point of migration is that the target enumeration differs. The map is also parsed before any state transition, so a typo cannot leave the target process locked or half-restored, and a map that does not cover every visible GPU is refused here rather than producing a considerably less specific failure from the driver. Requires CUDA 13 headers: 12.x declares CUcheckpointRestoreArgs as an opaque reserved[8] with no gpuPairs member, so the builder stage moves to 13.0.3-devel-ubuntu22.04. The ubuntu22.04 base is kept deliberately, because the binary runs against the criu bundle's glibc 2.35 and a 24.04 builder would raise the floor and break that contract. Tested against stubs, so no GPU is needed: 18 cases covering both UUID spellings, uppercase, index resolution, ordering, and the malformed and miscounted maps. Three mutations were checked rather than assumed -- not skipping dashes, dropping the count check, and transposing old and new in a pair -- and each turns cases red. The transposition matters most: it parses cleanly and would migrate onto the wrong device. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
…ogether The CRIU plugin runs cuda-checkpoint one pid per exec, so an N-rank job means N separate processes. On a single-GPU workload that only costs a driver attach per spawn. On a multi-GPU job the gaps between those processes are where collective traffic resumes: rank 0 is locked while rank 1 is still running, rank 1 posts a collective rank 0 will never service, and rank 1's own lock then waits on an operation that cannot complete. That is the shape of the "hangs on 2nd rank (lock timeout)" failure recorded the last time multi-GPU capture was attempted. --pid now repeats and accepts comma-separated lists, and the actions apply to every pid from the one process. A single --pid behaves exactly as before, including exit codes, so the plugin's existing calls are untouched. Adds --action save, which locks every rank before checkpointing any of them. The two-phase order is the hypothesis being tested, not an implementation detail. A failed lock rolls back the ranks already locked, because a rank left LOCKED is a hung workload and worse than a failed capture; a failed checkpoint stops immediately and leaves locks in place, since the ranks already checkpointed are recoverable only by restoring them, which the caller drives. This is deliberately a mechanism test rather than a daemon. If lock ordering is what multi-GPU capture is failing on, one invocation against a live TP=2 pod shows it. If it hangs identically, the ordering is not the problem and a resident daemon built on the same idea would not have helped either. 13 stub-driven cases covering the ordering, both rollback paths, single-pid equivalence, and pid parsing. The ordering assertion was mutation-checked: degrading save to interleaved lock+checkpoint turns three cases red. That mattered to verify, because the degraded form still succeeds on every single-GPU workload and would only resurface as a multi-GPU hang. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
The agent enumerates every process holding GPU state, uses the first one to pick a dump target, and then never looks at the list again. Nothing checks that those processes actually reached the images, so a capture can lose the ranks it exists for and still be published with a checkpoint ID and a zero exit code. Reproduced on dev1: quiescing NCCL on a live TP=2 vLLM took the executor down with it, both TP workers exited, and CRIU dumped the surviving API server correctly in 21 seconds. Four GPU processes were enumerated; one was captured. The result was a well-formed 1.5G checkpoint of a workload whose weights and KV cache were never in it -- about 3 percent of the expected size, and the size was the only symptom. dumpV2 now resolves each GPU pid to its in-namespace pid before dumping, and afterwards requires a core image for each. Missing any fails the capture and names the pids. The resolve has to happen first because a dump that is not leave-running kills the tree, so /proc is gone by the time the images exist. A pid that cannot be resolved is skipped rather than fatal -- it has usually just exited on its own, and failing a healthy capture because one short-lived helper raced us would be worse than the gap. A workload with no GPU processes is unaffected. This is the capture-side counterpart to the restore guards: those stop a cold start being published as a restore, this stops an incomplete capture being published as a checkpoint. Both exist because a green result that is quietly wrong costs more than a failure. Six tests, and the two load-bearing assertions were mutation-checked rather than assumed: accepting a partial capture when any process is present, and resolving the outer NSpid field instead of the innermost, each turn cases red. The second matters because CRIU names images after the in-namespace pid, so using the host pid would look for files that never exist and fail every capture instead. Note: the internal/agent suite has a pre-existing flake, seen once in about ten runs and not reproducible in eight further attempts. It is unrelated to this change, which adds no concurrency or timing, but it is worth chasing separately. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
Multi-GPU CRIU was refused outright in the agent, on the reasoning that
cuda-checkpoint blocks on peer state and the D2H path could never reconstruct
CUDA context state on restore, so multi-GPU had to use rootfs/cachedir. The
first half is true. The conclusion was too strong: with peer state absent,
criu-v2 captures and restores a tensor-parallel workload intact.
Measured on 8x H100 80GB, driver 580.126.16, TinyLlama at TP=2 and TP=4:
TP=2 checkpoint 2m34s 56G restore 1m00s PASS x3
TP=4 checkpoint 4m16s 110G restore 1m14s PASS
The captures are complete, not merely successful: every Worker_TP rank appears
in the dumped tree, cuda_plugin pauses every GPU pid, and there is one ~28.4G
image per rank, matching the per-GPU memory budget. The restored pods served
live inference. Restore scaled far better than the size did, 2x the data for
23 percent more time, which suggests the per-rank GPU restores overlap rather
than serialising.
NVSNAP_MULTI_GPU_CRIU=1 lifts the refusal and keeps the ordinary criu-v2
engine: in-namespace dump, cuda_plugin driving cuda-checkpoint per rank, no
interception library and no D2H. It inherits the CRIU-layer fixes that made
the injection stack unnecessary for single GPU rather than reviving it.
NVSNAP_LEGACY_MULTI_GPU_D2H=1 selects the old quiesce plus D2H path instead;
the two were previously one switch, which sent runs down the path nobody meant
to test. Neither is on by default.
Also fixes GPU device counting, which fed nvidia-smi a LD_LIBRARY_PATH with the
driver tree ahead of the container's. The driver tree ships its own libc, so
nvidia-smi aborted on a symbol mismatch, the query "failed", and every workload
was reported single-GPU. That failure was logged at warning and returned 1,
which is fail-open in the worst way: it both skipped multi-GPU handling and hid
that it had been skipped.
The workload is a new manifest rather than a change to vllm-tp2, which stays on
the rootfs path. Its restore placeholder is generated from the nvsnap.io/path
annotation, so the two paths no longer share a file.
What this does not do is remove the configuration constraint. Every cross-GPU
transport must be off before capture, and a bisect showed that set is close to
irreducible: --enforce-eager and NCCL_P2P_DISABLE each break it individually,
and one of NCCL_SHM_DISABLE / VLLM_ALLREDUCE_USE_SYMM_MEM is required as well.
NCCL reaches a peer several independent ways and vLLM adds its own, so closing
one door leaves the others open. --enforce-eager is the exception worth
tracking: it stands in for post-restore CUDA graph re-capture, which only the
engine can do, rather than for a mechanism nobody has.
docs/proposals/multi-gpu-criu-v2.md records the measurements, the bisect, and
what remains open.
Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
The multi-GPU result so far was TinyLlama, which proves the topology but not the
scale anyone cares about. This adds the production-shaped case and it works.
Llama-3.1-70B, TP=4, --gpu-memory-utilization 0.85
checkpoint 9m56s 290G
restore 2m02s
restored pod served a completion correctly
Complete, not merely successful: seven tasks dumped including all four
Worker_TP ranks, cuda_plugin paused six GPU pids, and four 76.5G images, one per
rank, matching the per-GPU budget.
Restore scales much better than size. Across TinyLlama TP=2 (56G, 60s),
TinyLlama TP=4 (110G, 74s) and 70B TP=4 (290G, 122s), five times the data costs
twice the time, and the 70B restore moved 290G at about 2.4 GB/s against roughly
0.9 GB/s measured on single GPU. The per-rank GPU restores are overlapping
rather than serialising. That contradicts the premise behind the deferred
per-pid context parallelisation work, which assumed they serialise, and is worth
re-examining before anyone invests there.
The workload is a new manifest rather than a change to vllm-70b, which stays on
the rootfs path, matching how vllm-tp2-criu was split from vllm-tp2.
Two harness constants blocked this and neither is a mechanism limit. The
checkpoint step defaults to a 600s timeout and the capture alone took 596s, so
CHECKPOINT_TIMEOUT has to be raised. More seriously, any workload matching *70b*
had POD_READY_TIMEOUT pinned to 1800s, and a cold run needs more than 32 minutes
just to pull and load ~140G of weights; the first attempt failed there and
reported "Pod ready FAIL", which reads like a capture problem and is not one.
Raised to 4200s with an override hook. A warm cache finishes far inside it.
Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
…s not transfer vLLM multi-GPU works on criu-v2 once every cross-GPU transport is off. The obvious next question is whether that recipe is about peer mappings in general or about vLLM specifically. It is about vLLM specifically. SGLang TP=2 (Llama-3.1-8B) with the equivalent set - --disable-cuda-graph in place of --enforce-eager, plus --disable-custom-all-reduce and the same three NCCL_*_DISABLE vars, which are engine-independent - hangs at capture. Not slowly. A criu and a cuda-checkpoint were still wedged three hours later, both blocked in anon_pipe_read: CRIU waiting on a cuda-checkpoint that never returns. That is worth distinguishing from a slow capture, because the harness reports both the same way. The first attempt returned an empty response after its 600s curl timeout, which looks like "too slow"; the retry then failed with "text file busy" on the cuda-checkpoint binary, which is the tell that the original process is still holding it. So the flag names transfer and the coverage does not. Something in SGLang still establishes cross-GPU mappings these flags leave open. Do not assume the recipe generalises to another engine without measuring. Adds the workload and its generated placeholder so the failure is reproducible rather than a note, and records in the proposal that a wedged capture leaves host processes behind which block later attempts on that node. NIM is not covered here. nim-qwen3-32b defines no command or args, so it runs the image entrypoint, and the criu-v2 generator requires the setsid stdio-redirect convention; it needs a command wrapper written before it can be tested at all. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
…g dead ends Two corrections and two eliminated hypotheses, all measured. The proposal claimed the sever flags work by leaving "no peer mappings at all". That is wrong. A vLLM rank that captures cleanly still holds mappings to the peer GPU's device node, 3 to its own and 2 to the peer's. Whatever the flags remove is not visible as device mappings, and the mechanism is not established. Stating it that confidently was not supported by anything measured. For the SGLang hang, two hypotheses were tested and both died. SGLang processes hold fds on all eight GPUs despite CUDA_VISIBLE_DEVICES naming two, which looked conclusive until vLLM turned out to have an identical per-device count. Peer device mappings looked like the next candidate until vLLM turned out to have those too. The only measured difference left is that SGLang holds about 40 /dev/nvidiactl mappings per rank against vLLM's 28, which is a lead rather than a cause. Flag coverage is now verified instead of assumed. --disable-cuda-graph and --disable-custom-all-reduce both exist and were accepted. SGLang's peer features are all opt-in (--enable-nccl-nvls, --enable-symm-mem, --enable-torch-symm-mem, --enable-p2p-check) and were never enabled, so the NCCL env vars were suppressing things that were already inactive. --disable-piecewise-cuda-graph and --disable-decode-cuda-graph exist and were not set; they are unlikely to affect a capture hang, since graphs fail late at restore-inference rather than at capture, but that is untested. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
…sign Every measurement in this proposal captures an engine that has just served a request and gone idle, which is what NVCF does: a pod is checkpointed when it has no traffic and later pods restore from that same artifact. Recorded because capture under live traffic behaves differently and someone would otherwise rediscover it from the data and file it as a defect. With eight concurrent requests against a TP=2 pod, rank 0 locks in about 3 seconds and rank 1 times out after 60. That is a deadlock, not slowness: locking rank 0 stops it servicing the collective rank 1 waits on, so locking the ranks together does not help either. It does not apply to the NVCF flow. The model also decides which numbers matter. Capture is amortised over every pod restoring from the artifact, so the 70B's ten minutes is paid once; restore is paid per pod, which makes the sub-linear restore scaling worth more than the capture cost. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
…confirm the hang NVSNAP_SECCOMP_ENABLED, USE_LIBUV and HF_HUB_ENABLE_HF_TRANSFER were all set to "0,1" instead of "0". They were collateral from deriving this manifest with a value-scoped substitution (replace value "0" with "0,1" for CUDA_VISIBLE_DEVICES), which matched every variable holding that value rather than the intended one. USE_LIBUV mattered: the single-GPU SGLang manifest sets it "0" deliberately, and libuv plus io_uring is the C/R hazard the patched-library stack exists for. A non-zero string may well read as enabled, so the run that produced the recorded SGLang hang could not be trusted. Repeated with the values corrected and a 30 minute timeout. Identical hang, and the same wedged-process signature: criu and cuda-checkpoint both blocked in anon_pipe_read. The confound was real and was not the cause, so the conclusion stands, now on clean inputs. Recorded in the proposal so the result is not re-litigated later. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
…pid bump works Qwen3-32B on TRT-LLM captures and restores through criu-v2: checkpoint 3m03s at 103G, restore 1m55s, restored pod served inference. Its process shape differs from vLLM's, with start_server.sh, an orted MPI daemon, four python3 ranks, and asymmetric GPU images of 83.5G and 24.7G rather than an even per-rank split. That makes multi-GPU criu-v2 work on two engines, and makes the SGLang failure specific to SGLang rather than a property of anything except vLLM. Three fixes were needed and only the first is NIM-shaped. The stock image has no command: it runs /opt/nvidia/nvidia_entrypoint.sh with cmd `bash -c $SERVER_START_SCRIPT_PATH`. The manifest reproduces that startup inside the setsid convention rather than replacing it, so the entrypoint still runs and still execs the script the image names. Stdio cannot go to /tmp. isRuntimeGeneratedPath treats /tmp as runtime-generated and drops it from the rootfs diff, so the placeholder restores an empty file and CRIU refuses with "File tmp/nim.out has bad size 0 (expect 21443)". It cannot go to the container root either, because the image runs as uid 1000. /opt/nim satisfies both and matches none of the excluded patterns. The placeholder must run as root, and this is the general one. It writes /proc/sys/kernel/ns_last_pid, privileged does not confer root, and an image defaulting to a non-root uid fails that write silently, leaves its pid range unreserved, and the restore dies with "Can't fork for 336: File exists" - the exact failure the pid reservation exists to prevent. The generator now emits runAsUser 0 for every placeholder. The restored workload's own uid comes from the checkpoint, so this does not change what it runs as. This was previously known only as a hand-written note on one manifest; it is a property of any non-root image and belonged in the generator. The bump's failure message was also parenthetical, which is why an unreserved range presented as a confusing restore error rather than as itself. It now says what will happen. Only the criu-v2 placeholders are regenerated here. The others would also pick up the same correction, but they carry hand-edits that #965 already fixes, and rewriting them now would collide with it. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
gazelle had not been re-run after checkpoint_v2_gpuguard_test.go was added, so the file was missing from the go_test srcs. The bazel target built and passed without ever compiling those six cases, and the check-gazelle CI step failed. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
The sever set carried NCCL_NVLS_DISABLE=1, which is not an NCCL variable. Grepping the shipped libnccl.so.2 finds NCCL_P2P_DISABLE, NCCL_SHM_DISABLE, NCCL_CUMEM_ENABLE and NCCL_NVLS_ENABLE, but no NCCL_NVLS_DISABLE. It had been a no-op for as long as it had been set, which is why the earlier bisect could never attribute anything to it. NVLS was off in practice anyway: with P2P and SHM disabled NCCL treats the ranks as separate nodes and builds no NVLS channels (confirmed by NCCL_DEBUG=INFO, "0 nvls channels"). So this corrects the spelling rather than the behaviour, but the set now says what it does. Re-ran vllm-tp2-criu end to end on the corrected flag: checkpoint 2m35s / 56G, restore 1m00s, post-restore inference OK. Matches the run this recipe was originally validated against. Co-Authored-By: Balaji Ganesan <bganesan@nvidia.com>
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
2 participants
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.
Draft. The mechanism works and is measured. The reason it is draft is the configuration compromise below, which a reviewer should push back on before this is more than opt-in.
Why
Multi-GPU workloads fall back to the cachedir path, which captures the container filesystem and no GPU state, so a restored pod reloads weights from scratch. This makes criu-v2 work for tensor-parallel workloads, with weights and KV cache in the checkpoint.
The compromise
cuda-checkpoint refuses to checkpoint a process holding cross-GPU state, and every high-performance TP transport creates exactly that. So capture only succeeds once all of them are off, and the workload has to be started that way. This is the cost of the feature, not a tuning detail:
vLLM
NIM / TRT-LLM
What that costs:
--enforce-eageris worth separating from the rest, because it is the one item with an exit. A captured CUDA graph holds GPU resources the checkpoint destroys, so replaying it after restore drives dead handles. It fails late and misleadingly: capture succeeds, restore succeeds,/v1/modelsanswers, and the first real inference hangs. Nothing outside the engine can fix it, since NCCL has no record of which graphs captured its kernels and cannot invalidate them. An engine that re-captured graphs after restore would remove this flag. The NCCL knobs need either a transparent sever or newer driver support.NIM is the interesting exception: TRT-LLM configures graphs at engine build time and the image exposes no knob, so it ran with graphs and both capture and post-restore inference still worked. That is observed, not explained, and I would not generalise from it.
A bisect showed the vLLM set is close to irreducible:
--enforce-eagerandNCCL_P2P_DISABLEeach break capture on their own, and one of SHM/SYMM_MEM is also required.--disable-custom-all-reducewas not isolated individually, so treat the set as known-good, not minimal.The NVLS entry was previously spelled
NCCL_NVLS_DISABLE=1, which is not an NCCL variable at all - the shippedlibnccl.so.2containsNCCL_NVLS_ENABLEbut noNCCL_NVLS_DISABLE. It had been a no-op the whole time, which is why the bisect could never attribute anything to it. NVLS is off regardless (NCCL_DEBUG=INFOreports0 nvls channels, since P2P and SHM disabled make NCCL treat the ranks as separate nodes), so this corrects the spelling rather than the behaviour. Re-validated end to end on the corrected flag: checkpoint 2m35s / 56G, restore 1m00s, inference OK - identical to the original run.What changed
Multi-GPU capture on criu-v2, opt-in via
NVSNAP_MULTI_GPU_CRIU=1. Ordinary criu-v2 engine throughout: in-namespace dump, cuda_plugin driving cuda-checkpoint per rank, no interception library and no D2H.NVSNAP_LEGACY_MULTI_GPU_D2H=1selects the old path. These were previously one switch, which silently routed runs down the path nobody meant to test. Neither is on by default.New workloads rather than edits to existing ones, so the two engines never share a manifest:
vllm-tp2-criu,vllm-70b-criu,nim-qwen3-32b-criu,sglang-tp2-criu.Two agent defects fixed on the way:
LD_LIBRARY_PATHwith the driver tree ahead of the container's. The driver tree ships its own libc, nvidia-smi aborted on the symbol mismatch, the query "failed", and every workload was counted as single-GPU. Fail-open, and it hid that it had failed./proc/sys/kernel/ns_last_pid, privileged does not confer root, and an image defaulting to a non-root uid fails that write silently, so the pid range is never reserved and restore dies withCan't fork for 336: File exists. This was known only as a hand-written note on one manifest; it applies to any non-root image and belongs in the generator.Plus
--gpu-map, exposing the r580 GPU-migrationgpuPairson restore so a checkpoint is not pinned to the GPU slot it came from. Needs CUDA 13 headers, so the builder stage moves to13.0.3-devel-ubuntu22.04, staying on ubuntu22.04 to match the criu bundle's glibc 2.35.Customer Release Notes
Not customer visible. Off by default, and the configuration compromise above has to be resolved before it is offered.
Plan Summary
Not applicable.
Usage
# agent: agent.extraEnv NVSNAP_MULTI_GPU_CRIU=1 CAPTURE_PATH=criu-v2 ./scripts/test-e2e.sh vllm-tp2-criu CAPTURE_PATH=criu-v2 CHECKPOINT_TIMEOUT=2400 ./scripts/test-e2e.sh vllm-70b-criuTesting
8x H100 80GB (p5.48xlarge), driver 580.126.16:
Passes were checked for completeness, not exit status: a partial capture still exits zero and still writes a well-formed checkpoint, just a smaller one. Verified every rank present in the dumped tree, cuda_plugin pausing every GPU pid, one GPU image per rank matching the per-GPU memory budget, and the restored pod serving live inference.
Restore scales better than size: 5x the data costs 2x the time, and the 70B restore moved 290G at roughly 2.4 GB/s against 0.9 GB/s on single GPU. Per-rank restores overlap rather than serialise, which contradicts the premise behind the deferred parallelisation work in #437.
Unit tests: 18 cases for
--gpu-mapparsing, 13 for multi-pid save ordering, 6 for the capture guard. Load-bearing assertions were mutation-checked, including a transposed GPU pair, which parses cleanly and would migrate onto the wrong device.Notes
SGLang hangs at capture with criu and cuda-checkpoint both wedged in a plain pipe read. Four hypotheses are dead: device fds and peer mappings match a passing vLLM rank,
--disable-overlap-scheduledoes not help, and its peer features are opt-in and were never on. The workload ships so the failure is reproducible rather than a note.This extends the custom cuda-checkpoint CLI that #730 proposes removing. Upstream's CLI does not expose GPU migration, so adopting it would drop
--gpu-mapunless NVIDIA adds it. Worth a decision, not a merge conflict discovered later.Only the four criu-v2 restore placeholders are regenerated. The rest need the same
runAsUsercorrection but carry hand-edits that #965 fixes, so regenerating now would collide.References
Relates to #1183
Relates to #437
Relates to #730
Related Pull Requests
#965 (touches the same generated placeholders; land first)
Dependencies
None. CUDA builder image moves 12.8.1 to 13.0.3, both
nvidia/cuda:*-devel-ubuntu22.04.