Repository navigation
[bug] Fix LoRA disk reloading and sleep issues - #1606
Conversation
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
![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 a uses_lora_weight_sync flag to the RemoteInferenceClient to manage VRAM usage during sleep cycles, specifically forcing a lower sleep level to preserve the base model when only LoRA adapters are synced. It also refactors update_lora_from_disk to perform an unload-then-load sequence, bypassing a vLLM performance issue related to load_inplace. The review feedback suggests hardening error handling when parsing JSON responses from the inference server to avoid masking HTTP errors with content type exceptions.
| if resp.status >= 400: | ||
| body = await resp.json() | ||
| raise_for_status(resp, body) |
There was a problem hiding this comment.
The call to resp.json() is risky here. If the server returns a non-JSON error response (e.g., a 502 Bad Gateway from a proxy or a 500 Internal Server Error with a plain text stack trace), await resp.json() will raise an aiohttp.ContentTypeError. This masks the actual HTTP status code and error message, making debugging harder. It is safer to wrap the JSON parsing in a try-except block or check the content type before parsing, especially since raise_for_status is already designed to handle a None body by falling back to the standard resp.raise_for_status().
| if resp.status >= 400: | |
| body = await resp.json() | |
| raise_for_status(resp, body) | |
| if resp.status >= 400: | |
| try: | |
| body = await resp.json() | |
| except Exception: | |
| body = None | |
| raise_for_status(resp, body) |
| if resp.status >= 400: | ||
| body = await resp.json() | ||
| raise_for_status(resp, body) |
There was a problem hiding this comment.
Similar to the _unload_on_server implementation, calling resp.json() directly on a non-2xx response can lead to a ContentTypeError if the error body is not valid JSON. Wrapping this in a try-except block ensures that the actual HTTP error is surfaced via raise_for_status even if the error response body is malformed or plain text.
| if resp.status >= 400: | |
| body = await resp.json() | |
| raise_for_status(resp, body) | |
| if resp.status >= 400: | |
| try: | |
| body = await resp.json() | |
| except Exception: | |
| body = None | |
| raise_for_status(resp, body) |
| active_lora_name: Optional[str] = None | ||
| """Name of the active LoRA adapter. If set, generation requests use this adapter instead of the base model.""" | ||
|
|
||
| uses_lora_weight_sync: bool = False |
There was a problem hiding this comment.
Why do we need a separate flag here? Can't we use just active_lora_name is not None as the condition?
There was a problem hiding this comment.
planning to get rid of active_lora_name in multi-lora, https://github.com/NovaSky-AI/SkyRL/pull/1579/changes#diff-e098add986cd397d2cb2fa04042e96a468524e733f54be101ee04d400e4c4270L200
| Implementation: unload then load (with load_inplace=False), rather than a single | ||
| load with load_inplace=True. Reason: vLLM's /v1/load_lora_adapter caches the | ||
| LoRARequest object server-side keyed by lora_name (see OpenAIServingModels.lora_requests). | ||
| If load_inplace=True is set on that cached request, every subsequent generate request | ||
| causes LRUCacheWorkerLoRAManager.add_adapter to re-read the adapter from disk on every | ||
| model step (line 299-307 of vllm/lora/worker_manager.py checks `load_inplace` per-step). | ||
| The cleanest fix is upstream (clear load_inplace before caching); until that lands, | ||
| unloading first lets us reuse the same lora_name without ever needing load_inplace=True. |
There was a problem hiding this comment.
Nit: Docstring stylke is not consistent with the rest of the codebase. Can you clean this up and also convert this into a TODO:
The cleanest fix is upstream (clear load_inplace before caching); until that lands,
unloading first lets us reuse the same lora_name without ever needing load_inplace=True.
Bug fixes found in #1616 While testing out the tinker CI, found a port collision bug. async def _run_server(self) -> None: """Internal method to run the HTTP server.""" # Release the port reservation right before vLLM rebinds. if self._port_reservation is not None: self._port_reservation.close() # ← step A self._port_reservation = None ... sock_addr = (self._cli_args.host, self._cli_args.port) sock = create_server_socket(sock_addr) There exists a window of time where we free port so vllm can take it, but other actors can also. Fixed this by allocating a specific, non-overlapping port range for each actor. Additionally, changed the disk reloading workaround #1606 created here to more closely mirror the upstream change vllm-project/vllm#41482, adding a new custom endpoint with the load_inplace override. --------- Signed-off-by: hao-aaron <ahao@anyscale.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This PR introduces a tinker ci, closing #1553. It uses tinker_cookbook.recipes.math_rl.train. Currently the `REWARD_MIN_VALUE` is set to 0, but will be tightened once we have a couple of runs. Added both colocated and fully async tests. https://wandb.ai/sky-posttraining-uc-berkeley/gsm8k_tinker_ci?nw=nwuserahao https://wandb.ai/sky-posttraining-uc-berkeley/gsm8k_tinker_fully_async_ci?nw=nwuserahao While testing out the tinker CI, found a port collision bug. ``` async def _run_server(self) -> None: """Internal method to run the HTTP server.""" # Release the port reservation right before vLLM rebinds. if self._port_reservation is not None: self._port_reservation.close() # ← step A self._port_reservation = None ... sock_addr = (self._cli_args.host, self._cli_args.port) sock = create_server_socket(sock_addr) ``` There exists a window of time where we free port so vllm can take it, but other actors can also. Fixed this by allocating a specific, non-overlapping port range for each actor. Additionally, changed the disk reloading workaround #1606 created here to more closely mirror the upstream change vllm-project/vllm#41482, adding a new custom endpoint with the load_inplace override. --------- Signed-off-by: hao-aaron <ahao@anyscale.com> Signed-off-by: ahao-anyscale <ahao@anyscale.com>
Two bugs fixed here:
load_inplace=Truecauses each step to trigger load from disk causing massive slowdowns. This requires an upstream vllm fix, but I provide a temporary workaround by unloading the lora first then loading withload_inplace=False. This could be the issue behind _SKYRL_USE_NEW_INFERENCE=1 (default) causes ~15x performance regression vs legacy path #1582Validated by running previously slow lora run and ensuring nonzero reward: https://wandb.ai/sky-posttraining-uc-berkeley/gsm8k_0.5b_lora/runs/f44fhwwl