Add RL-Kernel 8GPU runbook, metrics, and config support#4
Add RL-Kernel 8GPU runbook, metrics, and config support#4inaniloquentee wants to merge 8 commits into
Conversation
Signed-off-by: inaniloquentee <3051000145@qq.com>
Signed-off-by: inaniloquentee <3051000145@qq.com>
Signed-off-by: inaniloquentee <3051000145@qq.com>
|
Warning Review limit reached
Next review available in: 30 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThis PR introduces an RL-Kernel ChangesRL-Kernel linear_logp Feature
Profiling, Offload, and Weight-Sync Infrastructure
RL-Kernel Documentation
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Model as model.py forward_step
participant LossFn as loss.get_log_probs_and_entropy
participant RLK as rl_kernel.maybe_compute_linear_logp
participant Op as RL-Kernel linear_logp op
participant Timer as baseline_timer/CudaEventTimerQueue
Model->>Model: build LinearLogpContext from output layer weights
Model->>LossFn: forward hidden states + rl_kernel_linear_logp_context
LossFn->>RLK: hidden_states, target_ids, context
alt RL-Kernel available and constraints satisfied
RLK->>Op: apply(weight, bias, hidden_states)
Op-->>RLK: log_probs
RLK->>Timer: record dispatch/token/call metrics
RLK-->>LossFn: log_probs
else fallback required
RLK-->>LossFn: None
LossFn->>LossFn: materialize logits and compute native logprob
LossFn->>Timer: record baseline runtime/CUDA-event metrics
end
LossFn-->>Model: log_probs, entropy, fallback_count
Estimated code review effort: 5 (Critical) | ~120 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/_unit_stubs.py (2)
169-206: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRestore the removed vLLM submodule stubs
tests/_unit_stubs.py:169-206still needsvllm.utils.system_utils,vllm.entrypoints.cli, andvllm.entrypoints.openaiforvime.backends.vllm_utils.vllm_engineandvime.backends.vllm_utils.arguments; without them, those imports fail in a vllm-free test environment.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/_unit_stubs.py` around lines 169 - 206, The vLLM test stubs in install_vllm_cli_stubs are missing required submodules, causing imports from vime.backends.vllm_utils.vllm_engine and vime.backends.vllm_utils.arguments to fail when vllm is not installed. Extend the existing stub setup to register vllm.utils.system_utils, vllm.entrypoints.cli, and vllm.entrypoints.openai in sys.modules, and attach them to the vllm module hierarchy alongside the existing vllm.utils and vllm.engine stubs so the import tree resolves cleanly in the test environment.
56-61: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winRestore the
vllm_router.launch_router.RouterArgsstub
tests/utils/test_megatron_role_config.pycallsinstall_rollout_optional_stubs()and then importsvime.utils.arguments, but that module doesfrom vllm_router.launch_router import RouterArgsat import time. An emptyModuleType("vllm_router")is not enough here, so CPU-only runs without the real package will fail on import.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/_unit_stubs.py` around lines 56 - 61, The rollout stubs currently only register a bare vllm_router module, but vime.utils.arguments imports RouterArgs from vllm_router.launch_router at import time. Update install_rollout_optional_stubs() in tests/_unit_stubs.py to also stub the vllm_router.launch_router submodule and provide a RouterArgs symbol so tests that import vime.utils.arguments can run without the real package.
🧹 Nitpick comments (12)
vime/backends/megatron_utils/arguments.py (1)
163-165: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStale comment no longer matches the conditional default.
The comment
# always use zero optimizerat Line 164 no longer reflects reality now thatargs.use_distributed_optimizeris env-driven and can be disabled viaVIME_USE_DISTRIBUTED_OPTIMIZER=0.✏️ Update comment
- # always use zero optimizer + # use distributed (ZeRO-style) optimizer by default; disable via VIME_USE_DISTRIBUTED_OPTIMIZER=0 args.use_distributed_optimizer = os.environ.get("VIME_USE_DISTRIBUTED_OPTIMIZER", "1") != "0"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/backends/megatron_utils/arguments.py` around lines 163 - 165, The comment in _set_default_megatron_args is stale because use_distributed_optimizer is no longer always enabled; update or remove the comment so it matches the env-driven default set from VIME_USE_DISTRIBUTED_OPTIMIZER. Keep the wording aligned with the actual behavior in arguments.py and the _set_default_megatron_args function.vime/backends/megatron_utils/model_provider.py (1)
275-307: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value
matched_output_layersubstring check can diverge fromwants_output_layer's literal match, silently skipping intended gradient enabling.
wants_output_layerchecks patterns against the literal string"output_layer.weight", butmatched_output_layeronly checks whether"output_layer"is a substring of any matched parameter name (Line 287). If a pattern happens to match some otheroutput_layer-related parameter (e.g., a bias or norm containing "output_layer" in its path) without matchingoutput_layer.weightitself,matched_output_layerbecomesTrueand the shared-embedding fallback (Lines 288-307) is skipped — silently leaving the tied output/embedding weight untrained even though the user requested it.Tighten the check to specifically look for the
output_layer.weightname rather than a loose substring.🛠️ Proposed fix
- matched_output_layer = any("output_layer" in name for name in matched_names) + matched_output_layer = any(name.endswith("output_layer.weight") for name in matched_names)Also, Ruff flags the blind
except Exceptionat Line 294 (BLE001); consider narrowing to the expected exception types (e.g.,AttributeError,RuntimeError) raised byshared_embedding_or_output_weight()for clearer failure diagnostics.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/backends/megatron_utils/model_provider.py` around lines 275 - 307, The output-layer fallback in model_provider.py is too broad: matched_output_layer currently uses a loose "output_layer" substring check, which can incorrectly suppress the shared-weight training path in the same block that uses wants_output_layer. Update the logic around the matched_names scan and the shared_embedding_or_output_weight fallback to detect the exact output_layer.weight parameter name instead of any output_layer-related substring, so the tied embedding/output weight is enabled whenever the user explicitly requested that target. While you are there, narrow the blanket except around shared_embedding_or_output_weight() to the specific exceptions it can raise, keeping the logger debug/warning behavior intact.Source: Linters/SAST tools
vime/utils/train_metric_utils.py (1)
16-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePotential key collision between timer-derived and metric-derived
perf/entries.
log_dictis first populated withperf/{key}_timefromlog_dict_raw, then updated withperf/{key}frommetric_dict_raw. If a metric is ever named e.g."actor_train_time", itsperf/actor_train_timekey would silently collide with (and overwrite) the timer-derivedperf/actor_train_timefrom a timer named"actor_train". Not currently triggered by any visible caller, but worth a naming convention/assertion to avoid silent overwrites as more metrics are added.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/utils/train_metric_utils.py` around lines 16 - 25, The `train_metric_utils` merge logic can silently overwrite `perf/` entries when a metric name matches a timer-derived key. Update the `timer_instance.log_dict()` / `timer_instance.metric_dict()` handling so `log_dict` construction in `train_metric_utils` avoids collisions, either by using a stricter naming convention for `perf/{key}_time` versus `perf/{key}` or by asserting that the two key sets are disjoint before merging. Keep the fix localized around the `Timer`-based dictionary assembly and ensure the final `log_dict` cannot contain duplicate `perf/` keys.vime/backends/megatron_utils/actor.py (1)
49-156: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated profiler/env helpers across
actor.pyandrollout.py.
_parse_capture_rollouts,_should_nsys_capture, and_NsysCudaProfilerCapturehere are near-verbatim copies ofvime/ray/rollout.py:43-119(that copy additionally guards ontorch.cuda.is_available()). The two are already diverging. Consider extracting these into a shared helper (e.g.,vime/utils/nsys_capture.py) so gating semantics and the CUDA-availability guard stay consistent.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/backends/megatron_utils/actor.py` around lines 49 - 156, The profiler/env helper logic in _parse_capture_rollouts, _should_nsys_capture, and _NsysCudaProfilerCapture is duplicated and already diverging from the matching implementation in rollout.py. Move this shared Nsight Systems capture logic into a common helper module and have actor.py and rollout.py import it so the CUDA-availability check and rollout gating behavior stay consistent in one place.vime/backends/vllm_utils/vllm_engine.py (1)
454-461: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a timeout to the profiler control POSTs.
start_profile/stop_profileare invoked viaray.get(...)inside_VLLMCudaProfilerCapturearound rollout generation. Without a request timeout, a stuck engine would blockray.getand stall the wholegenerate()path indefinitely (Ruff S113 also flags line 459).♻️ Suggested change
- response = requests.post(f"http://{self.server_host}:{self.server_port}/start_profile", json={}) + response = requests.post( + f"http://{self.server_host}:{self.server_port}/start_profile", json={}, timeout=30 + ) response.raise_for_status() return {"ok": True} def stop_profile(self): - response = requests.post(f"http://{self.server_host}:{self.server_port}/stop_profile", json={}) + response = requests.post( + f"http://{self.server_host}:{self.server_port}/stop_profile", json={}, timeout=30 + ) response.raise_for_status() return {"ok": True}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/backends/vllm_utils/vllm_engine.py` around lines 454 - 461, The profiler control requests in start_profile and stop_profile currently post without a timeout, which can leave ray.get waiting forever if the vLLM engine stalls. Update the requests.post calls in vllm_engine.py to include a finite timeout for both start_profile and stop_profile, keeping the existing response.raise_for_status flow and the returned {"ok": True} behavior unchanged.Source: Linters/SAST tools
scripts/run-qwen3-30B-A3B.sh (1)
492-503: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winQuote the array expansions in the
train.pyinvocation.
${MODEL_ARGS[@]}…${RLK_ARGS[@]}are expanded unquoted (Shellcheck SC2068). This re-splits/globs each element; notably--moe-layer-freq "[1,1,…]"is a glob bracket expression that could expand against filenames in the CWD, and any element containing spaces would be split into multiple argv entries. Use quoted expansions.♻️ Suggested change
- ${MODEL_ARGS[@]} \ - ${CKPT_ARGS[@]} \ - ${ROLLOUT_ARGS[@]} \ - ${OPTIMIZER_ARGS[@]} \ - ${GRPO_ARGS[@]} \ - ${WANDB_ARGS[@]} \ - ${TB_ARGS[@]} \ - ${PERF_ARGS[@]} \ - ${EVAL_ARGS[@]} \ - ${VLLM_ARGS[@]} \ - ${MISC_ARGS[@]} \ - ${RLK_ARGS[@]} + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${TB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${VLLM_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${RLK_ARGS[@]}"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/run-qwen3-30B-A3B.sh` around lines 492 - 503, The train.py argument assembly in the run-qwen3-30B-A3B.sh script is expanding several arrays unquoted, which can trigger shell re-splitting and glob expansion. Update the invocation that uses MODEL_ARGS, CKPT_ARGS, ROLLOUT_ARGS, OPTIMIZER_ARGS, GRPO_ARGS, WANDB_ARGS, TB_ARGS, PERF_ARGS, EVAL_ARGS, VLLM_ARGS, MISC_ARGS, and RLK_ARGS so each array expansion is quoted, preserving each element as a single argv entry even when values contain spaces or bracket characters.Source: Linters/SAST tools
tests/test_rl_kernel_linear_logp_integration.py (1)
30-31: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueStatic analysis: mutable class attribute default (RUF012).
calls: list[dict] = []is a shared mutable class-level default. It's used intentionally here as a spy accumulator across instantiations (reset via_FakeLinearLogpOp.calls.clear()), so it's not a functional bug, but consider annotating withClassVar[list[dict]]to signal intent to the linter and future readers.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_rl_kernel_linear_logp_integration.py` around lines 30 - 31, The _FakeLinearLogpOp spy uses a shared mutable class attribute for calls, which triggers the mutable default warning; update the class attribute declaration in _FakeLinearLogpOp to be explicitly marked as a class variable so the intended shared accumulator is clear to the linter and readers, while keeping the existing clear/reset behavior intact.Source: Linters/SAST tools
vime/utils/rl_kernel.py (1)
6-7: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTwo-tier op validation is currently dead code.
RL_KERNEL_SUPPORTED_OPSandRL_KERNEL_INTEGRATED_OPSare identical today, so any op that survivesparse_rl_kernel_ops'sRL_KERNEL_SUPPORTED_OPScheck (Line 29) is guaranteed to already be inRL_KERNEL_INTEGRATED_OPS, making theunsupported = [...]branch innormalize_rl_kernel_args(Lines 66-73) unreachable. This is presumably intentional scaffolding for a future op that is "supported" (parseable) but not yet "integrated" (backend-wired) — consider a brief comment clarifying the intended distinction so future contributors don't add new ops to the wrong constant.Also applies to: 66-73
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/utils/rl_kernel.py` around lines 6 - 7, The two-tier validation in rl_kernel is currently indistinguishable because RL_KERNEL_SUPPORTED_OPS and RL_KERNEL_INTEGRATED_OPS contain the same value, making the unsupported branch in normalize_rl_kernel_args effectively dead. Update the rl_kernel module by adding a brief clarifying comment near RL_KERNEL_SUPPORTED_OPS / RL_KERNEL_INTEGRATED_OPS and parse_rl_kernel_ops / normalize_rl_kernel_args explaining the intended difference between “parseable” and “backend-integrated” ops so future additions go to the correct constant.vime/backends/megatron_utils/megatron_to_hf/qwen2.py (1)
60-67: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLayer-norm mapping expansion is duplicated verbatim in
qwen3moe.py.The identical set-membership logic (
input_layernorm.weight/pre_mlp_layernorm.weight/post_attention_layernorm.weighthandling) is repeated in bothqwen2.pyandqwen3moe.py. Consider extracting a shared helper (e.g._map_layernorm_rest(rest, layer_idx)) used by both converters to avoid future drift between the two.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/backends/megatron_utils/megatron_to_hf/qwen2.py` around lines 60 - 67, The layer-norm rest-to-target mapping in Qwen2 is duplicated in Qwen3MoE, so centralize this shared logic to prevent drift. Extract the repeated set-membership handling from the converter method in qwen2.py into a common helper such as _map_layernorm_rest(rest, layer_idx), then have both qwen2.py and qwen3moe.py call that helper for the input_layernorm/post_attention_layernorm cases.vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py (2)
137-175: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate tp_size-selection logic between
_bucket_size_bytesand_get_megatron_local_param_info_buckets.Both functions independently re-derive
param_size = info.size * tp_sizeusing the same.experts.substring check. Extract a shared helper (e.g._param_full_size_bytes(info)) to avoid the two implementations drifting apart.♻️ Suggested refactor
+def _param_full_size_bytes(info: ParamInfo) -> int: + tp_size = ( + mpu.get_expert_tensor_parallel_world_size() + if ".experts." in info.name + else mpu.get_tensor_model_parallel_world_size() + ) + return info.size * tp_size + + def _bucket_size_bytes(param_infos: Sequence[ParamInfo]) -> int: - total = 0 - for info in param_infos: - if ".experts." in info.name: - tp_size = mpu.get_expert_tensor_parallel_world_size() - else: - tp_size = mpu.get_tensor_model_parallel_world_size() - total += info.size * tp_size - return total + return sum(_param_full_size_bytes(info) for info in param_infos)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py` around lines 137 - 175, The TP-size selection and full parameter size calculation are duplicated in _bucket_size_bytes and _get_megatron_local_param_info_buckets, which can drift over time. Extract the shared logic into a single helper used by both functions, ideally one that takes a ParamInfo and returns its full byte size based on the .experts. check and the appropriate mpu world-size accessor, then update both call sites to use that helper.
26-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueVerbose per-bucket logging also fires during checkpoint saves.
This iterator is shared with
hf_checkpoint_saver.py's full-checkpoint save path, so every checkpoint save will now emit the same[weight-sync]-prefixed bucket timing prints, which is a slightly misleading label outside the rollout weight-sync context. Not incorrect, just a minor observability nit; consider routing through a proper logger with a context-appropriate prefix instead of bare🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py` around lines 26 - 61, The per-bucket timing output in hf_weight_iterator_direct.py is using a weight-sync-specific print prefix even when this iterator is reused by the full checkpoint save path, so the logs are misleading. Update the logging in the iterator method that loops over megatron_local_param_info_buckets to use a proper logger or a context-aware prefix that can reflect checkpoint-save usage as well as rollout sync, and keep the bucket timing messages consistent with the call site.vime_plugins/mbridge/__init__.py (1)
10-23: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMonkey-patching private
_OTHER_MAPPING/methods on third-party classes is fragile.This relies on
mbridge's internal (underscore-prefixed) attributes/methods (_OTHER_MAPPING,_weight_name_mapping_mlp,_weight_to_mcore_format), which are not part of any public API contract and could change/rename acrossmbridgereleases, silently breaking this override without an import error. Consider pinning thembridgedependency version and/or adding an assertion that these attributes exist at import time to fail fast if the upstream library changes shape.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@vime_plugins/mbridge/__init__.py` around lines 10 - 23, The current import-time override in mbridge/__init__.py monkey-patches underscore-prefixed internals on Qwen2MoEBridge and Qwen3MoEBridge, which is fragile against upstream mbridge changes. Update this patching logic to fail fast by asserting the expected private symbols exist before modifying them, and add a version pin or compatibility guard around the Qwen2MoEBridge/Qwen3MoEBridge setup so changes to _OTHER_MAPPING or related internals are caught immediately instead of silently breaking behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh`:
- Line 36: The NSYS default in the benchmark script points to the Nsight Compute
installation instead of the Nsight Systems binary used by nsys profile. Update
the NSYS fallback in the script’s environment setup to a valid Nsight Systems
path (for example, a bin/nsys location), so the existing checks in the profiling
flow succeed without requiring an override.
In `@vime-RLK-8gpu-single-op-runbook.md`:
- Line 12: Clarify the benchmark tracing guidance in the runbook so the official
no-trace path is explicit: in the section that mentions TRACE_MODE and the
benchmark flow, state that benchmark numbers should be collected with
TRACE_MODE=none, while TRACE_MODE=train|rollout|all are only for tracing/debug
runs. Update the wording around the TRACE_MODE setting to avoid implying all is
the default benchmark mode, and keep the guidance anchored to the TRACE_MODE
benchmark description in the runbook.
In `@vime-RLK-final-metrics-config-matrix.md`:
- Line 49: The matrix row for the trace setting is breaking because the raw pipe
characters in the `TRACE_MODE=train|rollout|all` cell are being parsed as column
separators. Update that cell in the markdown table to escape or rephrase the
pipes so the row stays a single cell and the rest of the matrix renders
correctly. Use the existing `trace` row as the target when making the edit.
In `@vime/backends/megatron_utils/update_weight/common.py`:
- Around line 130-132: The trainable-only filtering is incorrectly applied
inside the shared named_params_and_buffers helper, which also affects HF
checkpoint export paths. Move this logic out of the shared helper and add a
trainable_only flag to the weight-sync call sites in the update_weight flow so
only the sync path filters by requires_grad, while hf_checkpoint_saver.py
continues to receive all tensors.
---
Outside diff comments:
In `@tests/_unit_stubs.py`:
- Around line 169-206: The vLLM test stubs in install_vllm_cli_stubs are missing
required submodules, causing imports from vime.backends.vllm_utils.vllm_engine
and vime.backends.vllm_utils.arguments to fail when vllm is not installed.
Extend the existing stub setup to register vllm.utils.system_utils,
vllm.entrypoints.cli, and vllm.entrypoints.openai in sys.modules, and attach
them to the vllm module hierarchy alongside the existing vllm.utils and
vllm.engine stubs so the import tree resolves cleanly in the test environment.
- Around line 56-61: The rollout stubs currently only register a bare
vllm_router module, but vime.utils.arguments imports RouterArgs from
vllm_router.launch_router at import time. Update
install_rollout_optional_stubs() in tests/_unit_stubs.py to also stub the
vllm_router.launch_router submodule and provide a RouterArgs symbol so tests
that import vime.utils.arguments can run without the real package.
---
Nitpick comments:
In `@scripts/run-qwen3-30B-A3B.sh`:
- Around line 492-503: The train.py argument assembly in the
run-qwen3-30B-A3B.sh script is expanding several arrays unquoted, which can
trigger shell re-splitting and glob expansion. Update the invocation that uses
MODEL_ARGS, CKPT_ARGS, ROLLOUT_ARGS, OPTIMIZER_ARGS, GRPO_ARGS, WANDB_ARGS,
TB_ARGS, PERF_ARGS, EVAL_ARGS, VLLM_ARGS, MISC_ARGS, and RLK_ARGS so each array
expansion is quoted, preserving each element as a single argv entry even when
values contain spaces or bracket characters.
In `@tests/test_rl_kernel_linear_logp_integration.py`:
- Around line 30-31: The _FakeLinearLogpOp spy uses a shared mutable class
attribute for calls, which triggers the mutable default warning; update the
class attribute declaration in _FakeLinearLogpOp to be explicitly marked as a
class variable so the intended shared accumulator is clear to the linter and
readers, while keeping the existing clear/reset behavior intact.
In `@vime_plugins/mbridge/__init__.py`:
- Around line 10-23: The current import-time override in mbridge/__init__.py
monkey-patches underscore-prefixed internals on Qwen2MoEBridge and
Qwen3MoEBridge, which is fragile against upstream mbridge changes. Update this
patching logic to fail fast by asserting the expected private symbols exist
before modifying them, and add a version pin or compatibility guard around the
Qwen2MoEBridge/Qwen3MoEBridge setup so changes to _OTHER_MAPPING or related
internals are caught immediately instead of silently breaking behavior.
In `@vime/backends/megatron_utils/actor.py`:
- Around line 49-156: The profiler/env helper logic in _parse_capture_rollouts,
_should_nsys_capture, and _NsysCudaProfilerCapture is duplicated and already
diverging from the matching implementation in rollout.py. Move this shared
Nsight Systems capture logic into a common helper module and have actor.py and
rollout.py import it so the CUDA-availability check and rollout gating behavior
stay consistent in one place.
In `@vime/backends/megatron_utils/arguments.py`:
- Around line 163-165: The comment in _set_default_megatron_args is stale
because use_distributed_optimizer is no longer always enabled; update or remove
the comment so it matches the env-driven default set from
VIME_USE_DISTRIBUTED_OPTIMIZER. Keep the wording aligned with the actual
behavior in arguments.py and the _set_default_megatron_args function.
In `@vime/backends/megatron_utils/megatron_to_hf/qwen2.py`:
- Around line 60-67: The layer-norm rest-to-target mapping in Qwen2 is
duplicated in Qwen3MoE, so centralize this shared logic to prevent drift.
Extract the repeated set-membership handling from the converter method in
qwen2.py into a common helper such as _map_layernorm_rest(rest, layer_idx), then
have both qwen2.py and qwen3moe.py call that helper for the
input_layernorm/post_attention_layernorm cases.
In `@vime/backends/megatron_utils/model_provider.py`:
- Around line 275-307: The output-layer fallback in model_provider.py is too
broad: matched_output_layer currently uses a loose "output_layer" substring
check, which can incorrectly suppress the shared-weight training path in the
same block that uses wants_output_layer. Update the logic around the
matched_names scan and the shared_embedding_or_output_weight fallback to detect
the exact output_layer.weight parameter name instead of any output_layer-related
substring, so the tied embedding/output weight is enabled whenever the user
explicitly requested that target. While you are there, narrow the blanket except
around shared_embedding_or_output_weight() to the specific exceptions it can
raise, keeping the logger debug/warning behavior intact.
In `@vime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py`:
- Around line 137-175: The TP-size selection and full parameter size calculation
are duplicated in _bucket_size_bytes and _get_megatron_local_param_info_buckets,
which can drift over time. Extract the shared logic into a single helper used by
both functions, ideally one that takes a ParamInfo and returns its full byte
size based on the .experts. check and the appropriate mpu world-size accessor,
then update both call sites to use that helper.
- Around line 26-61: The per-bucket timing output in
hf_weight_iterator_direct.py is using a weight-sync-specific print prefix even
when this iterator is reused by the full checkpoint save path, so the logs are
misleading. Update the logging in the iterator method that loops over
megatron_local_param_info_buckets to use a proper logger or a context-aware
prefix that can reflect checkpoint-save usage as well as rollout sync, and keep
the bucket timing messages consistent with the call site.
In `@vime/backends/vllm_utils/vllm_engine.py`:
- Around line 454-461: The profiler control requests in start_profile and
stop_profile currently post without a timeout, which can leave ray.get waiting
forever if the vLLM engine stalls. Update the requests.post calls in
vllm_engine.py to include a finite timeout for both start_profile and
stop_profile, keeping the existing response.raise_for_status flow and the
returned {"ok": True} behavior unchanged.
In `@vime/utils/rl_kernel.py`:
- Around line 6-7: The two-tier validation in rl_kernel is currently
indistinguishable because RL_KERNEL_SUPPORTED_OPS and RL_KERNEL_INTEGRATED_OPS
contain the same value, making the unsupported branch in
normalize_rl_kernel_args effectively dead. Update the rl_kernel module by adding
a brief clarifying comment near RL_KERNEL_SUPPORTED_OPS /
RL_KERNEL_INTEGRATED_OPS and parse_rl_kernel_ops / normalize_rl_kernel_args
explaining the intended difference between “parseable” and “backend-integrated”
ops so future additions go to the correct constant.
In `@vime/utils/train_metric_utils.py`:
- Around line 16-25: The `train_metric_utils` merge logic can silently overwrite
`perf/` entries when a metric name matches a timer-derived key. Update the
`timer_instance.log_dict()` / `timer_instance.metric_dict()` handling so
`log_dict` construction in `train_metric_utils` avoids collisions, either by
using a stricter naming convention for `perf/{key}_time` versus `perf/{key}` or
by asserting that the two key sets are disjoint before merging. Keep the fix
localized around the `Timer`-based dictionary assembly and ensure the final
`log_dict` cannot contain duplicate `perf/` keys.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d2c712a6-04fb-4248-93ed-dc9bb113800d
⛔ Files ignored due to path filters (1)
vime-RLK-single-op-performance-visualization.svgis excluded by!**/*.svg
📒 Files selected for processing (36)
scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.shscripts/models/qwen3-30B-A3B.shscripts/run-qwen3-30B-A3B.shtests/_unit_stubs.pytests/test_cuda_event_timer.pytests/test_rl_kernel_args.pytests/test_rl_kernel_linear_logp_integration.pytests/test_rl_kernel_logp_integration.pytests/utils/test_vllm_engine.pyvime-RLK-8gpu-single-op-runbook.mdvime-RLK-final-metrics-config-matrix.mdvime-RLK-long-run-practice-report.mdvime-RLK.mdvime/backends/megatron_utils/actor.pyvime/backends/megatron_utils/arguments.pyvime/backends/megatron_utils/baseline_timer.pyvime/backends/megatron_utils/cuda_event_timer.pyvime/backends/megatron_utils/data.pyvime/backends/megatron_utils/loss.pyvime/backends/megatron_utils/megatron_to_hf/qwen2.pyvime/backends/megatron_utils/megatron_to_hf/qwen3moe.pyvime/backends/megatron_utils/model.pyvime/backends/megatron_utils/model_provider.pyvime/backends/megatron_utils/rl_kernel.pyvime/backends/megatron_utils/update_weight/common.pyvime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.pyvime/backends/megatron_utils/update_weight/update_weight_from_tensor.pyvime/backends/vllm_utils/vllm_engine.pyvime/ray/actor_group.pyvime/ray/rollout.pyvime/utils/arguments.pyvime/utils/memory_utils.pyvime/utils/rl_kernel.pyvime/utils/timer.pyvime/utils/train_metric_utils.pyvime_plugins/mbridge/__init__.py
| SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" | ||
| VIME_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" | ||
| VENV="${VIME_PYTHON_ENV:-${WORKSPACE_ROOT}/vime-rlk-env}" | ||
| NSYS="${NSYS:-/opt/nvidia/nsight-compute/2024.1.1/host/target-linux-x64/nsys}" |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
NSYS default points at the Nsight Compute tree, not Nsight Systems.
The script runs nsys profile, but the nsys binary ships with Nsight Systems (e.g. /opt/nvidia/nsight-systems/<ver>/bin/nsys or /usr/local/cuda/bin/nsys); the nsight-compute directory ships ncu. With TRACE_MODE != none, this default fails the -x check at Line 293 unless NSYS is overridden. Point the default at a Nsight Systems path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh` at line 36, The
NSYS default in the benchmark script points to the Nsight Compute installation
instead of the Nsight Systems binary used by nsys profile. Update the NSYS
fallback in the script’s environment setup to a valid Nsight Systems path (for
example, a bin/nsys location), so the existing checks in the profiling flow
succeed without requiring an override.
| - 不设置 `LOAD_DEBUG_ROLLOUT_DATA`，不走 debug train-only。 | ||
| - 先跑 candidate full-gradient 找最大不 OOM，再补 baseline。 | ||
| - 每个成功配置必须完成 rollout 0-11；取 rollout 3-11 的均值填表。 | ||
| - trace 用 `TRACE_MODE=train|rollout|all`；默认 `all`，需要少 overhead 时才改成 `train` 或 `none`。 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Clarify the no-trace benchmark path.
TRACE_MODE=all reads like the default run mode, but the benchmark flow uses TRACE_MODE=none for official numbers. Make that explicit so trace overhead doesn’t get mixed into performance comparisons.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vime-RLK-8gpu-single-op-runbook.md` at line 12, Clarify the benchmark tracing
guidance in the runbook so the official no-trace path is explicit: in the
section that mentions TRACE_MODE and the benchmark flow, state that benchmark
numbers should be collected with TRACE_MODE=none, while
TRACE_MODE=train|rollout|all are only for tracing/debug runs. Update the wording
around the TRACE_MODE setting to avoid implying all is the default benchmark
mode, and keep the guidance anchored to the TRACE_MODE benchmark description in
the runbook.
| | prompt data | `/workspace/dapo-math-17k/dapo-math-17k.jsonl` | | ||
| | baseline timer | `VIME_BASELINE_LINEAR_LOGP_TIMER=1`, `VIME_BASELINE_CUDA_EVENT_TIMER=1` | | ||
| | candidate timer | `VIME_RL_KERNEL=1`, `VIME_RL_KERNEL_LINEAR_LOGP_BACKEND=cuda`, `VIME_RL_KERNEL_CUDA_EVENT_TIMER=1` | | ||
| | trace | `TRACE_MODE=train|rollout|all`，默认 `all`；actor train 用 CUDA profiler range，rollout 用 vLLM CUDA profiler endpoint | |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Escape the | characters in this table cell.
TRACE_MODE=train|rollout|all splits the row into extra columns, so the matrix renders incorrectly and drops the rest of the data. Replace the raw pipes or rephrase the cell.
🛠️ Suggested fix
-| trace | `TRACE_MODE=train|rollout|all`，默认 `all`；actor train 用 CUDA profiler range，rollout 用 vLLM CUDA profiler endpoint |
+| trace | `TRACE_MODE={train,rollout,all}`，默认 `all`；actor train 用 CUDA profiler range，rollout 用 vLLM CUDA profiler endpoint |📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| | trace | `TRACE_MODE=train|rollout|all`，默认 `all`；actor train 用 CUDA profiler range，rollout 用 vLLM CUDA profiler endpoint | | |
| | trace | `TRACE_MODE={train,rollout,all}`，默认 `all`；actor train 用 CUDA profiler range，rollout 用 vLLM CUDA profiler endpoint | |
🧰 Tools
🪛 markdownlint-cli2 (0.22.1)
[warning] 49-49: Spaces inside code span elements
(MD038, no-space-in-code)
[warning] 49-49: Table column count
Expected: 2; Actual: 4; Too many cells, extra data will be missing
(MD056, table-column-count)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vime-RLK-final-metrics-config-matrix.md` at line 49, The matrix row for the
trace setting is breaking because the raw pipe characters in the
`TRACE_MODE=train|rollout|all` cell are being parsed as column separators.
Update that cell in the markdown table to escape or rephrase the pipes so the
row stays a single cell and the rest of the matrix renders correctly. Use the
existing `trace` row as the target when making the edit.
Source: Linters/SAST tools
| if os.environ.get("VIME_SYNC_TRAINABLE_WEIGHTS_ONLY", "0") == "1": | ||
| ans = ((name, tensor) for name, tensor in ans if getattr(tensor, "requires_grad", False)) | ||
|
|
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🔴 Critical | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the shared helper and its call sites.
rg -n "named_params_and_buffers|VIME_SYNC_TRAINABLE_WEIGHTS_ONLY|megatron_local_weights" vime -S
# Show the relevant files with line numbers if present.
for f in \
vime/backends/megatron_utils/update_weight/common.py \
vime/backends/megatron_utils/hf_checkpoint_saver.py
do
if [ -f "$f" ]; then
echo "===== $f ====="
cat -n "$f" | sed -n '1,260p'
fi
doneRepository: RL-Align/vime
Length of output: 25552
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect only the relevant slices if the files are large.
for f in \
vime/backends/megatron_utils/update_weight/common.py \
vime/backends/megatron_utils/hf_checkpoint_saver.py
do
if [ -f "$f" ]; then
echo "===== outline: $f ====="
ast-grep outline "$f" --view expanded || true
fi
doneRepository: RL-Align/vime
Length of output: 1710
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# If checkpoint saver is elsewhere, locate all callers of the helper.
rg -n "named_params_and_buffers\(" vime -SRepository: RL-Align/vime
Length of output: 1164
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the relevant code slices only.
for f in \
vime/backends/megatron_utils/update_weight/common.py \
vime/backends/megatron_utils/hf_checkpoint_saver.py \
scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh
do
if [ -f "$f" ]; then
echo "===== $f ====="
wc -l "$f"
fi
done
echo "===== common.py slice ====="
sed -n '1,240p' vime/backends/megatron_utils/update_weight/common.py | cat -n
echo "===== hf_checkpoint_saver.py slice ====="
sed -n '1,260p' vime/backends/megatron_utils/hf_checkpoint_saver.py | cat -n
echo "===== benchmark script slice ====="
sed -n '1,260p' scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh | cat -nRepository: RL-Align/vime
Length of output: 33794
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
files = [
Path("vime/backends/megatron_utils/update_weight/common.py"),
Path("vime/backends/megatron_utils/hf_checkpoint_saver.py"),
Path("scripts/benchmarks/run-qwen3-30B-A3B-8gpu-rlk-12rollout.sh"),
]
for p in files:
print(f"=== {p} exists={p.exists()} ===")
if p.exists():
text = p.read_text()
for needle in ["named_params_and_buffers", "VIME_SYNC_TRAINABLE_WEIGHTS_ONLY", "megatron_local_weights", "TRAIN_SCOPE", "output_layer"]:
if needle in text:
print(f"found {needle}")
PYRepository: RL-Align/vime
Length of output: 581
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find where the HF saver is invoked, to see whether it runs in the same process as the sync paths.
rg -n "save_hf_model_direct\(|hf_checkpoint_saver|save_hf" vime -SRepository: RL-Align/vime
Length of output: 1362
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the actor call site that uses named_params_and_buffers.
sed -n '200,260p' vime/backends/megatron_utils/actor.py | cat -nRepository: RL-Align/vime
Length of output: 3144
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect any direct references in scripts/config for raw HF saves.
rg -n "save_hf_model_direct|--save-hf|save_hf" . -SRepository: RL-Align/vime
Length of output: 1940
Scope the trainable-only filter to the sync path
named_params_and_buffers also feeds hf_checkpoint_saver.py, so VIME_SYNC_TRAINABLE_WEIGHTS_ONLY=1 will drop frozen backbone/expert tensors from raw HF checkpoint exports too. Pass a trainable_only flag from the weight-sync call sites instead of gating this shared helper globally.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@vime/backends/megatron_utils/update_weight/common.py` around lines 130 - 132,
The trainable-only filtering is incorrectly applied inside the shared
named_params_and_buffers helper, which also affects HF checkpoint export paths.
Move this logic out of the shared helper and add a trainable_only flag to the
weight-sync call sites in the update_weight flow so only the sync path filters
by requires_grad, while hf_checkpoint_saver.py continues to receive all tensors.
2 participants
Summary
This PR collects the RL-Kernel
linear_logp8xH100 experiment support and documentation from the T01/T03/T06 runs.Included intentionally, preferring not to miss related files:
T07 status
T07 was checked before opening this PR, but it cannot safely run right now. NVML reports external
[Not Found]compute apps occupying about 41GB on every H100:2195449-2195456kill -9from the container returnsNo such process.nvidia-smi --gpu-reset -i 0,1,2,3,4,5,6,7returnsNot Supported.Starting T07 in this state would almost certainly OOM, so no new T07 metrics are included.
Validation
python -m py_compileon changed Python filesbash -non changed benchmark/run scriptsvime-RLK-single-op-performance-visualization.svgGPU/integration validation is blocked by the external GPU memory holders above.
Summary by CodeRabbit
New Features
linear_logp, including new runtime controls, profiling hooks, and performance/memory tracking.Bug Fixes
Tests
Documentation