[tinker] 9/n towards Kimi K2.6: colocated engine wake/offload for cold sample paths

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:

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: