[fix] Avoid inference engine initialization for the SFT code path#1878
Merged
Conversation
_to_training_batch in the Tinker SkyRL-Train backend unconditionally initialized the vLLM inference engines just to construct a VLLMRenderer, so pure SFT workloads (forward_backward + optim_step, no sampling) paid for full inference-engine startup and the GPU memory it reserves. Only image chunks actually need vLLM's render endpoint; text-only batches now render locally via render_model_input and never touch the inference engines. Image-bearing (VLM) batches keep the existing lazy engine init. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
…n training path Image-bearing training batches previously initialized the full vLLM inference engines just to reach /v1/chat/completions/render. Rendering is pure CPU preprocessing (chat template + HF processor), so run vLLM's GPU-less render server (`vllm launch render` equivalent) in a child process instead and point VLLMRenderer at it. The stock CLI infers the device from the current platform, which fails on GPU-less nodes (e.g. Ray head nodes) with the CUDA wheel; the child process forces the CPU platform and an explicit cpu DeviceConfig so the server runs identically everywhere. With this, forward/forward_backward never initialize inference engines for any workload -- engines start only for sample and save_weights_for_sampler. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
![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 introduces a CPU-only render server (CPURenderServer) and client (RenderServerClient) to handle multi-modal preprocessing (such as rendering image chunks) on CPU without spinning up a full GPU-based inference engine. Text-only batches are rendered locally, ensuring that the training path avoids inference engine startup overhead. Comprehensive tests have been added to verify these behaviors. The review feedback highlights an opportunity to optimize RenderServerClient by reusing a single httpx.AsyncClient instance across requests to leverage connection pooling and avoid the overhead of repeatedly establishing TCP/SSL connections.
SumanthRH
requested changes
Jul 13, 2026
SumanthRH
left a comment
SumanthRH
left a comment
Member
There was a problem hiding this comment.
Looking good!
I left a comment on connection pooling and about the client used for rendering with VLM RL .
Per review on NovaSky-AI#1878: RenderServerClient created a fresh httpx.AsyncClient per request, defeating connection pooling for the concurrent per-image render calls within a batch. Cache one pooled client, mirroring RemoteInferenceClient._get_session: rebuild it when the running event loop changes (callers drive the renderer via asyncio.run per batch, so a client bound to the previous loop is unusable), cap the pool at SKYRL_HTTP_CONNECTION_LIMIT, and keep keepalive_expiry below uvicorn's 5s server-side keep-alive to avoid reusing connections the server already closed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
… engines are up Per review on NovaSky-AI#1878: - Add RendererClientProtocol in renderer.py and retype VLLMRenderer.__init__ to accept it, making the "anything serving /v1/chat/completions/render" contract explicit. Satisfied by both RemoteInferenceClient and RenderServerClient. - _create_render_client now branches: when inference engines are already initialized (e.g. VLM RL, where sampling brought them up), reuse RemoteInferenceClient so render requests fan out across all engine API servers; otherwise fall back to the driver-local CPU-only render server. SFT still never initializes engines. - When engines come up, drop any CPU-render-backed renderer and shut down the CPU render server so the next image batch upgrades to the engine-backed client. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
SumanthRH
approved these changes
Jul 15, 2026
Member
|
Looking good! It would be good to do an E2E test for VLM SFT against SkyRL's Tinker API and ensure that the new local render server is used as expected |
Member
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
2 participants
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.
What does this PR do?
Avoids inference engine initialization for the SFT code path in the Tinker SkyRL-Train backend (P0).
_to_training_batch— the batch-conversion step behind everyforward/forward_backwardrequest — unconditionally called_ensure_inference_engines()just to construct aVLLMRenderer. As a result, pure SFT workloads (forward_backward+optim_step, never sampling) still booted the full vLLM inference engine fleet and weight-sync state, wasting GPU memory and startup time.Changes
skyrl/backends/skyrl_train_backend.py:_to_training_batchnow scans the batch forImageChunk/ImageAssetPointerChunk:render_model_input— no inference engines, no HTTP.VLLMRenderer, which needs vLLM's/renderendpoint for image placeholder tokens._ensure_inference_engines()call sites (sample,save_sampler_checkpoint) genuinely need engines and are unchanged.For text-only inputs
VLLMRendereris equivalent torender_model_input(already pinned bytest_matches_render_model_input), so rendered output is identical.Tests
tests/tinker/skyrl_train/test_text_only_batch_no_inference.py:_ensure_inference_enginesstubbed to raise (proving it is never called)tests/backends/skyrl_train/test_renderer.py(5/5) andtests/tinker/skyrl_train/test_sample_session_routing.pypass.Follow-up: CPU-only render server for VLM batches (commit 2)
Image-bearing training batches previously still initialized the inference engines just to reach vLLM's
/v1/chat/completions/render. Rendering is pure CPU preprocessing, so the backend now runs vLLM's GPU-less render server (vllm launch renderequivalent) in a child process and pointsVLLMRendererat it via a minimalRenderServerClient.CPURenderServerforces vLLM's CPU platform + explicitcpudevice in the child so it also works on GPU-less nodes (e.g. Ray head nodes) where the CUDA wheel cannot infer a device.delete_model.forward/forward_backwardnever initialize inference engines for any workload — engines start only forsampleandsave_weights_for_sampler.tests/tinker/skyrl_train/test_render_server.pyrenders a real image through the server and asserts placeholder splicing plus decodedpixel_values/image_grid_thw.🤖 Generated with Claude Code