refactor(inference): extract reusable RemoteInferenceGenerator#1911
refactor(inference): extract reusable RemoteInferenceGenerator#1911dyurk-lila wants to merge 4 commits into
Conversation
Fix routed-expert replay (R3) correctness for RL training and introduce a shared token-metadata layout that later routed-expert and sampler-support work builds on. - Scope global RouterReplay state to one Megatron pipeline schedule so a forward-only logprob pass can no longer leak backward replay state into the next training schedule (clear before the schedule and in `finally`). - Keep `rollout_expert_indices` ragged and treat its length as the captured-prefix length. Derive a `router_padding_mask` after left padding that marks alignment padding and the uncaptured trajectory suffix, and carry it through the training data, replay experiences, microbatch padding, and the Megatron model call. - Build one `TokenMetadataLayout` per microbatch and apply it to both routes and the padding mask. Generic construction, alignment, next-token shifting, and packed-output restoration live in `skyrl/utils/token_metadata.py`. - Pass Megatron's `padding_mask` through the model and apply a narrow compatibility shim so `[tokens]` masks broadcast over experts in expert-bias accounting. - Slice every per-trajectory generator field generically during dynamic-sampling replacement and filtering so route metadata stays attached to its trajectory. Synthetic padding rows use distinct dummy experts `[0, ..., topk - 1]`; the mask excludes them from expert-bias accounting while preserving Megatron's dropless `tokens * topk` dispatcher invariant.
Store routed-expert (R3) generation data as compact NumPy arrays instead of large nested Python lists, and send it over the network base64-encoded alongside its shape and dtype. Expert IDs are compacted to the smallest safe uint8/int16/int32 dtype, vLLM responses and client responses use orjson, and preprocessing accepts the decoded NumPy route arrays directly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Under pipeline parallelism each rank replays only its local router layers, but the Megatron worker eagerly expanded the full global-layer routed-expert tensor to int32 before replay setup, allocating a large device temporary for unused layers. Keep routed-expert IDs in their compact dtype through whole-batch device movement, index_select the current PP stage's router layers before metadata alignment, and perform the single int32 conversion inside _split_replay_indices so only the bounded PP-local slice is materialized as int32. Also validate the 4D replay-indices shape up front. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Extract the single-request HTTP generation path out of RemoteInferenceClient into a standalone RemoteInferenceGenerator and a RemoteGenerateResult dataclass. RemoteInferenceClient now owns an internal generator and delegates session management, _post, and _generate_single to it. This is a pure refactor with no functionality change: endpoint routing, retry/backoff, cache_salt handling, serialization, and lifecycle behavior are all preserved. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
|
/gemini review |
![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 refactors the MoE router replay mechanism and optimizes inference server communication by introducing a dedicated RemoteInferenceGenerator and a compact base64-encoded payload format for routed expert indices. It also adds a router_padding_mask to training batches and refactors token metadata alignment to correctly handle padding and context-parallel sharding. The reviewer feedback highlights critical improvements: ensuring the _generator field is properly cleared and reset during serialization/deserialization in RemoteInferenceClient, falling back to a default value instead of raising a ValueError for non-finite logprobs in vllm_server_actor.py, and handling list-type unwrapped models when virtual pipeline parallelism is enabled in replay_utils.py to prevent runtime crashes.
| state = self.__dict__.copy() | ||
| state["_session"] = None | ||
| state["_gen_sem"] = None | ||
| state["_detok_sem"] = None | ||
| state["_sem_loop"] = None |
There was a problem hiding this comment.
The _generator field is not cleared during serialization in __getstate__. Since the test test_serialization asserts that restored._generator is None after deserialization, this will cause the test to fail if the generator was initialized before pickling. Additionally, clearing it ensures lazy re-initialization with a fresh event loop upon deserialization.
| state = self.__dict__.copy() | |
| state["_session"] = None | |
| state["_gen_sem"] = None | |
| state["_detok_sem"] = None | |
| state["_sem_loop"] = None | |
| state = self.__dict__.copy() | |
| state["_generator"] = None | |
| state["_gen_sem"] = None | |
| state["_detok_sem"] = None | |
| state["_sem_loop"] = None |
| self.__dict__.update(state) | ||
| self._session = None | ||
| self._gen_sem = None | ||
| self._detok_sem = None | ||
| self._sem_loop = None |
There was a problem hiding this comment.
The _generator field should be explicitly reset to None in __setstate__ to ensure it is lazily re-initialized on the correct event loop after deserialization.
| self.__dict__.update(state) | |
| self._session = None | |
| self._gen_sem = None | |
| self._detok_sem = None | |
| self._sem_loop = None | |
| self.__dict__.update(state) | |
| self._generator = None | |
| self._gen_sem = None | |
| self._detok_sem = None | |
| self._sem_loop = None |
| logprob = lp_dict[tid].logprob | ||
| if not math.isfinite(logprob): | ||
| raise ValueError("Out of range float values are not JSON compliant") | ||
| content.append({"logprob": logprob}) |
There was a problem hiding this comment.
Raising a ValueError when a logprob is not finite (e.g., -inf) will crash the entire generation request and fail the training run. It is much safer to fall back to a compliant default value like -9999.0 (which is already used as the default for missing logprobs) to prevent unexpected crashes.
| logprob = lp_dict[tid].logprob | |
| if not math.isfinite(logprob): | |
| raise ValueError("Out of range float values are not JSON compliant") | |
| content.append({"logprob": logprob}) | |
| logprob = lp_dict[tid].logprob | |
| if not math.isfinite(logprob): | |
| logprob = -9999.0 | |
| content.append({"logprob": logprob}) |
| unwrapped_model = unwrap_model(model) | ||
| # GPTModel scatters its mask beside the embedding on the first PP stage. HybridModel | ||
| # scatters only the embedding, so its mask must always be scattered here. | ||
| if not isinstance(unwrapped_model, HybridModel) and unwrapped_model.pre_process: | ||
| return router_padding_mask |
There was a problem hiding this comment.
When virtual pipeline parallelism (VPP) is enabled, unwrap_model(model) can return a list of models. In this case, calling unwrapped_model.pre_process will raise an AttributeError and crash the training run. We should handle the list case by extracting the first model instance.
unwrapped_model = unwrap_model(model)
if isinstance(unwrapped_model, list):
unwrapped_model = unwrapped_model[0]
# GPTModel scatters its mask beside the embedding on the first PP stage. HybridModel
# scatters only the embedding, so its mask must always be scattered here.
if not isinstance(unwrapped_model, HybridModel) and unwrapped_model.pre_process:
return router_padding_mask2 participants
Summary
Pure refactor (no functionality change) that extracts the single-request HTTP
generation path out of
RemoteInferenceClientso it can be reused on its own:RemoteInferenceGeneratorand aRemoteGenerateResultdataclass.RemoteInferenceClientown an internalRemoteInferenceGeneratoranddelegate session management,
_post, and_generate_singleto it.the full inference/control-plane client.
cache_salthandling,serialization, and lifecycle behavior.
Testing
uv run --isolated --extra dev --extra fsdp pytest tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py(58 passed)
ruffandblackclean on the changed files.