From 3e71c92f9716c3c5ae5a734aed6ef1d672065037 Mon Sep 17 00:00:00 2001 From: casper-hansen Date: Thu, 13 Aug 2026 08:11:40 +0000 Subject: [PATCH 1/3] [tinker] 9/n towards Kimi K2.6: base-model sampling and colocated engine wake/offload Four fixes that make sampling work reliably on the SkyRL-Train backend outside the train->save_weights->sample happy path: - create_sampling_client(base_model=...) maps to model_id "" on the API side, but sample() validated every model_id against the registered adapters and rejected "" as unknown. Treat falsy model_ids as base-model requests. - Under LoRA weight sync (megatron + merge_lora=false), resolve_policy_model_name() returns the skyrl-lora adapter alias, so base-model sampling 404'd on vLLM: the alias only exists after the first sampler-weight save, and applying adapter deltas to a base-model request would be wrong anyway. Resolve falsy model_ids to generator.inference_engine.served_model_name / the policy model path. - Colocated engines are slept right after init and around every training op, and only save_weights_for_sampler woke them -- so a cold sample (base model, or an already-synced adapter) queued against sleeping engines and hung forever. Track engine sleep state on the backend, wake (weights + KV cache) on the sample path after offloading any GPU-resident trainer via the new WorkerDispatch.offload_for_sampling, and normalize to the asleep state before save_weights_for_sampler's wake->broadcast->wake dance. - Lazy engine bring-up runs on the first sampling-related call, which in a multi-tenant service can land right after another tenant's forward/forward_backward left the trainer GPU-resident; under colocate_all the engines' startup allocation then fails ("Engine core initialization failed"). Offload the trainer first, matching the build path's build -> offload -> engines order. Co-authored-by: Cursor --- .../skyrl_train/workers/worker_dispatch.py | 15 ++++ skyrl/backends/skyrl_train_backend.py | 76 ++++++++++++++++--- 2 files changed, 82 insertions(+), 9 deletions(-) diff --git a/skyrl/backends/skyrl_train/workers/worker_dispatch.py b/skyrl/backends/skyrl_train/workers/worker_dispatch.py index 52b8f71a3c..506755a644 100644 --- a/skyrl/backends/skyrl_train/workers/worker_dispatch.py +++ b/skyrl/backends/skyrl_train/workers/worker_dispatch.py @@ -193,6 +193,21 @@ def _offload(self, model: str, offload_optimizer: bool = True, offload_model: bo if offload_optimizer: self._gpu_state[model].optimizer_on_gpu = False + def offload_for_sampling(self, model: str = "policy") -> None: + """Fully offload a colocated trainer so inference engines can reclaim VRAM. + + Used by cold sample paths (no preceding weight sync): the engines are + woken directly, so the trainer left GPU-resident by a forward/optim op + must first move to CPU. No-op when nothing is on the GPU. + """ + if not self.colocate_all: + return + state = self._gpu_state.get(model) + if state is None: + return + if state.model_on_gpu or state.optimizer_on_gpu: + self._offload(model, offload_optimizer=True, offload_model=True) + def mark_all_offloaded(self) -> None: """Mark all models as offloaded (call after build_models when colocate_all).""" for model in self._actor_groups: diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index ca67f2ed5e..a5587eb87c 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -142,6 +142,10 @@ def __init__(self, base_model: str, config: SkyRLTrainBackendOverrides): # New inference infrastructure self._server_groups: list = [] self._inference_router = None + # Colocated engines are slept after init and around training ops; + # sample paths must wake them (tracked here so wakes are not issued + # against already-awake engines). + self._engines_asleep = False # Optional hook invoked on inference-engine state changes (after # _create_new_inference_client, on delete_model teardown). The host @@ -378,6 +382,7 @@ def _create_new_inference_client(self): # LoRA weight sync is in use, since level 2 would discard the base model). if is_colocated: asyncio.run(client.sleep()) + self._engines_asleep = True def _create_render_client(self) -> RendererClientProtocol: """Return a client for vLLM's ``/v1/chat/completions/render``. @@ -414,6 +419,14 @@ def _ensure_inference_engines(self): if self._inference_engines_initialized: return + # A preceding training op (another tenant's forward/forward_backward) + # may have left the trainer GPU-resident; under colocate_all the + # engines' startup allocation (gpu_memory_utilization of each GPU) + # then OOMs. Offload the trainer before bringing the engines up -- + # the same order the build path uses (build -> offload -> engines). + if self._dispatch is not None: + self._dispatch.offload_for_sampling("policy") + self._create_new_inference_client() self._dispatch.set_inference_engine_client(self._inference_engine_client) @@ -783,11 +796,34 @@ def _extract_metrics(self, data: dict) -> dict[str, float]: def _sleep_inference_engines(self): """Sleep inference engines to free GPU memory for training.""" if self._inference_engines_initialized and self._cfg.trainer.placement.colocate_all: + if self._engines_asleep: + return lora_cfg = self._cfg.trainer.policy.model.lora # TODO(team): remove once vllm fixes this # otherwise waking it up will output gibberish: https://github.com/vllm-project/vllm/issues/17103 sleep_level = 1 if lora_cfg and lora_cfg.rank > 0 else 2 asyncio.run(self._inference_engine_client.sleep(level=sleep_level)) + self._engines_asleep = True + + def _wake_inference_engines_for_sampling(self): + """Wake colocated engines before serving sample requests. + + Inverse of :meth:`_sleep_inference_engines`. A cold sample -- base + model, or an already-synced adapter, with no interleaved training op + -- must not rely on ``save_weights_for_sampler`` having woken the + engines: without this, requests queue against sleeping engines and + hang. The trainer may be GPU-resident from a preceding forward / + optim op, so it is offloaded first to give the engines their VRAM + back. + """ + if not (self._inference_engines_initialized and self._cfg.trainer.placement.colocate_all): + return + if not self._engines_asleep: + return + self._dispatch.offload_for_sampling("policy") + asyncio.run(self._inference_engine_client.wake_up(tags=["weights"])) + asyncio.run(self._inference_engine_client.wake_up(tags=["kv_cache"])) + self._engines_asleep = False def _validate_batch_role_and_loss(self, role: str, loss_fn: str): if role == "critic" and loss_fn not in {"ppo", "ppo_critic"}: @@ -1009,21 +1045,25 @@ def sample( save_weights_for_sampler() explicitly before calling sample() if weights have been updated. """ - # 1. Ensure inference engines are initialized + # 1. Ensure inference engines are initialized and awake self._ensure_inference_engines() + self._wake_inference_engines_for_sampling() # 2. Validate every model_id in the batch is a known policy. Multi-LoRA # mixes adapters in one batched sample call (the engine batches across # model_ids in find_batchable_sample); we route each request via the - # `model` field in _sample_with_remote_client below. + # `model` field in _sample_with_remote_client below. An empty model_id + # is base-model sampling (create_sampling_client(base_model=...): the + # API maps it to model_id "") and must not be treated as unknown -- + # _sample_with_remote_client routes it to the served base model name. unique_models = set(prepared_batch.all_model_ids) - unknown = [mid for mid in unique_models if mid not in self._model_ids_to_role] + unknown = [mid for mid in unique_models if mid and mid not in self._model_ids_to_role] if unknown: error = types.ErrorResponse( error=f"Sampling requested for unknown model_id(s): {sorted(unknown)}", status="error" ) return {req_id: error for req_id, *_ in prepared_batch.request_batch_slices} - non_policy = [mid for mid in unique_models if self._model_ids_to_role.get(mid) != "policy"] + non_policy = [mid for mid in unique_models if mid and self._model_ids_to_role.get(mid) != "policy"] if non_policy: error = types.ErrorResponse( error=f"Sampling is only supported for policy models, got non-policy: {sorted(non_policy)}", @@ -1043,12 +1083,24 @@ def _sample_with_remote_client( # Resolve the inference-engine model name per request. With multi-LoRA # the adapter name on vLLM IS the Tinker model_id (registered by # save_sampler_checkpoint via load_lora_adapter). Single-tenant / - # FFT path falls back to resolve_policy_model_name(cfg). + # FFT path falls back to resolve_policy_model_name(cfg). An empty + # model_id is base-model sampling and must target the served base + # model directly: resolve_policy_model_name would return the LoRA + # adapter alias under LoRA weight sync, which (a) does not exist on + # the engines until the first sampler-weight save and (b) would wrongly + # apply adapter deltas to a base-model request. fallback_model_name = resolve_policy_model_name(self._cfg) - per_request_models = [ - mid if (self._base_lora_signature is not None and mid in self._model_ids_to_role) else fallback_model_name - for mid in prepared_batch.all_model_ids - ] + base_model_name = ( + self._cfg.generator.inference_engine.served_model_name or self._cfg.trainer.policy.model.path + ) + per_request_models = [] + for mid in prepared_batch.all_model_ids: + if not mid: + per_request_models.append(base_model_name) + elif self._base_lora_signature is not None and mid in self._model_ids_to_role: + per_request_models.append(mid) + else: + per_request_models.append(fallback_model_name) # Prompt logprobs are a property of the prompt, and all `num_samples` # samples of a request share one prompt, so only ask for them on the @@ -1269,11 +1321,17 @@ def save_sampler_checkpoint(self, output_path, model_id: str, persist: bool = Tr # Lazily create inference engines on first sampling-related call self._ensure_inference_engines() + # The colocated sync dance (wake weights -> broadcast -> wake KV cache) + # assumes engines start asleep; a preceding sample leaves them awake. + self._sleep_inference_engines() + # Multi-LoRA: pass model_id so the dispatch swaps the right adapter in # before broadcasting and the worker registers it on vLLM under that # name. None for the FFT / single-tenant path uses legacy behavior. sync_id = model_id if self._base_lora_signature is not None else None asyncio.run(self._dispatch.save_weights_for_sampler(model_id=sync_id)) + # The colocated sync path leaves the engines awake (weights + KV cache). + self._engines_asleep = False logger.info(f"Synced weights for {model_id} to inference engines via NCCL") if persist: From 9c720282cf1bb27ee845108880598a5b116c6b91 Mon Sep 17 00:00:00 2001 From: casper-hansen Date: Thu, 13 Aug 2026 09:27:33 +0000 Subject: [PATCH 2/3] [tinker] address review: offload every colocated model before waking engines offload_for_sampling only offloaded the named role (callers passed "policy"), so a preceding critic forward/forward_backward left the critic GPU-resident on the cold-sample and lazy engine bring-up paths and could OOM the engines' startup allocation. Offload every tracked GPU-resident model instead. Co-authored-by: Cursor --- .../skyrl_train/workers/worker_dispatch.py | 17 ++++++++--------- skyrl/backends/skyrl_train_backend.py | 4 ++-- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/skyrl/backends/skyrl_train/workers/worker_dispatch.py b/skyrl/backends/skyrl_train/workers/worker_dispatch.py index 506755a644..d2ea6b60d8 100644 --- a/skyrl/backends/skyrl_train/workers/worker_dispatch.py +++ b/skyrl/backends/skyrl_train/workers/worker_dispatch.py @@ -193,20 +193,19 @@ def _offload(self, model: str, offload_optimizer: bool = True, offload_model: bo if offload_optimizer: self._gpu_state[model].optimizer_on_gpu = False - def offload_for_sampling(self, model: str = "policy") -> None: - """Fully offload a colocated trainer so inference engines can reclaim VRAM. + def offload_for_sampling(self) -> None: + """Fully offload every colocated trainer model so inference engines can reclaim VRAM. Used by cold sample paths (no preceding weight sync): the engines are - woken directly, so the trainer left GPU-resident by a forward/optim op - must first move to CPU. No-op when nothing is on the GPU. + woken directly, so any model left GPU-resident by a forward/optim op + (policy, critic, ...) must first move to CPU. No-op when nothing is on + the GPU. """ if not self.colocate_all: return - state = self._gpu_state.get(model) - if state is None: - return - if state.model_on_gpu or state.optimizer_on_gpu: - self._offload(model, offload_optimizer=True, offload_model=True) + for model, state in self._gpu_state.items(): + if state.model_on_gpu or state.optimizer_on_gpu: + self._offload(model, offload_optimizer=True, offload_model=True) def mark_all_offloaded(self) -> None: """Mark all models as offloaded (call after build_models when colocate_all).""" diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index a5587eb87c..ab7b0981b0 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -425,7 +425,7 @@ def _ensure_inference_engines(self): # then OOMs. Offload the trainer before bringing the engines up -- # the same order the build path uses (build -> offload -> engines). if self._dispatch is not None: - self._dispatch.offload_for_sampling("policy") + self._dispatch.offload_for_sampling() self._create_new_inference_client() @@ -820,7 +820,7 @@ def _wake_inference_engines_for_sampling(self): return if not self._engines_asleep: return - self._dispatch.offload_for_sampling("policy") + self._dispatch.offload_for_sampling() asyncio.run(self._inference_engine_client.wake_up(tags=["weights"])) asyncio.run(self._inference_engine_client.wake_up(tags=["kv_cache"])) self._engines_asleep = False From 39dd0672e5aac30273812dbc2681993c1d8206e1 Mon Sep 17 00:00:00 2001 From: Casper Hansen Date: Sat, 5 Sep 2026 07:06:39 +0000 Subject: [PATCH 3/3] [tinker] Drop base-model sampling from the cold-sample fixes Per maintainer review: keep the SkyRL-Train backend's contract that every sampled model_id is a registered policy. sample() rejects an empty model_id as unknown again and _sample_with_remote_client resolves names through resolve_policy_model_name only, as on main. The engine sleep-state tracking, the wake on the sample path and the trainer offload before engine bring-up are unchanged; their docstrings no longer cite base-model sampling as a motivating case. Made with Cursor --- skyrl/backends/skyrl_train_backend.py | 46 ++++++++++----------------- 1 file changed, 16 insertions(+), 30 deletions(-) diff --git a/skyrl/backends/skyrl_train_backend.py b/skyrl/backends/skyrl_train_backend.py index be45a4378e..826d196d70 100644 --- a/skyrl/backends/skyrl_train_backend.py +++ b/skyrl/backends/skyrl_train_backend.py @@ -850,13 +850,14 @@ def _sleep_inference_engines(self): def _wake_inference_engines_for_sampling(self): """Wake colocated engines before serving sample requests. - Inverse of :meth:`_sleep_inference_engines`. A cold sample -- base - model, or an already-synced adapter, with no interleaved training op - -- must not rely on ``save_weights_for_sampler`` having woken the - engines: without this, requests queue against sleeping engines and - hang. The trainer may be GPU-resident from a preceding forward / - optim op, so it is offloaded first to give the engines their VRAM - back. + Inverse of :meth:`_sleep_inference_engines`. A cold sample -- a + request against an already-synced adapter with no + ``save_weights_for_sampler`` of its own in between (a training op, + this tenant's or another's, slept the engines since the last sync) + -- must not rely on the sync having woken the engines: without this, + requests queue against sleeping engines and hang. The trainer may be + GPU-resident from a preceding forward / optim op, so it is offloaded + first to give the engines their VRAM back. """ if not (self._inference_engines_initialized and self._cfg.trainer.placement.colocate_all): return @@ -1094,18 +1095,15 @@ def sample( # 2. Validate every model_id in the batch is a known policy. Multi-LoRA # mixes adapters in one batched sample call (the engine batches across # model_ids in find_batchable_sample); we route each request via the - # `model` field in _sample_with_remote_client below. An empty model_id - # is base-model sampling (create_sampling_client(base_model=...): the - # API maps it to model_id "") and must not be treated as unknown -- - # _sample_with_remote_client routes it to the served base model name. + # `model` field in _sample_with_remote_client below. unique_models = set(prepared_batch.all_model_ids) - unknown = [mid for mid in unique_models if mid and mid not in self._model_ids_to_role] + unknown = [mid for mid in unique_models if mid not in self._model_ids_to_role] if unknown: error = types.ErrorResponse( error=f"Sampling requested for unknown model_id(s): {sorted(unknown)}", status="error" ) return {req_id: error for req_id, *_ in prepared_batch.request_batch_slices} - non_policy = [mid for mid in unique_models if mid and self._model_ids_to_role.get(mid) != "policy"] + non_policy = [mid for mid in unique_models if self._model_ids_to_role.get(mid) != "policy"] if non_policy: error = types.ErrorResponse( error=f"Sampling is only supported for policy models, got non-policy: {sorted(non_policy)}", @@ -1125,24 +1123,12 @@ def _sample_with_remote_client( # Resolve the inference-engine model name per request. With multi-LoRA # the adapter name on vLLM IS the Tinker model_id (registered by # save_sampler_checkpoint via load_lora_adapter). Single-tenant / - # FFT path falls back to resolve_policy_model_name(cfg). An empty - # model_id is base-model sampling and must target the served base - # model directly: resolve_policy_model_name would return the LoRA - # adapter alias under LoRA weight sync, which (a) does not exist on - # the engines until the first sampler-weight save and (b) would wrongly - # apply adapter deltas to a base-model request. + # FFT path falls back to resolve_policy_model_name(cfg). fallback_model_name = resolve_policy_model_name(self._cfg) - base_model_name = ( - self._cfg.generator.inference_engine.served_model_name or self._cfg.trainer.policy.model.path - ) - per_request_models = [] - for mid in prepared_batch.all_model_ids: - if not mid: - per_request_models.append(base_model_name) - elif self._base_lora_signature is not None and mid in self._model_ids_to_role: - per_request_models.append(mid) - else: - per_request_models.append(fallback_model_name) + per_request_models = [ + mid if (self._base_lora_signature is not None and mid in self._model_ids_to_role) else fallback_model_name + for mid in prepared_batch.all_model_ids + ] # Prompt logprobs are a property of the prompt, and all `num_samples` # samples of a request share one prompt, so only ask for them on the