Skip to content

[feat] Support Colocated Weight Sync via CUDA IPC for vime - #22

Merged
aoshen02 merged 4 commits into
mainfrom
update_weights_tensor
May 27, 2026
Merged

[feat] Support Colocated Weight Sync via CUDA IPC for vime#22
aoshen02 merged 4 commits into
mainfrom
update_weights_tensor

Conversation

@knlnguyen1802

@knlnguyen1802 knlnguyen1802 commented May 21, 2026

Copy link
Copy Markdown
Collaborator

co-author: @SamitHuang

Purpose

Enable colocated weight synchronization between the Megatron trainer and vLLM rollout engines running on the same GPU(s), using CUDA IPC (Inter-Process Communication) via Ray instead of NCCL distributed broadcast. This avoids cross-node traffic and reduces weight sync latency when trainer and inference engine share the same physical GPUs.


What is Changed

New: UpdateVLLMWeightFromTensor (slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py)

A new weight-update class following the vLLM RLHF IPC approach. It handles the full colocated sync lifecycle:

  1. Megatron parameters are converted to HF format via HfWeightIteratorBase.
  2. All trainer ranks call IPCWeightTransferEngine.trainer_send_weights(send_mode="ray"), creating a CUDA IPC handle per GPU.
  3. Handles are all-gathered and merged so every vLLM worker can pick the handle matching its physical GPU UUID.
  4. For non-colocated overflow engines, the existing NCCL distributed broadcast path (update_weights_from_distributed) is preserved unchanged.

The per-update_weights call lifecycle is:

release_memory_occupation(level=0)   # free KV cache + model weights
init_weight_transfer_engine          # first call only
start_weight_update
  [for each HF weight chunk]
    trainer_send_weights (IPC)
    update_weights_from_distributed  # overflow/distributed engines
finish_weight_update
resume_memory_occupation

Modified: slime/backends/vllm_utils/vllm_engine.py

  • update_weights(update_info) — new public Ray-callable entry point. Since ipc_handles are Python callables (closures from monkey_patch_torch_reductions) that cannot be JSON-serialized, they are serialized with cloudpickle and base64-encoded into ipc_handles_pickled, which the vLLM IPCWeightTransferUpdateInfo accepts when VLLM_ALLOW_INSECURE_SERIALIZATION=1 is set.
  • release_memory_occupation(level=1) — now accepts a level parameter. level=0 releases both KV cache and model weights (required before IPC tensor injection); level=1 (default, unchanged) releases KV cache only.
  • init_weight_transfer_engine(payload) — new method, posts to /init_weight_transfer_engine with retry logic (3 attempts with back-off).
  • start_weight_update(is_checkpoint_format) — new method, posts to /start_weight_update to enter IPC weight-update mode.
  • finish_weight_update() — new method, posts to /finish_weight_update to exit IPC weight-update mode.
  • Auto-config when colocate=True: automatically injects --weight-transfer-config '{"backend":"ipc"}' and --worker-extension-cls slime.backends.vllm_utils.vllm_worker_extension.vLLMColocateWorkerExtension into the vLLM serve command, unless the user has already overridden them.

New: slime/backends/vllm_utils/vllm_worker_extension.py

Introduces vLLMColocateWorkerExtension, passed to vllm serve via --worker-extension-cls. On instantiation inside each vLLM worker process, it applies _VLLMHijack.hijack(), which monkey-patches IPCWeightTransferEngine.receive_weights to call monkey_patch_torch_reductions() before deserializing CUDA IPC handles. The patch is idempotent and applied automatically without requiring explicit patching from the trainer side.

Modified: slime/backends/megatron_utils/actor.py

Minor update to wire UpdateVLLMWeightFromTensor into the actor's weight-update dispatch.


Test

A comprehensive unit test suite is added in tests/test_update_weight_from_tensor_vllm.py :

  • All tests run without GPU, CUDA, or any heavy training framework (Megatron, vLLM, Ray). All heavy dependencies are injected as lightweight stubs before the module under test is loaded.
  • UpdateVLLMWeightFromTensor is instantiated via a helper that directly sets instance attributes, bypassing GPU-requiring initialization paths (HfWeightIteratorBase.create(), etc.).
  • Remote Ray calls are intercepted by RecordingRemoteMethod / RecordingEngine / RecordingIPCEngine objects, enabling assertions on call arguments without touching any CUDA primitives.
  • Tests cover: IPC handle serialization, colocated vs. distributed engine dispatch, lifecycle ordering (release → init → start → send → finish → resume), and retry behavior on init_weight_transfer_engine failures.

Result validate for run Qwen3-4B under 4 card (TP2, 4 vllm engine and 4 actor train)
eval_trends
raw_reward_trend
train_trends

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for updating colocated vLLM rollout engines using CUDA IPC (Ray mode), primarily through the new UpdateVLLMWeightFromTensor class and associated vLLM engine configurations. The implementation handles both colocated IPC transfers and distributed NCCL fallbacks, supported by a new worker extension for internal patching. Feedback focuses on performance optimizations, specifically recommending batching engine handles in trainer_send_weights to minimize synchronization overhead and parallelizing remote initialization calls. Other suggestions include removing redundant environment variable settings and cleaning up debug code.

Comment thread slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py Outdated
Comment thread slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py Outdated
Comment thread slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py Outdated
Comment thread slime/backends/vllm_utils/vllm_worker_extension.py Outdated
Comment thread tests/test_update_weight_from_tensor.py Outdated
Comment thread slime/backends/vllm_utils/vllm_worker_extension.py Outdated
Comment thread slime/backends/vllm_utils/vllm_worker_extension.py Outdated
Comment thread slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py Outdated
@CalvinXKY

Copy link
Copy Markdown
Collaborator

If we only call start_weight_update / finish_weight_update from rank 0 and do not barrier after each HF chunk, faster ranks can finish an engine while slower ranks are still sending update_weights, which triggers RuntimeError: start_weight_update must be called before update_weights (HTTP 500) on later training steps.

^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/executor/uniproc_executor.py", line 93, in collective_rpc [repeated 2x across cluster]
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]     return self.model_executor.collective_rpc(method, timeout, args, kwargs)
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]     result = run_method(self.driver_worker, method, args, kwargs)
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/serial_utils.py", line 510, in run_method
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]     return func(*args, **kwargs)
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]            ^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/worker/gpu_worker.py", line 1040, in update_weights
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341]     raise RuntimeError(
(VLLMEngine pid=1256031) (EngineCore pid=1256895) ERROR 05-22 07:20:17 [core.py:1341] RuntimeError: start_weight_update must be called before update_weights.
(VLLMEngine pid=1256031) (APIServer pid=1256559) ERROR 05-22 07:20:17 [server_utils.py:377] Exception caught. Request id: None
(VLLMEngine pid=1256031) (APIServer pid=1256559) INFO:     10.155.96.7:36304 - "POST /update_weights HTTP/1.1" 500 Internal Server Error
(MegatronTrainRayActor pid=1258985) [2026-05-22 07:20:14] memory_utils.py:47 - [Rank 1] Memory-Usage after offload model: {'gpu': '1', 'total_GB': 79.25, 'free_GB': 72.67, 'used_GB': 6.58, 'allocated_GB': 5.46, 'reserved_GB': 6.86, 'host_total_GB': 1007.16, 'host_available_GB': 846.98, 'host_used_GB': 160.19, 'host_free_GB': 464.33} [repeated 3x across cluster]
(MegatronTrainRayActor pid=1258984) [2026-05-22 07:20:17] reloadable_process_group.py:165 - Reloading 6 process groups in pid 1258984 [repeated 3x across cluster]
(MegatronTrainRayActor pid=1258984) [2026-05-22 07:20:17] memory_utils.py:47 - [Rank 3] Memory-Usage before update_weights: {'gpu': '3', 'total_GB': 79.25, 'free_GB': 71.64, 'used_GB': 7.62, 'allocated_GB': 5.45, 'reserved_GB': 6.49, 'host_total_GB': 1007.16, 'host_available_GB': 846.99, 'host_used_GB': 160.17, 'host_free_GB': 464.35} [repeated 3x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559) ERROR:    Exception in ASGI application
(VLLMEngine pid=1256031) (APIServer pid=1256559) Traceback (most recent call last):
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/uvicorn/protocols/http/httptools_impl.py", line 421, in run_asgi
(VLLMEngine pid=1256031) (APIServer pid=1256559)     result = await app(  # type: ignore[func-returns-value]
(VLLMEngine pid=1256031) (APIServer pid=1256559)              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/starlette/routing.py", line 716, in __call__ [repeated 11x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)     return await self.app(scope, receive, send)
(VLLMEngine pid=1256031) (APIServer pid=1256559)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await super().__call__(scope, receive, send)
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await self.middleware_stack(scope, receive, send) [repeated 2x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)     raise exc [repeated 4x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await self.app(scope, receive, _send)
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await self.app(scope, receive, send) [repeated 3x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await self.app(scope, receive, send_wrapper)
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await wrap_app_handling_exceptions(self.app, conn)(scope, receive, send)
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/starlette/_exception_handler.py", line 42, in wrapped_app [repeated 4x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await app(scope, receive, sender) [repeated 2x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/fastapi/routing.py", line 674, in app [repeated 4x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await route.handle(scope, receive, send)
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/starlette/routing.py", line 290, in handle
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await wrap_app_handling_exceptions(app, request)(scope, receive, send)
(VLLMEngine pid=1256031) (APIServer pid=1256559)     response = await f(request)
(VLLMEngine pid=1256031) (APIServer pid=1256559)                ^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (APIServer pid=1256559)     raw_response = await run_endpoint_function(
(VLLMEngine pid=1256031) (APIServer pid=1256559)                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/fastapi/routing.py", line 328, in run_endpoint_function
(VLLMEngine pid=1256031) (APIServer pid=1256559)     return await dependant.call(**values)
(VLLMEngine pid=1256031) (APIServer pid=1256559)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ [repeated 2x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/async_llm.py", line 1101, in update_weights [repeated 2x across cluster]
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await engine_client(raw_request).update_weights(
(VLLMEngine pid=1256031) (APIServer pid=1256559)     await self.collective_rpc(
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/async_llm.py", line 970, in collective_rpc
(VLLMEngine pid=1256031) (APIServer pid=1256559)     return await self.engine_core.collective_rpc_async(
(VLLMEngine pid=1256031) (APIServer pid=1256559)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core_client.py", line 1132, in collective_rpc_async
(VLLMEngine pid=1256031) (APIServer pid=1256559)     return await self.call_utility_async(
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core_client.py", line 1039, in call_utility_async
(VLLMEngine pid=1256031) (APIServer pid=1256559)     return await self._call_utility_async(method, *args, engine=self.core_engine)
(VLLMEngine pid=1256031) (APIServer pid=1256559)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(VLLMEngine pid=1256031) (APIServer pid=1256559)   File "/usr/local/lib/python3.12/dist-packages/vllm/v1/engine/core_client.py", line 1053, in _call_utility_async
(VLLMEngine pid=1256031) (APIServer pid=1256559)     return await future
(VLLMEngine pid=1256031) (APIServer pid=1256559)            ^^^^^^^^^^^^
(VLLMEngine pid=1256031) (APIServer pid=1256559) Exception: Call to collective_rpc method failed: start_weight_update must be called before update_weights.

@CalvinXKY

Copy link
Copy Markdown
Collaborator

After I made a few modifications, the issue with step1 was resolved:

image

# engines (not per update call).
self._ipc_initialized: bool = False
# vLLM IPC handle payloads may use cloudpickle on the Ray/HTTP bridge.
os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's already set in slime/backends/vllm_utils/vllm_engine.py when colocate=True

# engines (not per update call).
self._ipc_initialized: bool = False
# vLLM IPC handle payloads may use cloudpickle on the Ray/HTTP bridge.
os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's already set in slime/backends/vllm_utils/vllm_engine.py when colocate=True

Comment thread slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py Outdated
trainer_args=trainer_args,
)

if self._distributed_engines and self._is_distributed_src_rank:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's weird to put distributed related code in ipc file, I think it's the legacy design from slime. This case does exist, do you think it's better to move the code to distributed files?

@knlnguyen1802 knlnguyen1802 May 25, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should keep it, it's the case that run muti node but within 1 node the weight sync can still be colocated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But I thought we were talking about a case with both colocate and distributed elements. Wouldn't it be better to separate the distributed code into its own file rather than mixing everything together, and then simply invoke it here?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, it's the case of mixing colocated and distributed where weight within same node will be send and update first, then weight is send to distributed node.
But the problem is that slime will only use 1 mode either non-colocated and colocated, where non-colocated can handle this case smoothly and for colocated mode, it'll need to invoke both kind of code

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why it would need to invoke both?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why it would need to invoke both?

I believe it's for multi nodes training but we havent test it yet

aoshen02 added a commit that referenced this pull request May 26, 2026
Brings in vLLMColocateWorkerExtension + IPC weight-transfer-config backend
+ full start/finish_weight_update colocate bracket. Supersedes the
session-bracket-only fix in b59d4d8 (kept locally for history; the
PR #22 path is the load-bearing one for actual colocate IPC).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

# Conflicts:
#	requirements.txt
#	slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
#	slime/backends/vllm_utils/vllm_engine.py
aoshen02 added a commit that referenced this pull request May 26, 2026
…T_CONFIG_YAML

Round-2 rename sweep after PR #18 (and the earlier `7aead3d` round):

Docs:
  docs/{en,zh}/advanced/sglang-config.md   → vllm-config.md (git mv)
  docs/_static/image/sglang_config.png    → vllm_config.png (regenerated:
    redrawn with matplotlib so the three embedded labels — "data generation
    with arbitary sglang deployment", and "sglang router" twice — are now
    "data generation with arbitrary vllm deployment" and "vllm router";
    "arbitary" typo also fixed)
  Update all references in docs/{en,zh}/index.rst, usage.md, megatron-config.md,
  zh/advanced/slime_vllm_backend_design_v1.md.
  Inside vllm-config.md: --sglang-config→--vllm-config, YAML key sglang:→vllm:,
  args.sglang_model_routers→args.vllm_model_routers, slime.rollout.sglang_rollout
  →slime.rollout.vllm_rollout, sglang_basic.yaml→vllm_basic.yaml (and the
  three other sample filenames), "SGLang ServerArgs"→"vLLM EngineArgs",
  "SGLang Model Gateway (sgl-router)"→"vllm-router",
  "python -m sglang.launch_server"→"vllm serve".

docs/{en,zh}/advanced/reproducibility.md:
  # sglang config              →  # vLLM config
  --sglang-enable-deterministic-inference  →  --vllm-enable-deterministic-inference
  --sglang-attention-backend flashinfer    →  --vllm-attention-backend flashinfer

docs/en/blogs/introducing_slime.md: "slime exclusively integrates SGLang"
  → "slime/vime exclusively integrates vLLM". (Same edit in zh blog.)

Tests:
  4 files (test_qwen2.5_0.5B_vllm_config{,_distributed}.py,
  test_vllm_config_mixed_offload{,_ft}.py): ROLLOUT_CONFIG_YAML→VLLM_CONFIG_YAML
  + "Inline rollout config"→"Inline vLLM config".

slime/utils/arguments.py:
  10 help strings on top-level rollout flags (--colocate description,
  --rollout-temperature/top-p/top-k, max-context/prompt/response-len, stop
  words/token-ids) said "the inference engine" generically — now say "vLLM"
  to match the vLLM-only state of the repo.

slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py:
  one docstring line "all rollout engines" → "all vLLM engines" (this is the
  colocate IPC path, vLLM-specific; PR #22 already merged).

examples/search-r1/generate_with_search.py: one comment about strict token
  alignment of "the inference engine" → "vLLM".

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request May 26, 2026
The README B200 paragraph told users to pass `--sglang-mm-attention-backend
sdpa` to work around FA3 missing on Blackwell. Two problems:

1. The `--sglang-*` namespace was removed by PR #18, so following this
   advice now fails at parse_args.
2. The advice is unnecessary for vllm: vllm auto-dispatches the ViT encoder
   on Blackwell to FA4 (or FA2 fallback) without any user intervention.
   - vllm/v1/attention/backends/fa_utils.py:77-104 picks fa_version=4 when
     device_capability.major == 10 and explicitly rejects fa_version=3 on
     Blackwell with a warning.
   - vllm/platforms/cuda.py:404-440 (get_vit_attn_backend) iterates
     [FLASH_ATTN, TRITON_ATTN, TORCH_SDPA, FLASHINFER] and selects the
     first whose supports_compute_capability accepts (10, 0). That's
     FLASH_ATTN, which internally uses FA4.

Rewrote the paragraph to explain the vllm-native default and document
`--vllm-mm-encoder-attn-backend TORCH_SDPA` only as a manual escape hatch.
The HF-side `--attn-implementation flash_attention_2` note is kept since
it's still relevant when the model is loaded via Hugging Face Transformers.

Also adds pr18-post-fix-audit.md cataloging the remaining sglang residue
in gcl/clean-sglang at 6abfdbc, separating PR #18 fixes from PR #22
territory.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request May 26, 2026
PR #22's colocate-IPC weight transfer path calls
_apply_monkey_patch_torch_reductions (update_weight_from_tensor.py:247,319),
which lazy-imports `monkey_patch_torch_reductions` from
`slime.backends.megatron_utils.sglang`. PR #18 deleted that module —
the function now lives at
`slime.backends.megatron_utils.update_weight.torch_patch` (vendored from
sglang's patch_torch utility; same code, new home).

On PR #22 standalone this worked because the deletion hasn't landed on
main. On `gcl/clean-sglang` (= main + PR #18 + PR #22 merge), the
import raises `ModuleNotFoundError: No module named 'slime.backends.
megatron_utils.sglang'` the first time the colocate path is hit. This
shows up in every smoke test that uses `--colocate` (gb10-smoke
reproduces it within ~2 minutes of startup, right after Megatron loads
weights and tries to sync to vLLM).

Pointed both the production call site (update_weight_from_tensor.py:52)
and the test stub (test_update_weight_from_tensor.py:48-50) at the new
torch_patch module. The sibling
`hf_weight_iterator_direct.py:16` already uses
`from .torch_patch import monkey_patch_torch_reductions` post-PR-18, so
this just brings the second consumer into line.

The PR #22 author cannot land this fix on their own branch because
`torch_patch.py` does not exist on `main` yet — the rewire only makes
sense in the integrated PR #18 + PR #22 state, which is exactly
`gcl/clean-sglang`.

Audit §1A in pr18-post-fix-audit.md tracked this.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
@knlnguyen1802
knlnguyen1802 force-pushed the update_weights_tensor branch from 1d431bc to 132066f Compare May 26, 2026 09:04
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
@aoshen02

Copy link
Copy Markdown
Collaborator

Hi @knlnguyen1802 — found a test/production drift that should be cleaned up in this PR before merge.

What happened

Commit 47158c1 ("feat(colocate): vLLM 0.21 IPC weight sync for --colocate mode") in this PR deleted the leader-skip path that PR #21 (f7ac630) had introduced on VLLMEngine:

  • _SKIP_NON_LEADER constant — deleted
  • _skipped_if_not_leader() method — deleted
  • The PR-Sync to vLLM 0.21.0 #21 version of start_weight_update (was at line 569, with if skipped := self._skipped_if_not_leader(): return skipped) — replaced with this PR's version at line 797 (no leader-skip)
  • Same for finish_weight_update

That's a deliberate design choice — fine on its own. But PR #21 also added two tests in tests/unit/backends/vllm_utils/test_vllm_engine.py (commit blob acda0879) that assert leader-skip behavior:

tests/unit/backends/vllm_utils/test_vllm_engine.py::test_start_weight_update_skipped_when_node_rank_nonzero
tests/unit/backends/vllm_utils/test_vllm_engine.py::test_finish_weight_update_skipped_when_node_rank_nonzero

Both set engine.node_rank = 1 and monkeypatch.setattr(engine, "_post_json", lambda *a, **k: pytest.fail("should not POST")), then assert engine.start_weight_update() == {"ok": True, "skipped": True}. After 47158c1's production change, these tests fail with Failed: should not POST (the new start_weight_update posts unconditionally).

Verification

tests/unit/backends/vllm_utils/test_vllm_engine.py is the same blob (acda087) on both this PR's HEAD and main, and _skipped_if_not_leader is gone from this PR's slime/backends/vllm_utils/vllm_engine.py. Running pytest tests/unit/backends/vllm_utils/test_vllm_engine.py against this branch reproduces:

========================= 2 failed, 26 passed in 1.59s =========================
FAILED tests/unit/backends/vllm_utils/test_vllm_engine.py::test_start_weight_update_skipped_when_node_rank_nonzero
FAILED tests/unit/backends/vllm_utils/test_vllm_engine.py::test_finish_weight_update_skipped_when_node_rank_nonzero

Fix

Delete those two test functions (~lines 106–114 and 334–342 of test_vllm_engine.py) as part of this PR, with a one-line comment explaining the removal — e.g.:

# (PR #21's leader-skip path was removed in commit 47158c1 of this PR;
# vime engines are always single-node per actor, so node_rank is always 0
# and the skip branch was unreachable. Tests for that branch deleted.)

That keeps the test surface honest about what production actually does on this branch.

(I'd send a small PR against update_weights_tensor if you'd prefer that to a manual edit — let me know.)

@knlnguyen1802

Copy link
Copy Markdown
Collaborator Author

Hi @knlnguyen1802 — found a test/production drift that should be cleaned up in this PR before merge.

What happened

Commit 47158c1 ("feat(colocate): vLLM 0.21 IPC weight sync for --colocate mode") in this PR deleted the leader-skip path that PR #21 (f7ac630) had introduced on VLLMEngine:

  • _SKIP_NON_LEADER constant — deleted
  • _skipped_if_not_leader() method — deleted
  • The PR-Sync to vLLM 0.21.0 #21 version of start_weight_update (was at line 569, with if skipped := self._skipped_if_not_leader(): return skipped) — replaced with this PR's version at line 797 (no leader-skip)
  • Same for finish_weight_update

That's a deliberate design choice — fine on its own. But PR #21 also added two tests in tests/unit/backends/vllm_utils/test_vllm_engine.py (commit blob acda0879) that assert leader-skip behavior:

tests/unit/backends/vllm_utils/test_vllm_engine.py::test_start_weight_update_skipped_when_node_rank_nonzero
tests/unit/backends/vllm_utils/test_vllm_engine.py::test_finish_weight_update_skipped_when_node_rank_nonzero

Both set engine.node_rank = 1 and monkeypatch.setattr(engine, "_post_json", lambda *a, **k: pytest.fail("should not POST")), then assert engine.start_weight_update() == {"ok": True, "skipped": True}. After 47158c1's production change, these tests fail with Failed: should not POST (the new start_weight_update posts unconditionally).

Verification

tests/unit/backends/vllm_utils/test_vllm_engine.py is the same blob (acda087) on both this PR's HEAD and main, and _skipped_if_not_leader is gone from this PR's slime/backends/vllm_utils/vllm_engine.py. Running pytest tests/unit/backends/vllm_utils/test_vllm_engine.py against this branch reproduces:

========================= 2 failed, 26 passed in 1.59s =========================
FAILED tests/unit/backends/vllm_utils/test_vllm_engine.py::test_start_weight_update_skipped_when_node_rank_nonzero
FAILED tests/unit/backends/vllm_utils/test_vllm_engine.py::test_finish_weight_update_skipped_when_node_rank_nonzero

Fix

Delete those two test functions (~lines 106–114 and 334–342 of test_vllm_engine.py) as part of this PR, with a one-line comment explaining the removal — e.g.:

# (PR #21's leader-skip path was removed in commit 47158c1 of this PR;
# vime engines are always single-node per actor, so node_rank is always 0
# and the skip branch was unreachable. Tests for that branch deleted.)

That keeps the test surface honest about what production actually does on this branch.

(I'd send a small PR against update_weights_tensor if you'd prefer that to a manual edit — let me know.)

Hi , I think it's better if you can do a small PR again this. Thanks

aoshen02 added a commit that referenced this pull request May 26, 2026
Two regressions surfaced after PR #18 merged PR #22 (update_weights_tensor)
into gcl/clean-sglang:

1. ``slime/backends/vllm_utils/vllm_engine.py`` had duplicate
   ``def start_weight_update`` / ``def finish_weight_update`` definitions.
   PR #21 (#f7ac630, "Sync to vLLM 0.21.0") added a leader-aware version at
   ~line 569/581 with ``_skipped_if_not_leader()`` self-protection. PR #22's
   #47158c1 commit added its own no-skip version at ~line 820/834 without
   removing PR #21's. Python class-body semantics shadow the earlier
   definitions, so the no-skip version was the one running at runtime.

   The merge artifact didn't cause runtime crashes (vime engines always
   have ``self.node_rank = 0`` — only assignment in the codebase, line 459),
   but it left two defensive code paths dead-coded and confused review.

   Resolution: keep PR #22's design (drop leader-skip path entirely since
   ``node_rank`` is structurally 0 in current vime), remove the duplicates,
   and also remove the now-orphan ``_SKIP_NON_LEADER`` constant and
   ``_skipped_if_not_leader`` method. PR #22 author has been pinged
   (#22 comment) to also drop the corresponding
   ``test_*_skipped_when_node_rank_nonzero`` tests on their branch.

2. ``tests/test_update_weight_from_tensor.py::_make_instance`` (added on
   gcl/pr18-tests-ci via #b6a357a) bypasses ``__init__`` with
   ``object.__new__`` and hand-sets attributes. It was missing
   ``obj._ipc_engine = None`` and 12 tests that mutate
   ``obj._colocated_engines = [...]`` after ``_make_instance`` returned
   never wired ``_ipc_engine`` either. Production code at
   ``update_weight_from_tensor.py:237`` gates IPC lifecycle on
   ``self._ipc_engine is not None`` and the helper sets ``_ipc_engine`` to
   one designated engine in ``connect_rollout_engines`` — the tests need to
   replicate that designation.

   Resolution: add ``obj._ipc_engine = None`` to the ``_make_instance``
   base, plus ``obj._ipc_engine = obj._colocated_engines[0] if obj._colocated_engines else None``
   at every non-empty ``_colocated_engines`` mutation site (12 occurrences,
   batch-applied). Updated ``test_multiple_colocated_engines_all_get_lifecycle_calls``
   to match production: release/resume/init are fanned out to all engines,
   but start/finish_weight_update is sent only to the designated
   ``_ipc_engine`` (PR #22's coordinator-per-rank design).

Verified on h200-1 with vime-vllm-cu129-latest image (single GPU,
Qwen3-4B compatible):
- ``pytest tests/test_update_weight_from_tensor.py``: 29 passed
- ``pytest tests/unit/backends/vllm_utils/test_vllm_engine.py``: 26 passed,
  2 failed (``test_*_skipped_when_node_rank_nonzero`` — owner-pinged via
  #22).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request May 26, 2026
Commit 3b12c77 on this branch consolidated start_weight_update /
finish_weight_update to PR #22's no-leader-skip version (dropping the
PR #21 defensive guard that vime's node_rank=0 architecture never
exercises). The two tests below were left behind asserting the now-
removed self-skip behavior:

- test_start_weight_update_skipped_when_node_rank_nonzero
- test_finish_weight_update_skipped_when_node_rank_nonzero

Both set node_rank=1 and assert _post_json is never called; after the
3b12c77 dedup, production posts unconditionally (engines are always
single-node-per-actor, node_rank is structurally 0), so these tests
fail with "should not POST". Their assertion no longer reflects the
production contract.

Deleting them. The remaining 26 tests in this file still verify the
HTTP shapes for /start_weight_update, /finish_weight_update,
/update_weights, /sleep, /wake_up, etc.

PR #22 author was pinged about this in
#22 (comment)
but has not acted; doing the cleanup on PR #18 since these tests are
unblocking the unit-all batch run.

After this commit:
  pytest tests/unit/backends/vllm_utils/test_vllm_engine.py
  → 26 passed (was 26 passed, 2 failed)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
@aoshen02
aoshen02 merged commit 43f922c into main May 27, 2026
10 of 16 checks passed
aoshen02 added a commit that referenced this pull request May 27, 2026
After PR #22 added the colocated CUDA IPC weight sync path, the IPC branch
bypasses ``VLLMEngine.update_weights_from_tensor`` entirely (data moves
through ``IPCWeightTransferEngine.trainer_send_weights`` directly, not the
``/update_weights`` RPC). That RPC is the normal place where the engine's
``self._weight_version`` gets recorded; the IPC path leaves it at the
initial ``None``.

The fail mode shows up the first time IPC weight sync runs with
``--ci-test``: ``actor.update_weights`` at
slime/backends/megatron_utils/actor.py picks a random rollout engine,
calls ``get_weight_version()``, and compares against the trainer-side
``weight_updater.weight_version``. Because the engine's
``_weight_version`` is still ``None``, ``get_weight_version`` falls back
to ``GET /v1/models`` and returns the model path string
(e.g. ``/root/models/Qwen2.5-0.5B-Instruct``); the updater returns the
integer version string (``"1"``, ``"2"``, ...). They never match:

    RuntimeError: Weight version mismatch! Engine:
    /root/models/Qwen2.5-0.5B-Instruct, Updater: 1

Reproducer: ``tests/test_qwen2.5_0.5B_short.py`` and
``tests/test_qwen2.5_0.5B_ppo_critic_only_short.py`` once they reach
``update_weights`` — both colocate + IPC, both have ``--ci-test`` set.

Thread the version through:

- ``VLLMEngine.finish_weight_update(weight_version: str | None = None)``
  now sets ``self._weight_version = str(weight_version)`` before issuing
  the ``POST /finish_weight_update``.
- ``UpdateWeightFromTensor.update_weights`` (step 5) passes
  ``str(self.weight_version)`` when the coordinator rank fires the
  per-engine ``finish_weight_update`` RPC.

Each colocated engine's coordinator rank fires its own
``finish_weight_update``, so every engine ends up with the new version
recorded; the distributed path is unchanged.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
CalvinXKY pushed a commit that referenced this pull request May 27, 2026
…ses (#45)

After PR #22 added the colocated CUDA IPC weight sync path, the IPC branch
bypasses ``VLLMEngine.update_weights_from_tensor`` entirely (data moves
through ``IPCWeightTransferEngine.trainer_send_weights`` directly, not the
``/update_weights`` RPC). That RPC is the normal place where the engine's
``self._weight_version`` gets recorded; the IPC path leaves it at the
initial ``None``.

The fail mode shows up the first time IPC weight sync runs with
``--ci-test``: ``actor.update_weights`` at
slime/backends/megatron_utils/actor.py picks a random rollout engine,
calls ``get_weight_version()``, and compares against the trainer-side
``weight_updater.weight_version``. Because the engine's
``_weight_version`` is still ``None``, ``get_weight_version`` falls back
to ``GET /v1/models`` and returns the model path string
(e.g. ``/root/models/Qwen2.5-0.5B-Instruct``); the updater returns the
integer version string (``"1"``, ``"2"``, ...). They never match:

    RuntimeError: Weight version mismatch! Engine:
    /root/models/Qwen2.5-0.5B-Instruct, Updater: 1

Reproducer: ``tests/test_qwen2.5_0.5B_short.py`` and
``tests/test_qwen2.5_0.5B_ppo_critic_only_short.py`` once they reach
``update_weights`` — both colocate + IPC, both have ``--ci-test`` set.

Thread the version through:

- ``VLLMEngine.finish_weight_update(weight_version: str | None = None)``
  now sets ``self._weight_version = str(weight_version)`` before issuing
  the ``POST /finish_weight_update``.
- ``UpdateWeightFromTensor.update_weights`` (step 5) passes
  ``str(self.weight_version)`` when the coordinator rank fires the
  per-engine ``finish_weight_update`` RPC.

Each colocated engine's coordinator rank fires its own
``finish_weight_update``, so every engine ends up with the new version
recorded; the distributed path is unchanged.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 added a commit that referenced this pull request May 27, 2026
…version-with-data)

Background
----------
After PR #22 introduced the colocated CUDA IPC path, vime ended up with three
``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was
inconsistent:

- ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version``
  inside the RPC, version travels with data — slime-style.
- ``update_weights`` (IPC path, called from
  ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's
  ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never
  recorded ``_weight_version`` — vLLM's payload schema does not carry it.
- ``update_weights_from_tensor`` (PR #18 legacy entry): kept a
  SGLang-ish ``serialized_named_tensors`` payload and wrote
  ``_weight_version``, but had no callers in main.

The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine:
/root/models/<...>, Updater: N" failure on every colocated test with
``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto
``finish_weight_update``).

slime's design avoids this entirely: both IPC and distributed call
``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same
RPC name across both repos, with version travelling alongside the data
in the same RPC.

This PR
-------
Rewire vime's IPC path to match slime's interface:

1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point.
   Signature ``(update_info: dict, weight_version: str | None, flush_cache)``;
   payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names /
   shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from
   ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after
   the POST succeeds — mirrors ``update_weights_from_distributed``'s
   post-POST ordering so a failed transfer never advances the engine's
   tracked version.

2. Delete ``vllm_engine.update_weights`` — was the vLLM
   ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used
   after step 4 below.

3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only
   ``update_weights_from_tensor``'s old SGLang-ish path called.

4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version``
   now lives in the data-carrying RPC, so the bookend no longer needs to
   piggyback a kwarg. (Undoes the kwarg added in #45.)

5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in
   ``_send_hf_chunk_via_ipc`` with direct
   ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)``
   for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing
   vLLM's ``reduce_tensor`` for IPC handle creation
   (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is
   ours — so we don't fork the vLLM IPC protocol, just route through our own
   RPC surface.

Why not just keep #45's piggyback?
- #45 worked but coupled version bookkeeping to the lifecycle hook
  ``finish_weight_update`` instead of the data RPC. The wire shape doesn't
  match slime's, and a new IPC-style entry point added later would have to
  remember to also write ``_weight_version`` — exactly the trap PR #22 fell
  into. Centralising the write inside the data RPC removes the trap.

Unit tests
----------
- ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC
  call recording stays complete.
- Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` ->
  ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``,
  asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and
  that ``finish_weight_update`` is now stateless (no kwargs).
- Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``:
  asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata
  fields pass through, ``_weight_version`` advances on POST success.
- Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``:
  asserts POST failure does not advance ``_weight_version`` (matches
  the same post-POST ordering review note from #45).

Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py
(``_weight_transfer_http_timeout``, ``_response_json_or_fallback``,
``server_host``) are unchanged from main — main has 8 failed / 24 passed,
this PR has 8 failed / 26 passed (the two added tests). Not in scope here.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 added a commit that referenced this pull request May 27, 2026
Split from PR #18 (gcl/clean-sglang). One of 4 PRs splitting the
original PR #18 by content area: docs (#38) / examples (#39) /
**tests+CI** / core runtime. 42 files / +~750 / -~700.

These are bundled in a single PR because the CI workflows reference
test file names by string — splitting them would create a window where
either tests are renamed but CI still points at the old names, or vice
versa, breaking CI mid-roll.

What this PR does:

(A) tests/ (38 files):

- Mechanical CLI-flag rename: --sglang-* → --vllm-* equivalents in all
  test scripts (matches the table now used in scripts/ and examples/).
- Variable rename: SGLANG_ARGS → VLLM_ARGS where present.
- 4 file renames (R086-R091, all >85% similarity):
    test_qwen2.5_0.5B_opd_sglang.py        → test_qwen2.5_0.5B_opd_vllm.py
    test_qwen2.5_0.5B_sglang_config.py     → test_qwen2.5_0.5B_vllm_config.py
    test_qwen2.5_0.5B_sglang_config_distributed.py
                                           → test_qwen2.5_0.5B_vllm_config_distributed.py
    test_sglang_config_mixed_offload.py    → test_vllm_config_mixed_offload.py
    test_sglang_config_mixed_offload_ft.py → test_vllm_config_mixed_offload_ft.py
    tests/utils/test_sglang_config.py      → tests/utils/test_vllm_config.py
- 2 new tests for the IPC weight-transfer path landed in PR #18:
    tests/test_update_weight_from_tensor.py
    tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py
  (These are PR #22 / colocate-IPC test coverage; the production code
  the slim PR #18 ships will rely on the same code from PR #22.)

(B) .github/ (4 files):

- workflows/conda-ci.yml: container image lmsysorg/sglang → vime
  (inferactinc/public:vime-vllm-cu129-latest).
- workflows/pr-test.yml + pr-test.yml.j2 (template):
    * Container images (slimerl/slime[-test]:latest → vime image) on
      every job that ran on the sglang-era base.
    * e2e-test-sglang-config job → e2e-test-vllm-config job (renamed
      label `run-ci-sglang-config` → `run-ci-vllm-config`; matrix
      `test_file` entries updated to point at the renamed test files
      in (A)).
    * e2e-test-megatron + e2e-test-image matrices: `_opd_sglang.py`
      entries → `_opd_vllm.py`.
- ISSUE_TEMPLATE/bug_report.yml: drop the "SGLang version (if
  relevant):" environment field, add "vLLM version:" and
  "vllm-router version:" lines. (PR #36 already changed
  "CUDA/ROCm version" → "CUDA version" earlier; that change is
  preserved.)

Sgl residue intentionally kept (4 hits — all anti-regression
assertions that prove sglang code paths are gone, not residual
references to bring back):

- tests/test_update_weight_from_tensor.py:753 — comment "The vLLM IPC
  implementation must NOT contain sglang-style Gloo gather code".
- tests/unit/backends/vllm_utils/test_arguments.py:233-237 — three
  assertions that --sglang-router-ip, --sglang-router-port, and
  sglang_router_ip are NOT present in the argument parser.

Tests + CI must land together; splitting them risks a window where
the CI matrix references test files by names that don't exist yet
(or no longer exist). After this lands, the test_file string in CI
matches the test files on disk.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
aoshen02 added a commit that referenced this pull request May 27, 2026
…version-with-data)

Background
----------
After PR #22 introduced the colocated CUDA IPC path, vime ended up with three
``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was
inconsistent:

- ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version``
  inside the RPC, version travels with data — slime-style.
- ``update_weights`` (IPC path, called from
  ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's
  ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never
  recorded ``_weight_version`` — vLLM's payload schema does not carry it.
- ``update_weights_from_tensor`` (PR #18 legacy entry): kept a
  SGLang-ish ``serialized_named_tensors`` payload and wrote
  ``_weight_version``, but had no callers in main.

The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine:
/root/models/<...>, Updater: N" failure on every colocated test with
``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto
``finish_weight_update``).

slime's design avoids this entirely: both IPC and distributed call
``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same
RPC name across both repos, with version travelling alongside the data
in the same RPC.

This PR
-------
Rewire vime's IPC path to match slime's interface:

1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point.
   Signature ``(update_info: dict, weight_version: str | None, flush_cache)``;
   payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names /
   shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from
   ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after
   the POST succeeds — mirrors ``update_weights_from_distributed``'s
   post-POST ordering so a failed transfer never advances the engine's
   tracked version.

2. Delete ``vllm_engine.update_weights`` — was the vLLM
   ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used
   after step 4 below.

3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only
   ``update_weights_from_tensor``'s old SGLang-ish path called.

4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version``
   now lives in the data-carrying RPC, so the bookend no longer needs to
   piggyback a kwarg. (Undoes the kwarg added in #45.)

5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in
   ``_send_hf_chunk_via_ipc`` with direct
   ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)``
   for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing
   vLLM's ``reduce_tensor`` for IPC handle creation
   (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is
   ours — so we don't fork the vLLM IPC protocol, just route through our own
   RPC surface.

Why not just keep #45's piggyback?
- #45 worked but coupled version bookkeeping to the lifecycle hook
  ``finish_weight_update`` instead of the data RPC. The wire shape doesn't
  match slime's, and a new IPC-style entry point added later would have to
  remember to also write ``_weight_version`` — exactly the trap PR #22 fell
  into. Centralising the write inside the data RPC removes the trap.

Unit tests
----------
- ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC
  call recording stays complete.
- Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` ->
  ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``,
  asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and
  that ``finish_weight_update`` is now stateless (no kwargs).
- Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``:
  asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata
  fields pass through, ``_weight_version`` advances on POST success.
- Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``:
  asserts POST failure does not advance ``_weight_version`` (matches
  the same post-POST ordering review note from #45).

Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py
(``_weight_transfer_http_timeout``, ``_response_json_or_fallback``,
``server_host``) are unchanged from main — main has 8 failed / 24 passed,
this PR has 8 failed / 26 passed (the two added tests). Not in scope here.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 added a commit that referenced this pull request May 27, 2026
…version-with-data)

Background
----------
After PR #22 introduced the colocated CUDA IPC path, vime ended up with three
``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was
inconsistent:

- ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version``
  inside the RPC, version travels with data — slime-style.
- ``update_weights`` (IPC path, called from
  ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's
  ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never
  recorded ``_weight_version`` — vLLM's payload schema does not carry it.
- ``update_weights_from_tensor`` (PR #18 legacy entry): kept a
  SGLang-ish ``serialized_named_tensors`` payload and wrote
  ``_weight_version``, but had no callers in main.

The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine:
/root/models/<...>, Updater: N" failure on every colocated test with
``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto
``finish_weight_update``).

slime's design avoids this entirely: both IPC and distributed call
``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same
RPC name across both repos, with version travelling alongside the data
in the same RPC.

This PR
-------
Rewire vime's IPC path to match slime's interface:

1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point.
   Signature ``(update_info: dict, weight_version: str | None, flush_cache)``;
   payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names /
   shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from
   ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after
   the POST succeeds — mirrors ``update_weights_from_distributed``'s
   post-POST ordering so a failed transfer never advances the engine's
   tracked version.

2. Delete ``vllm_engine.update_weights`` — was the vLLM
   ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used
   after step 4 below.

3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only
   ``update_weights_from_tensor``'s old SGLang-ish path called.

4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version``
   now lives in the data-carrying RPC, so the bookend no longer needs to
   piggyback a kwarg. (Undoes the kwarg added in #45.)

5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in
   ``_send_hf_chunk_via_ipc`` with direct
   ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)``
   for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing
   vLLM's ``reduce_tensor`` for IPC handle creation
   (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is
   ours — so we don't fork the vLLM IPC protocol, just route through our own
   RPC surface.

Why not just keep #45's piggyback?
- #45 worked but coupled version bookkeeping to the lifecycle hook
  ``finish_weight_update`` instead of the data RPC. The wire shape doesn't
  match slime's, and a new IPC-style entry point added later would have to
  remember to also write ``_weight_version`` — exactly the trap PR #22 fell
  into. Centralising the write inside the data RPC removes the trap.

Unit tests
----------
- ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC
  call recording stays complete.
- Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` ->
  ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``,
  asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and
  that ``finish_weight_update`` is now stateless (no kwargs).
- Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``:
  asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata
  fields pass through, ``_weight_version`` advances on POST success.
- Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``:
  asserts POST failure does not advance ``_weight_version`` (matches
  the same post-POST ordering review note from #45).

Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py
(``_weight_transfer_http_timeout``, ``_response_json_or_fallback``,
``server_host``) are unchanged from main — main has 8 failed / 24 passed,
this PR has 8 failed / 26 passed (the two added tests). Not in scope here.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
CalvinXKY pushed a commit that referenced this pull request May 27, 2026
…version-with-data)

Background
----------
After PR #22 introduced the colocated CUDA IPC path, vime ended up with three
``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was
inconsistent:

- ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version``
  inside the RPC, version travels with data — slime-style.
- ``update_weights`` (IPC path, called from
  ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's
  ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never
  recorded ``_weight_version`` — vLLM's payload schema does not carry it.
- ``update_weights_from_tensor`` (PR #18 legacy entry): kept a
  SGLang-ish ``serialized_named_tensors`` payload and wrote
  ``_weight_version``, but had no callers in main.

The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine:
/root/models/<...>, Updater: N" failure on every colocated test with
``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto
``finish_weight_update``).

slime's design avoids this entirely: both IPC and distributed call
``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same
RPC name across both repos, with version travelling alongside the data
in the same RPC.

This PR
-------
Rewire vime's IPC path to match slime's interface:

1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point.
   Signature ``(update_info: dict, weight_version: str | None, flush_cache)``;
   payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names /
   shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from
   ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after
   the POST succeeds — mirrors ``update_weights_from_distributed``'s
   post-POST ordering so a failed transfer never advances the engine's
   tracked version.

2. Delete ``vllm_engine.update_weights`` — was the vLLM
   ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used
   after step 4 below.

3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only
   ``update_weights_from_tensor``'s old SGLang-ish path called.

4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version``
   now lives in the data-carrying RPC, so the bookend no longer needs to
   piggyback a kwarg. (Undoes the kwarg added in #45.)

5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in
   ``_send_hf_chunk_via_ipc`` with direct
   ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)``
   for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing
   vLLM's ``reduce_tensor`` for IPC handle creation
   (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is
   ours — so we don't fork the vLLM IPC protocol, just route through our own
   RPC surface.

Why not just keep #45's piggyback?
- #45 worked but coupled version bookkeeping to the lifecycle hook
  ``finish_weight_update`` instead of the data RPC. The wire shape doesn't
  match slime's, and a new IPC-style entry point added later would have to
  remember to also write ``_weight_version`` — exactly the trap PR #22 fell
  into. Centralising the write inside the data RPC removes the trap.

Unit tests
----------
- ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC
  call recording stays complete.
- Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` ->
  ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``,
  asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and
  that ``finish_weight_update`` is now stateless (no kwargs).
- Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``:
  asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata
  fields pass through, ``_weight_version`` advances on POST success.
- Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``:
  asserts POST failure does not advance ``_weight_version`` (matches
  the same post-POST ordering review note from #45).

Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py
(``_weight_transfer_http_timeout``, ``_response_json_or_fallback``,
``server_host``) are unchanged from main — main has 8 failed / 24 passed,
this PR has 8 failed / 26 passed (the two added tests). Not in scope here.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
CalvinXKY pushed a commit that referenced this pull request May 28, 2026
…version-with-data) (#48)

* refactor(weight-sync): align IPC RPC contract with slime (single-RPC version-with-data)

Background
----------
After PR #22 introduced the colocated CUDA IPC path, vime ended up with three
``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was
inconsistent:

- ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version``
  inside the RPC, version travels with data — slime-style.
- ``update_weights`` (IPC path, called from
  ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's
  ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never
  recorded ``_weight_version`` — vLLM's payload schema does not carry it.
- ``update_weights_from_tensor`` (PR #18 legacy entry): kept a
  SGLang-ish ``serialized_named_tensors`` payload and wrote
  ``_weight_version``, but had no callers in main.

The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine:
/root/models/<...>, Updater: N" failure on every colocated test with
``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto
``finish_weight_update``).

slime's design avoids this entirely: both IPC and distributed call
``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same
RPC name across both repos, with version travelling alongside the data
in the same RPC.

This PR
-------
Rewire vime's IPC path to match slime's interface:

1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point.
   Signature ``(update_info: dict, weight_version: str | None, flush_cache)``;
   payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names /
   shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from
   ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after
   the POST succeeds — mirrors ``update_weights_from_distributed``'s
   post-POST ordering so a failed transfer never advances the engine's
   tracked version.

2. Delete ``vllm_engine.update_weights`` — was the vLLM
   ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used
   after step 4 below.

3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only
   ``update_weights_from_tensor``'s old SGLang-ish path called.

4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version``
   now lives in the data-carrying RPC, so the bookend no longer needs to
   piggyback a kwarg. (Undoes the kwarg added in #45.)

5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in
   ``_send_hf_chunk_via_ipc`` with direct
   ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)``
   for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing
   vLLM's ``reduce_tensor`` for IPC handle creation
   (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is
   ours — so we don't fork the vLLM IPC protocol, just route through our own
   RPC surface.

Why not just keep #45's piggyback?
- #45 worked but coupled version bookkeeping to the lifecycle hook
  ``finish_weight_update`` instead of the data RPC. The wire shape doesn't
  match slime's, and a new IPC-style entry point added later would have to
  remember to also write ``_weight_version`` — exactly the trap PR #22 fell
  into. Centralising the write inside the data RPC removes the trap.

Unit tests
----------
- ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC
  call recording stays complete.
- Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` ->
  ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``,
  asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and
  that ``finish_weight_update`` is now stateless (no kwargs).
- Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``:
  asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata
  fields pass through, ``_weight_version`` advances on POST success.
- Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``:
  asserts POST failure does not advance ``_weight_version`` (matches
  the same post-POST ordering review note from #45).

Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py
(``_weight_transfer_http_timeout``, ``_response_json_or_fallback``,
``server_host``) are unchanged from main — main has 8 failed / 24 passed,
this PR has 8 failed / 26 passed (the two added tests). Not in scope here.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(weight-sync): correct IPC slot leader gating and gather group

Two regressions in the previous commit only fire when Megatron TP !=
rollout-num-gpus-per-engine (e.g. parallel-check sweeps Megatron TP=1
with rollout TP=2). Both surface via
``tests/test_qwen3_0.6B_parallel_check.py``.

Bug 1: leader gating wrong reference group
------------------------------------------
``connect_rollout_engines`` used::

    if mpu.get_tensor_model_parallel_rank() == 0:
        self._ipc_engine_coordinator = True

The intent was "TP rank 0 within the engine GPU slot", but
``mpu.get_tensor_model_parallel_rank()`` is the Megatron TP rank, not
the engine-slot rank. When Megatron TP=1, every trainer rank sees
``tp_rank=0`` and becomes a coordinator. For slot_size > 1, both ranks
in the slot then call ``start_weight_update`` → the second call
explodes::

    Worker failed with error 'start_weight_update called while a weight
    update is already active. Call finish_weight_update first.'

Fix: gate on ``rank == start`` (lowest trainer rank in the engine GPU
range). Unique per slot regardless of Megatron parallelism.

Bug 2: gather group wrong scope
-------------------------------
``_send_hf_chunk_via_ipc`` used ``mpu.get_tensor_model_parallel_group()``
to all_gather IPC payloads from peers in the engine slot. Again Megatron
TP group ≠ engine slot when their world sizes differ. The merge then
only had the coordinator's own UUID; downstream workers reading a
different physical GPU got::

    ValueError: IPC handle not found for GPU UUID <peer>.
    Available UUIDs: ['<coordinator>']

Fix: build per-slot process groups in ``connect_rollout_engines``
collectively (every trainer rank calls ``dist.new_group(slot_ranks)``
for every engine slot, keeps the one it belongs to). Use that group
instead of Megatron's TP group for the gather and the trailing barrier.

Validation
----------
Ran ``tests/test_qwen3_0.6B_parallel_check.py`` on 8×H200 with
``--num-rollout 2``. The test sweeps tp_size ∈ {1, 2, 4, 8} × pp_size
∈ {1, 2, 4} × cp_size ∈ {1, 2, 4, 8} for num_gpus ∈ {8, 4, 2} —
every leg also uses ``--rollout-num-gpus-per-engine 2``, so the
Megatron-TP=1 cases now exercise the new slot-group path. Pre-fix:
fails on the very first Megatron-TP=1 config with bug 1 above; after
bug 1 is patched, the next iteration fails with bug 2. Post-fix: the
entire ~2-hour sweep completes with "Job succeeded" on every leg.

Other tests already exercising IPC at Megatron-TP=2 == rollout-TP=2 are
unaffected by this fix (``rank == start`` is equivalent to
``tp_rank == 0`` when the slot fits a single Megatron TP group). The
following also pass post-fix as a sanity check:
``test_qwen3_4B_ppo``, ``test_qwen3_4B_ppo_train_critic_only``,
``test_qwen3_4B_ppo_disaggregate``, ``test_mimo_7B_mtp_only_grad``,
``test_moonlight_16B_A3B``, ``test_quick_start_glm4_9B``,
``test_qwen2.5_0.5B_{short,async_short,debug_rollout_then_train,
ppo_critic_only_short}``, ``test_qwen3.5_0.8B_gsm8k_{short,async_short}``.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(weight-sync): use gloo backend for ipc payload gather, add multi-gpu test

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: SamitHuang <285365963@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* refactor(weight-sync): drop dead _apply_monkey_patch_torch_reductions calls

The two _apply_monkey_patch_torch_reductions() call sites in this file
(trainer send + vLLM worker hijack) are no-ops on the IPC path here:

1. We route IPC handles by physical GPU UUID dict key (set on the trainer
   side at _build_ipc_update_info_from_named_tensors via _current_gpu_uuid()
   from torch.cuda.get_device_properties().uuid). The receiver looks up by
   its own UUID, independent of args[6].

2. vLLM's IPCWeightTransferEngine.receive_weights unconditionally overwrites
   args[6] with the receiver's local device_index before calling
   rebuild_cuda_tensor. Whatever a torch reductions patch encodes into
   args[6] is therefore discarded.

The patch was the historical mechanism (sglang upstream) for translating
device indices across CUDA_VISIBLE_DEVICES boundaries by stuffing UUID
strings into args[6]. Our UUID-keyed dict + vLLM's explicit device_index
override accomplish the same thing without the global torch reductions
mutation.

Also expand the _build_ipc_update_info_from_named_tensors docstring to
spell out the UUID-keyed routing contract so future readers don't have
to chase this through git history.

Side effect: hf_weight_iterator_direct.py also calls
monkey_patch_torch_reductions() at module-collective time. That call site
is similarly decorative (only NCCL broadcast / all_gather collectives run
there, no cross-process pickling) but lives outside this file's scope and
is not touched here. Tracked alongside #29.

* refactor(weight-sync): finish removing monkey_patch_torch_reductions dead code

Follow-up to 39bf899 (deleted _apply_monkey_patch_torch_reductions from
update_weight_from_tensor.py). With that helper gone, two more references
are now dead in the vime IPC weight-transfer path:

1. hf_weight_iterator_direct.py:48 called monkey_patch_torch_reductions()
   at the top of _get_megatron_full_params(). On vime this never has effect
   on the IPC handle path: _get_megatron_full_params only runs NCCL
   broadcast/all_gather collectives (no cross-process pickling), and the
   chunks it returns are subsequently sent via PR #48's UUID-keyed
   {gpu_uuid: reduce_tensor(weight)} dict that vLLM's receiver routes by
   physical UUID + explicit args[6] overwrite. The call survives in slime/
   miles upstream because their downstream path pickles tensors through
   sglang's MultiprocessingSerializer.serialize (ForkingPickler →
   reduce_tensor), where the patched encoding/decoding does real work; PR
   #48 does not use that pipeline, so the call here was incidentally
   inherited rather than functionally required.

2. slime/backends/megatron_utils/sglang.py's monkey_patch_torch_reductions
   re-export + __all__ entry now have no remaining importers in vime.
   Remove them.

Also (this commit, B):

3. vllm_engine.py's update_weights_from_tensor docstring referred to
   "closures injected by _apply_monkey_patch_torch_reductions" as the
   reason for cloudpickle. That helper is gone; the cloudpickle is still
   correct because reduce_tensor returns a (rebuild_fn, args) tuple where
   the rebuild_fn is a module-level callable that JSON can't serialise.
   Update the docstring to reflect the actual reason.

Net behaviour: identical — the deletions remove dead code paths. The
patch_torch shim itself is still importable for callers outside vime
(none currently in this tree).

Cross-references:
- #29 — issue documenting the no-op stub
- 39bf899 — prior commit that deleted the helper from update_weight_from_tensor.py

* docs(weight-sync): drop sglang refs + add IPC call-stack notes

- update_weight_from_tensor.py module docstring: rewrite from a pure
  vLLM perspective. Step (2) now spells out the merged
  {uuid_G0: handle, uuid_G1: handle, ...} dict + collective_rpc fan-out
  inside the vLLM server, and step (3) makes version-with-data atomicity
  explicit. Drop the "match slime's sglang_engine signature" framing.
- vllm_engine.update_weights_from_tensor: collapse the long sglang-vs-vLLM
  compare block into a focused docstring describing what the POST does
  and why ipc_handles needs cloudpickle. Add a Chinese call-stack
  walkthrough (trainer → ★this method★ → server collective_rpc → per-TP
  worker receive_weights) and note that `node_rank != 0` is a dead branch
  since VLLMEngine pins node_rank=0 (see PR #48 review comment for the
  follow-up cleanup).
- Hoist base64/cloudpickle imports to module scope so the hot path no
  longer pays the per-call import overhead.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

---------

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: SamitHuang <285365963@qq.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: SamitHuang <285365963@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
aoshen02 added a commit that referenced this pull request May 28, 2026
The 786-line tests/test_update_weight_from_tensor.py is a stale rebase
leftover from the original PR #18 branch — it predates the IPC test
file PR #22 landed at the canonical unit-test path
(tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py)
and predates PR #48's single-RPC weight-version contract.

Comparing the two:
* Both stub sys.modules / torch.distributed at module import time, so
  having two files compounds the test-isolation issue Gemini raised
  (PR #40 comment #1).
* Coverage overlaps materially (e.g. test_ipc_init_called_on_first_update_only
  ≈ test_ipc_init_runs_once — same invariant, different wording).
* The nested file is up-to-date with PR #48's RPC contract
  (update_weights_from_tensor.remote(**fields, weight_version=...));
  the top-level file still uses the pre-#48 lifecycle shape and does
  not exercise the coordinator slot fields.
* The nested path matches repo convention: tests/unit/ for mock-only
  unit tests, tests/ top level for e2e scripts.

Closes Gemini comment #1 on PR #40. Gemini comment #2 (the same stub
pattern in the surviving nested file) is a pre-existing issue from
PR #22 / #48 and out of scope for this rename PR — to be addressed
in a follow-up that converts _install_stubs() to an autouse
module-scoped fixture with save/restore.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
CalvinXKY pushed a commit that referenced this pull request May 30, 2026
* tests + CI: complete sglang→vllm rename across tests/ and .github/

Split from PR #18 (gcl/clean-sglang). One of 4 PRs splitting the
original PR #18 by content area: docs (#38) / examples (#39) /
**tests+CI** / core runtime. 42 files / +~750 / -~700.

These are bundled in a single PR because the CI workflows reference
test file names by string — splitting them would create a window where
either tests are renamed but CI still points at the old names, or vice
versa, breaking CI mid-roll.

What this PR does:

(A) tests/ (38 files):

- Mechanical CLI-flag rename: --sglang-* → --vllm-* equivalents in all
  test scripts (matches the table now used in scripts/ and examples/).
- Variable rename: SGLANG_ARGS → VLLM_ARGS where present.
- 4 file renames (R086-R091, all >85% similarity):
    test_qwen2.5_0.5B_opd_sglang.py        → test_qwen2.5_0.5B_opd_vllm.py
    test_qwen2.5_0.5B_sglang_config.py     → test_qwen2.5_0.5B_vllm_config.py
    test_qwen2.5_0.5B_sglang_config_distributed.py
                                           → test_qwen2.5_0.5B_vllm_config_distributed.py
    test_sglang_config_mixed_offload.py    → test_vllm_config_mixed_offload.py
    test_sglang_config_mixed_offload_ft.py → test_vllm_config_mixed_offload_ft.py
    tests/utils/test_sglang_config.py      → tests/utils/test_vllm_config.py
- 2 new tests for the IPC weight-transfer path landed in PR #18:
    tests/test_update_weight_from_tensor.py
    tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py
  (These are PR #22 / colocate-IPC test coverage; the production code
  the slim PR #18 ships will rely on the same code from PR #22.)

(B) .github/ (4 files):

- workflows/conda-ci.yml: container image lmsysorg/sglang → vime
  (inferactinc/public:vime-vllm-cu129-latest).
- workflows/pr-test.yml + pr-test.yml.j2 (template):
    * Container images (slimerl/slime[-test]:latest → vime image) on
      every job that ran on the sglang-era base.
    * e2e-test-sglang-config job → e2e-test-vllm-config job (renamed
      label `run-ci-sglang-config` → `run-ci-vllm-config`; matrix
      `test_file` entries updated to point at the renamed test files
      in (A)).
    * e2e-test-megatron + e2e-test-image matrices: `_opd_sglang.py`
      entries → `_opd_vllm.py`.
- ISSUE_TEMPLATE/bug_report.yml: drop the "SGLang version (if
  relevant):" environment field, add "vLLM version:" and
  "vllm-router version:" lines. (PR #36 already changed
  "CUDA/ROCm version" → "CUDA version" earlier; that change is
  preserved.)

Sgl residue intentionally kept (4 hits — all anti-regression
assertions that prove sglang code paths are gone, not residual
references to bring back):

- tests/test_update_weight_from_tensor.py:753 — comment "The vLLM IPC
  implementation must NOT contain sglang-style Gloo gather code".
- tests/unit/backends/vllm_utils/test_arguments.py:233-237 — three
  assertions that --sglang-router-ip, --sglang-router-port, and
  sglang_router_ip are NOT present in the argument parser.

Tests + CI must land together; splitting them risks a window where
the CI matrix references test files by names that don't exist yet
(or no longer exist). After this lands, the test_file string in CI
matches the test files on disk.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

* on_policy_distillation: port from SGLang to vLLM /v1/completions

Follow-up on the test rename in this PR:
test_qwen2.5_0.5B_opd_sglang.py → test_qwen2.5_0.5B_opd_vllm.py.

The test only spawns a vLLM teacher and exercises the OPD pipeline; the
real broken piece was slime/rollout/on_policy_distillation.py, which
PR #18 left in SGLang request/response shape:

  request fields:
    "max_new_tokens": 0          (vLLM: "max_tokens")
    "return_logprob": True       (sglang-only)
    "logprob_start_len": 0       (sglang-only)
  response parsing:
    reward["meta_info"]["input_token_logprobs"]   (sglang shape)

vLLM 0.21 supports the same workflow natively via `prompt_logprobs`:

  request to POST /v1/completions:
    {
        "model": <teacher>,
        "prompt_token_ids": sample.tokens,
        "max_tokens": 1,
        "temperature": 0,
        "prompt_logprobs": 1,
        "logprobs": 0,
        "skip_special_tokens": False,
    }
  response:
    response["choices"][0]["prompt_logprobs"]   # list[dict[int, Logprob] | None]
      where Logprob is {"logprob": float, "rank": int, "decoded_token": str}

References checked against vllm source:
  - reference/vllm/vllm/entrypoints/openai/completion/protocol.py:91
    (request: prompt_logprobs: int | None)
  - reference/vllm/vllm/entrypoints/openai/completion/protocol.py:487
    (response: prompt_logprobs: list[dict[int, Logprob] | None] | None)
  - reference/vllm/vllm/logprobs.py:13
    (Logprob dataclass: logprob/rank/decoded_token)

Implementation notes:

1. JSON serializes int dict keys as strings, so `_logprob_for_token`
   tries both `pos_entry.get(token_id)` and `pos_entry.get(str(token_id))`.

2. `pos_entry` is `None` at position 0 (no prior context) — handled
   explicitly. We also gracefully degrade if a token at position `i` is
   not in the top-1 logprob dict (falls back to 0.0, same as the prior
   sglang code would do).

3. The Logprob dataclass `decoded_token` field is unused; we only read
   `.logprob`. Both dict and `Logprob` shapes are accepted in case the
   server uses a flatter serialization toggle.

4. `args.opd_teacher_model` is the new model-name arg; falls back to
   `args.hf_checkpoint` if not set, mirroring how vime's other rollout
   paths derive the model name.

Smoke-tested `_logprob_for_token` locally:
  - None entry → 0.0
  - int key + dict value → logprob
  - str key (JSON shape) → logprob
  - missing token → 0.0
  - flattened float value → float

Also drops 3 lines from
tests/unit/backends/vllm_utils/test_arguments.py: the
`--sglang-router-ip`/`--sglang-router-port`/`sglang_router_ip` anti-
regression assertions. Once the slim PR #18 lands and sglang is gone
from the runtime, those assertions are vacuous; treating sglang as
non-existent per the project policy.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

* tests: drop duplicate smoke test updates from PR40

* test(update_weight_from_tensor): drop stale _apply_monkey_patch_torch_reductions patch

The inner ``with patch(f"{MODULE_PATH}._apply_monkey_patch_torch_reductions"):``
context in _run_update suppressed a helper call that PR #48 has since deleted
from update_weight_from_tensor.py (commit 39bf899 on aoshen/align-ipc-rpc-with-slime).
After that PR lands the patched attribute won't exist and this line raises
AttributeError. Remove it now so the test survives PR #48 merge.

The ``sglang_mod.monkey_patch_torch_reductions = MagicMock()`` stub on the
fake sglang module is intentionally kept: on this branch the production code
still imports it via ``from ..sglang import monkey_patch_torch_reductions``
(both update_weight_from_tensor._apply_monkey_patch_torch_reductions on
PR #40's view of main, and hf_weight_iterator_direct.py at module level).
Removing the stub here would break the test on PR #40 alone; it can be
dropped in a follow-up once PR #48 finishes removing every import site.

Tests: ``tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py``
all 6 pass with this change applied to gcl/pr18-tests-ci HEAD.

* tests: drop duplicate top-level test_update_weight_from_tensor.py

The 786-line tests/test_update_weight_from_tensor.py is a stale rebase
leftover from the original PR #18 branch — it predates the IPC test
file PR #22 landed at the canonical unit-test path
(tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py)
and predates PR #48's single-RPC weight-version contract.

Comparing the two:
* Both stub sys.modules / torch.distributed at module import time, so
  having two files compounds the test-isolation issue Gemini raised
  (PR #40 comment #1).
* Coverage overlaps materially (e.g. test_ipc_init_called_on_first_update_only
  ≈ test_ipc_init_runs_once — same invariant, different wording).
* The nested file is up-to-date with PR #48's RPC contract
  (update_weights_from_tensor.remote(**fields, weight_version=...));
  the top-level file still uses the pre-#48 lifecycle shape and does
  not exercise the coordinator slot fields.
* The nested path matches repo convention: tests/unit/ for mock-only
  unit tests, tests/ top level for e2e scripts.

Closes Gemini comment #1 on PR #40. Gemini comment #2 (the same stub
pattern in the surviving nested file) is a pre-existing issue from
PR #22 / #48 and out of scope for this rename PR — to be addressed
in a follow-up that converts _install_stubs() to an autouse
module-scoped fixture with save/restore.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* test(vllm_config): use real get_model_url default endpoint /inference/v1/generate

get_model_url defaults to /inference/v1/generate (PR #18), not /v1/completions.
Aligns this test with PR #18's test_vllm_config.py so the two PRs no longer
conflict on this file and the assertion matches the actual runtime default.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Drop on_policy_distillation.py from tests+CI PR (now owned by runtime PR #18)

The OPD vLLM /v1/completions migration is a runtime change; it was folded into
the core-runtime PR (#18). Restore this file to main here so the two PRs no longer
overlap on it. #18 merges first, so this lands via #18.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop unit-test files now owned by runtime PR #18

test_vllm_config.py + the plugin_contracts tests are coupled to #18's runtime
rename (they import vllm_config / vllm_rollout, which #18 creates). They live in
#18; remove them here so the two PRs don't overlap. #18 merges first.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: restore vLLM rollout args dropped during sglang→vllm rename

The mechanical sglang→vllm rename dropped several rollout knobs instead of
mapping them to their vLLM equivalents, weakening CI coverage (cuda-graph
capture caps, speculative decoding, expert parallel). Restore them using the
mapping established by the converted production scripts on main
(run-glm4.7-30B-A3B.sh / run-glm5-744B-A40B.sh), verified against vLLM
AsyncEngineArgs:

  --sglang-cuda-graph-max-bs N            -> --vllm-max-cudagraph-capture-size N
  --sglang-cuda-graph-bs a b c            -> --vllm-cudagraph-capture-sizes a b c
  --sglang-ep-size N                      -> --vllm-enable-expert-parallel
  --sglang-speculative-* (eagle)          -> --vllm-speculative-config '{"method":"eagle","num_speculative_tokens":K}'

Also:
- glm4.7 pd: fix --vllm-max-num-seqs (was 8, taken from cuda-graph-max-bs;
  --sglang-max-running-requests was 16) and split out cuda-graph capture.
- fix sglang→rollout mis-renames in temp-file prefixes (→ vllm_*).
- test_vllm_config: rename test_update_weights_default_true →
  test_update_weights_defaults_to_none (it asserts `is None`).

Dropped sglang flags with no vLLM equivalent (enable-dp-lm-head,
moe-dense-tp-size, watchdog-timeout, mamba-scheduler-strategy,
disaggregation-transfer-backend, enable-metrics) stay dropped; PD KV-transfer
is driven by --prefill-num-servers + the --vllm-config prefill/decode topology.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(plugin_contracts): migrate from sglang_rollout to vllm_rollout

The three plugin-contract tests still imported slime.rollout.sglang_rollout
and called install_stubs(with_sglang_router=True), but _shared.install_stubs
already dropped that parameter — so all three failed at collection
(TypeError: unexpected keyword 'with_sglang_router'). Complete the migration:

- install_stubs(with_sglang_router=True, ...) -> install_stubs(...)
- import generate_and_rm / generate_rollout from slime.rollout.vllm_rollout
- default rollout/eval path string -> slime.rollout.vllm_rollout.generate_rollout
  (matches runtime default at slime/utils/arguments.py:233)
- FakeGenerateState: sglang_enable_deterministic_inference ->
  vllm_enable_deterministic_inference, with group_sampling_seeds defaulting to
  None and gated on the flag (mirrors the already-migrated
  tests/unit/rollout/test_vllm_rollout.py).

All 34 plugin-contract cases pass (were 3 collection errors before).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(update_weight_from_tensor): drop stale slime…megatron_utils.sglang mock

The test pre-registered a sys.modules mock for
slime.backends.megatron_utils.sglang (monkey_patch_torch_reductions), left over
from when update_weight_from_tensor imported it. The module under test no longer
imports that module (its real deps are get_gloo_group / HfWeightIteratorBase /
update_weight_from_distributed), so the mock is dead. Removing it makes tests/
and .github/ fully sglang-free. Test still passes (7/7).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* [Clean] Remove SGLang runtime code

Rebuilt against current main so the PR contains only the SGLang runtime
removal -- the docs / tests-ci / examples / scripts / docker portions were
split into separate PRs that have since merged.

- Delete dead SGLang server/runtime code: sglang_utils/{arguments,sglang_engine}.py,
  rollout/sglang_rollout.py, the megatron_utils/sglang.py re-export shim, and all
  docker/**/sglang.patch files.
- Rename the rollout config module sglang_utils/sglang_config.py ->
  vllm_utils/vllm_config.py (SglangConfig -> VllmConfig, _resolve_sglang_config ->
  _resolve_vllm_config, --sglang-config -> --vllm-config); inline the
  GPU_MEMORY_TYPE_* constants in rollout.py.
- Add megatron_utils/fp8_helpers.py for the UE8M0 fp8 helpers formerly re-exported
  through the sglang shim; repoint quantizer_fp8 to it.
- Swap sglang_router -> vllm_router in http_utils/wandb_utils; drop the dead
  sglang-router dependency from requirements.txt.
- Finish the SGLang->vLLM rename in the runtime so it is internally consistent and
  matches the tests landing in the tests/CI PR:
  * router args --router-* -> --vllm-router-* (vllm_router_ip/port/timeout);
  * get_model_url reads vllm_model_routers (aligning with rollout.py);
  * --opd-type sglang -> vllm; engine_overrides rename;
  * sglang_enable_deterministic_inference -> vllm_enable_deterministic_inference,
    wired to a real --vllm-enable-deterministic-inference flag (exports
    VLLM_BATCH_INVARIANT=1);
  * consistent_hash session-id routing uses vllm-router's x-session-id header;
  * drop dead trace helper build_sglang_meta_trace_attrs; de-SGLang comments/docstrings.
- Rename test_sglang_config.py -> test_vllm_config.py and de-SGLang the
  plugin-contract tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Address review: finish de-SGLang + fold OPD/router-policy into runtime

- naming: replace residual generic "rollout engine"/"engine"/"comm" wording
  with concrete vLLM (engine_overrides -> vllm_overrides; arguments help text;
  http_utils comments; rollout.py "inference workers"). sglang->vllm is correct,
  sglang->generic is not.
- megatron_to_hf: drop the q_a_proj/kv_a_proj_with_mqa pairing + _cached_tensors
  global. That was sglang-only: sglang's loader torch.cat's both shards within a
  single load_weights call (needs them co-bucketed), whereas vLLM loads each shard
  independently via stacked_params_mapping into fused_qkv_a_proj. Also fix the
  misleading "merge into single fused name" comment.
- docker/Dockerfile: remove now-dead sglang/sglang-router --no-deps stubs + the
  build-time `import sglang` smoke check (slime no longer imports sglang_router).
- OPD: migrate on_policy_distillation.py teacher logprobs to vLLM /v1/completions
  (prompt_logprobs) instead of sglang return_logprob / meta_info.input_token_logprobs.
- routing replay: register --vllm-router-policy (dest=router_policy) so the
  consistent_hash x-session-id session-affinity path is actually wired (was dead).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Review follow-ups: mirror slime vllm_config parsing + restore vLLM process cleanup

- vllm_config.from_yaml: drop the needless `models_raw` intermediate and iterate
  `data["vllm"]` directly, restoring the "Accept both server_groups / legacy
  engine_groups" comment -- mirrors slime's sglang_config.from_yaml line-for-line.
- command_utils.execute_train: re-add a process kill for leftover rollout engines
  as `pkill -9 -f "vllm serve"` (the old `pkill -9 sglang` was dropped with no
  vLLM equivalent), so stale engines don't hold GPUs/ports across runs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop test changes from runtime PR; tests live in the tests+CI PR (#40)

The plugin_contracts tests and the test_sglang_config -> test_vllm_config rename
are coupled to the test/CI rename effort and are owned by #40. Restore them to
main here so #18 is purely the SGLang runtime removal. #18 merges first; #40
rebases and re-lands the vLLM test versions.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fp8_helpers: copy SGLang verbatim (fix import crash); pin vLLM deep_gemm env

- fp8_helpers.py: replace the bespoke rewrite with SGLang's exact implementations
  of quant_weight_ue8m0 / transform_scale_ue8m0 and their DeepGEMM helpers
  (per_block_cast_to_fp8, ceil_to_ue8m0, ceil_div, ceil_align, the torch-impl
  packer). deep_gemm is imported lazily inside the functions (as SGLang does), so
  module import no longer requires deep_gemm. This fixes the module-level
  `NameError: _get_tma_aligned_size` that crashed `import megatron_to_hf` on any
  deep_gemm image, and drops the invented sf-stride fixup block that was not in
  upstream. Only should_deepgemm_weight_requant_ue8m0 stays vLLM-adapted
  (is_deep_gemm_e8m0_used) since SGLang's reads SGLang-internal deep_gemm_wrapper.
- vllm_engine.launch_server_process: set VLLM_USE_DEEP_GEMM=1 +
  VLLM_DEEP_GEMM_WARMUP=relax explicitly (setdefault) alongside VLLM_BATCH_INVARIANT,
  replacing SGLang's removed deep_gemm precompile/warmup envs. All vLLM engine env
  now lives in the subprocess env builder (single source of truth).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fp8_helpers: revert to vLLM impl + fix the import NameError; add vLLM trace attrs

- fp8_helpers.py: keep the vLLM-based implementation (uses vllm.utils.deep_gemm,
  consistent with the vLLM runtime) rather than the SGLang verbatim copy. Fix the
  module-level crash: the `try` block referenced `_get_tma_aligned_size` before it
  was bound (the "pre-imported with fallback" import was never written), which
  raised NameError whenever deep_gemm imported successfully -- and NameError is not
  caught by `except ImportError`, so `import megatron_to_hf` crashed on any
  deep_gemm image. Replace the bogus self-assignment with the real import:
  `from vllm.utils.deep_gemm import get_tma_aligned_size as _get_tma_aligned_size`.
- trace_utils/vllm_rollout: add build_vllm_meta_trace_attrs and attach finish_reason
  + token usage to the vllm_inference_generate span (mirrors SGLang's
  build_sglang_meta_trace_attrs; vLLM responses lack the pd_* timing, which lives
  in vLLM's own OTLP traces).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(opd): score teacher via /inference/v1/generate with prompt_logprobs

Move the vllm OPD teacher path off the OpenAI /v1/completions endpoint onto
vime's native /inference/v1/generate (the same endpoint the rollout engines
use), and fix three latent issues:

1. model field: /inference/v1/generate takes `model` as OPTIONAL. Stop
   defaulting to args.hf_checkpoint (the *student* name, which mis-names a
   teacher!=student server). Add --opd-teacher-model; send `model` only when
   set, otherwise omit it (single-model teacher servers use their loaded model).
2. multimodal: the old code sent image_data to a token-only endpoint, which is
   invalid. Raise NotImplementedError until the
   /v1/chat/completions/render -> /inference/v1/generate flow is wired (mirrors
   slime.rollout.vllm_rollout.generate).
3. logprob robustness: read top-level GenerateResponse.prompt_logprobs, assert
   it is present and length-aligned with token_ids, assert the per-sample tensor
   covers response_length, and raise (not silently return 0.0) on a missing
   token logprob. vLLM always includes the actual prompt token in
   prompt_logprobs, so a miss is a real error.

Alignment is unchanged (plp[i] <-> tokens[i], skip pos 0, take [-response_length:]).

Follow-up (separate, in the tests PR): the OPD e2e test must launch a teacher
that exposes /inference/v1/generate and point --rm-url at it.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* chore(clean-sglang): purge SGLang from tools, train scripts, and build infra

tools/: drop dead `args.sglang_enable_ep_moe` shim (read nowhere); reword
profile/replay helpers to vLLM and map analyzer hints to vLLM flags
(--enforce-eager, --gpu-memory-utilization). train{,_async}.py: comments
SGLang -> vLLM.

build infra: remove build_conda.sh (SGLang-only conda path); drop the GB300
sgl-kernel install from the Dockerfile; delete docker/npu_patch/ wholesale.

docker base image: bump to vLLM v0.22.0. justfile ARM recipes now pin the real
multi-arch vLLM base images instead of the dead SGLANG_IMAGE_TAG/
ENABLE_SGLANG_PATCH build-args -- cu129-arm64 -> v0.22.0-cu129-ubuntu2404
(CUDA 12.9), cu13-arm64 -> v0.22.0-ubuntu2404 (the default-CUDA tag is already
CUDA 13.0) + ENABLE_CUDA_13=1. vLLM tags are multi-arch manifests, so docker
selects the arm64 image automatically on an ARM host.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* ci(clean-sglang): drop conda-build workflow (ran deleted build_conda.sh on SGLang image)

The single build-conda job ran `bash build_conda.sh` (removed in the previous
commit) inside an lmsysorg/sglang container. With the SGLang-only conda path
gone, the whole workflow is dead.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* chore(clean-sglang): fix stale SGLang refs in docs/skills; tidy comments

docs/conf.py: point the "edit on GitHub" links at vllm-project/vime instead of
the inherited sgl-project.github.io repo. .claude/skills/*: update the dead
`slime/rollout/sglang_rollout.py` references to `vllm_rollout.py` (the real
default is slime.rollout.vllm_rollout.generate_rollout).

justfile: drop the redundant BASE_IMAGE override on release-cu129-arm64 (it
equalled the Dockerfile default; the multi-arch manifest already resolves
arm64). train{,_async}.py: drop stray "the" in the W&B comment.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(tests): use method=mtp (not eagle) in vllm speculative config

The migrated speculative configs pass no draft `model`, so method=eagle
raises "num_speculative_tokens was provided but without speculative model"
in vLLM's SpeculativeConfig. These models carry embedded MTP layers, so
method=mtp is correct and unblocks the mimo MTP-only-grad test (#19).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(cleanup): target renamed vLLM subprocesses in pkill so VRAM is freed

vLLM's set_process_title() renames the VRAM-holding subprocesses
(VLLM::EngineCore, VLLM::Worker_TP*, vllm::router), so their cmdline no
longer contains "vllm serve". The previous `pkill -9 -f "vllm serve"`
matched only the launcher and left engine/worker children holding GPU
memory, leaking it into the next run — masked only by the indiscriminate
`pkill -9 python`, which is unsafe on colocate/shared nodes.

Match both the launcher and the renamed children with
`pkill -9 -f '[v]llm serve|VLL[M]::'`; the [v]/[M] bracket trick keeps the
pattern from matching pkill's own cmdline. This makes the broad python
kill unnecessary, so its already-commented-out lines are removed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(speculative): use method=mtp (not eagle) for embedded-MTP models

vLLM's SpeculativeConfig requires an explicit draft `model` for
method=eagle; with only num_speculative_tokens set it raises
"num_speculative_tokens was provided but without speculative model".
The migrated configs in scripts/examples/docs pass no model, so they must
use method=mtp, which reuses the target checkpoint's embedded MTP layer
(DeepSeek-R1, GLM-4.x-MoE, MiMo, Qwen3-Next/3.5).

The two docs examples that pass an explicit "model" are genuine eagle
usage and are left unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(vllm): launch each rollout engine with its ServerGroup's per-group TP

launch_server_process / _init_normal derived tensor-parallel size and
CUDA_VISIBLE_DEVICES from the global --rollout-num-gpus-per-engine, ignoring
the per-engine num_gpus_per_engine already carried on the VLLMEngine actor.

A ServerGroup configured with num_gpus_per_engine greater than the global
flag (e.g. tp=2) therefore launched as tp=1, while the NCCL weight-sync
rendezvous sized world_size from engine_gpu_counts (the per-group value).
The two disagreed: the trainer waited for a rank the under-sized engine
never started, so init_weight_transfer_engine hung for 300s
("3/4 clients joined") and the job failed.

Honor the per-engine num_gpus_per_engine at launch, falling back to the
global flag when unset (matches the SGLang path and PR #66's
_compute_server_args).

Verified on H200: tests/test_qwen2.5_0.5B_vllm_config_distributed now
launches engine0 tp=2 / engine1 tp=1, update_weights completes in 1.1s
(was a 301s timeout), and rollout+eval proceed.

AI assistance (Claude Code) was used for this change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(ckpt): add --dist-ckpt-optim-fully-reshardable for PAO+offload save/load

test_qwen3_4B_ckpt.py uses precision-aware optimizer + cpu-offload
(HybridDeviceOptimizer). Under the default dp_reshardable (bucket-centric)
optimizer sharding, save/load produce unequal-length param_state lists, so
dist-ckpt load fails with
"Cannot merge two lists with different lengths (81 and 79)".

fully_reshardable is model-centric and immune to bucket-layout changes.
Verified on the r3 image (Megatron-LM 0.16.0rc0 @ 1dcf0da): save+load both
succeed, and source review confirms master_param / step / HybridDeviceOptimizer
sync are handled on this path. This is the flag described in PR #50 that was
never actually merged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* refactor(router-args): hybrid naming — vllm_ for ip/port, bare router_ for timeout

vllm-router's RouterArgs.from_cli_args only supports prefix "" or "router_"
(never "vllm_router_"), and excludes host/port from its CLI via
exclude_host_port=True. So:

- --vllm-router-ip / --vllm-router-port keep the vllm_ prefix: RouterArgs does
  not own these CLI flags, vime does (populated via _start_router's manual
  router_args.host/port assignment), so the vllm_ prefix is free and marks them
  as vime-owned endpoint config.
- --router-request-timeout-secs goes bare (dest router_request_timeout_secs): it
  is a genuine RouterArgs field, so it shares the --router-* namespace with
  policy / cache_threshold / retries / … and flows through from_cli_args like
  the other knobs.
- --vllm-router-policy keeps dest=router_policy (unchanged).

Also fixes conftest fixture to seed vllm_router_ip/port (was bare router_ip/port,
which never matched the vllm_engine reader) and updates README/README_zh prose.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
momo609 pushed a commit that referenced this pull request Jun 8, 2026
* Fix colocated mode weight sync

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

* Fix pre-commit

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

* Fix pre-commit

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

* Fix test

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>

---------

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
momo609 pushed a commit that referenced this pull request Jun 8, 2026
…ses (#45)

After PR #22 added the colocated CUDA IPC weight sync path, the IPC branch
bypasses ``VLLMEngine.update_weights_from_tensor`` entirely (data moves
through ``IPCWeightTransferEngine.trainer_send_weights`` directly, not the
``/update_weights`` RPC). That RPC is the normal place where the engine's
``self._weight_version`` gets recorded; the IPC path leaves it at the
initial ``None``.

The fail mode shows up the first time IPC weight sync runs with
``--ci-test``: ``actor.update_weights`` at
slime/backends/megatron_utils/actor.py picks a random rollout engine,
calls ``get_weight_version()``, and compares against the trainer-side
``weight_updater.weight_version``. Because the engine's
``_weight_version`` is still ``None``, ``get_weight_version`` falls back
to ``GET /v1/models`` and returns the model path string
(e.g. ``/root/models/Qwen2.5-0.5B-Instruct``); the updater returns the
integer version string (``"1"``, ``"2"``, ...). They never match:

    RuntimeError: Weight version mismatch! Engine:
    /root/models/Qwen2.5-0.5B-Instruct, Updater: 1

Reproducer: ``tests/test_qwen2.5_0.5B_short.py`` and
``tests/test_qwen2.5_0.5B_ppo_critic_only_short.py`` once they reach
``update_weights`` — both colocate + IPC, both have ``--ci-test`` set.

Thread the version through:

- ``VLLMEngine.finish_weight_update(weight_version: str | None = None)``
  now sets ``self._weight_version = str(weight_version)`` before issuing
  the ``POST /finish_weight_update``.
- ``UpdateWeightFromTensor.update_weights`` (step 5) passes
  ``str(self.weight_version)`` when the coordinator rank fires the
  per-engine ``finish_weight_update`` RPC.

Each colocated engine's coordinator rank fires its own
``finish_weight_update``, so every engine ends up with the new version
recorded; the distributed path is unchanged.

Signed-off-by: aoshen02 <aoshen@inferact.ai>
momo609 pushed a commit that referenced this pull request Jun 8, 2026
…version-with-data) (#48)

* refactor(weight-sync): align IPC RPC contract with slime (single-RPC version-with-data)

Background
----------
After PR #22 introduced the colocated CUDA IPC path, vime ended up with three
``update_weights*`` RPC entry points whose ``_weight_version`` bookkeeping was
inconsistent:

- ``update_weights_from_distributed`` (NCCL path): writes ``_weight_version``
  inside the RPC, version travels with data — slime-style.
- ``update_weights`` (IPC path, called from
  ``IPCWeightTransferEngine.trainer_send_weights``): forwarded vLLM's
  ``IPCWeightTransferUpdateInfo`` to ``/update_weights`` over HTTP but never
  recorded ``_weight_version`` — vLLM's payload schema does not carry it.
- ``update_weights_from_tensor`` (PR #18 legacy entry): kept a
  SGLang-ish ``serialized_named_tensors`` payload and wrote
  ``_weight_version``, but had no callers in main.

The IPC gap was the root cause of #41-era's "Weight version mismatch! Engine:
/root/models/<...>, Updater: N" failure on every colocated test with
``--ci-test`` (fixed in #45 by piggybacking ``weight_version`` onto
``finish_weight_update``).

slime's design avoids this entirely: both IPC and distributed call
``engine.update_weights_from_tensor.remote(..., weight_version=N)`` — same
RPC name across both repos, with version travelling alongside the data
in the same RPC.

This PR
-------
Rewire vime's IPC path to match slime's interface:

1. ``vllm_engine.update_weights_from_tensor`` is now the IPC entry point.
   Signature ``(update_info: dict, weight_version: str | None, flush_cache)``;
   payload carries vLLM's ``IPCWeightTransferUpdateInfo`` (names / dtype_names /
   shapes / ipc_handles), the trainer constructs it with ``reduce_tensor`` from
   ``torch.multiprocessing.reductions``. Records ``_weight_version`` only after
   the POST succeeds — mirrors ``update_weights_from_distributed``'s
   post-POST ordering so a failed transfer never advances the engine's
   tracked version.

2. Delete ``vllm_engine.update_weights`` — was the vLLM
   ``IPCWeightTransferEngine.trainer_send_weights`` entry, no longer used
   after step 4 below.

3. Delete ``vllm_engine._run_vllm_weight_update`` — dead helper that only
   ``update_weights_from_tensor``'s old SGLang-ish path called.

4. Revert ``finish_weight_update`` to a stateless POST — ``_weight_version``
   now lives in the data-carrying RPC, so the bookend no longer needs to
   piggyback a kwarg. (Undoes the kwarg added in #45.)

5. Replace ``IPCWeightTransferEngine.trainer_send_weights(...)`` calls in
   ``_send_hf_chunk_via_ipc`` with direct
   ``engine.update_weights_from_tensor.remote(update_info=..., weight_version=...)``
   for both slot_size paths (slot_size==1 and slot_size>1). vime keeps reusing
   vLLM's ``reduce_tensor`` for IPC handle creation
   (via ``_build_ipc_update_info_from_named_tensors``) — only the dispatch is
   ours — so we don't fork the vLLM IPC protocol, just route through our own
   RPC surface.

Why not just keep #45's piggyback?
- #45 worked but coupled version bookkeeping to the lifecycle hook
  ``finish_weight_update`` instead of the data RPC. The wire shape doesn't
  match slime's, and a new IPC-style entry point added later would have to
  remember to also write ``_weight_version`` — exactly the trap PR #22 fell
  into. Centralising the write inside the data RPC removes the trap.

Unit tests
----------
- ``RecordingVLLMEngine`` learns ``update_weights_from_tensor`` so engine RPC
  call recording stays complete.
- Renamed ``test_trainer_send_weights_uses_single_llm_handle_per_rank`` ->
  ``test_send_via_ipc_dispatches_update_weights_from_tensor_with_version``,
  asserts the new RPC name + kwargs (``update_info``, ``weight_version``) and
  that ``finish_weight_update`` is now stateless (no kwargs).
- Added ``test_update_weights_from_tensor_posts_ipc_update_info_and_records_version``:
  asserts ipc_handles get cloudpickle'd into ipc_handles_pickled, metadata
  fields pass through, ``_weight_version`` advances on POST success.
- Added ``test_update_weights_from_tensor_does_not_advance_version_on_failure``:
  asserts POST failure does not advance ``_weight_version`` (matches
  the same post-POST ordering review note from #45).

Pre-existing test failures in tests/unit/backends/vllm_utils/test_vllm_engine.py
(``_weight_transfer_http_timeout``, ``_response_json_or_fallback``,
``server_host``) are unchanged from main — main has 8 failed / 24 passed,
this PR has 8 failed / 26 passed (the two added tests). Not in scope here.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(weight-sync): correct IPC slot leader gating and gather group

Two regressions in the previous commit only fire when Megatron TP !=
rollout-num-gpus-per-engine (e.g. parallel-check sweeps Megatron TP=1
with rollout TP=2). Both surface via
``tests/test_qwen3_0.6B_parallel_check.py``.

Bug 1: leader gating wrong reference group
------------------------------------------
``connect_rollout_engines`` used::

    if mpu.get_tensor_model_parallel_rank() == 0:
        self._ipc_engine_coordinator = True

The intent was "TP rank 0 within the engine GPU slot", but
``mpu.get_tensor_model_parallel_rank()`` is the Megatron TP rank, not
the engine-slot rank. When Megatron TP=1, every trainer rank sees
``tp_rank=0`` and becomes a coordinator. For slot_size > 1, both ranks
in the slot then call ``start_weight_update`` → the second call
explodes::

    Worker failed with error 'start_weight_update called while a weight
    update is already active. Call finish_weight_update first.'

Fix: gate on ``rank == start`` (lowest trainer rank in the engine GPU
range). Unique per slot regardless of Megatron parallelism.

Bug 2: gather group wrong scope
-------------------------------
``_send_hf_chunk_via_ipc`` used ``mpu.get_tensor_model_parallel_group()``
to all_gather IPC payloads from peers in the engine slot. Again Megatron
TP group ≠ engine slot when their world sizes differ. The merge then
only had the coordinator's own UUID; downstream workers reading a
different physical GPU got::

    ValueError: IPC handle not found for GPU UUID <peer>.
    Available UUIDs: ['<coordinator>']

Fix: build per-slot process groups in ``connect_rollout_engines``
collectively (every trainer rank calls ``dist.new_group(slot_ranks)``
for every engine slot, keeps the one it belongs to). Use that group
instead of Megatron's TP group for the gather and the trailing barrier.

Validation
----------
Ran ``tests/test_qwen3_0.6B_parallel_check.py`` on 8×H200 with
``--num-rollout 2``. The test sweeps tp_size ∈ {1, 2, 4, 8} × pp_size
∈ {1, 2, 4} × cp_size ∈ {1, 2, 4, 8} for num_gpus ∈ {8, 4, 2} —
every leg also uses ``--rollout-num-gpus-per-engine 2``, so the
Megatron-TP=1 cases now exercise the new slot-group path. Pre-fix:
fails on the very first Megatron-TP=1 config with bug 1 above; after
bug 1 is patched, the next iteration fails with bug 2. Post-fix: the
entire ~2-hour sweep completes with "Job succeeded" on every leg.

Other tests already exercising IPC at Megatron-TP=2 == rollout-TP=2 are
unaffected by this fix (``rank == start`` is equivalent to
``tp_rank == 0`` when the slot fits a single Megatron TP group). The
following also pass post-fix as a sanity check:
``test_qwen3_4B_ppo``, ``test_qwen3_4B_ppo_train_critic_only``,
``test_qwen3_4B_ppo_disaggregate``, ``test_mimo_7B_mtp_only_grad``,
``test_moonlight_16B_A3B``, ``test_quick_start_glm4_9B``,
``test_qwen2.5_0.5B_{short,async_short,debug_rollout_then_train,
ppo_critic_only_short}``, ``test_qwen3.5_0.8B_gsm8k_{short,async_short}``.

Signed-off-by: aoshen02 <aoshen@inferact.ai>

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(weight-sync): use gloo backend for ipc payload gather, add multi-gpu test

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: SamitHuang <285365963@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

* refactor(weight-sync): drop dead _apply_monkey_patch_torch_reductions calls

The two _apply_monkey_patch_torch_reductions() call sites in this file
(trainer send + vLLM worker hijack) are no-ops on the IPC path here:

1. We route IPC handles by physical GPU UUID dict key (set on the trainer
   side at _build_ipc_update_info_from_named_tensors via _current_gpu_uuid()
   from torch.cuda.get_device_properties().uuid). The receiver looks up by
   its own UUID, independent of args[6].

2. vLLM's IPCWeightTransferEngine.receive_weights unconditionally overwrites
   args[6] with the receiver's local device_index before calling
   rebuild_cuda_tensor. Whatever a torch reductions patch encodes into
   args[6] is therefore discarded.

The patch was the historical mechanism (sglang upstream) for translating
device indices across CUDA_VISIBLE_DEVICES boundaries by stuffing UUID
strings into args[6]. Our UUID-keyed dict + vLLM's explicit device_index
override accomplish the same thing without the global torch reductions
mutation.

Also expand the _build_ipc_update_info_from_named_tensors docstring to
spell out the UUID-keyed routing contract so future readers don't have
to chase this through git history.

Side effect: hf_weight_iterator_direct.py also calls
monkey_patch_torch_reductions() at module-collective time. That call site
is similarly decorative (only NCCL broadcast / all_gather collectives run
there, no cross-process pickling) but lives outside this file's scope and
is not touched here. Tracked alongside #29.

* refactor(weight-sync): finish removing monkey_patch_torch_reductions dead code

Follow-up to 39bf899 (deleted _apply_monkey_patch_torch_reductions from
update_weight_from_tensor.py). With that helper gone, two more references
are now dead in the vime IPC weight-transfer path:

1. hf_weight_iterator_direct.py:48 called monkey_patch_torch_reductions()
   at the top of _get_megatron_full_params(). On vime this never has effect
   on the IPC handle path: _get_megatron_full_params only runs NCCL
   broadcast/all_gather collectives (no cross-process pickling), and the
   chunks it returns are subsequently sent via PR #48's UUID-keyed
   {gpu_uuid: reduce_tensor(weight)} dict that vLLM's receiver routes by
   physical UUID + explicit args[6] overwrite. The call survives in slime/
   miles upstream because their downstream path pickles tensors through
   sglang's MultiprocessingSerializer.serialize (ForkingPickler →
   reduce_tensor), where the patched encoding/decoding does real work; PR
   #48 does not use that pipeline, so the call here was incidentally
   inherited rather than functionally required.

2. slime/backends/megatron_utils/sglang.py's monkey_patch_torch_reductions
   re-export + __all__ entry now have no remaining importers in vime.
   Remove them.

Also (this commit, B):

3. vllm_engine.py's update_weights_from_tensor docstring referred to
   "closures injected by _apply_monkey_patch_torch_reductions" as the
   reason for cloudpickle. That helper is gone; the cloudpickle is still
   correct because reduce_tensor returns a (rebuild_fn, args) tuple where
   the rebuild_fn is a module-level callable that JSON can't serialise.
   Update the docstring to reflect the actual reason.

Net behaviour: identical — the deletions remove dead code paths. The
patch_torch shim itself is still importable for callers outside vime
(none currently in this tree).

Cross-references:
- #29 — issue documenting the no-op stub
- 39bf899 — prior commit that deleted the helper from update_weight_from_tensor.py

* docs(weight-sync): drop sglang refs + add IPC call-stack notes

- update_weight_from_tensor.py module docstring: rewrite from a pure
  vLLM perspective. Step (2) now spells out the merged
  {uuid_G0: handle, uuid_G1: handle, ...} dict + collective_rpc fan-out
  inside the vLLM server, and step (3) makes version-with-data atomicity
  explicit. Drop the "match slime's sglang_engine signature" framing.
- vllm_engine.update_weights_from_tensor: collapse the long sglang-vs-vLLM
  compare block into a focused docstring describing what the POST does
  and why ipc_handles needs cloudpickle. Add a Chinese call-stack
  walkthrough (trainer → ★this method★ → server collective_rpc → per-TP
  worker receive_weights) and note that `node_rank != 0` is a dead branch
  since VLLMEngine pins node_rank=0 (see PR #48 review comment for the
  follow-up cleanup).
- Hoist base64/cloudpickle imports to module scope so the hot path no
  longer pays the per-call import overhead.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

---------

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: SamitHuang <285365963@qq.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: SamitHuang <285365963@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
momo609 pushed a commit that referenced this pull request Jun 8, 2026
* tests + CI: complete sglang→vllm rename across tests/ and .github/

Split from PR #18 (gcl/clean-sglang). One of 4 PRs splitting the
original PR #18 by content area: docs (#38) / examples (#39) /
**tests+CI** / core runtime. 42 files / +~750 / -~700.

These are bundled in a single PR because the CI workflows reference
test file names by string — splitting them would create a window where
either tests are renamed but CI still points at the old names, or vice
versa, breaking CI mid-roll.

What this PR does:

(A) tests/ (38 files):

- Mechanical CLI-flag rename: --sglang-* → --vllm-* equivalents in all
  test scripts (matches the table now used in scripts/ and examples/).
- Variable rename: SGLANG_ARGS → VLLM_ARGS where present.
- 4 file renames (R086-R091, all >85% similarity):
    test_qwen2.5_0.5B_opd_sglang.py        → test_qwen2.5_0.5B_opd_vllm.py
    test_qwen2.5_0.5B_sglang_config.py     → test_qwen2.5_0.5B_vllm_config.py
    test_qwen2.5_0.5B_sglang_config_distributed.py
                                           → test_qwen2.5_0.5B_vllm_config_distributed.py
    test_sglang_config_mixed_offload.py    → test_vllm_config_mixed_offload.py
    test_sglang_config_mixed_offload_ft.py → test_vllm_config_mixed_offload_ft.py
    tests/utils/test_sglang_config.py      → tests/utils/test_vllm_config.py
- 2 new tests for the IPC weight-transfer path landed in PR #18:
    tests/test_update_weight_from_tensor.py
    tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py
  (These are PR #22 / colocate-IPC test coverage; the production code
  the slim PR #18 ships will rely on the same code from PR #22.)

(B) .github/ (4 files):

- workflows/conda-ci.yml: container image lmsysorg/sglang → vime
  (inferactinc/public:vime-vllm-cu129-latest).
- workflows/pr-test.yml + pr-test.yml.j2 (template):
    * Container images (slimerl/slime[-test]:latest → vime image) on
      every job that ran on the sglang-era base.
    * e2e-test-sglang-config job → e2e-test-vllm-config job (renamed
      label `run-ci-sglang-config` → `run-ci-vllm-config`; matrix
      `test_file` entries updated to point at the renamed test files
      in (A)).
    * e2e-test-megatron + e2e-test-image matrices: `_opd_sglang.py`
      entries → `_opd_vllm.py`.
- ISSUE_TEMPLATE/bug_report.yml: drop the "SGLang version (if
  relevant):" environment field, add "vLLM version:" and
  "vllm-router version:" lines. (PR #36 already changed
  "CUDA/ROCm version" → "CUDA version" earlier; that change is
  preserved.)

Sgl residue intentionally kept (4 hits — all anti-regression
assertions that prove sglang code paths are gone, not residual
references to bring back):

- tests/test_update_weight_from_tensor.py:753 — comment "The vLLM IPC
  implementation must NOT contain sglang-style Gloo gather code".
- tests/unit/backends/vllm_utils/test_arguments.py:233-237 — three
  assertions that --sglang-router-ip, --sglang-router-port, and
  sglang_router_ip are NOT present in the argument parser.

Tests + CI must land together; splitting them risks a window where
the CI matrix references test files by names that don't exist yet
(or no longer exist). After this lands, the test_file string in CI
matches the test files on disk.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

* on_policy_distillation: port from SGLang to vLLM /v1/completions

Follow-up on the test rename in this PR:
test_qwen2.5_0.5B_opd_sglang.py → test_qwen2.5_0.5B_opd_vllm.py.

The test only spawns a vLLM teacher and exercises the OPD pipeline; the
real broken piece was slime/rollout/on_policy_distillation.py, which
PR #18 left in SGLang request/response shape:

  request fields:
    "max_new_tokens": 0          (vLLM: "max_tokens")
    "return_logprob": True       (sglang-only)
    "logprob_start_len": 0       (sglang-only)
  response parsing:
    reward["meta_info"]["input_token_logprobs"]   (sglang shape)

vLLM 0.21 supports the same workflow natively via `prompt_logprobs`:

  request to POST /v1/completions:
    {
        "model": <teacher>,
        "prompt_token_ids": sample.tokens,
        "max_tokens": 1,
        "temperature": 0,
        "prompt_logprobs": 1,
        "logprobs": 0,
        "skip_special_tokens": False,
    }
  response:
    response["choices"][0]["prompt_logprobs"]   # list[dict[int, Logprob] | None]
      where Logprob is {"logprob": float, "rank": int, "decoded_token": str}

References checked against vllm source:
  - reference/vllm/vllm/entrypoints/openai/completion/protocol.py:91
    (request: prompt_logprobs: int | None)
  - reference/vllm/vllm/entrypoints/openai/completion/protocol.py:487
    (response: prompt_logprobs: list[dict[int, Logprob] | None] | None)
  - reference/vllm/vllm/logprobs.py:13
    (Logprob dataclass: logprob/rank/decoded_token)

Implementation notes:

1. JSON serializes int dict keys as strings, so `_logprob_for_token`
   tries both `pos_entry.get(token_id)` and `pos_entry.get(str(token_id))`.

2. `pos_entry` is `None` at position 0 (no prior context) — handled
   explicitly. We also gracefully degrade if a token at position `i` is
   not in the top-1 logprob dict (falls back to 0.0, same as the prior
   sglang code would do).

3. The Logprob dataclass `decoded_token` field is unused; we only read
   `.logprob`. Both dict and `Logprob` shapes are accepted in case the
   server uses a flatter serialization toggle.

4. `args.opd_teacher_model` is the new model-name arg; falls back to
   `args.hf_checkpoint` if not set, mirroring how vime's other rollout
   paths derive the model name.

Smoke-tested `_logprob_for_token` locally:
  - None entry → 0.0
  - int key + dict value → logprob
  - str key (JSON shape) → logprob
  - missing token → 0.0
  - flattened float value → float

Also drops 3 lines from
tests/unit/backends/vllm_utils/test_arguments.py: the
`--sglang-router-ip`/`--sglang-router-port`/`sglang_router_ip` anti-
regression assertions. Once the slim PR #18 lands and sglang is gone
from the runtime, those assertions are vacuous; treating sglang as
non-existent per the project policy.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

* tests: drop duplicate smoke test updates from PR40

* test(update_weight_from_tensor): drop stale _apply_monkey_patch_torch_reductions patch

The inner ``with patch(f"{MODULE_PATH}._apply_monkey_patch_torch_reductions"):``
context in _run_update suppressed a helper call that PR #48 has since deleted
from update_weight_from_tensor.py (commit 39bf899 on aoshen/align-ipc-rpc-with-slime).
After that PR lands the patched attribute won't exist and this line raises
AttributeError. Remove it now so the test survives PR #48 merge.

The ``sglang_mod.monkey_patch_torch_reductions = MagicMock()`` stub on the
fake sglang module is intentionally kept: on this branch the production code
still imports it via ``from ..sglang import monkey_patch_torch_reductions``
(both update_weight_from_tensor._apply_monkey_patch_torch_reductions on
PR #40's view of main, and hf_weight_iterator_direct.py at module level).
Removing the stub here would break the test on PR #40 alone; it can be
dropped in a follow-up once PR #48 finishes removing every import site.

Tests: ``tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py``
all 6 pass with this change applied to gcl/pr18-tests-ci HEAD.

* tests: drop duplicate top-level test_update_weight_from_tensor.py

The 786-line tests/test_update_weight_from_tensor.py is a stale rebase
leftover from the original PR #18 branch — it predates the IPC test
file PR #22 landed at the canonical unit-test path
(tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor.py)
and predates PR #48's single-RPC weight-version contract.

Comparing the two:
* Both stub sys.modules / torch.distributed at module import time, so
  having two files compounds the test-isolation issue Gemini raised
  (PR #40 comment #1).
* Coverage overlaps materially (e.g. test_ipc_init_called_on_first_update_only
  ≈ test_ipc_init_runs_once — same invariant, different wording).
* The nested file is up-to-date with PR #48's RPC contract
  (update_weights_from_tensor.remote(**fields, weight_version=...));
  the top-level file still uses the pre-#48 lifecycle shape and does
  not exercise the coordinator slot fields.
* The nested path matches repo convention: tests/unit/ for mock-only
  unit tests, tests/ top level for e2e scripts.

Closes Gemini comment #1 on PR #40. Gemini comment #2 (the same stub
pattern in the surviving nested file) is a pre-existing issue from
PR #22 / #48 and out of scope for this rename PR — to be addressed
in a follow-up that converts _install_stubs() to an autouse
module-scoped fixture with save/restore.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* test(vllm_config): use real get_model_url default endpoint /inference/v1/generate

get_model_url defaults to /inference/v1/generate (PR #18), not /v1/completions.
Aligns this test with PR #18's test_vllm_config.py so the two PRs no longer
conflict on this file and the assertion matches the actual runtime default.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Drop on_policy_distillation.py from tests+CI PR (now owned by runtime PR #18)

The OPD vLLM /v1/completions migration is a runtime change; it was folded into
the core-runtime PR (#18). Restore this file to main here so the two PRs no longer
overlap on it. #18 merges first, so this lands via #18.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop unit-test files now owned by runtime PR #18

test_vllm_config.py + the plugin_contracts tests are coupled to #18's runtime
rename (they import vllm_config / vllm_rollout, which #18 creates). They live in
#18; remove them here so the two PRs don't overlap. #18 merges first.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test: restore vLLM rollout args dropped during sglang→vllm rename

The mechanical sglang→vllm rename dropped several rollout knobs instead of
mapping them to their vLLM equivalents, weakening CI coverage (cuda-graph
capture caps, speculative decoding, expert parallel). Restore them using the
mapping established by the converted production scripts on main
(run-glm4.7-30B-A3B.sh / run-glm5-744B-A40B.sh), verified against vLLM
AsyncEngineArgs:

  --sglang-cuda-graph-max-bs N            -> --vllm-max-cudagraph-capture-size N
  --sglang-cuda-graph-bs a b c            -> --vllm-cudagraph-capture-sizes a b c
  --sglang-ep-size N                      -> --vllm-enable-expert-parallel
  --sglang-speculative-* (eagle)          -> --vllm-speculative-config '{"method":"eagle","num_speculative_tokens":K}'

Also:
- glm4.7 pd: fix --vllm-max-num-seqs (was 8, taken from cuda-graph-max-bs;
  --sglang-max-running-requests was 16) and split out cuda-graph capture.
- fix sglang→rollout mis-renames in temp-file prefixes (→ vllm_*).
- test_vllm_config: rename test_update_weights_default_true →
  test_update_weights_defaults_to_none (it asserts `is None`).

Dropped sglang flags with no vLLM equivalent (enable-dp-lm-head,
moe-dense-tp-size, watchdog-timeout, mamba-scheduler-strategy,
disaggregation-transfer-backend, enable-metrics) stay dropped; PD KV-transfer
is driven by --prefill-num-servers + the --vllm-config prefill/decode topology.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(plugin_contracts): migrate from sglang_rollout to vllm_rollout

The three plugin-contract tests still imported slime.rollout.sglang_rollout
and called install_stubs(with_sglang_router=True), but _shared.install_stubs
already dropped that parameter — so all three failed at collection
(TypeError: unexpected keyword 'with_sglang_router'). Complete the migration:

- install_stubs(with_sglang_router=True, ...) -> install_stubs(...)
- import generate_and_rm / generate_rollout from slime.rollout.vllm_rollout
- default rollout/eval path string -> slime.rollout.vllm_rollout.generate_rollout
  (matches runtime default at slime/utils/arguments.py:233)
- FakeGenerateState: sglang_enable_deterministic_inference ->
  vllm_enable_deterministic_inference, with group_sampling_seeds defaulting to
  None and gated on the flag (mirrors the already-migrated
  tests/unit/rollout/test_vllm_rollout.py).

All 34 plugin-contract cases pass (were 3 collection errors before).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(update_weight_from_tensor): drop stale slime…megatron_utils.sglang mock

The test pre-registered a sys.modules mock for
slime.backends.megatron_utils.sglang (monkey_patch_torch_reductions), left over
from when update_weight_from_tensor imported it. The module under test no longer
imports that module (its real deps are get_gloo_group / HfWeightIteratorBase /
update_weight_from_distributed), so the mock is dead. Removing it makes tests/
and .github/ fully sglang-free. Test still passes (7/7).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* [Clean] Remove SGLang runtime code

Rebuilt against current main so the PR contains only the SGLang runtime
removal -- the docs / tests-ci / examples / scripts / docker portions were
split into separate PRs that have since merged.

- Delete dead SGLang server/runtime code: sglang_utils/{arguments,sglang_engine}.py,
  rollout/sglang_rollout.py, the megatron_utils/sglang.py re-export shim, and all
  docker/**/sglang.patch files.
- Rename the rollout config module sglang_utils/sglang_config.py ->
  vllm_utils/vllm_config.py (SglangConfig -> VllmConfig, _resolve_sglang_config ->
  _resolve_vllm_config, --sglang-config -> --vllm-config); inline the
  GPU_MEMORY_TYPE_* constants in rollout.py.
- Add megatron_utils/fp8_helpers.py for the UE8M0 fp8 helpers formerly re-exported
  through the sglang shim; repoint quantizer_fp8 to it.
- Swap sglang_router -> vllm_router in http_utils/wandb_utils; drop the dead
  sglang-router dependency from requirements.txt.
- Finish the SGLang->vLLM rename in the runtime so it is internally consistent and
  matches the tests landing in the tests/CI PR:
  * router args --router-* -> --vllm-router-* (vllm_router_ip/port/timeout);
  * get_model_url reads vllm_model_routers (aligning with rollout.py);
  * --opd-type sglang -> vllm; engine_overrides rename;
  * sglang_enable_deterministic_inference -> vllm_enable_deterministic_inference,
    wired to a real --vllm-enable-deterministic-inference flag (exports
    VLLM_BATCH_INVARIANT=1);
  * consistent_hash session-id routing uses vllm-router's x-session-id header;
  * drop dead trace helper build_sglang_meta_trace_attrs; de-SGLang comments/docstrings.
- Rename test_sglang_config.py -> test_vllm_config.py and de-SGLang the
  plugin-contract tests.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* Address review: finish de-SGLang + fold OPD/router-policy into runtime

- naming: replace residual generic "rollout engine"/"engine"/"comm" wording
  with concrete vLLM (engine_overrides -> vllm_overrides; arguments help text;
  http_utils comments; rollout.py "inference workers"). sglang->vllm is correct,
  sglang->generic is not.
- megatron_to_hf: drop the q_a_proj/kv_a_proj_with_mqa pairing + _cached_tensors
  global. That was sglang-only: sglang's loader torch.cat's both shards within a
  single load_weights call (needs them co-bucketed), whereas vLLM loads each shard
  independently via stacked_params_mapping into fused_qkv_a_proj. Also fix the
  misleading "merge into single fused name" comment.
- docker/Dockerfile: remove now-dead sglang/sglang-router --no-deps stubs + the
  build-time `import sglang` smoke check (slime no longer imports sglang_router).
- OPD: migrate on_policy_distillation.py teacher logprobs to vLLM /v1/completions
  (prompt_logprobs) instead of sglang return_logprob / meta_info.input_token_logprobs.
- routing replay: register --vllm-router-policy (dest=router_policy) so the
  consistent_hash x-session-id session-affinity path is actually wired (was dead).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Review follow-ups: mirror slime vllm_config parsing + restore vLLM process cleanup

- vllm_config.from_yaml: drop the needless `models_raw` intermediate and iterate
  `data["vllm"]` directly, restoring the "Accept both server_groups / legacy
  engine_groups" comment -- mirrors slime's sglang_config.from_yaml line-for-line.
- command_utils.execute_train: re-add a process kill for leftover rollout engines
  as `pkill -9 -f "vllm serve"` (the old `pkill -9 sglang` was dropped with no
  vLLM equivalent), so stale engines don't hold GPUs/ports across runs.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Drop test changes from runtime PR; tests live in the tests+CI PR (#40)

The plugin_contracts tests and the test_sglang_config -> test_vllm_config rename
are coupled to the test/CI rename effort and are owned by #40. Restore them to
main here so #18 is purely the SGLang runtime removal. #18 merges first; #40
rebases and re-lands the vLLM test versions.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fp8_helpers: copy SGLang verbatim (fix import crash); pin vLLM deep_gemm env

- fp8_helpers.py: replace the bespoke rewrite with SGLang's exact implementations
  of quant_weight_ue8m0 / transform_scale_ue8m0 and their DeepGEMM helpers
  (per_block_cast_to_fp8, ceil_to_ue8m0, ceil_div, ceil_align, the torch-impl
  packer). deep_gemm is imported lazily inside the functions (as SGLang does), so
  module import no longer requires deep_gemm. This fixes the module-level
  `NameError: _get_tma_aligned_size` that crashed `import megatron_to_hf` on any
  deep_gemm image, and drops the invented sf-stride fixup block that was not in
  upstream. Only should_deepgemm_weight_requant_ue8m0 stays vLLM-adapted
  (is_deep_gemm_e8m0_used) since SGLang's reads SGLang-internal deep_gemm_wrapper.
- vllm_engine.launch_server_process: set VLLM_USE_DEEP_GEMM=1 +
  VLLM_DEEP_GEMM_WARMUP=relax explicitly (setdefault) alongside VLLM_BATCH_INVARIANT,
  replacing SGLang's removed deep_gemm precompile/warmup envs. All vLLM engine env
  now lives in the subprocess env builder (single source of truth).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fp8_helpers: revert to vLLM impl + fix the import NameError; add vLLM trace attrs

- fp8_helpers.py: keep the vLLM-based implementation (uses vllm.utils.deep_gemm,
  consistent with the vLLM runtime) rather than the SGLang verbatim copy. Fix the
  module-level crash: the `try` block referenced `_get_tma_aligned_size` before it
  was bound (the "pre-imported with fallback" import was never written), which
  raised NameError whenever deep_gemm imported successfully -- and NameError is not
  caught by `except ImportError`, so `import megatron_to_hf` crashed on any
  deep_gemm image. Replace the bogus self-assignment with the real import:
  `from vllm.utils.deep_gemm import get_tma_aligned_size as _get_tma_aligned_size`.
- trace_utils/vllm_rollout: add build_vllm_meta_trace_attrs and attach finish_reason
  + token usage to the vllm_inference_generate span (mirrors SGLang's
  build_sglang_meta_trace_attrs; vLLM responses lack the pd_* timing, which lives
  in vLLM's own OTLP traces).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(opd): score teacher via /inference/v1/generate with prompt_logprobs

Move the vllm OPD teacher path off the OpenAI /v1/completions endpoint onto
vime's native /inference/v1/generate (the same endpoint the rollout engines
use), and fix three latent issues:

1. model field: /inference/v1/generate takes `model` as OPTIONAL. Stop
   defaulting to args.hf_checkpoint (the *student* name, which mis-names a
   teacher!=student server). Add --opd-teacher-model; send `model` only when
   set, otherwise omit it (single-model teacher servers use their loaded model).
2. multimodal: the old code sent image_data to a token-only endpoint, which is
   invalid. Raise NotImplementedError until the
   /v1/chat/completions/render -> /inference/v1/generate flow is wired (mirrors
   slime.rollout.vllm_rollout.generate).
3. logprob robustness: read top-level GenerateResponse.prompt_logprobs, assert
   it is present and length-aligned with token_ids, assert the per-sample tensor
   covers response_length, and raise (not silently return 0.0) on a missing
   token logprob. vLLM always includes the actual prompt token in
   prompt_logprobs, so a miss is a real error.

Alignment is unchanged (plp[i] <-> tokens[i], skip pos 0, take [-response_length:]).

Follow-up (separate, in the tests PR): the OPD e2e test must launch a teacher
that exposes /inference/v1/generate and point --rm-url at it.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* chore(clean-sglang): purge SGLang from tools, train scripts, and build infra

tools/: drop dead `args.sglang_enable_ep_moe` shim (read nowhere); reword
profile/replay helpers to vLLM and map analyzer hints to vLLM flags
(--enforce-eager, --gpu-memory-utilization). train{,_async}.py: comments
SGLang -> vLLM.

build infra: remove build_conda.sh (SGLang-only conda path); drop the GB300
sgl-kernel install from the Dockerfile; delete docker/npu_patch/ wholesale.

docker base image: bump to vLLM v0.22.0. justfile ARM recipes now pin the real
multi-arch vLLM base images instead of the dead SGLANG_IMAGE_TAG/
ENABLE_SGLANG_PATCH build-args -- cu129-arm64 -> v0.22.0-cu129-ubuntu2404
(CUDA 12.9), cu13-arm64 -> v0.22.0-ubuntu2404 (the default-CUDA tag is already
CUDA 13.0) + ENABLE_CUDA_13=1. vLLM tags are multi-arch manifests, so docker
selects the arm64 image automatically on an ARM host.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* ci(clean-sglang): drop conda-build workflow (ran deleted build_conda.sh on SGLang image)

The single build-conda job ran `bash build_conda.sh` (removed in the previous
commit) inside an lmsysorg/sglang container. With the SGLang-only conda path
gone, the whole workflow is dead.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* chore(clean-sglang): fix stale SGLang refs in docs/skills; tidy comments

docs/conf.py: point the "edit on GitHub" links at vllm-project/vime instead of
the inherited sgl-project.github.io repo. .claude/skills/*: update the dead
`slime/rollout/sglang_rollout.py` references to `vllm_rollout.py` (the real
default is slime.rollout.vllm_rollout.generate_rollout).

justfile: drop the redundant BASE_IMAGE override on release-cu129-arm64 (it
equalled the Dockerfile default; the multi-arch manifest already resolves
arm64). train{,_async}.py: drop stray "the" in the W&B comment.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* fix(tests): use method=mtp (not eagle) in vllm speculative config

The migrated speculative configs pass no draft `model`, so method=eagle
raises "num_speculative_tokens was provided but without speculative model"
in vLLM's SpeculativeConfig. These models carry embedded MTP layers, so
method=mtp is correct and unblocks the mimo MTP-only-grad test (#19).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(cleanup): target renamed vLLM subprocesses in pkill so VRAM is freed

vLLM's set_process_title() renames the VRAM-holding subprocesses
(VLLM::EngineCore, VLLM::Worker_TP*, vllm::router), so their cmdline no
longer contains "vllm serve". The previous `pkill -9 -f "vllm serve"`
matched only the launcher and left engine/worker children holding GPU
memory, leaking it into the next run — masked only by the indiscriminate
`pkill -9 python`, which is unsafe on colocate/shared nodes.

Match both the launcher and the renamed children with
`pkill -9 -f '[v]llm serve|VLL[M]::'`; the [v]/[M] bracket trick keeps the
pattern from matching pkill's own cmdline. This makes the broad python
kill unnecessary, so its already-commented-out lines are removed.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(speculative): use method=mtp (not eagle) for embedded-MTP models

vLLM's SpeculativeConfig requires an explicit draft `model` for
method=eagle; with only num_speculative_tokens set it raises
"num_speculative_tokens was provided but without speculative model".
The migrated configs in scripts/examples/docs pass no model, so they must
use method=mtp, which reuses the target checkpoint's embedded MTP layer
(DeepSeek-R1, GLM-4.x-MoE, MiMo, Qwen3-Next/3.5).

The two docs examples that pass an explicit "model" are genuine eagle
usage and are left unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* fix(vllm): launch each rollout engine with its ServerGroup's per-group TP

launch_server_process / _init_normal derived tensor-parallel size and
CUDA_VISIBLE_DEVICES from the global --rollout-num-gpus-per-engine, ignoring
the per-engine num_gpus_per_engine already carried on the VLLMEngine actor.

A ServerGroup configured with num_gpus_per_engine greater than the global
flag (e.g. tp=2) therefore launched as tp=1, while the NCCL weight-sync
rendezvous sized world_size from engine_gpu_counts (the per-group value).
The two disagreed: the trainer waited for a rank the under-sized engine
never started, so init_weight_transfer_engine hung for 300s
("3/4 clients joined") and the job failed.

Honor the per-engine num_gpus_per_engine at launch, falling back to the
global flag when unset (matches the SGLang path and PR #66's
_compute_server_args).

Verified on H200: tests/test_qwen2.5_0.5B_vllm_config_distributed now
launches engine0 tp=2 / engine1 tp=1, update_weights completes in 1.1s
(was a 301s timeout), and rollout+eval proceed.

AI assistance (Claude Code) was used for this change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(ckpt): add --dist-ckpt-optim-fully-reshardable for PAO+offload save/load

test_qwen3_4B_ckpt.py uses precision-aware optimizer + cpu-offload
(HybridDeviceOptimizer). Under the default dp_reshardable (bucket-centric)
optimizer sharding, save/load produce unequal-length param_state lists, so
dist-ckpt load fails with
"Cannot merge two lists with different lengths (81 and 79)".

fully_reshardable is model-centric and immune to bucket-layout changes.
Verified on the r3 image (Megatron-LM 0.16.0rc0 @ 1dcf0da): save+load both
succeed, and source review confirms master_param / step / HybridDeviceOptimizer
sync are handled on this path. This is the flag described in PR #50 that was
never actually merged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* refactor(router-args): hybrid naming — vllm_ for ip/port, bare router_ for timeout

vllm-router's RouterArgs.from_cli_args only supports prefix "" or "router_"
(never "vllm_router_"), and excludes host/port from its CLI via
exclude_host_port=True. So:

- --vllm-router-ip / --vllm-router-port keep the vllm_ prefix: RouterArgs does
  not own these CLI flags, vime does (populated via _start_router's manual
  router_args.host/port assignment), so the vllm_ prefix is free and marks them
  as vime-owned endpoint config.
- --router-request-timeout-secs goes bare (dest router_request_timeout_secs): it
  is a genuine RouterArgs field, so it shares the --router-* namespace with
  policy / cache_threshold / retries / … and flows through from_cli_args like
  the other knobs.
- --vllm-router-policy keeps dest=router_policy (unchanged).

Also fixes conftest fixture to seed vllm_router_ip/port (was bare router_ip/port,
which never matched the vllm_engine reader) and updates README/README_zh prose.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
@aoshen02
aoshen02 deleted the update_weights_tensor branch June 8, 2026 14:17
aoshen02 added a commit that referenced this pull request Jun 10, 2026
…bject patch

arguments.py (361→346 lines):
- Import FlexibleArgumentParser from vllm.utils.argparse_utils; use it in
  vllm_parse_args() and get_vllm_cli_action_table() so vLLM's deprecated
  kwarg is handled natively on Python 3.12 without a shim
- Remove _ARGPARSE_UNSUPPORTED_KWARGS + _strip_unsupported_argparse_kwargs
- Remove import logging / logger (unused)

reloadable_process_group.py:
- Drop dist.gather_object monkey-patch (added by PR #22, never in slime;
  callers explicitly use all_gather_object to stay on the patched path)

Architecture note: subprocess vllm serve kept intentionally — run_server()
is not in vllm.__all__ and changes across minor versions; AReaL uses the
same Popen pattern for the same reason.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
aoshen02 added a commit that referenced this pull request Jun 11, 2026
…bject patch

arguments.py (361→346 lines):
- Import FlexibleArgumentParser from vllm.utils.argparse_utils; use it in
  vllm_parse_args() and get_vllm_cli_action_table() so vLLM's deprecated
  kwarg is handled natively on Python 3.12 without a shim
- Remove _ARGPARSE_UNSUPPORTED_KWARGS + _strip_unsupported_argparse_kwargs
- Remove import logging / logger (unused)

reloadable_process_group.py:
- Drop dist.gather_object monkey-patch (added by PR #22, never in slime;
  callers explicitly use all_gather_object to stay on the patched path)

Architecture note: subprocess vllm serve kept intentionally — run_server()
is not in vllm.__all__ and changes across minor versions; AReaL uses the
same Popen pattern for the same reason.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>
CalvinXKY pushed a commit that referenced this pull request Jun 11, 2026
* sync: complete slime #1920, #1967, #1985 — lint, fully_async, PYTHONUNBUFFERED

Traced 3 individual diffs back to their source slime PRs (以点带面) and
synced all remaining changes from each:

slime #1985 (make tests shorter):
  - Wrap NamedTemporaryFile across 3 lines in test_vllm_config_mixed_offload_ft.py
  - Remove extra blank line in test_vllm_config_mixed_offload.py
  (parameter shortening already synced in vime PR #218)

slime #1920 (move fully_async example to main codebase):
  - Rewrite README.md to match upstream (qwen2.5-0.5B, not qwen3-4b)
  - Add run-qwen2.5-0.5B-fully_async.sh with proper vLLM translations
  - Delete run-qwen3-4b-fully_async.sh (upstream removed it)

slime #1967 (fix PYTHONBUFFERED typo):
  - Fix PYTHONBUFFERED=16 → PYTHONUNBUFFERED=1 in 3 scripts:
    run-glm4.7-30B-A3B.sh, run-glm4.7-355B-A32B.sh, run-minimax-m2.sh
  (command_utils.py already fixed in prior sync)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: correct §2.4 flag translations in analyze_profile.py

--enforce-eager → --vllm-enforce-eager
--gpu-memory-utilization → --vllm-gpu-memory-utilization

These are diagnostic hint strings, not CLI invocations, but should still
use the canonical vime flag names (§2.4 translation rules).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix: mirror slime underline + revert group_id comment to rollout_id

- analyze_profile.py: match slime's 30-char underline (was 28)
- run_qwen36_35b_a3b_swe_8nodes.sh: revert group_id→rollout_id in
  comment (partial sync of slime #2013, target commit 44d29ee5)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* refactor(arguments): use FlexibleArgumentParser; revert dist.gather_object patch

arguments.py (361→346 lines):
- Import FlexibleArgumentParser from vllm.utils.argparse_utils; use it in
  vllm_parse_args() and get_vllm_cli_action_table() so vLLM's deprecated
  kwarg is handled natively on Python 3.12 without a shim
- Remove _ARGPARSE_UNSUPPORTED_KWARGS + _strip_unsupported_argparse_kwargs
- Remove import logging / logger (unused)

reloadable_process_group.py:
- Drop dist.gather_object monkey-patch (added by PR #22, never in slime;
  callers explicitly use all_gather_object to stay on the patched path)

Architecture note: subprocess vllm serve kept intentionally — run_server()
is not in vllm.__all__ and changes across minor versions; AReaL uses the
same Popen pattern for the same reason.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* ci: revert aoshen02 CI additions; align J2 template with slime

Reverts PR #26 (pre-commit gate) + PR #110 (e2e-test-unit).
Syncs from slime: opened/reopened trigger types, --pull=always,
pip deps (requests ray safetensors), cpu-unittest rename.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* docs(docker): translate Chinese to English in README

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(ci): stub vllm_router in plugin_contracts for cpu-unittest

Mirrors slime's with_sglang_router stub pattern: add with_vllm_router
kwarg to install_stubs() and pass it from test_plugin_generate_contracts.
vllm_rollout.py imports vllm_router at module level; without the stub
the cpu-unittest job fails with ModuleNotFoundError.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* fix(ci): add missing with_vllm_router=True stub to remaining plugin contracts

test_plugin_rollout_contracts and test_plugin_path_loading_contracts both
import vllm_rollout (which has bare `import vllm_router` at module level)
but were not passing with_vllm_router=True to install_stubs — mirroring
the same gap fixed in test_plugin_generate_contracts.

Mirrors slime: all three tests pass with_sglang_router=True.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* docs(coding_agent_rl): align generate.py and README.md wording with slime

- generate.py: drop backtick-wrapping around /inference/v1/generate in module docstring; collapse to one line matching slime style
- README.md: condense rollout-max-*-len paragraph (remove verbose "sampling-params" / "generation length" verbiage, restore `max_tokens` inline like slime's `max_new_tokens` form)
- README.md: trim vLLM response-structure detail (choices[0].token_ids / choices[0].logprobs.content[i].logprob) from token-out bullet — matches slime's abstraction level
- README.md: add missing 3-line unit-test sentence after provenance paragraph (slime has it, vime was missing it)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* refactor(arguments): remove dist_ckpt_optim_fully_reshardable warning block

Not present in slime; drop to maintain mirror parity.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* refactor(rollout): remove redundant assertions and finalization re-assignments

on_policy_distillation.py: drop two assertions not present in slime
(len(plp)==len(sample.tokens) and len(t_log_prob)>=response_length).

vllm_streaming_rollout.py: drop 9-line finalization block that re-set
sample.tokens/response/response_length/rollout_log_probs/loss_mask after
the streaming loop had already written the same values.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

* style: apply black formatting to fix pre-commit CI

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: aoshen02 <aoshen@inferact.ai>

---------

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants