[RFC]: DP pause/sleep correctness: coordinator latch, request admission, device-idle completion, and sleep/pause layering

[aoshen02]

Motivation.

Pause and sleep are the control plane RL post-training stands on: every weight sync is a
pause → (trainer work) → sleep → wake → resume cycle. This RFC tracks four separable problems
on that path. They are one topic: what each call claims versus what it actually guarantees.

#	Problem	Status
1	Coordinator's engines_running latches True; drain can only time out	Fixed — #51481 (merged)
2	Per-mode request admission during pause was specified (#32103) but never implemented	#51488 (draft)
3	Pause completion is CPU-only: the future resolves while the device is still busy — and in DP the pause path itself manufactures the GPU work it fails to wait for	#52914, #52957 (in review)
4	sleep embeds request-fate policy (mode, default abort) that belongs to pause, and the pause→sleep flow runs the quiescence consensus twice	#53082 (draft, the end-state implementation) — seeking consensus

---

1. The coordinator latch (fixed by #51481).

After a pause, a single request arriving from a front-end set the DP coordinator's
engines_running to True permanently: True was a prediction made when forwarding a wake,
while False was an observation requiring the engines to actually step — and paused engines
never step, so wave_complete was never produced and True had no path back.
wait_for_requests_to_drain could then only reach its 300 s timeout, and the stale True
published to front-ends made them skip the FIRST_REQ notification, so after resume nothing
woke the engines. All pause modes were affected.

The merged fix reports both edges from the engines (rank 0 reports the False→True transition,
guarded by not pending_pause so the pause kick-start is not misreported) and stops predicting
in the coordinator. No new fields or message types; the decision stays in the single
coordinator process, so ranks cannot diverge (the failure mode that sank #38009).

2. Per-mode admission: specified, never implemented (#51488).

#32103 specified "Frontend blocks new requests" for all modes; the implementation instead
queues late arrivals into the scheduler's waiting queue regardless of mode — mode only
describes requests already in flight. One rule generates the intended behavior:

mode says who keeps custody of a request during the pause.
keep — the engine carries requests across, so late arrivals are held too.
abort / wait — the pause is a generation boundary, so late arrivals are rejected.

Today abort says "custody is yours" and then silently keeps late arrivals: a request accepted
under old weights executes under new ones; reschedule-style partial rollout gets a second
rollout for the same sample; and none of it is observable, because the request is excluded
from get_num_unfinished_requests() under PAUSED_NEW and a paused engine publishes no stats.
#51488 implements the custody rule at the front-end (EnginePausedError, admission closed
before the pause is requested, reopened if the pause fails; keep still accepts).

A front-end guard cannot be airtight — each front-end's view of "paused" is a lagging replica —
which is why (1) had to land first: with it, the residue is untidy rather than fatal.

3. Pause completion is not a device barrier — and in DP, pause creates GPU work (#52914, #52957).

Reported from production: pause + wait_for_requests_to_drain before sleep does not
guarantee an idle GPU; the flow survives on the incidental ~0.5 s between pause and sleep.
Verified against main:

• The pause future resolves off has_work() — a pure CPU predicate — and the DP pause
  consensus is a gloo CPU all-reduce. Nothing on the pause path touches a CUDA stream.
• In DP the situation is stronger than "no barrier": DPEngineCoreProc._pause_complete
  unconditionally kick-starts idle engines so every rank reaches the consensus all-reduce, and
  the fixed every-32-step sync cadence makes each pause call manufacture 32 dummy forwards
  per rank. The consensus iteration launches its final dummy batch before the all-reduce
  and nothing ever waits on it (not in batch_queue, invisible to has_work()); the future
  resolves one iteration later while that forward can still be executing.
• Drain cannot compensate: the coordinator only tracks request-driven waves, so during the
  burst it publishes "not running" and drain returns immediately — a structural false
  positive, not a race.
• sleep() itself is safe only because GPUWorker.sleep synchronizes the device before
  unmapping. Bare pause is sleep minus the synchronizing half; callers that run their own
  kernels or NCCL between pause and sleep race the in-flight dummy-batch collectives.

Measured (DP=2 + EP, PowerMoE-3b, dummy batch inflated to ~110 ms to model a large-MoE step):
pause_generation() on an idle engine takes 3.4 s and returns with the device still
busy; drain returns in 0.00 s during the burst while pause_done=False.

Fixes, split into two independent PRs (repro scripts in the #52914 description):

• [Bugfix][DP] Synchronize the device on pause completion #52914 — resolve the pause future only after collective_rpc("synchronize_device") on
  every worker, in all three pause completion paths (in-proc, fast path, deferred idle
  callback). Makes pause completion honest regardless of in-flight work.
• [Core] Sync DP state on the first step of a wave #52957 — sync DP state on the first step of a wave as well as every 32nd, so an
  idle-engine pause reaches consensus after one dummy forward instead of 32 — measured
  3.4 s → 0.17 s. Also caps the cost of the re-pause sleep() runs internally (problem 4).

4. Sleep's semantics are entangled with pause's.

sleep was born pure: #12987 introduced it as a memory-state transition only
(executor.sleep(level) — offload/discard, wake_up restores). #33195 routed level 0 through
pause_scheduler(), and #34528 completed today's shape:

sleep(level, mode) = pause_scheduler(mode, clear_cache=level>=1)  # request fate + quiescence
                     + executor.sleep(level)                      # memory release

The memory layer is cleanly unaware of requests (executor.sleep has no mode). The
entanglement is at the engine API: sleep answers "what happens to the requests" — a question
that belongs to pause — through a duplicated mode vocabulary with a policy default
(abort), and sleep(level=0) is simply a pause under another name.

Concrete costs of the entanglement, observed in review and production flows:

• The standard RL flow pause → sync → sleep runs the DP quiescence consensus twice; the
  second run kick-starts all ranks again for up to 32 dummy forwards (~3.4 s measured, reduced
  to ~0.2 s by [Core] Sync DP state on the first step of a wave #52957 — mitigated, not removed).
• The internal re-pause cannot be guarded with "already paused → skip": the guard would read a
  local snapshot of a global in-flight property (pending_pause / engines_running /
  ignore_start_dp_wave all flip inside the consensus iteration, while the sleep RPC arrives
  at each rank at an unsynchronized time), and ranks landing on opposite sides of the flip
  deadlock the gloo rendezvous — one rank enters the all-reduce, the other never does. The
  unconditional kick-start is what makes the current protocol safe.
• sleep(mode="keep", level>=1) has compound semantics that surprise callers: requests survive
  but their KV is preempted and recomputed from scratch after wake — "keep the request, drop
  the compute" — expressible, but not discoverable from a memory API.

Proposed Change.

Problems 1–3 have concrete fixes (merged / in review, linked above). The new proposal is for
problem 4: restore the layering so that pause owns request fate and quiescence, and sleep
owns memory state — each call guaranteeing exactly what its name claims.

Target end state:

• pause(mode) — request fate (custody rule from problem 2) + group quiescence + device
  barrier ([Bugfix][DP] Synchronize the device on pause completion #52914). Its completion means: no request needs the GPU, and the GPU is idle.
• sleep(level) — memory release only. Precondition: the engine is paused; raises
  (EngineNotPausedError) otherwise. No mode parameter.
• wake_up(tags) — memory restore only; full wake resumes the scheduler as today, partial
  wake keeps it paused.

Why a raise is distributively safe where a skip is not: a failed precondition check never
enters a collective — every rank either proceeds into the same protocol or errors out to the
client, so local-snapshot divergence produces a retryable error instead of a hung all-reduce.
The fan-out does need to be all-or-nothing (check on all ranks before any rank releases
memory) to avoid partial sleep; that check-then-commit shape is an open design point below.

Migration (non-breaking, phased):

1. Now: document the layering (sleep = pause + release; level 0 is an alias of pause),
   land [Bugfix][DP] Synchronize the device on pause completion #52914/[Core] Sync DP state on the first step of a wave #52957/[Core] Reject new requests while generation is paused #51488.
2. Deprecation period: sleep called on an un-paused engine keeps today's behavior
   (internal pause, mode honored) but emits a DeprecationWarning steering callers to
   explicit pause(mode) → sleep(level); mode on sleep is marked deprecated.
3. After the period: un-paused sleep raises; mode is removed; sleep(level=0) is
   removed or kept as a documented alias of pause.

Open questions.

1. Direct-sleep users (e.g. SPMD colocate frameworks call llm.sleep() without a prior
   pause): is a deprecation cycle across two releases enough runway?
2. Should the sleep fan-out be made check-then-commit (verify all ranks paused, then release),
   or is raise-and-retry on partial failure acceptable?
3. sleep(level=0): deprecate in favor of pause, or keep as an alias?
4. Does anyone depend on wait queuing new requests during a pause (problem 2's custody rule
   aligns it with abort)?
5. For keep, what shape should the client-visible "paused, not finished" marker take — a
   RequestOutput field, or a protocol-level event in the OpenAI-compatible layer?

Feedback Period.

One week.

CC List.

@njhill @markmc (DP coordinator) @hao-aaron (#37024, #39366) @kouroshHakha (#32103)
@junjzhang (#36594)

Any Other Things.

Background findings that constrain the fixes, kept for the record:

• Wave coordination is MoE-only (enable_wave_coordination = model_config.is_moe); dense DP
  never reaches this code.
• The coordinator cannot infer a pause: pause is a client→engine utility RPC that bypasses it,
  and a paused engine publishes no stats — observability is lowest exactly in the state drain
  cares about.
• The scheduler keeps two disagreeing ledgers for requests queued during a pause: visible in
  len(self.waiting) for stats, hidden from get_num_unfinished_requests() under
  PAUSED_NEW. Count-based drain predicates are therefore not viable.
• keep + clear_cache=True already preempts running requests and recomputes after resume;
  what is missing is the client-visible pause signal (open question 5).
• Fix history before this RFC was engine-side only ([bug] Fix deadlock with pause resume and collective_rpc #37024, [BUG] Two phase pause to prevent deadlock #39366); [bug] Fix remaining START_DP_WAVE pause race in _handle_client_request #38009 was rejected
  because per-engine checks make ranks diverge and break the all-reduce rendezvous.

AI assistance was used while investigating and drafting this RFC; all code references were
traced and verified by hand against main.