[tinker] 14/n towards Kimi K2.6: sleep colocated engines around Tinker API checkpoint ops - #2065
Open
casper-hansen wants to merge 4 commits into
Open
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>
…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>
save_checkpoint backloaded the trainer against awake colocated engines (218GiB/GPU at gpu_memory_utilization=0.8), OOMing the first save_every checkpoint of a training run: the client refreshes its sampling session (waking the engines) right before saving. Sleep the engines first -- the same idiom forward/forward_backward already use -- in save_checkpoint and load_checkpoint; the next sampling call or weight sync wakes them. The _engines_asleep guard makes this a no-op when the engines are already asleep. Co-authored-by: Cursor <cursoragent@cursor.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request improves VRAM management and state tracking for colocated inference engines and trainers. It introduces tracking for the sleep state of inference engines, offloads trainer models to CPU before waking or initializing engines to prevent OOM errors, and correctly handles base-model sampling requests with empty model IDs. The review feedback suggests raising a clear RuntimeError if self._dispatch is None during engine initialization to avoid a subsequent AttributeError.
Comment on lines
+430
to
+431
| if self._dispatch is not None: | ||
| self._dispatch.offload_for_sampling() |
Contributor
There was a problem hiding this comment.
If self._dispatch is None, the check if self._dispatch is not None: prevents a crash on offload_for_sampling(), but the code will immediately crash on line 435 with AttributeError: 'NoneType' object has no attribute 'set_inference_engine_client'. To prevent this and provide a clear error message, we should explicitly raise a RuntimeError if self._dispatch is None at the start of _ensure_inference_engines.
Suggested change
| if self._dispatch is not None: | |
| self._dispatch.offload_for_sampling() | |
| if self._dispatch is None: | |
| raise RuntimeError("Model not initialized") | |
| self._dispatch.offload_for_sampling() |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part of the Kimi K2.6/K2.7 series. Stacked on #2031 (9/n) — its commits appear in this diff; review the last commit ("Sleep colocated engines around checkpoint ops") for this PR's own change.
What
save_checkpointandload_checkpointsleep the colocated inference engines before asking the dispatch to backload the trainer, using the same_sleep_inference_engines()idiomforward/forward_backwardalready use. The_engines_asleeptracking from #2031 makes the call a no-op when the engines are already asleep; the next sampling call or weight sync wakes them.Why
Hit on our Kimi K2.7-Code LoRA runs (2x8xB300, colocated vLLM at
gpu_memory_utilization=0.8, i.e. ~218GiB/GPU held by the engines while awake) through the Tinker API: the firstsave_everycheckpoint of a run OOMed reproducibly. The trigger is a client-side pattern, not a rare race — the training client refreshes its sampling session right before saving, which wakes the engines, and the checkpoint's trainer backload (masters + optimizer) then cannot fit. The same hazard applies on the restore side, soload_checkpointgets the identical guard.Made with Cursor