CUDA | 木叶吟

GPU Pause, Resume, and Migration: The Missing Primitive in Cluster Scheduling

Sun, 17 May 2026 15:00:00 +0800

GPU cluster scheduling would be much easier if a running GPU job behaved like an ordinary CPU process. Pause it. Move it. Resume it somewhere else. Reclaim the device when a higher-priority job arrives. Repair fragmentation without killing user work.

In practice, this is exactly where GPU scheduling gets stuck.

A CPU process can be checkpointed by saving its address space, file descriptors, and kernel-visible state. A GPU task has an extra half of its life outside the normal process abstraction: CUDA contexts, device allocations, streams, events, library handles, kernels in flight, and data resident in GPU memory. The operating system does not naturally know how to serialize that state. The scheduler can stop the host process, but that is not the same thing as having a correct, portable checkpoint of the GPU computation.

FlowGPU is about turning GPU checkpoint/restore into a system primitive. Before FlowGPU became a full system, I wrote cudaw as the first version of the codebase: a CUDA wrapper prototype for interposing on runtime calls, tracking GPU objects, translating application-visible addresses, and making pause/resume/migration possible above an unmodified CUDA application.

Why Schedulers Want This Primitive

Pause/resume and migration change what a scheduler can do.

Without GPU checkpoint/restore, preemption is often blunt. A scheduler can kill a job, ask the framework to checkpoint at a pre-defined training boundary, or wait until the user code cooperates. That is acceptable for some training loops, but it is poorly aligned with cluster events. A high-priority job may arrive now. A GPU may fail now. A fragmented placement may need repair now. Framework-level checkpoints are usually placed for application convenience, not scheduler control.

With a transparent GPU checkpoint, the scheduler gets stronger operations:

pause a job and release its GPU memory;
resume it later on the same GPU;
migrate it to another GPU or node;
checkpoint periodically for fault tolerance;
defragment the cluster by moving jobs away from awkward placements;
support elastic scaling and priority scheduling with less user code involvement.

This is the missing link between scheduling policy and GPU execution. A scheduler may know the right decision, but without a safe migration primitive, it cannot act on that decision cheaply.

The CUDA Wrapper View

The basic idea behind my cudaw prototype is to place a wrapper between the application and CUDA runtime. Instead of letting the application talk directly to libcudart, the wrapper intercepts CUDA calls such as allocation, memory copy, and kernel launch. From the scheduler’s perspective, this creates an execution log and a shadow view of GPU state.

This wrapper layer can record which device memory regions exist, what host-side pointers correspond to them, how data moves between CPU and GPU, and which kernels are launched with which arguments. It can also maintain virtual GPU addresses: the application sees stable logical addresses, while the wrapper maps them to real CUDA allocations underneath. That indirection is what makes restore and migration plausible, because the restored task may receive different physical GPU addresses on the target device.

In a simplified checkpoint flow, the wrapper reaches a safe point, synchronizes GPU work, copies live GPU memory into a checkpoint image, saves enough metadata to reconstruct CUDA state, and releases the device. Restore reverses the process: allocate memory on the target GPU, rebuild mappings, copy data back, replay necessary CUDA setup calls, and continue execution.

This early prototype captured the central intuition that later shaped FlowGPU. GPU migration is not magic; it is state reconstruction. The hard part is making the reconstructed world indistinguishable from the original one.

Where Wrapper-Only Designs Struggle

The wrapper idea is powerful, but the edge cases are brutal.

First, CUDA state is larger than cudaMalloc and cudaMemcpy. Real applications use streams, events, cuBLAS, cuDNN, NCCL, memory pools, unified memory, graph execution, and framework allocators. Many of these objects are opaque: CUDA exposes handles, not serializable internals. A checkpoint system must record and replay the operations that created or mutated them.

Second, address identity matters. A pointer value may be stored inside application data structures, kernel arguments, framework metadata, or library state. If restore gives the program a different GPU virtual address, the application can become subtly wrong even if the bytes were copied correctly.

Third, deep learning frameworks hide memory behavior. PyTorch and TensorFlow often reserve large GPU memory blocks and keep them for reuse. Much of that reserved memory may be inactive at a given moment. A naive checkpoint that saves everything allocated by the runtime can produce enormous checkpoint images, even when the useful live state is much smaller.

Fourth, distributed training is a synchronization problem. A consistent checkpoint of a multi-GPU job requires pausing all participating ranks safely. With NCCL communication, pausing one side of a blocking send/receive pair while the other side waits can deadlock the checkpoint protocol itself.

These are the problems FlowGPU is designed to handle systematically.

FlowGPU’s Core Move

FlowGPU’s key insight is that prior system-level GPU checkpoint/restore designs coupled C/R with API forwarding. In API forwarding, all GPU operations pass through a privileged central process. That makes interception and state separation easier, but it imposes runtime overhead, creates GPU address conflicts under sharing, and blocks some GPU features.

FlowGPU decouples checkpoint/restore from virtualization.

During normal execution, each task uses a per-task intercept library. GPU operations stay private to that task and go directly to the GPU, avoiding the IPC overhead of a central forwarding process. When checkpointing is needed, FlowGPU creates a ghost process. The ghost process temporarily takes over GPU state, while the original process becomes a conventional CPU process that can be checkpointed with CRIU. GPU state and CPU state are saved in parallel, then recombined during restore.

This design keeps the useful part of interception without forcing every GPU operation through a virtualization server during normal execution.

Making Checkpoints Small and Correct

FlowGPU adds several mechanisms that are especially important for deep learning workloads.

Active memory identification avoids saving the whole framework-reserved memory pool. FlowGPU inserts a memory stub at stable DL framework backend allocation/free interfaces, tracking the memory regions that are actually active. It can also wait briefly for active memory to reach a low point in the training iteration before checkpointing. This matters because active memory in training can fluctuate dramatically between the end of an iteration and the activation-heavy middle of forward/backward execution.

Virtual memory management preserves GPU address identity. FlowGPU intercepts GPU allocations and uses CUDA VMM APIs such as cuMemAddressReserve, cuMemCreate, and cuMemMap to reserve and remap the same virtual addresses on restore. That removes a major source of correctness bugs for pointer-rich GPU applications.

Record/replay handles opaque runtime objects. Since CUDA streams, events, contexts, and library handles cannot simply be read out as bytes, FlowGPU records operations that create or modify them and replays those operations during recovery.

The pause mechanism is refined for distributed tasks. FlowGPU coordinates pausing across ranks, but avoids a known NCCL deadlock pattern by resuming all instances after a timeout if a complete pause cannot be achieved. This is a small detail with a large consequence: checkpointing must not introduce a failure mode worse than the one it tries to solve.

For multi-GPU tasks, FlowGPU also performs fine-grained deduplication. Replicated model parameters may appear on multiple GPUs, but runtime memory blocks rarely match exactly. FlowGPU deduplicates fixed-size regions, reducing checkpoint image size for distributed jobs.

What This Means for Scheduling

Once GPU pause/resume becomes practical, several scheduling policies become more realistic.

Priority scheduling can preempt a low-priority GPU job without throwing away all its progress. Fairness scheduling can redistribute service over time with lower disruption. Fragmentation-aware schedulers can migrate jobs to rebuild contiguous placements for gang-scheduled workloads. Fault-tolerance systems can checkpoint at scheduler-controlled intervals instead of relying only on framework checkpoints. Elastic schedulers can shrink, expand, or relocate jobs with a clearer recovery path.

The primitive also changes the economics of GPU sharing. If a job can be paused and restored quickly, a cluster can take more aggressive actions under bursty demand. Online inference, training, and HPO workloads no longer need to live in completely isolated resource islands; the scheduler has a better way to move work when priorities change.

FlowGPU’s evaluation shows why the details matter. It reports no runtime overhead during normal single-GPU execution because tasks can access the GPU directly without API forwarding. For DL tasks, it reduces checkpoint pause time by 6.2x to 15x over POS and up to 10.4x over Singularity. Restore time drops by 12x to 18x over POS and up to 4.1x over Singularity. For migration, FlowGPU outperforms Singularity by up to 2.1x and PyTorch framework-level checkpointing by 1.7x to 4.5x.

Those numbers are not only checkpointing results. They are scheduling-enablement results. A slow checkpoint is a policy that the scheduler cannot afford to use often. A fast, transparent checkpoint becomes a real control knob.

The Takeaway

GPU scheduling is often discussed in terms of algorithms: fairness metrics, placement heuristics, bin packing, elastic allocation, and priority queues. But the scheduler is only as powerful as the execution primitives beneath it.

cudaw was my first working cut at the wrapper-level intuition: interpose on CUDA, virtualize what the application sees, and reconstruct GPU state when needed. FlowGPU pushes that intuition into a more complete system design: per-task interception for low overhead, ghost processes for state separation, active-memory tracking for small images, VMM for address correctness, and distributed pause logic for multi-GPU workloads.

The result is a cleaner boundary between policy and mechanism. The scheduler decides when a job should pause, resume, or move. The checkpoint/restore layer makes that decision safe enough to execute.

Paper: FlowGPU: Transparent and Efficient GPU Checkpointing and Restore
Early codebase: yzs981130/cudaw

GPU 任务的暂停、恢复与迁移：调度器一直缺的那块拼图

Sun, 17 May 2026 15:00:00 +0800

如果正在运行的 GPU 任务能像普通 CPU 进程一样行动，GPU 集群调度会容易很多。暂停它，移动它，在别处恢复它。当更高优先级任务到来时回收设备。通过迁移修复碎片化，而不杀掉用户工作。

现实中，GPU 调度恰恰卡在这里。

CPU 进程可以通过保存 address space、file descriptor 和 kernel-visible state 来 checkpoint。GPU 任务则有另一半生命在普通进程抽象之外：CUDA context、device allocation、stream、event、library handle、正在运行的 kernel，以及驻留在 GPU memory 中的数据。操作系统并不天然知道如何 serialize 这些状态。调度器可以停止 host process，但这不等于拥有一个正确、可迁移的 GPU computation checkpoint。

FlowGPU 的目标就是把 GPU checkpoint/restore 变成系统原语。在 FlowGPU 成为完整系统之前，我写过 cudaw 作为代码库的最早版本：一个 CUDA wrapper prototype，用来 interpose runtime call、跟踪 GPU object、翻译 application-visible address，并让 pause/resume/migration 在不修改 CUDA application 的情况下变得可能。

为什么调度器需要这个原语

Pause/resume 和 migration 会改变调度器能做什么。

没有 GPU checkpoint/restore 时，抢占往往很粗暴。调度器可以杀掉任务，要求 framework 在预定义 training boundary checkpoint，或者等用户代码主动配合。这对某些 training loop 可以接受，但和集群事件并不匹配。高优先级任务可能现在就到达，GPU 可能现在就故障，碎片化 placement 也可能现在就需要修复。Framework-level checkpoint 通常服务于应用便利性，而不是调度器控制。

有了透明 GPU checkpoint，scheduler 就拥有了更强的操作：

暂停任务并释放 GPU memory；
稍后在同一张 GPU 上恢复；
迁移到另一张 GPU 或另一个节点；
为 fault tolerance 做周期性 checkpoint；
通过迁移任务修复集群碎片化；
用更少用户代码介入支持 elastic scaling 和 priority scheduling。

这是调度策略和 GPU execution 之间缺失的连接。调度器也许知道正确决策是什么，但如果没有安全 migration primitive，它就无法低成本执行这个决策。

CUDA Wrapper 视角

我的 cudaw prototype 的基本想法，是在应用和 CUDA runtime 之间放一个 wrapper。应用不再直接和 libcudart 交互，而是由 wrapper 拦截 allocation、memory copy、kernel launch 等 CUDA call。从调度器的角度看，这会形成一份 execution log，以及 GPU state 的 shadow view。

这个 wrapper layer 可以记录哪些 device memory region 存在、哪些 host-side pointer 与之对应、数据如何在 CPU 和 GPU 之间移动、哪些 kernel 带着哪些参数被 launch。它也可以维护 virtual GPU address：应用看到稳定的 logical address，而 wrapper 在底层把它们映射到真实 CUDA allocation。这层 indirection 让 restore 和 migration 变得可能，因为恢复后的任务在目标设备上可能拿到不同的 physical GPU address。

在一个简化的 checkpoint flow 中，wrapper 到达 safe point，同步 GPU work，把 live GPU memory 拷贝到 checkpoint image，保存足以重建 CUDA state 的 metadata，然后释放设备。Restore 则反向执行：在目标 GPU 上分配 memory，重建 mapping，把数据拷贝回来，replay 必要的 CUDA setup call，并继续执行。

这个早期 prototype 捕捉到的核心直觉，后来也影响了 FlowGPU。GPU migration 不是魔法，而是状态重建。难点在于，让重建后的世界和原来的世界无法区分。

纯 Wrapper 设计难在哪里

Wrapper 思路很有力量，但边界情况非常残酷。

第一，CUDA state 远不止 cudaMalloc 和 cudaMemcpy。真实应用会使用 stream、event、cuBLAS、cuDNN、NCCL、memory pool、unified memory、graph execution 和 framework allocator。许多对象是 opaque 的：CUDA 暴露的是 handle，而不是可序列化的内部状态。Checkpoint system 必须记录并 replay 创建或修改它们的操作。

第二，address identity 很重要。Pointer value 可能存储在应用数据结构、kernel argument、framework metadata 或 library state 里。如果 restore 后程序看到不同的 GPU virtual address，即使 bytes 被正确拷贝，应用也可能出现非常隐蔽的错误。

第三，深度学习框架会隐藏 memory behavior。PyTorch 和 TensorFlow 通常会预留大块 GPU memory 并复用它们。在某个时刻，其中很多 reserved memory 可能并不活跃。Naive checkpoint 如果保存 runtime 分配过的一切，就会产生巨大的 checkpoint image，即使真正有用的 live state 小得多。

第四，distributed training 是一个同步问题。多 GPU 任务的一致 checkpoint 需要安全暂停所有参与 rank。遇到 NCCL communication 时，如果暂停 blocking send/receive 的一侧，而另一侧还在等待，checkpoint protocol 本身就可能 deadlock。

这些正是 FlowGPU 试图系统化解决的问题。

FlowGPU 的核心动作

FlowGPU 的关键 insight 是，以前的 system-level GPU checkpoint/restore 设计经常把 C/R 和 API forwarding 绑在一起。在 API forwarding 中，所有 GPU operation 都经过一个 privileged central process。这让 interception 和 state separation 更容易，但会引入 runtime overhead，在 sharing 场景下产生 GPU address conflict，并阻碍部分 GPU feature。

FlowGPU 把 checkpoint/restore 从 virtualization 中解耦出来。

正常执行时，每个任务使用 per-task intercept library。GPU operation 仍然属于该任务，并直接访问 GPU，避免 central forwarding process 带来的 IPC overhead。需要 checkpoint 时，FlowGPU 创建一个 ghost process。Ghost process 临时接管 GPU state，而原 process 变成一个传统 CPU process，可以用 CRIU checkpoint。GPU state 和 CPU state 并行保存，再在 restore 时重新组合。

这个设计保留了 interception 有用的部分，同时避免正常执行期间所有 GPU operation 都绕过 virtualization server。

让 Checkpoint 更小也更可靠

FlowGPU 增加了几项对深度学习任务尤其重要的机制。

Active memory identification 避免保存整个 framework-reserved memory pool。FlowGPU 在稳定的 DL framework backend allocation/free interface 上插入 memory stub，跟踪真正活跃的 memory region。它也可以在 checkpoint 前短暂等待 active memory 降到 training iteration 中较低的位置。原因是 training 中 active memory 可能在 iteration 末尾和 forward/backward 中 activation 最重的阶段之间剧烈波动。

Virtual memory management 用于保留 GPU address identity。FlowGPU 拦截 GPU allocation，并使用 cuMemAddressReserve、cuMemCreate、cuMemMap 等 CUDA VMM API，在 restore 时保留并 remap 相同的 virtual address。这移除了 pointer-rich GPU application 中一类主要 correctness bug。

Record/replay 用来处理 opaque runtime object。CUDA stream、event、context 和 library handle 无法简单读成 bytes，因此 FlowGPU 记录创建或修改它们的操作，并在恢复期间 replay。

Pause mechanism 也针对 distributed task 做了细化。FlowGPU 会协调多个 rank 的暂停，但为了避免一种已知 NCCL deadlock pattern，如果完整 pause 无法达成，它会在 timeout 后恢复所有 instance。这个细节很小，后果很大：checkpointing 不能引入一个比原问题更糟的 failure mode。

对于 multi-GPU task，FlowGPU 还做了细粒度 deduplication。Replicated model parameter 可能出现在多张 GPU 上，但 runtime memory block 很少完全一致。FlowGPU 对固定大小 region 做去重，降低分布式任务的 checkpoint image size。

这对调度意味着什么

一旦 GPU pause/resume 变得实用，很多调度策略就更现实了。

Priority scheduling 可以抢占低优先级 GPU 任务，而不丢掉它的全部进度。Fairness scheduling 可以用更低扰动在时间上重新分配 service。Fragmentation-aware scheduler 可以迁移任务，重建 gang-scheduled workload 需要的连续 placement。Fault-tolerance system 可以按调度器控制的间隔 checkpoint，而不只依赖 framework checkpoint。Elastic scheduler 可以更清晰地 shrink、expand 或 relocate 任务。

这个原语也改变了 GPU sharing 的经济性。如果一个任务可以快速 pause 和 restore，集群就能在 bursty demand 下采取更激进的动作。Online inference、training 和 HPO workload 不必完全生活在彼此隔离的资源孤岛里；当优先级变化时，调度器有了更好的移动工作方式。

FlowGPU 的评估展示了细节为什么重要。论文报告称，因为任务可以不经过 API forwarding 直接访问 GPU，它在正常 single-GPU execution 中没有 runtime overhead。对于 DL task，相比 POS，FlowGPU 将 checkpoint pause time 降低 6.2x 到 15x；相比 Singularity，最多降低 10.4x。Restore time 相比 POS 降低 12x 到 18x，相比 Singularity 最多降低 4.1x。Migration 方面，FlowGPU 最多比 Singularity 快 2.1x，比 PyTorch framework-level checkpointing 快 1.7x 到 4.5x。

这些数字不只是 checkpointing 结果，也是调度能力被释放出来的结果。慢 checkpoint 是调度器不敢频繁使用的 policy；快而透明的 checkpoint 才会变成真正的 control knob。

小结

GPU 调度经常被讨论成算法问题：fairness metric、placement heuristic、bin packing、elastic allocation 和 priority queue。但调度器的能力上限取决于下面的 execution primitive。

cudaw 是我对 wrapper-level 直觉的第一版实现：interpose CUDA，virtualize application 看到的东西，并在需要时重建 GPU state。FlowGPU 把这个直觉推进成更完整的系统设计：per-task interception 保证低开销，ghost process 实现 state separation，active-memory tracking 缩小 image，VMM 保证 address correctness，distributed pause logic 支撑 multi-GPU workload。

最终结果是 policy 和 mechanism 之间更清晰的边界。调度器决定一个任务什么时候应该 pause、resume 或 move；checkpoint/restore layer 让这个决策足够安全，可以真正执行。

Paper: FlowGPU: Transparent and Efficient GPU Checkpointing and Restore
Early codebase: yzs981130/cudaw