feat(gym): steer the tail of a rollout step to a faster server group - #8
Open
aoshen02 wants to merge 1 commit into
Open
feat(gym): steer the tail of a rollout step to a faster server group#8aoshen02 wants to merge 1 commit into
aoshen02 wants to merge 1 commit into
Conversation
A rollout step finishes when its slowest sample does, so the last few percent
of samples set the step's wall clock while most of the fleet sits idle. Once
a configurable fraction of a step has finished, steer what remains to a named
server group -- typically a small, low-latency pool -- and let the throughput
pool drain.
- `tail_route_threshold` (fraction, unset disables) and `tail_route_group`.
- Replicas are now registered at the router with `labels: {"group": <name>}`,
taken from the generation backend's URL-to-group map. Backends without
groups register unlabelled, as before.
- The signal to Gym is a sentinel file under the Gym log dir, exported as
NEMO_GYM_TAIL_ROUTE_SENTINEL. vllm_model runs several uvicorn workers with
no shared memory and the flip happens mid-step, so a file is the only thing
they all observe at once. It is removed at the start of every step.
- `max_concurrent_rollouts` caps rollouts in flight. Steering can only reach
requests that have not been dispatched yet, so without a cap the whole batch
is already in flight when the threshold trips and there is nothing left to
steer. This is a prerequisite, not an optimization.
- `_validate_tail_route` fails setup rather than degrading silently: threshold
outside (0,1], no router_url, no server groups, an unknown group name, or a
policy_model with no worker_group_header configured.
No unit tests: the mechanism only shows up across a real router, several
engine groups and a full step. Verified end to end instead (9B, 3 nodes,
4 TP1 engines, 128 sessions, consistent_hash, threshold 0.95):
- flip fired at 122/128 = ceil(0.95 * 128)
- after the flip the throughput group froze at 2428 requests while the
low-latency group absorbed the remaining 693
- spillover to the wrong group: 1 of 332 routing decisions (0.3%), on the
boundary second of the flip
- prefix cache hit rate 95.7%, unchanged from a run without steering
- HTTP 500s: 0, preemptions: 0
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen524@gmail.com>
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.
base 为什么是一条集成分支
这个特性同时依赖两个仍未合并的上游 PR:
feat/multi-server-group—— 提供server_groups与dp_openai_server_base_urls_by_group(URL→组 的映射)feat/external-vllm-router—— 提供向 router 注册 replica 的基础base 分支
integration/group-routing= NVIDIA-NeMo#3520 的 head + 取自 NVIDIA-NeMo#3663 的nemo_gym.py单个文件。不做完整 merge,是因为 main 自 NVIDIA-NeMo#3520 分叉以来对
grpo.py已漂移 422 行,合并等于替 NVIDIA-NeMo#3520 重做 rebase—— 那是 NVIDIA-NeMo#3520 自己的事,不该混进这个 PR。NVIDIA-NeMo#3520 不碰
nemo_gym.py,所以直接取文件是安全的。因此本 PR 的 diff 只包含长尾引导这一层,不重复任何已开 PR 的内容。
真正上游化需要等 NVIDIA-NeMo#3520 与 NVIDIA-NeMo#3663 先落地。
做什么
一个 rollout step 的墙钟由最慢的那个样本决定,所以最后百分之几的样本拖着整步,
而此时绝大多数引擎已经空闲。当某一比例的样本完成后,把剩下的引导到指定的 server group
(通常是小而快的低延迟池),让高吞吐池排空。
tail_route_threshold(比例,不设即关闭)与tail_route_grouplabels: {"group": <name>},来自生成后端的 URL→组 映射。没有分组的后端照旧不带标签注册
NEMO_GYM_TAIL_ROUTE_SENTINEL导出)。vllm_model 会起多个 uvicorn worker、彼此无共享内存,而翻转发生在 step 中途 ——
文件是唯一能让它们同时观察到的东西。每步开始时移除
max_concurrent_rollouts是前提而非优化。引导只能作用于尚未发出的请求;没有并发上限时,阈值触发的那一刻整批请求早已在途,没有任何东西可引导
_validate_tail_route让配置错误在 setup 阶段失败而不是静默降级:阈值不在 (0,1]、缺
router_url、没有 server groups、组名未知、或 policy_model 未配worker_group_header配套改动
需要同时具备:
x-worker-group分组过滤(否则标签无人读)vllm_model:读哨兵文件并在出站请求上打该 header测试
没有单元测试。 这个机制只在「真实 router + 多个引擎组 + 完整一步」的组合下才显现,
纯 Python 层能测的部分测不到它真正的失败模式。改为端到端验证
(9B / 3 节点 / 4 个 TP1 引擎 / 128 会话 / consistent_hash / 阈值 0.95):
作为对照,同一套配置在缺少分组过滤修复的旧 router上,翻转后外溢 31.5%,
且单个 session 在两组之间来回抖动。
静态检查:
ruff check全过、ruff format干净、pyrefly错误数与 base 持平(122,未新增)。AI 辅助完成。所有改动已逐行 review,验证方式与结果如上。