feat(sglang): support checkpoint-engine refit

[tianyi-zhang-02]

What does this PR do?

Adds checkpoint-engine refit for non-colocated SGLang:

• one NIXL receiver per SGLang rank, with weight-name alignment when rank-local bucket boundaries differ;
• dense and full-weight MoE refit;
• the full SGLang update session (pause → KV invalidation → begin_weight_update → transfer → end_weight_update → resume);
• setup-time guards for the supported envelope: one node per engine, dp_size=1, pp_size=1, no expert-weight sharding, and no in_place pause mode;
• real tp_size=1 and tp_size=2 functional coverage.

The receive loop stays outside the actor class, so this does not turn the existing synchronous SGLang worker into an asyncio Ray actor.

The legacy refit-metadata gather is skipped for SGLang: this path does not use it, and gathering every DTensor shard previously OOMed the TP2 test before the server could receive a weight. The functional test below then reached both refits and completed training.

This covers the SGLang half of #3288. It has no unmerged dependency; the Megatron-generation half remains out of scope.

Validation

Current head: 97252916e1aa6c9082cb2716d2d82bb05837a5a1, including main at fd7112c374af23b3bcc392637fb737495f8a3f31.

On 2026-09-21, both changed checkpoint-engine unit-test files passed on Runpod Secure Cloud: 50 passed, 0 failed/skipped. Environment: 1× H100 80 GB, nvcr.io/nvidia/nemo-rl:v0.7.0, driver 580.126.09, CUDA 13.0, PyTorch 2.11.0+cu130. The current-main test fixture required pytest==9.0.3 and pinned nemo-lens commit b85578f, installed only in the temporary Pod. The refresh also updates one factory-routing test for main's newly supported non-colocated Megatron route.

Earlier, at head 5925bcf5b07fd438e9f385fb2bd1e1d21531ce33, tests/functional/grpo_sglang_nixl_non_colocated_tp2.sh passed on a Runpod 4×H100 host with 2 SGLang TP ranks and 2 DTensor policy ranks. That run used NCCL_NVLS_ENABLE=0, matching the repository's H100 recipe workaround.

Step	NIXL refit	loss	train/token_mult_prob_error
1	0.383 s	0	1.0145671
2	0.186 s	4.77e-6	1.0133847

Both steps completed without OOM, traceback, or NCCL error; the metric stayed below the functional threshold of 1.05. The 4-GPU functional test was not rerun after the main-only refresh and test expectation change; this session provides 1-GPU unit coverage only.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[tianyi-zhang-02]

GPU follow-up

I don't have an environment where tests/functional/grpo_sglang_nixl_non_colocated.sh runs end to end, so I validated the transport layer with a focused 2-GPU probe. Ran against nixl 1.3.0, the version this PR pins.

Setup

Single node, 2 GPUs, torch 2.7.1+cu128, nixl 1.3.0 (UCX backend, which bound to real IB HCAs plus cuda_ipc — not a tcp/shm fallback).

Topology mirrors what the PR introduces:

• 2 policy-side senders, one NIXLCheckpointEngine each (cuda:0, cuda:1)
• one actor holding 2 rank-local receivers — the piece this PR adds
• the SGLang HTTP server is stubbed by a separate process that deserializes the flattened bucket from the CUDA IPC handle and compares tensors, so the IPC handoff is real and cross-process

To exercise _aligned_checkpoint_engine_batches the shards are deliberately asymmetric (rank 1 carries 2× rank 0's rows). The probe counts each rank's transport batches and fails if they come out equal, since in that case the alignment path was never exercised at all.

Result

create_backend(UCX): OK
refit returned in 2.91s
transport batches per rank : [5, 6] (diverged=True)
aligned batches consumed   : 9
cross-process IPC loads    : 9
rank 0: 9 weights bit-exact (6912 rows, device=cuda:0)
rank 1: 9 weights bit-exact (13824 rows, device=cuda:1)
kv cache invalidated       : True
engines finalized cleanly
RESULT: PASS

Every weight arrived bit-exact on both ranks, exactly once, per-rank content distinct, across genuinely different transport boundaries (5 vs 6 batches).

Caveats

1. The shard asymmetry is synthetic. Bucket boundaries are a function of each rank's (name, dtype, size) sequence, so ranks with identical shapes produce identical boundaries and the alignment path degenerates to a pass-through. This shows the mechanism is correct when fed divergence; it does not establish how often real workloads diverge. If uneven TP shards cannot actually occur in the supported configurations, the aligner could be replaced with an assert — I'd rather be told that than guess.
2. No GRPO loop, no real SGLang server, no train/token_mult_prob_error check. The functional script still needs a CI run. It is now registered with L1_Functional_Tests_SGLang.sh; previously nothing referenced it, so it had never run.
3. create_backend("UCX") failed intermittently with NIXL_ERR_BACKEND and cleared on re-run. Two hypotheses tested and both refuted: process-local configuration (a matrix over OpenSSL preload × CUDA init × nixl library paths came back clean, including the exact configuration that had failed the run before) and IB availability (several nodes sampled, all showing identical /sys/class/infiniband with the same ports ACTIVE, all passing). Cause unknown. Note this shares an error string with the UCX_TLS item in the environment notes below, but the failing runs had UCX_TLS unset, so I do not think they are the same thing. Flagging only in case it is a known 1.3.0 flake. Happy to give exact run counts if useful.
4. nixl 1.3.0 + Ray SIGSEGVs during interpreter teardown, after finalize() returns cleanly and after the result is known. The probe calls os._exit() once the verdict is in, so a teardown crash cannot masquerade as a transfer failure. That is in the probe only, not in shipped code.

One proposed change: drop an unreachable branch

In _aligned_checkpoint_engine_batches, by the time control reaches this every rank has either just been refilled with a non-empty batch (empty batches raise above) or is finished, and a finished-and-empty rank is already caught by the break/raise immediately above.

         if any(finished[rank] and not queue for rank, queue in enumerate(pending)):
             raise RuntimeError(
                 "Checkpoint-engine streams ended with different weights across "
                 "SGLang ranks."
             )
-        if any(not queue for queue in pending):
-            continue

         aligned = [[] for _queue in pending]

The reason to delete rather than keep it as a defensive guard: if the invariant ever did break, continue returns to the top, the empty rank is skipped by the refill because it is finished, and the same branch is taken again — a silent infinite loop inside a refit, instead of the loud RuntimeError directly above it. It does not guard anything; it would hide something. Happy to keep it if you disagree.

Environment notes, in case they save someone time

• UCX needs a network-capable transport even for an intra-node refit. UCX_TLS=sm,self,cuda_copy,cuda_ipc fails create_backend("UCX") with NIXL_ERR_BACKEND; adding tcp fixes it, as does leaving UCX_TLS unset.
• nixl 1.3.0 is a dispatcher wheel selecting nixl-cu{torch.version.cuda major}; it lists both nixl-cu12 and nixl-cu13 as unconditional Requires-Dist, so an air-gapped install needs --no-deps plus the one backend matching the image.
• Its binding links libssl.so.3 / libcrypto.so.3; UCX itself is vendored in the wheel. On an OpenSSL 1.1 image OpenSSL is the only gap, and it surfaces as a bare ImportError.
• NIXL requires the same bucket_size_bytes on sender and receiver; mismatched sizes fail with createXferReq: length mismatch at index 0. Consistent with _resolve_bucket_size_bytes computing one value for everyone — and it means divergent boundaries can only originate from differing shard shapes.
• A CUDA IPC handle cannot be re-imported by the process that exported it (CUDA error: invalid device context), so any test stubbing update_weights_from_tensor must decode in a separate process.
• Ray clears CUDA_VISIBLE_DEVICES on num_gpus=0 actors; set RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1 if the actor manages devices itself.

Also happy to push

A property-test module for the alignment logic: randomized per-rank bucket layouts asserting global weight order is preserved, plus a receiver that poisons its recycled buffer to pin down the double-buffer contract in _receive_weight_chunk_batches. That second one is the valuable half — a fake engine that allocates fresh tensors every batch will pass even if the code holds references into memory NIXL is about to overwrite, so without it the whole class of bug is untestable.

One thing worth confirming, since I could not verify it without a real server: SGLang 0.5.12.post1 asserts self._weight_update_in_progress in update_weights_from_tensor, set only by begin_weight_update, which appears nowhere in this tree. The colocated path posts to the same endpoint and presumably passes CI, so either _apply_sglang_compat_patches rewrites it or the dispatch differs. Does my path land on whatever makes the colocated one work?

[tianyi-zhang-02]

Pushed the property tests I mentioned, plus four more gaps. Each one was verified by mutation — the test fails when the behaviour is broken and passes when it is not — so none of them are assertions that cannot fail:

Property	Mutation it catches
Recycled receive buffers	making the aligner prefetch instead of advancing only empty ranks
Multi-dtype batches	posting only the first dtype group
weight_version across refits	bumping per POST instead of per refit
base_gpu_id remapping	using the physical id instead of the local one
Payload index ↔ SGLang rank	transposing the per-rank payload list

The buffer one is the reason I wanted this in the tree. nixl.py hands out views into a rotating pool, so a fake engine that allocates fresh tensors each batch makes that whole bug class invisible — every test passes while weights are silently corrupted. _RecyclingEngine poisons a buffer once it is recycled. The current aligner survives it because it only calls anext for a rank whose deque is already empty, so at most one batch per rank is outstanding; that is load-bearing and was previously unstated and untested.

The payload-index one is the property with the worst failure mode: SGLang indexes serialized_named_tensors by its own TP rank, so a transposed list loads every shard onto the wrong GPU and still returns success.

149 unit tests passing across tests/unit/models/generation/sglang/ and tests/unit/weight_sync/, ruff and pyrefly clean. Still no GRPO loop and nothing cross-node — the functional script needs a CI run.

[yuki-97]

#3330 and #3426 overlap with parts of this work. I am happy to keep this standalone or rebase it onto whichever direction you prefer.

hi @xiuhu17 @Kh4L, could you help to take a look on this?

[Kh4L]

Reviewed at 29e6c3d, following up on @yuki-97's ping.

The direction is right: reusing CheckpointEngineWeightSynchronizer with per-rank NIXL receivers and name-aligned rank-local streams is the correct shape for this, the envelope guards fail at setup as they should, and the mutation-verified unit tests (especially the recycled-buffer one) are genuinely good.

Two blockers, both invisible to the current tests (CI has only run DCO here, so the new functional test has never executed against a real engine):

1. The transfer never opens a begin_weight_update/end_weight_update session — the pinned SGLang scheduler hard-asserts on the first bucket, and skipping end_weight_update skips quantized-layout finalization (details inline).
2. asyncio.run() inside the worker: the mixin's inherited coroutine turns SGLangGenerationWorker into an asyncio Ray actor, so the entry point raises RuntimeError: asyncio.run() cannot be called from a running event loop — reproduced on the locked ray 2.56.1 (details + one-line fix inline).

Also inline: two test-coverage suggestions, two config-conventions items, a doc/comment accuracy fix for what refit_transport: null means when non-colocated, one piece of dead code, and a forward-compat note on the #3613 fault-tolerance interplay.

One nit with no diff line to attach to: the exemplar comment above refit_transport: null in examples/configs/grpo_math_1B.yaml still says nixl is vLLM-only ("Non-colocated vLLM also supports vllm_s3_sparse, vllm_zmq_sparse, and nixl.") — worth updating, and grpo_math_1B_sglang.yaml could carry a short documented refit_transport note since the keys are new on SGLangConfig.

On the overlap question from #3288: #3612 has merged, so this PR is the right vehicle for the SGLang checkpoint-engine path — keeping it standalone on top of the merged seam is the right call. The quantized follow-ups (#3614/#3615) will stack cleanly once the session envelope is in.

[Kh4L]

Blocker: the SGLang path never opens a weight-update session — the first bucket fails on a scheduler assert at the pinned sglang rev.

The sglang rev this repo pins (3003d70f, unchanged by this PR) gates tensor updates in scheduler_components/weight_updater.py:254-256:

assert (
    self._weight_update_in_progress
), "update_weights_from_tensor requires an open begin_weight_update session"

The flag is set only by begin_weight_update, and this path never calls it: prepare_for_generation(tags=[...]) is memory-onload (resume_memory_occupation), not session bracketing, and the mixin then POSTs /update_weights_from_tensor directly. Skipping end_weight_update is independently wrong too — it's what runs postprocess_weight to finalize quantized/MoE kernel layouts after the last bucket (the seam the quantized-refit follow-ups to #3612 rely on).

Fix: wrap the SGLang branch in the same envelope as the sibling _SGLangWeightSynchronizer._refit (sglang_weight_synchronizer.py:150-180): pause_generation → invalidate_kv_cache → begin_weight_update → transfer → end_weight_update in a finally → continue_generation in an outer finally. Both hooks already exist on SGLangGeneration and the worker. The pause also makes the multi-bucket update atomic against in-flight requests: each update_weights_from_tensor takes the server's model_update_lock per call, so requests admitted between buckets would otherwise run on a mixed old/new model, and invalidate_kv_cache()'s /flush_cache 400s while requests are pending.

(The unit tests stub the worker's update_weights_from_tensor, which is why this never surfaced — see the test-file comment.)

[tianyi-zhang-02]

You're right, there was no session here at all. Fixed in fbf667f: the SGLang branch now runs the same sequence as _SGLangWeightSynchronizer._refit, pause then flush then begin, transfer, end_weight_update in a finally, continue_generation in an outer one. c963c83 moves the kv_cache re-onload into that same finally, otherwise a failed refit resumes with the pool still released.

Dropped the worker's post-transfer invalidate_kv_cache() since the flush happens before the buckets now. Also took the in_place rejection from sglang_weight_synchronizer.py:80-83, this path pauses through the same hook and had no guard.

[Kh4L]

Blocker: this asyncio.run() raises RuntimeError: asyncio.run() cannot be called from a running event loop on the first refit.

Mixing this class into the @ray.remote SGLangGenerationWorker makes it an asyncio actor: Ray's has_async_methods uses inspect.getmembers(cls, predicate=is_async_func), which counts the inherited private _update_weights_from_checkpoint_engine_async (ray/actor.py:1806 at the locked ray 2.56.1). Asyncio actors run sync methods on the actor's event loop, so the first update_weights_from_checkpoint_engine.remote() crashes. Reproduced with a minimal actor of exactly this shape on ray 2.56.1 (a pure-sync control actor works).

The unit tests miss it because they asyncio.run the private coroutine on a non-actor object; tests/functional/grpo_sglang_nixl_non_colocated.sh will hit it at the step-1 refit.

Fix — make the entry point async and let Ray await it natively:

async def update_weights_from_checkpoint_engine(self) -> bool:
    return await self._update_weights_from_checkpoint_engine_async()

or keep the receive loop off the actor class entirely, like the vLLM analog (its asyncio.run runs inside vLLM's internal worker via collective_rpc, never as a Ray actor method). Either way, worth a visible note: any coroutine on this class flips all pre-existing worker RPCs (init/shutdown/colocated refit) from threaded to asyncio-actor semantics.

Minor, same function: if every receiver yields zero batches, this still bumps _checkpoint_engine_weight_version, invalidates the KV cache, prints "Loaded 0 rank-local tensors", and returns True. Consider raising when loaded_tensors == 0 so a no-op sender can't report a successful refit.

[tianyi-zhang-02]

Reproduced it, same error text. Fixed in 2034cb9, but with your second option rather than the one-liner: making the entry point async fixes the crash and keeps the flip, so every other worker RPC still gets asyncio-actor semantics out of a change that has nothing to do with them. The receive loop is a module-level coroutine now, so the mixin has no coroutine member and the actor stays as it was.

Took the minor point too, zero tensors raises instead of reporting a refit.

[Kh4L]

The suite stubs exactly the two contracts that are broken in production, so it stays green while the feature's only entry point fails:

1. _Worker overrides update_weights_from_tensor to return {"success": True}, hiding the server-side session assert at the pinned sglang rev ("update_weights_from_tensor requires an open begin_weight_update session").
2. Every test drives asyncio.run(worker._update_weights_from_checkpoint_engine_async()) on a plain object; nothing calls the public update_weights_from_checkpoint_engine through a Ray actor, which is where the asyncio-actor crash lives — that wrapper is the one mixin method with zero coverage.

Suggestions once the fixes are in: exercise update_weights_from_checkpoint_engine through a real @ray.remote actor (or at minimum switch one existing test to call the public wrapper), and assert the begin_weight_update → updates → end_weight_update ordering at the synchronizer level.

[tianyi-zhang-02]

Both right. The _Worker double now asserts an open session using the pinned rev's own message instead of returning success unconditionally, and the synchronizer tests pin begin/update/end ordering plus the failure paths: transfer raises, kv flush fails, begin_weight_update itself raises, pause_generation itself raises.

My first version of the no-asyncio-actor guard was wrong in a way worth flagging. I wrote iscoroutinefunction over the mixin; Ray uses iscoroutinefunction or isasyncgenfunction over the actor class. So an async generator moved back onto the mixin left it green, which is exactly the shape at hand since the receive loop is one. c963c83 calls Ray's has_async_methods instead.

Couldn't do the real @ray.remote actor in this file though. pytest names the module sglang.test_checkpoint_engine, so the directory shadows the real package and the actor fails to import in the worker. Can put it elsewhere if you want it.

[Kh4L]

Non-blocking: with tp_size=1 and a single one-GPU engine, the machinery that differentiates this PR degenerates to the trivial case — one stream (nothing to align), _validate_rank_batches compares nothing, and update_weights_from_tensor gets a one-payload list, so the payload-i → TP-rank-i scatter contract never runs against a real server. As test_checkpoint_engine_payload_index_matches_sglang_rank's own docstring notes, a transposed payload list "would load every shard onto the wrong GPU while still succeeding".

Suggest a tp_size=2 variant (e.g. cluster.gpus_per_node=3, generation on 2 GPUs) — there's precedent for 3-GPU functional tests (grpo_sc_gym_router_failover.sh), and the existing token_mult_prob_error < 1.05 gate would catch shard misplacement.

[tianyi-zhang-02]

Agreed, tp_size=1 doesn't exercise the scatter at all. Added grpo_sglang_nixl_non_colocated_tp2.sh in 6b83532 (gpus_per_node=3, engine on 2, trainer on 1), registered as full-mode in L1_Functional_Tests_SGLang.sh with the same token_mult_prob_error gate.

Worth saying plainly: I haven't run either functional test, so CI is the first real execution of both.

[Kh4L]

Call-site config defaults (config-conventions forbids these for TypedDict configs — defaults live only in the exemplar YAML):

• generation.cfg.get("sglang_cfg", {}) — sglang_cfg is a required key of SGLangConfig, and this branch already knows the backend is SGLang; the {} fallback lets a malformed config silently pass the dp/pp guards below instead of failing loudly.
• sglang_cfg.get("dp_size", 1) / .get("pp_size", 1) — non-None defaults at the access site; sibling code reads these keys bare (sglang_worker.py:449/475), and their defaults already live in grpo_math_1B_sglang.yaml:18-19.

sglang_cfg = generation.cfg["sglang_cfg"]
if sglang_cfg["dp_size"] != 1: ...
if sglang_cfg["pp_size"] != 1: ...

[tianyi-zhang-02]

Done in fbf667f, indexed the way you wrote it. shard_expert_weights lost its False literal too. Three fixtures were each missing a key the guards now need, so there's a _sglang_refit_cfg() helper.

[Kh4L]

Non-blocking: refit_cfg: NotRequired[Any] types a known-field block as Any (config-conventions: avoid Any for known-field configs). It's normalized by the same schema as vLLM's — checkpoint_engine_refit_config runs normalize_vllm_refit_config and writes the validated VllmRefitConfig back into config["refit_cfg"]. Mirror the vLLM analog (nemo_rl/models/generation/vllm/config.py:175):

refit_cfg: NotRequired[VllmRefitConfig | None]

vllm/config.py imports only typing/pydantic/interfaces, so this adds no import cycle and no vllm dependency to the SGLang venv (verified). If the vllm-named import feels off here, hoisting the schema to a shared module works too.

[tianyi-zhang-02]

Done in c3a1687. Checked the import cost, it's free: the vLLM config chain was already being pulled in before this PR.

If the vllm-named import here bothers you later I'm happy to hoist the three refit models into a shared module. weight_sync/checkpoint_engine_config.py already imports all of them.

[Kh4L]

This comment is inaccurate for non-colocated SGLang: with colocated.enabled: false, refit_transport: null selects SGLangDisaggregatedWeightSynchronizer (SGLang's NCCL weight-update group with a Megatron policy), not CUDA-IPC — exactly what this PR's own docs/guides/refit.md table says. Suggest:

# Null selects the default refit path: Ray CUDA-IPC when colocated,
# SGLang's NCCL weight-update group when non-colocated. ``nixl`` and
# custom checkpoint engines are supported for non-colocated generation.

Same wording nit in the new NotImplementedError in grpo.py ("or null for colocated CUDA-IPC refit") — that raise can fire in non-colocated runs where null means the NCCL path, steering users toward nixl as their only option.

[tianyi-zhang-02]

Fixed here and in the grpo.py raise (c3a1687), with one qualifier on your wording. Non-colocated null only reaches the NCCL group with a Megatron policy, factory.py:165 raises otherwise, so a DTensor user following that message hits a second NotImplementedError. Both now say it requires Megatron today. Drop it if you'd rather keep them short.

Also put the envelope and the in_place restriction in docs/guides/checkpoint-engine-refit.md.

[Kh4L]

Non-blocking: get_rollout_engine_urls() has zero callers anywhere in the tree — it looks like a leftover from the earlier HTTP-transport design. The checkpoint-engine path dispatches through run_checkpoint_engine_method() actor RPCs and never resolves engine URLs. Suggest dropping it here and landing it with whichever follow-up actually consumes it (the worker-side get_base_url() it wraps already exists on main).

[tianyi-zhang-02]

Gone in c3a1687. Checked for string dispatch too (getattr, run_checkpoint_engine_method, collective_rpc), nothing reaches it. Left get_base_url alone.

[Kh4L]

Forward-compat note (non-blocking), from the #3613 side of the stack: this early return makes the receiver set permanent for the actor's lifetime, and nothing can invalidate checkpoint_engines / _checkpoint_engine_target_devices / the NIXL process groups. #3613 adds engine fault tolerance where a dead SGLang engine is restarted and rejoined mid-run — after such a restart these cached receivers would hold stale NIXL registrations and rank state. No action needed in this PR beyond awareness; when both land we'll need a reset hook here (invalidate + re-init on recovery) rather than the early return. Flagging so the constraint is on record — happy to coordinate.

[tianyi-zhang-02]

Noted, and thanks for raising it now rather than after the fact. I left the constraint as a comment on the early return in 2034cb9 so it's at the site. Ping me when #3613 is close and I'll rebase around it.

[tianyi-zhang-02]

Thanks — this is a great review, and the framing was the useful part: both blockers were invisible to the tests I had, and that was the actual defect. I checked each point myself before changing anything; all nine hold. Fixes pushed, one section per comment.

Blocker 1 — no weight-update session. Confirmed. The SGLang branch of sync_weights only called prepare_for_generation(tags=[...]), which is resume_memory_occupation and nothing else, then went straight to update_weights_from_tensor. The sibling _SGLangWeightSynchronizer._refit wraps the same endpoint in the envelope, so the two SGLang transports were driving the engine through different contracts — that asymmetry is the giveaway, independent of the server-side assert.

Fixed by mirroring it: prepare_for_generation(weights) → pause_generation → invalidate_kv_cache → begin_weight_update → transfer → end_weight_update (inner finally) → prepare_for_generation(kv_cache) → continue_generation (outer finally). Everything from the pause onward is released on the failure path too, and the KV pool goes back before requests are readmitted rather than after. The mixin's post-transfer invalidate_kv_cache() is now redundant and gone — the flush happens before the buckets land instead of after.

While porting I also took the pause_generation_mode == "in_place" rejection from sglang_weight_synchronizer.py:80-83. This path pauses through the same hook and had no equivalent guard.

Blocker 2 — asyncio.run. Confirmed, and reproduced: a minimal Ray actor of exactly this shape raises RuntimeError: asyncio.run() cannot be called from a running event loop on the first call, verbatim.

I took your second option rather than the one-liner. Making the entry point async fixes the crash but leaves the flip in place — the mixin would still turn SGLangGenerationWorker into an asyncio actor for every pre-existing RPC, which is a concurrency change this PR has no business making. Moving the receive loop to a module-level coroutine (_load_checkpoint_engine_weights(worker)) removes the only coroutine member, so has_async_methods is False again and the actor keeps the sync semantics its other RPCs were written against. Verified side by side under Ray: has_async_methods is True for the old shape and False for the new one, and the same call that raises on the first returns cleanly on the second.

On the vLLM twin you flagged: it's safe, and for the reason you'd expect — VllmCheckpointEngineMixin lands on VllmInternalWorkerExtensionWithCheckpointEngine, a vLLM worker extension reached through collective_rpc, never a Ray actor class. Nothing to change there.

Took the minor point too: zero tensors received now raises instead of bumping the weight version and returning True.

Test blind spot. Both halves were right, and this was the part worth fixing properly.

• _Worker.update_weights_from_tensor now asserts an open session using the pinned rev's own message, so the double models the server instead of rubber-stamping it. Refits in that file go through a _refit() helper that opens and closes the session the way the synchronizer does.
• At the synchronizer: begin → update → end ordering, plus both failure paths — transfer raises (session still closed exactly once, generation resumed), and KV flush fails (session never opened, end_weight_update never called).
• The public wrapper now has coverage, and there's a guard asserting the actor never becomes an asyncio actor. That's the thing that actually regresses, so that's what I locked.
• Also covered: the engine-reported-failure branch, that end_weight_update never closes a session begin_weight_update failed to open, and that a pause which itself raises still resumes. Those three were asserted only in comments before.

All the new synchronizer tests are red on the previous head and green now. I also ran a mutation matrix over the envelope — begin moved after the transfer, end_weight_update lifted out of the inner finally, continue_generation moved inside the try, pause deleted, the KV-flush failure no longer raising, the KV re-acquire dropped, SGLang routed through the plain transfer, and the in_place guard removed — and every one of the eight is killed by a test that names it.

Worth reporting honestly: my first version of the asyncio guard was wrong in a way that would have mattered. It restated Ray's rule as iscoroutinefunction over the mixin, where Ray actually uses iscoroutinefunction or isasyncgenfunction over the actor class. An async generator moved back onto the mixin — which is exactly the shape here, since the receive loop is one — left the guard green while reinstating the bug. It now imports Ray's own has_async_methods and asks it about SGLangGenerationWorker, and I verified it goes red for both shapes.

One thing I could not do as suggested: a real @ray.remote actor inside test_checkpoint_engine.py. pytest names that module sglang.test_checkpoint_engine — the test directory shadows the real package — so the actor fails to import in the worker process. Happy to add it under a different path if you'd rather have the real thing.

tp_size=1 functional test. Fair, and the docstring you quoted is exactly the case that goes uncovered. Added grpo_sglang_nixl_non_colocated_tp2.sh (gpus_per_node=3, engine on 2, trainer on 1), registered as a full-mode test in L1_Functional_Tests_SGLang.sh. Flagging clearly: I have not run either functional test myself, so CI is the first real execution of both.

factory.py call-site defaults. Fixed as written — sglang_cfg indexed, dp_size/pp_size read bare. shard_expert_weights lost its False literal too. Three fixtures each omitted a key the guards now require, so there's a _sglang_refit_cfg() helper.

refit_cfg: NotRequired[Any]. Now NotRequired[VllmRefitConfig | None] with the direct import; you're right that it costs the SGLang venv nothing. If the vllm-named import in an sglang module bothers you later, I'm glad to hoist the three refit models into a shared home — weight_sync/checkpoint_engine_config.py already imports all of them.

refit_transport comment. Fixed, and the grpo.py raise with it — with one qualifier added to your suggested wording. refit_transport=null non-colocated only reaches SGLang's NCCL weight-update group with a Megatron policy: factory.py:165 raises NotImplementedError for anything else. Telling a DTensor/FSDP2 user that null gives them the NCCL path would walk them into a second raise, so both the comment and the message now say "which today requires a Megatron policy." Happy to drop the qualifier if you'd rather keep it terse.

Also wrote the envelope and the in_place restriction into docs/guides/checkpoint-engine-refit.md, since the constraint list there was the other place claiming to be complete.

get_rollout_engine_urls. Dropped. Left get_base_url alone.

#3613 / receiver reset. Understood, and thanks for putting it on record here rather than after the fact. I left the constraint as a comment on the early return so whoever adds the reset hook finds it at the site. Ping me when #3613 is close and I'll rebase around it rather than the other way round.

[tianyi-zhang-02]

Real-server follow-up on final SHA 5925bcf5b07fd438e9f385fb2bd1e1d21531ce33:

• fixed the TP2 pre-refit OOM by skipping the legacy metadata gather that SGLang never consumes;
• 31 synchronizer unit tests pass;
• NCCL_NVLS_ENABLE=0 bash tests/functional/grpo_sglang_nixl_non_colocated_tp2.sh passes on a Runpod 4×H100 80 GB host (driver 580.126.09, Python 3.13.14, torch 2.11.0+cu130; 2 SGLang TP ranks + 2 DTensor policy ranks);
• both training steps complete real NIXL weight transfers (0.383 s, 0.186 s), with no OOM/traceback/NCCL error;
• max(train/token_mult_prob_error) = 1.0145671 < 1.05.

cc @Kh4L — this is the real server path behind the two blockers from your review.