[tinker] 9/n towards Kimi K2.6: colocated engine wake/offload for cold sample paths - #2031
[tinker] 9/n towards Kimi K2.6: colocated engine wake/offload for cold sample paths#2031casper-hansen wants to merge 5 commits into
Conversation
…ine 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 <cursoragent@cursor.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces mechanisms to manage the state of colocated inference engines (sleeping/waking) and offload the trainer to CPU during sampling to prevent out-of-memory (OOM) errors. It also handles base-model sampling requests where the model ID is empty. The review feedback highlights a potential AttributeError in _wake_inference_engines_for_sampling due to a missing None check on self._dispatch, which should be guarded similarly to other methods.
| return | ||
| if not self._engines_asleep: | ||
| return | ||
| self._dispatch.offload_for_sampling("policy") |
There was a problem hiding this comment.
In _wake_inference_engines_for_sampling, self._dispatch is accessed directly without a None check. However, in _ensure_inference_engines (line 427), self._dispatch is explicitly guarded with if self._dispatch is not None:. To prevent a potential AttributeError: 'NoneType' object has no attribute 'offload_for_sampling' when the dispatch layer is not yet initialized, we should add a similar guard here.
| self._dispatch.offload_for_sampling("policy") | |
| if self._dispatch is not None: | |
| self._dispatch.offload_for_sampling("policy") |
There was a problem hiding this comment.
Declining this one — the None case is unreachable here by construction: _wake_inference_engines_for_sampling early-returns unless self._inference_engines_initialized, and that flag is only set True at the end of _ensure_inference_engines, two lines after it dereferences self._dispatch (set_inference_engine_client). So whenever the wake path runs, _dispatch is provably non-None. The guard you're pointing at in _ensure_inference_engines exists precisely because that code runs before this wiring is guaranteed.
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
…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 <cursoragent@cursor.com>
|
Hey Casper, thanks so much for putting this up! If my understanding is correct:
|
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.
Reviewed by Cursor Bugbot for commit 8b0a909. Configure here.
Resolve conflict in save_sampler_checkpoint: keep both the PR's _engines_asleep reset and main's _inference_adapter_ids tracking. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
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
…kpoint-sleep Picks up the base-model sampling removal from NovaSky-AI#2031: sample() rejects an empty model_id as unknown again and the per-request model name resolves through resolve_policy_model_name only. Kept this branch's _engines_sleep_level tracking (level-2 wake refusal, try/finally around the sampler sync) through the conflicts.
casper-hansen
changed the title
3 participants
What
Two fixes that make sampling work reliably on the SkyRL-Train backend outside the train -> save_weights -> sample happy path:
save_weights_for_samplerwoke them -- so a cold sample (a request against an already-synced adapter with no weight sync of its own in between, e.g. after a training op slept the engines) 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 newWorkerDispatch.offload_for_sampling, and normalize to the asleep state beforesave_weights_for_sampler's wake -> broadcast -> wake dance.Scope change
An earlier revision of this PR also made
create_sampling_client(base_model=...)(API-side model_id"") work on this backend. That was dropped after discussion with the maintainers (commit "Drop base-model sampling from the cold-sample fixes"):sample()keeps rejecting an empty model_id as unknown, and inference-engine model names resolve throughresolve_policy_model_nameonly, exactly as on main. The remaining changes are model-agnostic engine sleep/wake/offload bookkeeping.Part of the Kimi K2.x series (follow-up to #1862). Independent of the other PRs in the series.
Made with Cursor