From be93e71ec09d37d2fa9bdd494f5b6bde8b16b505 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Thu, 3 Sep 2026 17:33:54 +0000 Subject: [PATCH] fix(inference): reset requests when discarding KV cache Signed-off-by: bvolpato --- .claude/docs/weight_sync.md | 2 +- .../backends/skyrl_train/weight_sync/cache.py | 16 ++++++ .../skyrl_train/workers/fsdp/fsdp_worker.py | 9 +-- .../workers/megatron/megatron_worker.py | 9 +-- skyrl/train/config/config.py | 2 + .../distributed/test_worker_dispatch.py | 37 ++++++++++++ .../gpu/gpu_ci/test_offload_kv_weight_sync.py | 22 ++++--- .../gpu/gpu_ci/test_prefix_cache_reset.py | 25 ++++++++ .../skyrl_train/weight_sync/test_cache.py | 57 +++++++++++++++++++ tests/train/test_config.py | 5 +- 10 files changed, 165 insertions(+), 19 deletions(-) create mode 100644 skyrl/backends/skyrl_train/weight_sync/cache.py create mode 100644 tests/backends/skyrl_train/weight_sync/test_cache.py diff --git a/.claude/docs/weight_sync.md b/.claude/docs/weight_sync.md index 85b4baa25b..d5a97cae46 100644 --- a/.claude/docs/weight_sync.md +++ b/.claude/docs/weight_sync.md @@ -222,7 +222,7 @@ The `sharded_rdt_*.py` files are vendored from the vLLM PR: github.com/vllm-proj Non-colocated normally keeps the engine fully awake and does `pause_generation → broadcast → resume_generation`. The opt-in `generator.inference_engine.offload_kv_for_weight_sync` flag sleeps the engine (freeing the KV cache from GPU) *during* the sync so `gpu_memory_utilization` can be pushed higher (no need to keep KV cache resident alongside the weight-transfer scratch buffers). It turns on `enable_sleep_mode` (via `inference_servers/utils.py`). Requires non-colocated and non-LoRA. Orchestrated in `WorkerDispatch.save_weights_for_sampler`; the flow depends on the trainer: - **Synchronous trainer** (`fully_async.enabled=false`): generation is complete at sync time, so there are no in-flight requests. A plain `sleep() → wake_up(["weights"]) → broadcast → wake_up(["kv_cache"])` (the same three-phase pattern colocated uses) is enough — the standard `/sleep`+`/wake_up` endpoints discard the KV cache and free the memory. -- **Fully-async trainer** (`fully_async.enabled=true`): generation overlaps the sync, so `pause_generation` (KEEP) freezes in-flight requests, then the allocator is driven directly (see below) so the scheduler is **not** resumed on the weights wake. The KV cache is offloaded to CPU and restored so frozen requests resume with no abort or prefill recompute — **unless** `clear_kv_cache_on_weight_sync=true`, in which case the broadcast resets the prefix cache anyway, so the KV is discarded (skipping the CPU copy) rather than offloaded. +- **Fully-async trainer** (`fully_async.enabled=true`): generation overlaps the sync, so `pause_generation` (KEEP) freezes in-flight requests, then the allocator is driven directly (see below) so the scheduler is **not** resumed on the weights wake. The KV cache is offloaded to CPU and restored so frozen requests resume with no abort or prefill recompute — **unless** `clear_kv_cache_on_weight_sync=true`, in which case the broadcast resets running requests so their KV is recomputed after wake. This reset also runs when prefix caching is disabled: running requests still hold KV. The KV is discarded (skipping the CPU copy) rather than offloaded. The fully-async path is driven entirely from SkyRL — no vLLM patch. It deliberately avoids the `/sleep`+`/wake_up` HTTP endpoints (which route through `EngineCore.sleep`, force-clearing the prefix cache and preempting every running request at level ≥ 1). Instead it drives the per-worker `CuMemAllocator` directly via two `NewInferenceWorkerWrap` methods invoked over `/collective_rpc`: diff --git a/skyrl/backends/skyrl_train/weight_sync/cache.py b/skyrl/backends/skyrl_train/weight_sync/cache.py new file mode 100644 index 0000000000..c9f8aa0c74 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/cache.py @@ -0,0 +1,16 @@ +"""Cache invalidation policy shared by training backends during weight sync.""" + + +def should_reset_kv_cache( + *, + enable_prefix_caching: bool, + fully_async: bool, + clear_kv_cache_on_weight_sync: bool, +) -> bool: + """Whether weight sync must invalidate cached KV, including running requests. + + Synchronous training only needs to invalidate reusable prefix blocks. + Async training can keep requests in flight, whose KV exists even when + prefix caching is disabled. + """ + return clear_kv_cache_on_weight_sync if fully_async else enable_prefix_caching diff --git a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py index a74118d805..28d9c2cccd 100644 --- a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py +++ b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py @@ -30,6 +30,7 @@ WeightChunk, WeightExtractor, ) +from skyrl.backends.skyrl_train.weight_sync.cache import should_reset_kv_cache from skyrl.backends.skyrl_train.weight_sync.weight_extractor_utils import ( yield_module_grouped_chunks, ) @@ -295,13 +296,13 @@ async def broadcast_to_inference_engines( ): if inference_engine_client is None: inference_engine_client = self._weight_sync_inference_client - use_prefix_cache = inference_engine_cfg.enable_prefix_caching generator_dtype = str_to_torch_dtype(inference_engine_cfg.model_dtype) cache_reset_task = None sender_handles_prefix_cache_reset = self._weight_transfer_sender.handles_prefix_cache_reset - # Clear prefix cache for synchronous training or for async training if `clear_kv_cache_on_weight_sync` is set - reset_prefix_cache: bool = use_prefix_cache and ( - not self.cfg.fully_async.enabled or self.cfg.fully_async.clear_kv_cache_on_weight_sync + reset_prefix_cache = should_reset_kv_cache( + enable_prefix_caching=inference_engine_cfg.enable_prefix_caching, + fully_async=self.cfg.fully_async.enabled, + clear_kv_cache_on_weight_sync=self.cfg.fully_async.clear_kv_cache_on_weight_sync, ) send_chunks_kwargs = {"reset_prefix_cache": reset_prefix_cache} diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index d513f7fa5e..0696e544e1 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -58,6 +58,7 @@ WeightChunk, WeightExtractor, ) +from skyrl.backends.skyrl_train.weight_sync.cache import should_reset_kv_cache from skyrl.backends.skyrl_train.workers.megatron.adapter_store import ( AdapterStore, LoraSignature, @@ -1525,13 +1526,13 @@ async def broadcast_to_inference_engines( ): if inference_engine_client is None: inference_engine_client = self._weight_sync_inference_client - use_prefix_cache = inference_engine_cfg.enable_prefix_caching generator_dtype = str_to_torch_dtype(inference_engine_cfg.model_dtype) cache_reset_task = None sender_handles_prefix_cache_reset = self._weight_transfer_sender.handles_prefix_cache_reset - # Clear prefix cache for synchronous training or for async training if `clear_kv_cache_on_weight_sync` is set - reset_prefix_cache: bool = use_prefix_cache and ( - not self.cfg.fully_async.enabled or self.cfg.fully_async.clear_kv_cache_on_weight_sync + reset_prefix_cache = should_reset_kv_cache( + enable_prefix_caching=inference_engine_cfg.enable_prefix_caching, + fully_async=self.cfg.fully_async.enabled, + clear_kv_cache_on_weight_sync=self.cfg.fully_async.clear_kv_cache_on_weight_sync, ) send_chunks_kwargs = {"reset_prefix_cache": reset_prefix_cache} diff --git a/skyrl/train/config/config.py b/skyrl/train/config/config.py index dcbc0b2266..4ed1230645 100644 --- a/skyrl/train/config/config.py +++ b/skyrl/train/config/config.py @@ -967,6 +967,8 @@ class FullyAsyncConfig(BaseConfig): the epoch ends.""" clear_kv_cache_on_weight_sync: bool = False """Whether or not to clear the KV cache on weight sync. Defaults to False. + If True, running requests recompute their KV after the sync, including when + inference prefix caching is disabled. If False, we reuse KV cache from stale policies during generation (avoids recomputation at the cost of using slightly stale KV cache). """ diff --git a/tests/backends/skyrl_train/distributed/test_worker_dispatch.py b/tests/backends/skyrl_train/distributed/test_worker_dispatch.py index 6d862d973d..ef62f7f4d1 100644 --- a/tests/backends/skyrl_train/distributed/test_worker_dispatch.py +++ b/tests/backends/skyrl_train/distributed/test_worker_dispatch.py @@ -37,6 +37,43 @@ def _fft_dispatch_cfg(weight_sync_backend: str = "nccl") -> SimpleNamespace: class TestSaveWeights: """Tests for `WorkerDispatch.save_weights_for_sampler`""" + @pytest.mark.asyncio + @pytest.mark.parametrize("clear_kv_cache", [False, True]) + @pytest.mark.parametrize("enable_prefix_caching", [False, True]) + async def test_async_kv_offload_lifecycle(self, clear_kv_cache, enable_prefix_caching): + from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch + + dispatch = WorkerDispatch.__new__(WorkerDispatch) + dispatch.colocate_all = False + dispatch.cfg = _fft_dispatch_cfg() + dispatch.cfg.trainer.fully_async = SimpleNamespace(enabled=True, clear_kv_cache_on_weight_sync=clear_kv_cache) + ie_cfg = dispatch.cfg.generator.inference_engine + ie_cfg.offload_kv_for_weight_sync = True + ie_cfg.enable_prefix_caching = enable_prefix_caching + dispatch._prepare_for_weight_sync = MagicMock() + dispatch._finish_weight_sync = MagicMock() + dispatch.ensure_active_adapter = MagicMock() + events = [] + client = AsyncMock() + client.increment_weight_version = MagicMock() + client.pause_generation.side_effect = lambda: events.append("pause") + client.sleep_for_weight_sync.side_effect = lambda **kwargs: events.append(("sleep", kwargs["offload_kv"])) + client.wake_for_weight_sync.side_effect = lambda **kwargs: events.append(("wake", kwargs["tags"])) + client.resume_generation.side_effect = lambda: events.append("resume") + dispatch._inference_engine_client = client + dispatch._broadcast_to_inference_engines = MagicMock(side_effect=lambda *a, **kw: events.append("broadcast")) + + await dispatch.save_weights_for_sampler() + + assert events == [ + "pause", + ("sleep", not clear_kv_cache), + ("wake", ["weights"]), + "broadcast", + ("wake", ["kv_cache"]), + "resume", + ] + @pytest.mark.asyncio async def test_non_colocated_calls_pause_and_resume(self): from skyrl.backends.skyrl_train.workers.worker_dispatch import WorkerDispatch diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/test_offload_kv_weight_sync.py b/tests/backends/skyrl_train/gpu/gpu_ci/test_offload_kv_weight_sync.py index 51e7ad02b7..bc812b22fc 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/test_offload_kv_weight_sync.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/test_offload_kv_weight_sync.py @@ -3,8 +3,9 @@ With ``generator.inference_engine.offload_kv_for_weight_sync=True`` (fully-async), ``WorkerDispatch.save_weights_for_sampler`` freezes in-flight requests (KEEP pause), offloads the KV cache to CPU, re-syncs weights into the freed space, then restores -the KV cache. This asserts an in-flight request survives that offload/restore and -finishes cleanly. +the KV cache. With explicit KV clearing, it discards KV and resets running +requests instead, including when prefix caching is disabled. In-flight requests +must finish cleanly in either case. GPU Requirements: 2 GPUs (1 inference + 1 policy). @@ -39,7 +40,7 @@ LONG_PROMPT = "Tell me a very long, detailed story about a dragon who learns to code." -def _offload_kv_cfg() -> SkyRLTrainConfig: +def _offload_kv_cfg(clear_kv_cache: bool = False, enable_prefix_caching: bool = True) -> SkyRLTrainConfig: cfg = SkyRLTrainConfig() cfg.trainer.policy.model.path = MODEL cfg.trainer.critic.model.path = "" @@ -48,11 +49,12 @@ def _offload_kv_cfg() -> SkyRLTrainConfig: cfg.trainer.placement.policy_num_gpus_per_node = 1 cfg.trainer.remove_microbatch_padding = False cfg.trainer.logger = "console" - # Fully-async so the KV cache is offloaded to CPU and in-flight requests are - # frozen and resumed across the sync. + # Fully-async so requests remain in flight across the sync, either restoring + # their KV from CPU or recomputing it when explicit clearing is enabled. cfg.trainer.fully_async.enabled = True - cfg.trainer.fully_async.clear_kv_cache_on_weight_sync = False + cfg.trainer.fully_async.clear_kv_cache_on_weight_sync = clear_kv_cache ie = cfg.generator.inference_engine + ie.enable_prefix_caching = enable_prefix_caching ie.num_engines = 1 ie.tensor_parallel_size = 1 ie.run_engines_locally = True @@ -90,7 +92,10 @@ def _sample_payload(prompt_token_ids: List[int], model: str, max_tokens: int) -> @pytest.mark.asyncio -async def test_offload_kv_weight_sync_preserves_inflight_request(ray_init_fixture, monkeypatch): +@pytest.mark.parametrize("clear_kv_cache,enable_prefix_caching", [(False, True), (True, False)]) +async def test_offload_kv_weight_sync_preserves_inflight_request( + ray_init_fixture, monkeypatch, clear_kv_cache, enable_prefix_caching +): """An in-flight request must survive the KV offload/restore weight sync. While a long sample is mid-generation we run a real training step and @@ -102,7 +107,7 @@ async def test_offload_kv_weight_sync_preserves_inflight_request(ray_init_fixtur # ignore_eos isn't in the default Tinker->vLLM forwarding map; widen it. monkeypatch.setitem(_ric._TINKER_SAMPLE_TO_VLLM_PARAM_MAP, "ignore_eos", "ignore_eos") - cfg = _offload_kv_cfg() + cfg = _offload_kv_cfg(clear_kv_cache, enable_prefix_caching) tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) prompt_token_ids = _prompt_token_ids(tokenizer) @@ -174,6 +179,7 @@ async def spy_wake(*args, **kwargs): # Confirm the KV-offloading path was actually exercised. assert sleep_calls, "sleep_for_weight_sync was never called — offload_kv path not taken" + assert sleep_calls[0][1]["offload_kv"] is not clear_kv_cache assert any(tags and "weights" in tags for tags in wake_tags), "weights were never woken for the broadcast" assert any(tags and "kv_cache" in tags for tags in wake_tags), "KV cache was never restored" diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py b/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py index d10da48eb2..593b413b35 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py @@ -30,6 +30,7 @@ def _make_worker(worker_cls, handles_prefix_cache_reset: bool): worker._is_lora = False worker.cfg = SimpleNamespace( fully_async=SimpleNamespace(enabled=False, clear_kv_cache_on_weight_sync=False), + placement=SimpleNamespace(colocate_all=False), policy=SimpleNamespace(megatron_config=SimpleNamespace(lora_config=SimpleNamespace(merge_lora=False))), ) worker._weight_transfer_sender = AsyncMock() @@ -129,3 +130,27 @@ async def test_no_reset_when_prefix_caching_disabled(strategy, monkeypatch): client.reset_prefix_cache.assert_not_awaited() assert worker._weight_transfer_sender.send.await_args.kwargs["reset_prefix_cache"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)]) +@pytest.mark.parametrize("handles_reset", [False, True]) +@pytest.mark.parametrize("enable_prefix_caching", [False, True]) +@pytest.mark.parametrize("clear_kv_cache", [False, True]) +async def test_async_running_request_reset(strategy, handles_reset, enable_prefix_caching, clear_kv_cache, monkeypatch): + _patch_collectives(monkeypatch) + worker = _make_worker(get_worker_cls(strategy), handles_prefix_cache_reset=handles_reset) + worker.cfg.fully_async.enabled = True + worker.cfg.fully_async.clear_kv_cache_on_weight_sync = clear_kv_cache + client = AsyncMock() + ie_cfg = SimpleNamespace(enable_prefix_caching=enable_prefix_caching, model_dtype="bfloat16") + + await worker.broadcast_to_inference_engines(client, ie_cfg) + + # The sender gets the same policy regardless of who owns the reset, and + # running requests require invalidation even without reusable prefix blocks. + assert worker._weight_transfer_sender.send.await_args.kwargs["reset_prefix_cache"] is clear_kv_cache + if clear_kv_cache and not handles_reset: + client.reset_prefix_cache.assert_awaited_once_with(reset_running_requests=True) + else: + client.reset_prefix_cache.assert_not_awaited() diff --git a/tests/backends/skyrl_train/weight_sync/test_cache.py b/tests/backends/skyrl_train/weight_sync/test_cache.py new file mode 100644 index 0000000000..994409212c --- /dev/null +++ b/tests/backends/skyrl_train/weight_sync/test_cache.py @@ -0,0 +1,57 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from skyrl.backends.skyrl_train.weight_sync.cache import should_reset_kv_cache + + +@pytest.mark.parametrize( + "fully_async,enable_prefix_caching,clear_kv_cache_on_weight_sync,expected", + [ + (False, False, False, False), + (False, False, True, False), + (False, True, False, True), + (False, True, True, True), + (True, False, False, False), + (True, False, True, True), + (True, True, False, False), + (True, True, True, True), + ], +) +def test_weight_sync_cache_policy(fully_async, enable_prefix_caching, clear_kv_cache_on_weight_sync, expected): + assert ( + should_reset_kv_cache( + enable_prefix_caching=enable_prefix_caching, + fully_async=fully_async, + clear_kv_cache_on_weight_sync=clear_kv_cache_on_weight_sync, + ) + is expected + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enable_prefix_caching", [False, True]) +@pytest.mark.parametrize("clear_kv_cache", [False, True]) +async def test_sender_owned_async_reset(enable_prefix_caching, clear_kv_cache): + from skyrl.backends.skyrl_train.weight_sync.delta_strategy import ( + DeltaWeightTransferSender, + ) + + client = AsyncMock() + sender = DeltaWeightTransferSender(SimpleNamespace(sync_dir="/unused"), client) + reset = should_reset_kv_cache( + enable_prefix_caching=enable_prefix_caching, + fully_async=True, + clear_kv_cache_on_weight_sync=clear_kv_cache, + ) + await sender._apply_receiver_update({"target_version": 1}, rank=0, reset_prefix_cache=reset) + + if clear_kv_cache: + client.reset_prefix_cache.assert_awaited_once_with(reset_running_requests=True) + calls = [call[0] for call in client.mock_calls] + assert calls.index("pause_generation") < calls.index("reset_prefix_cache") < calls.index("start_weight_update") + else: + client.reset_prefix_cache.assert_not_awaited() + client.finish_weight_update.assert_awaited_once() + client.resume_generation.assert_awaited_once() diff --git a/tests/train/test_config.py b/tests/train/test_config.py index c567b49ec5..72d2b90055 100644 --- a/tests/train/test_config.py +++ b/tests/train/test_config.py @@ -420,13 +420,14 @@ def test_offload_kv_for_weight_sync_sync_trainer_ok(): @pytest.mark.parametrize("clear_kv_cache", [False, True]) -def test_offload_kv_for_weight_sync_async_ok(clear_kv_cache): +@pytest.mark.parametrize("enable_prefix_caching", [False, True]) +def test_offload_kv_for_weight_sync_async_ok(clear_kv_cache, enable_prefix_caching): cfg = SkyRLTrainConfig() cfg.trainer.placement.colocate_all = False cfg.generator.inference_engine.offload_kv_for_weight_sync = True cfg.trainer.fully_async.enabled = True cfg.trainer.fully_async.clear_kv_cache_on_weight_sync = clear_kv_cache - # Both clear_kv_cache settings are supported now. + cfg.generator.inference_engine.enable_prefix_caching = enable_prefix_caching validate_inference_engine_cfg(cfg)