perf(r3): route rollout experts to the trainer packed with cu_seqlens - #2080
perf(r3): route rollout experts to the trainer packed with cu_seqlens#2080dyurk-lila wants to merge 10 commits into
Conversation
Extract the single-request HTTP generation path out of RemoteInferenceClient into a standalone RemoteGenerateClient plus a RemoteGenerateResult dataclass, so routed-expert results can be obtained without constructing the full inference/control-plane client. RemoteInferenceClient now owns an internal RemoteGenerateClient and delegates session management, _post, and _generate_single to it. Endpoint routing, retry/backoff, cache_salt handling, serialization, and lifecycle behavior are unchanged.
For multi-turn rollouts, build routed-expert (R3) trace data incrementally as the conversation grows instead of re-gathering the whole conversation's routes on every turn. - Add `routed_experts_prompt_starts` to `InferenceEngineInput` and thread a per-request `routed_experts_prompt_start` through `RemoteInferenceClient` and `RemoteInferenceGenerator.generate()`, forwarded as a sampling param so the engine only returns routes for the newly generated suffix. - Introduce `TokenMetadataTrace` (token-aligned array accumulator) and `RoutedExpertTrace`, which records each generation's routes and finalizes a full per-token routed-expert array with loss-mask-aware terminal padding. - Track a `RoutedExpertTrace` on `AgentLoopState` and record routes per turn, replacing the previous whole-conversation re-gather in `SkyRLGymGenerator.agent_loop`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.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 optimizes MoE router replay by transitioning from a padded, batch-major representation to a packed tensor layout (PackedTensor) for ragged token-aligned batch fields. It also introduces a dedicated RemoteGenerateClient with optimized side-channel array serialization to reduce communication overhead, and adds CPU topology utilities to dynamically size thread pools based on cgroup quotas. The review feedback highlights a critical bug in list comparison within align_packed_token_metadata, suggests caching a CPU copy of cu_seqlens in PackedTensor to avoid expensive GPU-to-CPU synchronizations, and recommends explicitly clearing _generate_client during serialization in RemoteInferenceClient to prevent active sessions from being pickled.
| if segment_lengths != list(layout.sequence_lengths): | ||
| raise ValueError( | ||
| f"Packed metadata segments {segment_lengths} do not match " | ||
| f"trajectory lengths {list(layout.sequence_lengths)}" | ||
| ) |
There was a problem hiding this comment.
If layout.sequence_lengths is a PyTorch tensor, calling list(layout.sequence_lengths) will return a list of 0-D tensors. Comparing a list of Python ints (segment_lengths) to a list of 0-D tensors evaluates to False in Python, which will cause this validation to always fail and raise a ValueError incorrectly. Coerce layout.sequence_lengths to a list of Python ints using .tolist() if it is a tensor.
| if segment_lengths != list(layout.sequence_lengths): | |
| raise ValueError( | |
| f"Packed metadata segments {segment_lengths} do not match " | |
| f"trajectory lengths {list(layout.sequence_lengths)}" | |
| ) | |
| layout_lengths = layout.sequence_lengths.tolist() if hasattr(layout.sequence_lengths, "tolist") else list(layout.sequence_lengths) | |
| if segment_lengths != layout_lengths: | |
| raise ValueError( | |
| f"Packed metadata segments {segment_lengths} do not match " | |
| f"trajectory lengths {layout_lengths}" | |
| ) |
| if int(cu_seqlens[0]) != 0 or int(cu_seqlens[-1]) != values.shape[0]: | ||
| raise ValueError( | ||
| f"cu_seqlens must run from 0 to the {values.shape[0]} packed rows, " | ||
| f"got {int(cu_seqlens[0])} to {int(cu_seqlens[-1])}" | ||
| ) | ||
| self.values = values | ||
| self.cu_seqlens = cu_seqlens |
There was a problem hiding this comment.
To avoid expensive host-device synchronizations (CUDA syncs) when retrieving segments or slicing the PackedTensor on GPU, we should cache a CPU copy of cu_seqlens during initialization. Since cu_seqlens is extremely small (batch size + 1), keeping a CPU copy is very cheap and prevents blocking the GPU.
| if int(cu_seqlens[0]) != 0 or int(cu_seqlens[-1]) != values.shape[0]: | |
| raise ValueError( | |
| f"cu_seqlens must run from 0 to the {values.shape[0]} packed rows, " | |
| f"got {int(cu_seqlens[0])} to {int(cu_seqlens[-1])}" | |
| ) | |
| self.values = values | |
| self.cu_seqlens = cu_seqlens | |
| self.values = values | |
| self.cu_seqlens = cu_seqlens | |
| self.cu_seqlens_cpu = cu_seqlens.cpu() if cu_seqlens.device.type == "cuda" else cu_seqlens | |
| if int(self.cu_seqlens_cpu[0]) != 0 or int(self.cu_seqlens_cpu[-1]) != values.shape[0]: | |
| raise ValueError( | |
| f"cu_seqlens must run from 0 to the {values.shape[0]} packed rows, " | |
| f"got {int(self.cu_seqlens_cpu[0])} to {int(self.cu_seqlens_cpu[-1])}" | |
| ) |
| if isinstance(index, slice): | ||
| if index.step in (None, 1): | ||
| start, stop, _ = index.indices(len(self)) | ||
| stop = max(start, stop) | ||
| offsets = self.cu_seqlens[start : stop + 1] | ||
| return PackedTensor(self.values[int(offsets[0]) : int(offsets[-1])], offsets - offsets[0]) |
There was a problem hiding this comment.
Use the cached CPU copy of cu_seqlens (self.cu_seqlens_cpu) to avoid GPU-to-CPU synchronization when slicing the PackedTensor.
| if isinstance(index, slice): | |
| if index.step in (None, 1): | |
| start, stop, _ = index.indices(len(self)) | |
| stop = max(start, stop) | |
| offsets = self.cu_seqlens[start : stop + 1] | |
| return PackedTensor(self.values[int(offsets[0]) : int(offsets[-1])], offsets - offsets[0]) | |
| if isinstance(index, slice): | |
| if index.step in (None, 1): | |
| start, stop, _ = index.indices(len(self)) | |
| stop = max(start, stop) | |
| offsets = self.cu_seqlens[start : stop + 1] | |
| offsets_cpu = self.cu_seqlens_cpu[start : stop + 1] | |
| return PackedTensor(self.values[int(offsets_cpu[0]) : int(offsets_cpu[-1])], offsets - offsets[0]) |
| position = index + len(self) if index < 0 else index | ||
| if not 0 <= position < len(self): | ||
| raise IndexError(f"segment {index} is out of range for a packed batch of {len(self)}") | ||
| return self.values[int(self.cu_seqlens[position]) : int(self.cu_seqlens[position + 1])] |
There was a problem hiding this comment.
Use the cached CPU copy of cu_seqlens (self.cu_seqlens_cpu) to avoid GPU-to-CPU synchronization when indexing and retrieving segment views.
| position = index + len(self) if index < 0 else index | |
| if not 0 <= position < len(self): | |
| raise IndexError(f"segment {index} is out of range for a packed batch of {len(self)}") | |
| return self.values[int(self.cu_seqlens[position]) : int(self.cu_seqlens[position + 1])] | |
| position = index + len(self) if index < 0 else index | |
| if not 0 <= position < len(self): | |
| raise IndexError(f"segment {index} is out of range for a packed batch of {len(self)}") | |
| return self.values[int(self.cu_seqlens_cpu[position]) : int(self.cu_seqlens_cpu[position + 1])] |
| 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.
Explicitly set _generate_client to None in __getstate__ to ensure that any active client session is not serialized and is correctly re-initialized on demand after unpickling.
| state = self.__dict__.copy() | |
| state["_session"] = None | |
| state["_gen_sem"] = None | |
| state["_detok_sem"] = None | |
| state["_sem_loop"] = None | |
| state = self.__dict__.copy() | |
| state["_generate_client"] = 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.
Explicitly set _generate_client to None in __setstate__ to ensure it is correctly re-initialized on demand after unpickling.
| self.__dict__.update(state) | |
| self._session = None | |
| self._gen_sem = None | |
| self._detok_sem = None | |
| self._sem_loop = None | |
| self.__dict__.update(state) | |
| self._generate_client = None | |
| self._gen_sem = None | |
| self._detok_sem = None | |
| self._sem_loop = None |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.
Reviewed by Cursor Bugbot for commit c1e71a0. Configure here.
| rollout_expert_indices_out = agent_loop_state.routed_expert_trace.finalize( | ||
| token_count=len(prompt_ids) + len(response_ids), | ||
| loss_mask=[0] * len(prompt_ids) + loss_mask, | ||
| ) |
There was a problem hiding this comment.
Synthetic EOS breaks R3 finalize
High Severity
In single-turn mode the generator appends a synthetic EOS with loss_mask=1, then RoutedExpertTrace.finalize rejects any loss-active token past the captured route prefix. That makes R3 rollouts fail whenever an EOS is appended after stop.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit c1e71a0. Configure here.
| responses=responses, | ||
| rewards=rewards, | ||
| loss_masks=loss_masks, | ||
| ) |
There was a problem hiding this comment.
Broken GPU test kwargs
Medium Severity
The new variable-length replay test calls convert_prompts_responses_to_batch_tensors with tokenizer=..., but that helper expects pad_token_id. The test raises TypeError before exercising the packed-route path it was added to cover.
Reviewed by Cursor Bugbot for commit c1e71a0. Configure here.
2 participants
Problem
Routed-expert data was expanded into
[batch, max_sequence_length, layers, top_k]as soon as trajectories reached the trainer. For variable-length rollouts, left padding occupied a large fraction of this already-large side channel. Later batching and model-alignment steps copied the padded rectangle several more times.Packed representation
This change introduces
PackedTensor, represented by:valuestensor containing all real trajectory rows; andcu_seqlens, whose adjacent offsets identify each trajectory segment.Rollout routes therefore remain
[sum(sequence_lengths), layers, top_k]until the model-facing alignment step. Uncaptured suffix tokens still receive valid dummy routes, but batch-level left padding is no longer stored in the side channel.PackedTensorsupports the operations required by training batches and replay buffers: segment lookup, slicing, concatenation, repetition, device and dtype transfer, pinning, and padding-row construction. Offset builders derive segment boundaries from validated non-negative lengths, and the container checks the offset dtype, device, endpoints, and packed row count.Alignment contract
Packed metadata is expanded only against the shared
TokenMetadataLayout, which already describes batch padding, packed microbatches, and context-parallel sharding. Response-suffix fields use explicit segment starts rather than an ambiguous full-coverage flag: a trajectory’s side channel may begin after its prompt, while routes begin at the first real token.The model path then:
Performance
In the included four-arm collation benchmark, packed pooled collation reduced a representative global batch from 55.00 GiB to 29.99 GiB of driver memory and from 4.58 s to 1.95 s. The benefit follows padding density: highly ragged workloads improve most, while a uniform batch has no padding to remove and can be marginally slower because offsets still need to be carried.
Testing
PackedTensortests cover construction, offset validation, slicing, concatenation, repetition, padding, dtype/device conversion, and empty segments.Note
High Risk
Changes the training-batch layout, pickle path, and Megatron router-replay alignment for MoE expert indices. Bugs here can desync replayed routes from tokens or break distributed collectives.
Overview
Stops expanding routed-expert indices into a left-padded
[batch, max_seq, layers, topk]tensor. Routes now travel as aPackedTensor: one[sum(seq_len), layers, topk]buffer pluscu_seqlens, so batch padding is not stored in the side channel.TensorBatchtreats packed fields like tensors/TensorList(slice, cat, repeat, device move, pickle). Collation fills packed segments from a local thread pool; dummy distinct-expert rows still cover uncaptured suffixes and DP/microbatch padding.Megatron alignment expands packed segments only against
TokenMetadataLayout(optional per-segment starts for suffixes), then CP/TP-shards as before. Multi-turn generation accumulates routes inRoutedExpertTraceand can request a prompt start on/skyrl/v1/generate. Router replay is now rejected with virtual pipeline parallelism.Also generalizes the generate-wire ndarray envelope (splice base64 without materializing huge strings) and extracts
RemoteGenerateClientfor that path.Reviewed by Cursor Bugbot for commit c1e71a0. Bugbot is set up for automated code reviews on this repo. Configure here.