MOVZXandHinata
7f6055caa2
The pull-transport flag check is wired in (rank-gated, off by default, warns when enabled). Both diagnostic checks log that they ran. Test repairs for helpers v0.30.0 renamed. Co-Authored-By: Hinata <hinata@idihore.id> Signed-off-by: MOVZX <movzx@yahoo.com>
564 lines
23 KiB
Python
564 lines
23 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Base class for PLE layers that support cross-process CPU offload.
|
|
|
|
Synchronization between the offload process and the GPU main stream is
|
|
handled by :class:`CpuGpuSemaphore`, which uses CUDA's
|
|
``cuStreamWaitValue32`` / ``cuStreamWriteValue32`` stream memory operations.
|
|
|
|
Decode-step timeline, host-mapped PULL design
|
|
--------------------------------------------------------
|
|
Offload process (CPU, no CUDA) GPU main stream (forward)
|
|
------------------------------ ------------------------------
|
|
forward_impl() on CPU ... other GPU ops ...
|
|
wait host flag == 0 WaitValue32(host flag == 1) <- in graph
|
|
rows -> shared pinned host slot cuMemcpyHtoDAsync(slot -> GPU) <- in graph
|
|
publish fence; host flag = 1 return _gpu_output_buffer[:n]
|
|
... model consumes output ...
|
|
WriteValue32(host flag = 0) <- release_outputs
|
|
|
|
``release_outputs`` performs the sole reset. Safety requires serialized local
|
|
forwards on one ordered model stream: a new transaction cannot begin until the
|
|
prior release has been enqueued. If any prepared forward is abandoned or its
|
|
correspondence becomes uncertain, the worker must terminate; resetting,
|
|
retrying, or continuing could reuse stale rows or overwrite live input.
|
|
"""
|
|
|
|
import contextlib
|
|
import functools
|
|
import threading
|
|
from abc import ABC, abstractmethod
|
|
from collections.abc import Callable, Iterator
|
|
from typing import Any, cast
|
|
|
|
import torch
|
|
from cuda.bindings import driver as cuda_driver
|
|
from cuda.bindings.driver import CUstreamWaitValue_flags
|
|
from torch import nn
|
|
|
|
import vllm.envs as envs
|
|
from vllm.utils.torch_utils import direct_register_custom_op
|
|
|
|
# Module-level flag set to True inside the offload subprocess.
|
|
# Because the offload process and GPU worker processes are separate OS
|
|
# processes (spawned via multiprocessing), each has its own memory space.
|
|
# A plain module-level bool is sufficient -- no thread-local storage needed.
|
|
_offload_worker_flag = False
|
|
|
|
# Native FULL capture runs on one Python thread. The custom-op implementation
|
|
# appends the exact arguments from every wait-copy invocation while that thread
|
|
# is inside the graph manager's capture scope.
|
|
_ple_copy_capture = threading.local()
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def observe_ple_copy_geometry() -> Iterator[list[tuple[str, int, int]]]:
|
|
"""Collect ``(layer, rows, bytes)`` from actual wait-copy executions."""
|
|
if getattr(_ple_copy_capture, "records", None) is not None:
|
|
raise RuntimeError("nested PLE copy-geometry observation is unsupported")
|
|
records: list[tuple[str, int, int]] = []
|
|
_ple_copy_capture.records = records
|
|
try:
|
|
yield records
|
|
finally:
|
|
del _ple_copy_capture.records
|
|
|
|
|
|
def _record_ple_copy_geometry(
|
|
layer_name: str,
|
|
num_tokens: int,
|
|
row_bytes: int,
|
|
) -> None:
|
|
records = getattr(_ple_copy_capture, "records", None)
|
|
if records is not None:
|
|
rows = int(num_tokens)
|
|
records.append((str(layer_name), rows, rows * int(row_bytes)))
|
|
|
|
|
|
def is_offload_process() -> bool:
|
|
"""Return True inside the dedicated PLE CPU-offload subprocess."""
|
|
return _offload_worker_flag
|
|
|
|
|
|
def mark_as_offload_worker() -> None:
|
|
"""Mark the current process as the dedicated PLE offload worker."""
|
|
global _offload_worker_flag
|
|
_offload_worker_flag = True
|
|
|
|
|
|
def _cuda_check(result: Any, operation: str) -> Any:
|
|
"""Check the ``(CUresult, ...)`` tuple returned by cuda-python calls."""
|
|
error = result[0] if isinstance(result, tuple) else result
|
|
if error.value != 0:
|
|
raise RuntimeError(f"{operation} failed: {error}")
|
|
return result
|
|
|
|
|
|
class CpuGpuSemaphore:
|
|
"""Cross-process semaphore backed by a one-element int32 CUDA tensor.
|
|
|
|
The flag is stored in a regular GPU tensor that is shared through
|
|
PyTorch's CUDA IPC mechanism. Both processes obtain device pointers that
|
|
map to the same physical GPU memory, so a stream-memory wait in the GPU
|
|
worker observes the write issued by the offload process.
|
|
"""
|
|
|
|
RESET_VALUE = 0
|
|
DONE_VALUE = 1
|
|
|
|
def __init__(self, device: torch.device) -> None:
|
|
self._flag_tensor = torch.zeros(1, dtype=torch.int32, device=device)
|
|
self._flag_ptr: int | None = None
|
|
|
|
@classmethod
|
|
def from_ipc_tensor(cls, flag_tensor: torch.Tensor) -> "CpuGpuSemaphore":
|
|
"""Construct a semaphore from a CUDA tensor received through IPC."""
|
|
semaphore = cls.__new__(cls)
|
|
semaphore._flag_tensor = flag_tensor
|
|
semaphore._flag_ptr = None
|
|
return semaphore
|
|
|
|
@classmethod
|
|
def from_host_mapped(
|
|
cls, host_flag_tensor: torch.Tensor, device_ptr: int
|
|
) -> "CpuGpuSemaphore":
|
|
"""Semaphore on a page-locked, device-mapped shared host int32.
|
|
|
|
``host_flag_tensor`` keeps the shared storage alive; ``device_ptr`` is
|
|
its device address (``cuMemHostGetDevicePointer``). Stream memory ops
|
|
on it are executed by the GPU against host memory, so a CPU process
|
|
with no CUDA context can signal it with a plain store.
|
|
"""
|
|
semaphore = cls.__new__(cls)
|
|
semaphore._flag_tensor = host_flag_tensor
|
|
semaphore._flag_ptr = int(device_ptr)
|
|
return semaphore
|
|
|
|
@property
|
|
def flag_tensor(self) -> torch.Tensor:
|
|
"""Return the tensor that backs the semaphore flag."""
|
|
return self._flag_tensor
|
|
|
|
@property
|
|
def flag_ptr(self) -> int:
|
|
"""Device address used by the stream memory operations."""
|
|
if self._flag_ptr is not None:
|
|
return self._flag_ptr
|
|
return self._flag_tensor.data_ptr()
|
|
|
|
def reset(self, stream: torch.cuda.Stream | None = None) -> None:
|
|
"""Enqueue ``WriteValue32(flag=0)`` on ``stream``."""
|
|
if stream is None:
|
|
stream = torch.cuda.current_stream()
|
|
_cuda_check(
|
|
cuda_driver.cuStreamWriteValue32(
|
|
cuda_driver.CUstream(stream.cuda_stream),
|
|
cuda_driver.CUdeviceptr(self.flag_ptr),
|
|
self.RESET_VALUE,
|
|
0,
|
|
),
|
|
"CpuGpuSemaphore.reset",
|
|
)
|
|
|
|
def signal(self, stream: torch.cuda.Stream | None = None) -> None:
|
|
"""Enqueue ``WriteValue32(flag=1)`` on ``stream``."""
|
|
if stream is None:
|
|
stream = torch.cuda.current_stream()
|
|
_cuda_check(
|
|
cuda_driver.cuStreamWriteValue32(
|
|
cuda_driver.CUstream(stream.cuda_stream),
|
|
cuda_driver.CUdeviceptr(self.flag_ptr),
|
|
self.DONE_VALUE,
|
|
0,
|
|
),
|
|
"CpuGpuSemaphore.signal",
|
|
)
|
|
|
|
def wait_reset(self, stream: torch.cuda.Stream | None = None) -> None:
|
|
"""Enqueue ``WaitValue32(flag==0)`` on ``stream``."""
|
|
if stream is None:
|
|
stream = torch.cuda.current_stream()
|
|
_cuda_check(
|
|
cuda_driver.cuStreamWaitValue32(
|
|
cuda_driver.CUstream(stream.cuda_stream),
|
|
cuda_driver.CUdeviceptr(self.flag_ptr),
|
|
self.RESET_VALUE,
|
|
CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value,
|
|
),
|
|
"CpuGpuSemaphore.wait_reset",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Custom op: vllm::ple_offload_wait
|
|
# ---------------------------------------------------------------------------
|
|
# Wraps cuStreamWaitValue32 as a torch custom op so that
|
|
# torch.compile(fullgraph=True) treats the CUDA Driver call as an opaque node.
|
|
# The hidden_states argument creates a dynamo data-dependency edge that prevents
|
|
# the wait from being reordered before preceding GPU work.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _ple_offload_wait_impl(
|
|
sem_flag_tensor: torch.Tensor,
|
|
gpu_output_buffer: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
) -> None:
|
|
"""Wait for the CPU result without releasing its output buffer."""
|
|
stream = torch.cuda.current_stream()
|
|
cuda_stream = cuda_driver.CUstream(stream.cuda_stream)
|
|
dev_ptr = cuda_driver.CUdeviceptr(sem_flag_tensor.data_ptr())
|
|
_cuda_check(
|
|
cuda_driver.cuStreamWaitValue32(
|
|
cuda_stream,
|
|
dev_ptr,
|
|
CpuGpuSemaphore.DONE_VALUE,
|
|
CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value,
|
|
),
|
|
"cuStreamWaitValue32(done)",
|
|
)
|
|
|
|
|
|
def _ple_offload_wait_fake(
|
|
sem_flag_tensor: torch.Tensor,
|
|
gpu_output_buffer: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
) -> None:
|
|
"""Represent the side-effect-only wait during dynamo tracing."""
|
|
pass
|
|
|
|
|
|
direct_register_custom_op(
|
|
op_name="ple_offload_wait",
|
|
op_func=_ple_offload_wait_impl,
|
|
mutates_args=["gpu_output_buffer"],
|
|
fake_impl=_ple_offload_wait_fake,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Custom op: vllm::ple_offload_wait_copy (host-mapped pull design)
|
|
# ---------------------------------------------------------------------------
|
|
# The offload process writes the result into a page-locked shared host buffer
|
|
# and stores flag=1 in a page-locked shared host int32; it holds no CUDA
|
|
# context. On the CURRENT stream (captured into the graph like the wait above):
|
|
# cuStreamWaitValue32(host flag == 1) -> GPU polls host memory
|
|
# cuMemcpyAsync(gpu buffer <- host buffer, num_tokens * row_bytes)
|
|
# The copy engine reads host memory coherently after the flag was observed,
|
|
# and the consumers read device memory as before. The flag is NOT reset here:
|
|
# ``release_outputs`` resets it after the forward (unchanged semantics), so
|
|
# capture warm-ups and dummy runs behave exactly as with ``ple_offload_wait``.
|
|
# No CPU tensor and no raw address enters the compiled graph: the op resolves
|
|
# the live pointers by layer name (see _PULL_LAYERS).
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# Fork (found on the reference host): the host flag / host output pointers and
|
|
# the row size are NOT op arguments any more. As Python ints they were
|
|
# specialised into the compiled graph as constants, so a boot that LOADED the
|
|
# AOT compile cache (no Dynamo pass) waited on the previous process's addresses:
|
|
# cuStreamWaitValue32 -> CUDA_ERROR_INVALID_VALUE on the first forward. The op
|
|
# now carries only process-stable values (buffer, token count, layer name) and
|
|
# resolves the live pointers from this registry on every call.
|
|
_PULL_LAYERS: dict[str, "PleOffloadLayer"] = {}
|
|
|
|
|
|
def _ple_offload_wait_copy_impl(
|
|
gpu_output_buffer: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
num_tokens: int,
|
|
layer_name: str,
|
|
) -> None:
|
|
"""Resolve this process's live pointers for ``layer_name``, then wait+DMA."""
|
|
layer = _PULL_LAYERS.get(layer_name)
|
|
if layer is None or layer._host_output_ptr is None:
|
|
raise RuntimeError(
|
|
f"PLE wait_copy: no armed pull layer named {layer_name!r} in this "
|
|
f"process (armed: {sorted(_PULL_LAYERS)})"
|
|
)
|
|
_ple_offload_wait_copy_raw(
|
|
gpu_output_buffer,
|
|
hidden_states,
|
|
layer._sem.flag_ptr,
|
|
layer._host_output_ptr,
|
|
layer._host_row_bytes,
|
|
num_tokens,
|
|
layer_name,
|
|
)
|
|
|
|
|
|
def _ple_offload_wait_copy_raw(
|
|
gpu_output_buffer: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
flag_ptr: int,
|
|
host_ptr: int,
|
|
row_bytes: int,
|
|
num_tokens: int,
|
|
layer_name: str,
|
|
) -> None:
|
|
"""Wait for the host-side result, then DMA it into the GPU buffer."""
|
|
if not 0 <= int(num_tokens) <= gpu_output_buffer.shape[0] or int(row_bytes) != (
|
|
gpu_output_buffer.shape[1] * gpu_output_buffer.element_size()
|
|
):
|
|
raise RuntimeError(
|
|
f"PLE wait_copy out of bounds: {num_tokens} rows x {row_bytes} B vs "
|
|
f"GPU buffer {tuple(gpu_output_buffer.shape)} {gpu_output_buffer.dtype}"
|
|
)
|
|
nbytes = int(num_tokens) * int(row_bytes)
|
|
_record_ple_copy_geometry(layer_name, num_tokens, row_bytes)
|
|
stream = cuda_driver.CUstream(torch.cuda.current_stream().cuda_stream)
|
|
_cuda_check(
|
|
cuda_driver.cuStreamWaitValue32(
|
|
stream,
|
|
cuda_driver.CUdeviceptr(flag_ptr),
|
|
CpuGpuSemaphore.DONE_VALUE,
|
|
CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ.value,
|
|
),
|
|
"cuStreamWaitValue32(host flag done)",
|
|
)
|
|
if nbytes > 0:
|
|
# host_ptr is this worker's host address of page-locked (registered)
|
|
# memory, so the copy is a true async DMA and is capturable.
|
|
_cuda_check(
|
|
cuda_driver.cuMemcpyHtoDAsync(
|
|
cuda_driver.CUdeviceptr(gpu_output_buffer.data_ptr()),
|
|
host_ptr,
|
|
nbytes,
|
|
stream,
|
|
),
|
|
"cuMemcpyHtoDAsync(PLE host -> GPU)",
|
|
)
|
|
|
|
|
|
def _ple_offload_wait_copy_fake(
|
|
gpu_output_buffer: torch.Tensor,
|
|
hidden_states: torch.Tensor,
|
|
num_tokens: int,
|
|
layer_name: str,
|
|
) -> None:
|
|
"""Represent the side-effect-only wait + copy during dynamo tracing."""
|
|
pass
|
|
|
|
|
|
direct_register_custom_op(
|
|
op_name="ple_offload_wait_copy",
|
|
op_func=_ple_offload_wait_copy_impl,
|
|
mutates_args=["gpu_output_buffer"],
|
|
fake_impl=_ple_offload_wait_copy_fake,
|
|
)
|
|
|
|
|
|
class PleOffloadLayer(nn.Module, ABC):
|
|
"""Base class for embedding-like PLE layers that can run on CPU.
|
|
|
|
Subclasses implement :meth:`forward_impl`, which contains the regular
|
|
on-device computation. In a GPU worker with PLE offload enabled, subclass
|
|
constructors are skipped so their large weights are never allocated. The
|
|
CPU process constructs a complete copy and owns those weights instead.
|
|
"""
|
|
|
|
_is_cpu_offloaded = False
|
|
_gpu_output_buffer: torch.Tensor
|
|
_sem: CpuGpuSemaphore
|
|
# Host-mapped pull design: host address of the shared output buffer and
|
|
# the bytes per output row. ``None`` keeps the IPC push behaviour.
|
|
_host_output_ptr: int | None = None
|
|
_host_row_bytes: int = 0
|
|
_ple_offload_name: str = "<unbound>"
|
|
|
|
def __init_subclass__(cls, **kwargs: object) -> None:
|
|
"""Skip subclass initialization in cross-process GPU-worker mode."""
|
|
super().__init_subclass__(**kwargs)
|
|
original_init_obj = cls.__dict__.get("__init__")
|
|
if original_init_obj is None or not callable(original_init_obj):
|
|
return
|
|
original_init = cast(Callable[..., None], original_init_obj)
|
|
|
|
@functools.wraps(original_init)
|
|
def guarded_init(
|
|
self: "PleOffloadLayer", *args: object, **kwargs: object
|
|
) -> None:
|
|
if envs.VLLM_E2_PLE_PULL_TRANSPORT and not is_offload_process():
|
|
nn.Module.__init__(self)
|
|
return
|
|
original_init(self, *args, **kwargs)
|
|
|
|
cls.__init__ = guarded_init # type: ignore[method-assign, assignment]
|
|
|
|
@classmethod
|
|
def get_target_device(cls) -> torch.device:
|
|
"""Return CPU for the offload process and the active GPU otherwise."""
|
|
if envs.VLLM_E2_PLE_PULL_TRANSPORT:
|
|
return torch.device("cpu")
|
|
return torch.device("cuda", torch.accelerator.current_device_index())
|
|
|
|
@abstractmethod
|
|
def forward_impl(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
input_ids: torch.Tensor,
|
|
*args: object,
|
|
**kwargs: object,
|
|
) -> torch.Tensor:
|
|
"""Execute the actual embedding computation on the owning device."""
|
|
raise NotImplementedError
|
|
|
|
def get_offload_output_dtype(self, default_dtype: torch.dtype) -> torch.dtype:
|
|
"""Return the dtype used by cross-process output buffers."""
|
|
return default_dtype
|
|
|
|
def get_offload_output_dim(self, default_dim: int) -> int:
|
|
"""Return the last dimension used by cross-process output buffers."""
|
|
return default_dim
|
|
|
|
def initialize_dummy_offload_metadata(self, device: torch.device) -> None:
|
|
"""Initialize metadata needed by dummy offload execution."""
|
|
|
|
def setup_cross_process_offload(
|
|
self,
|
|
gpu_output_buffer: torch.Tensor,
|
|
semaphore: CpuGpuSemaphore,
|
|
host_output_ptr: int | None = None,
|
|
host_row_bytes: int = 0,
|
|
) -> None:
|
|
"""Configure the GPU-worker placeholder with its offload resources."""
|
|
self._is_cpu_offloaded = True
|
|
self._gpu_output_buffer = gpu_output_buffer
|
|
self._sem = semaphore
|
|
self._host_output_ptr = host_output_ptr
|
|
self._host_row_bytes = int(host_row_bytes)
|
|
if host_output_ptr is not None:
|
|
# Hardening: the registry key is the name the
|
|
# traced graph carries, so it must be a real, unique layer name.
|
|
name = self._ple_offload_name
|
|
if name == "<unbound>":
|
|
raise RuntimeError(
|
|
"PLE pull: setup_cross_process_offload on a layer with no "
|
|
"_ple_offload_name; the op would resolve nothing by name"
|
|
)
|
|
prior = _PULL_LAYERS.get(name)
|
|
if prior is not None and prior is not self:
|
|
raise RuntimeError(
|
|
f"PLE pull: two layers armed under the name {name!r}; the op "
|
|
"would resolve the wrong layer's pointers"
|
|
)
|
|
_PULL_LAYERS[name] = self
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
input_ids: torch.Tensor,
|
|
*args: object,
|
|
**kwargs: object,
|
|
) -> torch.Tensor:
|
|
"""Wait for an offloaded result or delegate to ``forward_impl``."""
|
|
if self._is_cpu_offloaded:
|
|
if self._host_output_ptr is not None:
|
|
torch.ops.vllm.ple_offload_wait_copy(
|
|
self._gpu_output_buffer,
|
|
hidden_states,
|
|
input_ids.shape[0],
|
|
self._ple_offload_name,
|
|
)
|
|
else:
|
|
torch.ops.vllm.ple_offload_wait(
|
|
self._sem.flag_tensor,
|
|
self._gpu_output_buffer,
|
|
hidden_states,
|
|
)
|
|
return self._gpu_output_buffer[: input_ids.shape[0]]
|
|
return self.forward_impl(hidden_states, input_ids, *args, **kwargs)
|
|
|
|
def release_offloaded_output(
|
|
self,
|
|
stream: torch.cuda.Stream | None = None,
|
|
) -> None:
|
|
"""Mark the cross-process output buffer reusable on ``stream``."""
|
|
if self._is_cpu_offloaded:
|
|
self._sem.reset(stream)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pre-launch staleness check for the pull transport (#53899 class)
|
|
# ---------------------------------------------------------------------------
|
|
# WHAT IT TESTS. The transport's correctness rests on one precondition: when a
|
|
# real request is published, every layer's semaphore must be CLEAR, so the
|
|
# ``cuStreamWaitValue32(flag == DONE_VALUE)`` inside the forward (line 279 for
|
|
# the pull path, line 215 for the IPC push path) can only be satisfied by the
|
|
# copy THIS step publishes. If a leftover ``DONE_VALUE`` survives into a real
|
|
# step, that wait is satisfied immediately by the PREVIOUS step's signal and
|
|
# the forward consumes the previous step's PLE rows. That is vllm#53899 on
|
|
# GB10/sm_121 exactly: "every forward consumes the previous step's PLE
|
|
# outputs", for the life of the server, with ``cudagraph_mode=NONE`` correct.
|
|
#
|
|
# In this tree the precondition is held structurally -- the dummy signal is
|
|
# consumed by the dummy forward, ``_release_outputs_impl`` resets every layer
|
|
# and synchronizes the model stream while ``_after_dummy`` is set, and the
|
|
# dummy->real transition takes a TP barrier -- but nothing ever READ the flag
|
|
# back, so the structural argument had no instrument. This is that instrument.
|
|
#
|
|
# WHY NOT AN EPOCH WORD. The stronger fix is a monotone epoch: the launch
|
|
# writes ``epoch = step``, the worker signals ``epoch``, and the placeholder
|
|
# waits ``EQ(epoch)``, so a stale signal carries the previous epoch and cannot
|
|
# satisfy the wait. It is not implementable on this transport as it stands.
|
|
# The compare value is an IMMEDIATE OPERAND of ``cuStreamWaitValue32``
|
|
# (``_ple_offload_wait_copy_impl``, this file, line 279), and that call is
|
|
# issued inside the model forward, which is captured: PIECEWISE and
|
|
# FULL_AND_PIECEWISE both bake the wait into the graph with the value frozen at
|
|
# capture time. Making the compare value change per replay needs the captured
|
|
# node's parameters rewritten before each launch
|
|
# (``cuGraphExecBatchMemOpNodeSetParams``); nothing in this tree's capture path
|
|
# exposes the graph exec handle to the connector, and the wait is reached
|
|
# through a torch custom op inside a compiled region. An epoch word would
|
|
# therefore work only with ``cudagraph_mode=NONE``, i.e. only in the
|
|
# configuration where #53899 does not reproduce. So: binary flag, read back.
|
|
#
|
|
# The check is HOST-SIDE and needs a stream synchronize to be ordered against
|
|
# a pending reset, so it is never run inside a capture and is off by default.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def flag_check_is_authoritative(tp_rank: int) -> bool:
|
|
"""Whether ``flag == RESET_VALUE`` before a launch is an INVARIANT here.
|
|
|
|
Only on the publishing rank (TP0). Why: the offload worker fans a
|
|
result out to every registered TP target, each gated only on that target's
|
|
own flag having been reset, so a non-publishing rank's flag can legitimately
|
|
carry THIS step's signal before that rank reaches its own pre-launch check.
|
|
Reading a set flag there and calling it stale terminates a healthy run.
|
|
|
|
On TP0 the invariant holds by construction: the worker learns of a request
|
|
only from ``_launch``, which runs after the check, and the progress monitor
|
|
refuses overlapping serving submissions, so the previous transaction has
|
|
been released and (after a stream synchronize) its reset has retired.
|
|
|
|
A non-publishing rank cannot tell "this step's fan-out arrived early" from
|
|
"last step's signal was never consumed", because the flag is BINARY and
|
|
carries no transaction identity -- the same representational limit that
|
|
rules out an epoch word. So those ranks are skipped, not checked weakly.
|
|
"""
|
|
return tp_rank == 0
|
|
|
|
|
|
def stale_offload_flags(
|
|
layers: "dict[str, PleOffloadLayer]",
|
|
) -> list[tuple[str, int]]:
|
|
"""Return ``(layer_name, flag_value)`` for every layer whose flag is set.
|
|
|
|
An empty list is the precondition holding. A non-empty list means at least
|
|
one semaphore still carries a signal that no forward consumed, so the next
|
|
``wait`` would be satisfied by the previous step's copy.
|
|
|
|
Pure and side-effect free: the caller owns the ordering (synchronize the
|
|
model stream first, and never call this inside a stream capture). Layers
|
|
that are not cross-process offloaded have no semaphore and are skipped.
|
|
"""
|
|
stale: list[tuple[str, int]] = []
|
|
for name, layer in layers.items():
|
|
if not getattr(layer, "_is_cpu_offloaded", False):
|
|
continue
|
|
value = int(layer._sem.flag_tensor.reshape(-1)[0].item())
|
|
if value != CpuGpuSemaphore.RESET_VALUE:
|
|
stale.append((name, value))
|
|
return stale
|