Merge samithuang/slime rebase branch into vime#8
Conversation
![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 adds vLLM as a rollout backend in Slime, providing backend-agnostic adapters and supporting weight synchronization via CUDA IPC and NCCL. It also introduces a new SlimeRouter and a Megatron patch for chunked gradient coalescing. Review feedback highlights a potential AttributeError in weight updates, a hardcoded override of the rollout offload flag, and the removal of important configuration validation checks. Furthermore, the feedback suggests improving the router's robustness against invalid JSON and division-by-zero errors, and fixing an ineffective variable reassignment in the weight distribution logic.
| ray.get(self.rollout_engine_lock.release.remote()) | ||
| if pbar is not None: | ||
| pbar.update(1) | ||
| pbar.update(1) |
There was a problem hiding this comment.
The explicit check if pbar is not None: was removed in this change. Since pbar defaults to None in the function signature and is not initialized in some calling paths (e.g., the use_vllm_packed branch in update_weights), this will cause an AttributeError when attempting to call .update(1).
| pbar.update(1) | |
| if pbar is not None: | |
| pbar.update(1) |
| args.offload_train = False | ||
| if args.offload_rollout is None: | ||
| args.offload_rollout = False | ||
| args.offload_rollout = False |
There was a problem hiding this comment.
This hardcoded assignment to False silently overrides the user-provided configuration for --offload-rollout. If the intention was to disable this feature temporarily, it should be handled by setting a default value or providing a warning to the user, rather than a silent override in the validation logic.
| # if hasattr(hf_config, hf_config_name) and hasattr(args, megatron_config_name): | ||
| # if not compare_fn(getattr(hf_config, hf_config_name), getattr(args, megatron_config_name)): | ||
| # errors.append( | ||
| # f"{hf_config_name} in hf config {getattr(hf_config, hf_config_name)} is not equal to " | ||
| # f"{megatron_config_name} {getattr(args, megatron_config_name)}, please check the config." | ||
| # ) |
There was a problem hiding this comment.
The validation logic for ensuring consistency between HuggingFace configuration and Megatron arguments has been commented out. This disables important safety checks that prevent model initialization with mismatched parameters (e.g., layer norm epsilon). These checks should be restored or properly integrated with the new vLLM backend logic.
| max_connections = getattr(args, "slime_router_max_connections", None) | ||
| if max_connections is None: | ||
| max_connections = ( | ||
| args.rollout_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine |
There was a problem hiding this comment.
There is a potential division by zero risk if args.rollout_num_gpus_per_engine is set to 0. While this might be an invalid configuration, adding a guard or a minimum value check during argument validation would improve robustness.
| # 2) Fallback to JSON body | ||
| if not worker_url: | ||
| body = await request.body() | ||
| payload = json.loads(body) if body else {} |
There was a problem hiding this comment.
Calling json.loads(body) without a try-except block will cause the router to crash (returning a 500 error) if the request body is not valid JSON. It is recommended to handle json.JSONDecodeError and return a 400 Bad Request response instead.
| if isinstance(model_update_groups, _NcclBridge): | ||
| model_update_groups.shutdown() | ||
| elif model_update_groups is not None: | ||
| model_update_groups = None |
There was a problem hiding this comment.
Reassigning the local variable model_update_groups to None has no effect on the caller's state. If the intention is to clear the reference in the calling object, this must be done explicitly by the caller or by returning None.
|
I recommend splitting it and landing Megatron ↔ vLLM weight sync first, then glue/args, then docs/scripts. Suggested PR split
Acceptance for P0: smoke SlimeRouter — recommend deferring (or dropping from this PR)Upstream slime has already deprecated The weight-sync work (especially non-colocate |
|
|
||
| # Resolve vLLM base URLs for colocated engines (blocking Ray call) | ||
| if dist.get_rank() == 0 and self._colocated_engines: | ||
| url_refs = [engine.get_vllm_url.remote() for engine in self._colocated_engines] |
There was a problem hiding this comment.
VLLMEngine exposes get_url(), not get_vllm_url(). Colocated IPC connect will raise AttributeError. Please use get_url.remote() or add an alias on VLLMEngine.
| logger = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| class UpdateVLLMWeightFromTensor: |
There was a problem hiding this comment.
MegatronTrainRayActor.sleep() calls disconnect_rollout_engines() when present. UpdateVLLMWeightFromTensor never implements it, so _NcclBridge / NCCL groups for overflow engines may leak across sleep/wake.
| # Bump weight version on sidecar for colocated engines | ||
| for engine in self._colocated_engines: | ||
| try: | ||
| ray.get(engine.set_weight_version.remote(self.weight_version)) |
There was a problem hiding this comment.
VLLMEngine has no set_weight_version (comment still says "sidecar"). This call always fails silently. Implement on VLLMEngine or remove.
|
Will break into smaller PRs starting from #22 |
aoshen02
force-pushed
the
merge/samithuang-rebase
branch
from
a38186e to
ebcc57c
Compare
|
Close this and will continuous update it with main branch |
4 participants
Summary
Merges samithuang/slime PR #12 (rebase branch, commit
9146f9f) into vime. Both branches forked from upstream slime8ef1fb4and independently added vLLM support; this PR reconciles them so vime retains its vLLM design (kaiyuan) while picking up samithuang's broader infrastructure (Jason & Long): non-colocate Megatron+vLLM weight sync, slime router, Megatron grad-coalesce patch, retool fix, e2e scriptsPer-file precedence
Kept verbatim from vime (kaiyuan vLLM core)
slime/backends/vllm_utils/vllm_engine.py(707-line monolithic engine)slime/backends/vllm_utils/__init__.py(empty)slime/rollout/vllm_rollout.py(992-line standalone rollout)slime/backends/sglang_utils/sglang_engine.pyslime/rollout/sglang_rollout.py(kept upstream; samithuang's backend-abstraction refactor dropped)Dropped from samithuang (incompatible with vime's design)
slime/backends/vllm_utils/vllm_translation_sidecar.pyslime/rollout/backends/{__init__,base_client,sglang_client,vllm_client}.pyTaken from samithuang (Jason & Long infrastructure)
slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py(colocate + non-colocate weight sync)slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.pyslime/backends/megatron_utils/weight_sync_utils.pyslime/backends/megatron_utils/megatron_patch/*(chunked grad coalesce)slime/backends/megatron_utils/{__init__,actor,arguments,sglang}.py,update_weight/common.pyslime/ray/actor_group.py,slime/rollout/base_types.py,slime/backends/sglang_utils/arguments.pyslime/router/router.py(SlimeRouter, gated by--use-slime-router)run-qwen2.5-0.5B-vllm.sh,run-qwen2.5-0.5B-reproducibility-noncolocate.sh,convert_qwen2.5_ckpt.sh,setup_for_vllm.mdexamples/retool+tests/plugin_contractsupdatesManual splices
slime/ray/rollout.py: based on samithuang. Vime's_sanitize_vllm_router_args,_vllm_router_args_from_cli, themodel_pathkwarg passed to engine actors, and the engine-shutdown loop inRolloutManager.dispose()are preserved._start_routerwas rewritten to support three modes:--use-slime-router(SlimeRouter),--rollout-backend vllm(vllm-router via http_utils), and sglang-router fallback. Samithuang's_start_vllm_rollout_servershelper and its early-exit instart_rollout_serverswere removed because they target a differentVLLMEnginesignature; vime'sServerGroup.start_enginesalready dispatches toVLLMEnginewhenrollout_backend == "vllm".slime/utils/arguments.py: union of both arg sets with duplicates removed (--rollout-backend,--vllm-gpu-memory-utilization,--vllm-weight-sync-packed). Kept samithuang's--rollout-backenddefault ofsglangfor backwards compat and vime's mutually-exclusive--vllm-weight-sync-packed/--no-vllm-weight-sync-packedgroup. Vime's auto-switch ofrollout_function_pathtovllm_rolloutwhen--rollout-backend vllmis preserved.slime/utils/http_utils.py: vime'srun_router((impl, router_args))payload contract kept; samithuang'ssglang_server_concurrency->rollout_server_concurrencyrename applied.slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py: vime's 3-line addition (import_is_vllm_backend, passuse_vllm=...andpacked=False) kept; it correctly wires the existing tensor-update path through samithuang's new vllm-aware signature.Test plan
python -m py_compileover all touched.pyfiles (already green locally).run-qwen2.5-0.5B-vllm.sh(--rollout-backend vllm --use-slime-router, dense Qwen2.5-0.5B on GSM8K).run-qwen2.5-0.5B-reproducibility-noncolocate.sh(non-colocate Megatron+vLLM weight sync).scripts/run-deepseek-r1.sh) to confirm no regression on the default backend.Known follow-ups (deferred)
docs/en/vllm/ROUTER_DESIGN.mdandrfc-rollout-backend-separation-plan.mdstill mention the droppedvllm_translation_sidecar.py; update once the design has settled.slime/router/router.pycoexist (selected by--rollout-backend vllmvs--use-slime-router); a future cleanup could unify them.