[AsyncGRPO] Match FSDP2 weight-transfer payloads to the vLLM manifest

[sfc-gh-truwase]

Fixes #7103.

What happens

On FSDP2 the vLLM weight manifest is collected in __init__, before Accelerate upcasts low-precision parameters to fp32. vLLM sizes each packed NCCL broadcast from that manifest, so the server expects bf16 bytes while the trainer streams fp32 ones. Measured on Qwen3-1.7B (2-trainer probe): the manifest describes 3,441,149,952 bytes and the stream carries 6,882,299,904. The two sides disagree on where each packed broadcast ends, and the transfer wedges.

Because mismatched NCCL counts are undefined, which packed broadcast stops is inferred, not observed.

What this changes

• Cast each streamed tensor back to the dtype recorded in the manifest, checking name, shape and count against it.
• Bound the wait for the sender after vLLM returns from /update_weights, so the failure observed here (server done, sender still blocked) raises on rank 0 after _CONTROL_TIMEOUT instead of hanging. A failure on the sending side, including the new validation errors, still waits for weight_sync_timeout.
• Drop _send_full_tensors_lockstep and restore the rank-0 send_weights path. Rank 0 still gathers on the send thread in the passing runs, so the cross-thread gather is not the cause and the lockstep approach is withdrawn.

Applied identically to AsyncGRPO and AsyncDistillation.

Reproduction

16 FSDP2 trainers against a 16-GPU vLLM server (tensor parallel 4, data parallel 4, /get_world_size 16), Qwen3-1.7B, train-begin sync, vLLM 0.26.

On unfixed parent f53b45f5e, rank 0's send thread returned from 241 of 310 parameter gathers and stopped. The last printed tensor was shape (2048, 6144), and every gathered tensor was fp32. Ranks 1–15 finished 310/310 on the training thread. /update_weights returned 200 while rank 0 was still sending. Those ranks then watchdog on the next FSDP all-gather (SeqNum=864, OpType=COALESCED, NumelIn=128, NumelOut=2048) after 1800s. NumelIn=128 is the 16-way shard of a 2048-element parameter. Sequence 864 is that all-gather, not the packed weight-transfer broadcast, and it is not the (2048, 6144) tensor. num_alloc_retries was 0 on this run. The trainer exited 1.

With this change (90a6b5593), the same 16+16 run gathered 310/310 on every rank. full_tensor prints were still fp32, because the cast is applied after the gather. Rank 0 still gathered on the send thread. The sync ended in 0.44s on rank 0 and about 6.6s on the other ranks, with no watchdog. The trainer exited 0.

An earlier 4+4 train-begin sync showed the same stop: rank 0 returned from 241 of 310 gathers, and ranks 1–3 watchdog on sequence 864 with NumelIn=512 for that same 2048-element parameter. With this change, that 4+4 run also finished 310/310. At 2 trainers, applying the same cast as a monkeypatch, not as this commit, let both the unfixed parent and the earlier lockstep revision complete 310/310. Without the cast, both hung.

Tests

TestManifestWeightStream in both experimental test files drives WeightTransferClient.send_weights with the vLLM producer stubbed out, and asserts the producer receives the manifest's names, dtypes, shapes and values starting from an fp32 stream, plus the mismatch, extra-parameter and short-stream rejections. The dtype test fails on the parent commit and passes here.

Scope

Validated with this commit: the 16+16 train-begin sync above, the 4+4 train-begin sync, and the unit tests. At 2 trainers, the cast was a monkeypatch rather than this commit.

Not run: a Qwen3-32B rerun, a sync after an optimizer step, repeated syncs, the PEFT merged path, a checksum of the weights vLLM stored, the bounded-join and validation-error paths on GPU, vLLM >= 0.28, and the distillation trainer on GPU. The original report also saw the stall on Qwen3-32B; an earlier bundle that is not this pull request synced that model in about 2s.

---

Note

Medium Risk
Changes the NCCL weight-transfer path used by AsyncGRPO and AsyncDistillation; incorrect casting could corrupt vLLM weights, but the change is narrowly aligned with the existing manifest.

Overview
Fixes FSDP2 async trainer→vLLM NCCL weight sync when the manifest dtypes (e.g. bf16) no longer match streamed tensors after Accelerate upcasts parameters to fp32.

WeightTransferClient.send_weights in both AsyncGRPO and AsyncDistillation now casts each tensor to the dtype listed in weight_update_info["dtype_names"] before the packed NCCL broadcast, so byte counts match what vLLM sized from the manifest.

Merged full-weight sync in async_grpo_trainer and async_distillation_trainer drops redundant if self.weight_transfer guards on pause, send, and resume on the main process.

Adds TestManifestWeightStream (with a stubbed NCCL producer) in both experimental test modules to assert an fp32 parameter stream is downcast to manifest dtypes before send.

^{Reviewed by Cursor Bugbot for commit 59dca9d. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 3 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 728fa05. Configure here.}

[albertvillanova]

Thanks for the detailed issue and for running this at 16T+16S scale, that kind of report is hard to come by. I went through the four changes one at a time against the vLLM sources, and I think only the first one is doing work. Details and evidence below.

1. Gather-then-send: right diagnosis, but it drops the memory bound

The underlying concern is real. On rank 0 the generator is consumed from the NCCL daemon thread, so default-PG full_tensor() all-gathers interleave with vLLM-PG broadcasts and with the blocking /update_weights HTTP call, while non-0 ranks run ahead. That is a plausible watchdog-timeout shape.

But _gather_weight_items() deletes exactly the property the replaced code was written to have. The comment it removes says so: "frees it once the generator advances, avoiding materializing the full model in memory". Every gathered parameter is now held in a list on rank 0 before a single byte is sent, so for the Qwen3-32B case in the description that is on the order of 64 GiB of bf16 on rank 0's device, on top of its shard, optimizer state and activations. It is asymmetric too: rank 0 OOMs while every other rank is comfortable. It fits on H200s; it will not fit on 80 GB parts, which is the common configuration.

The interleaving can be removed without giving that up, for instance by keeping the per-parameter gather loop in lockstep on all ranks and handing rank 0's tensors to the send thread through a small bounded queue. Only a few tensors are live at a time and the all-gathers stay ordered across ranks.

2. is_checkpoint_format=True -> False is incorrect, and inert on the wire

is_checkpoint_format does not mean "reload from a checkpoint path". From vllm/v1/worker/gpu_model_runner.py::reload_weights:

is_checkpoint_format: set to False if weights have already been processed into kernel format (repacking, renaming, etc.)

True routes the stream through model.load_weights(), which does the HF-to-vLLM name mapping and the fused qkv_proj / gate_up_proj packing. False takes the other branch:

for name, loaded_weight in weights_iterator:
    param = _get_parameter_for_reload(model, name)   # model.get_submodule(...)
    param.copy_(loaded_weight)

with a warning that the weights must be "in kernel format and already sharded". TRL streams HF-named, unsharded tensors (model.layers.0.self_attn.q_proj.weight), so False would look up a submodule vLLM does not have and copy TP-unsharded shapes. True is the correct value.

The reason your run still worked is that the flag never reaches the worker. In vLLM 0.26, 0.27 and current main the route is:

@router.post("/start_weight_update")
async def start_weight_update(raw_request: Request):
    await engine_client(raw_request).start_weight_update()   # body is never read

and Worker.start_weight_update(self) takes no such argument. The parameter was dropped from the route after 0.24. So on vLLM >= 0.26 this change is a no-op, and on 0.22 to 0.24, still inside TRL's >= 0.22.0 floor, it would select the direct-copy path and break the sync. reload_weights already defaults to True, which is what we want.

3. _sync_weight_dtype(args) reads the wrong knob

Model precision in this trainer is AsyncGRPOConfig.dtype ("auto" | "bfloat16" | "float16" | "float32", default "float32"), not the bf16 / fp16 autocast flags. The two are orthogonal, so the helper is wrong in both directions:

• dtype="bfloat16" with bf16=False: every parameter is upcast to fp32 before the send, the manifest advertises float32, the NCCL payload doubles, and the dtype-mismatch warning in init_weight_transfer fires against a bf16 server.
• dtype="float32" (the default) with bf16=True: fp32 master weights are silently downcast to bf16 on the wire, which is precisely the trainer-generator precision gap the float32 default exists to close (the config docstring cites https://huggingface.co/papers/2510.26788 for it).

If a cast is genuinely needed under FSDP2 mixed precision, it should target the dtype the manifest already records, so the two sites cannot disagree by construction. Separately, args: Any plus two getattr calls on a dataclass whose fields always exist runs against the repository's guidance on defensive hasattr / getattr.

4. _vllm_param_name is not a behaviour change

The description lists _checkpoint_wrapped_module stripping as part of the fix, but it is already on main, and the new two-step replace is equivalent to the existing one-liner on every name shape I tried, including nested wrappers and the leading-prefix case:

module.model.layers.0.self_attn.q_proj.weight                       same
model.layers.0._checkpoint_wrapped_module.self_attn.q_proj.weight   same
_checkpoint_wrapped_module.model.embed_tokens.weight                same
model.layers.0._ckpt_wrapped._ckpt_wrapped.x.weight                 same

So test_vllm_param_name_strips_wrappers passes against unmodified main. It would be worth checking that a new test fails without its fix. Note also that main now derives these names in one place, _iter_vllm_named_params, so a second helper in AsyncGRPO and a third in distillation moves away from that.

The 4 GiB pack buffer

This one is not in the description, and I do not think it is a correctness fix. vLLM's pack_tensors treats buffer_size_bytes as a soft threshold: an oversize tensor only triggers warnings.warn, and producer and consumer use the identical total_bytes > buffer_size_bytes break rule, with the consumer allocating the actual chunk size rather than the nominal buffer. A 1.56 GB Qwen3-32B embed_tokens row therefore transfers correctly at the 1 GiB default; it lands in its own chunk and logs a warning.

What the bump costs is real: chunks grow to about 4 GiB, and with packed_num_buffers=2 in flight that is roughly 8 GiB of transient allocation on the trainer and on every vLLM worker, against a server whose OOM headroom the existing error message already calls out.

It also breaks vLLM >= 0.28. with_packed_nccl_buffer() injects the key into the per-round update info, but 0.28 moved it to the init info. NCCLWeightTransferUpdateInfo now carries only names / dtype_names / shapes, and its docstring states that the buffer geometry "is a must-agree wire param carried on the init info ... not here". The server does self.update_info_cls(**update_dict) in parse_update_info, so the extra key surfaces as ValueError: Invalid update_info. Main had to stop sending packed in that same payload for this reason (weight_transfer.py, the _HAS_STATEFUL_TRAINER_ENGINE branch), so this reintroduces the same class of key.

Smaller points

• The wait_for_everyone() added after init_weight_transfer() cannot be load-bearing: _sync_weight does a rank-0-only pause() and then wait_for_everyone() before any gather, so non-0 ranks already block there. The justifying comment describes a race the existing barrier prevents.
• Code quality would fail CI as it stands. ruff check reports two import-sorting errors (the new vllm_client / weight_transfer imports sit before async_rollout_worker in both test files) and ruff format reformats one over-length assert in test_async_distillation_trainer.py.
• The two weight_transfer.py copies get different comments above the same constant, and the _weight_update_info assignment gets a comment in AsyncGRPO but not in distillation. Duplicated blocks in this repository are expected to match word for word.
• The branch is behind main in a way that matters here: #7017 split _sync_weight into _sync_weight_merged / _sync_weight_lora and centralised name derivation, and #7001 added the vLLM 0.28 path. The AsyncGRPO side will need reworking rather than a textual rebase.

Suggestion

The interleave you identified in #7103 is worth fixing and the evidence for it is good. I would keep that one change, reshape it so the gather stays bounded rather than materialising the whole model on rank 0, and drop the other four, since three of them are no-ops or regressions and the pack-buffer change breaks 0.28. That also makes the 16T+16S validation attributable to the change it actually exercised.

[sfc-gh-truwase]

@albertvillanova thanks for the review and flagging these problems. Will address them.

[sfc-gh-truwase]

@albertvillanova please take another look

[sfc-gh-truwase]

@albertvillanova gentle ping

[AmineDiro]

weight_transfer should never be None here

[sfc-gh-truwase]

Dropped the weight_transfer is not None guard — do_send is now just accelerator.is_main_process. Rank 0 always has a client when this helper runs.

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[albertvillanova]

Much better than the first version, and the lockstep shape is right: the streaming memory bound is preserved, the two copies are byte-identical, and all six checks are green. One thing to fix, plus two notes.

The post-failure drain costs 1 s per remaining parameter. Once send_weights stops consuming, the queue stays full, so every send_queue.put(..., timeout=1.0) blocks the full second before except queue.Full gets to check send_done. From this PR's own Tests (experimental) log:

::test_propagates_send_failure           12:16:41.88 -> 12:16:47.95   6.07s
::test_returns_when_send_does_not_drain  12:16:48.18 -> 12:16:55.18   7.00s

That is about 26 s of a 232 s suite for three assertions, and on a dense 32B (around 707 parameter tensors) a failed send would take about 12 minutes to surface, with every non-0 rank parked in wait_for_everyone(). Checking the flag before the put fixes both and is less code:

    for name, full in gathered_params:
        # Rank 0 keeps iterating after the send side is done: `full_tensor()` is collective.
        while do_send and not send_done.is_set():
            try:
                send_queue.put((name, full), timeout=1.0)
                break
            except queue.Full:
                pass
        accelerator.wait_for_everyone()

with the same shape for the sentinel. I ran the loop standalone with that change: 0.0 s in every case, error still propagated.

Two small notes: the WeightTransferClient docstring still says the HTTP request "stays on the main thread", which is now _run_send; and pause() / resume() still guard on self.weight_transfer while the send no longer does, so if that guard is dead those two are too.

Last, the description still describes the first version (gather everything first, dtype cast, name stripping, is_checkpoint_format). It becomes the squashed commit message, so it is worth rewriting around what the change now does. Same for #7103, which this closes.

[kashif]

Reviewed against current main. The diagnosis in #7103 holds, and I'd fix it this way too.

Why the hang happens on main. NCCL collectives return before the work finishes, so on ranks 1..N-1 for _ in self._streaming_iter(): pass enqueues an all-gather for every parameter almost immediately. Those all-gathers then wait on rank 0, whose gathers run on the send thread and are paced by the vLLM broadcast and /update_weights. On a 32B model that wait exceeds the watchdog, hence "Rank 0 last enqueued work << other ranks".

Why the fix works. Moving full_tensor() onto the training thread keeps rank 0's default-group collectives off the send thread. The per-parameter wait_for_everyone() is load-bearing: it blocks the other ranks on the host so they can't run ahead and enqueue the rest of the model while rank 0 is waiting on vLLM. The bounded queue caps what's held on rank 0 at two gathered tensors, instead of materialising the model.

Alternative I considered. Swapping the threads inside WeightTransferClient (broadcast on the main thread, HTTP in the background) would remove the queue. But it changes the client for every caller, and it gives up the property that an NCCL failure surfaces as an HTTP error rather than a hang. Keeping WeightTransferProtocol untouched is the better trade.

Edge cases checked.

• A send_weights that returns without reading the iterator: a no-op transfer, WeightTransferClient with no group, or a backend that syncs out of band (e.g. Arctic's) ends the put loop instead of blocking on a full queue. Covered by test_returns_when_send_does_not_drain.
• Send failure: rank 0 stops queueing but keeps gathering, so FSDP ranks stay in lockstep, and the error is re-raised after the final barrier.
• I ran the helper across 2 ranks (torchrun + gloo) with a slow sender, a non-draining sender, and a failing sender. All three complete on both ranks.
• The AsyncGRPO and AsyncDistillation copies are identical, and LoRA sync is untouched.

Minor, no change needed. That's one barrier per parameter (~700 on a 32B model). Given the ~2s sync time reported on 32B, that's fine.

LGTM.

[AmineDiro]

Thanks for looking into this! I’ve been running FSDP2 training without hitting this hang at all, so I’d like to understand the conditions that trigger it ??

I understand the fix but other ranks enqueueing gathers ahead of rank 0 doesn’t necessarily cause a timeout if rank 0 keeps making progress 🤔 Do we know whether rank 0 is progressing slowly here, or actually stuck in the vLLM update?

#7103 also reports checkpoint-format and name/dtype issues. Can we reproduce the hang after correcting those, then show that this change alone resolves it on NCCL?

[albertvillanova]

@AmineDiro's question is the right one and I do not think the thread has answered it. Two things do not add up for me.

The timing. The watchdog fires on a collective pending longer than the process group timeout, which is ddp_timeout, default 1800 s. On a non-0 rank the all-gather for parameter k completes as soon as rank 0 issues its matching one, so the longest any work item can pend is bounded by the total duration of the transfer. That transfer is reported at around 2.3 s on 32B, and this change moves no bytes: same parameters, same NCCL, same packing, only a different thread calling full_tensor() plus the barriers. For the old path to reach 1800 s it would have to be roughly 800 times slower. @kashif's account of why the other ranks run ahead looks right to me, but running ahead only becomes a timeout if the transfer itself outlasts the timeout, and I cannot see where that comes from.

The traceback. The block in #7103 reports OpType=COALESCED, but full_tensor() goes through Shard._to_replicate_tensor and funcol.all_gather_tensor, which is an all-gather. And Rank 0 last enqueued work << other ranks is not what torch prints; the line is PG status: last enqueued work: N, last completed work: M. So the one quoted artifact names a different collective than the diagnosis does, and it is summarised rather than pasted.

What would close the gap. Every in-flight full_tensor() holds its full-size output buffer until the collective completes, so the old path leans on the caching allocator in a way that scales with model size and stays invisible at small scale. That is the only thing here I can see turning a 2 s transfer into a 30 minute one, and it would explain why @AmineDiro does not hit it. If that is what happened then this change is still the right fix, since the bounded queue is exactly what caps it, but the reasoning here and in #7103 would need rewriting. It is also cheap to test: torch.cuda.memory_stats()["num_alloc_retries"] either side of _sync_weight, on main, at any scale.

Could you post the unedited watchdog output from the original run, with the PG status lines and the real op type, and the ddp_timeout it used? Was the old path slow, or stuck?

To be clear, I have no objection to the code: the lockstep helper is clean and the drain fix landed correctly. The concern is that #7103 is the document everyone here is reasoning from, and merging this as its fix writes the mechanism into the history.

[sfc-gh-truwase]

@AmineDiro and @albertvillanova thanks for the feedback. I will look deeper.

[sfc-gh-truwase]

@albertvillanova @AmineDiro the code comments from the earlier review are already on the branch (69c08b4: check send_done before put, and the weight_transfer is not None guards are gone from the helper, pause, and resume). This reply is about the mechanism.

You are right that the #7103 quote cannot support the diagnosis. It was paraphrased. full_tensor() is an all-gather, not OpType=COALESCED, and the torch line is PG status: last enqueued work: N, last completed work: M. I do not have the unedited watchdog output from the 16T+16S run, or the ddp_timeout it used, so I cannot say whether rank 0 was slow or stuck.

I also agree that other ranks enqueueing ahead of rank 0 does not, by itself, turn a ~2s transfer into an 1800s timeout. Each all-gather completes once rank 0 issues the matching one. The old loop can still hold one full-size full_tensor() output per parameter until that collective finishes. The queue of size 2 is the bound on that. We have not recorded num_alloc_retries on main.

The 16T+16S completions were an earlier gather-everything-then-send build that also changed checkpoint format, dtype, names, and the pack buffer. Those four are not in this PR. The helper here is covered by TestLockstepWeightSend and a 2-rank gloo torchrun, and it has not been rerun at 16T+16S.

I rewrote the PR description and #7103 so they no longer state that quote, or that running ahead is what trips the watchdog. is_checkpoint_format, the dtype cast, name stripping, and the 4GiB buffer stay out.

[sfc-gh-truwase]

@albertvillanova @AmineDiro you were both right that running ahead cannot by itself turn a ~2s transfer into a 1800s timeout. I ran the A/B and the lockstep explanation does not hold, so I have withdrawn it and pushed a different fix.

The lockstep revision also hung. Two 2-trainer runs of this PR as it stood hung at 244/244 parameters on the per-parameter barrier, and two runs of the parent commit hung with rank 0 at 241 of 310 gathers against 310 on the other rank. num_alloc_retries was 0 in every run, so the allocator-retry idea is not supported either.

What the evidence points to. The manifest is collected in __init__, before Accelerate upcasts the bf16 parameters to fp32, and vLLM sizes each packed NCCL broadcast from that manifest. On Qwen3-1.7B the manifest describes 3,441,149,952 bytes while the stream carries 6,882,299,904, so the two sides disagree on where each packed broadcast ends.

Reproduced at 4 FSDP2 trainers against a tensor-parallel-4 server. On the unfixed parent, rank 0's send thread stopped after 241 of 310 gathers; ranks 1-3 queued all 310 and their watchdog fired on the next FSDP all-gather (SeqNum=864, the (2048,) parameter, NumelIn=512 NumelOut=2048) after 1800s. Rank 0's own FSDP process group was idle at 863/863, so rank 0 was not inside a gather. /update_weights returned 200 while rank 0's sender was still blocked and /finish_weight_update was never called, so the server thought it was done. Because mismatched NCCL counts are undefined, which packed broadcast stops is inferred, not observed.

The fix casts each streamed tensor back to the manifest dtype while checking name, shape and count, and bounds the wait for the sender after the server returns 200. With it, the same 4+4 run gathered 310/310 on every rank with no watchdog. At 2 trainers, the same cast as a monkeypatch let both the parent and the lockstep revision complete 310/310, two runs each. Rank 0 still gathers on the send thread in the passing runs, which is why I no longer think the cross-thread gather is the cause.

16 trainers and 16 samplers. The same stall reproduced on Qwen3-1.7B at the layout in #7103: 16 FSDP2 trainers and 16 sampler GPUs (vLLM 0.26, tensor parallel 4, data parallel 4, /get_world_size 16), train-begin sync. On parent f53b45f5e, rank 0 stopped after 241 of 310 fp32 gathers, ranks 1-15 finished 310/310, /update_weights returned 200 while rank 0 was still sending, and their watchdog fired on the next FSDP all-gather (SeqNum=864, NumelIn=128, NumelOut=2048) after 1800s. NumelIn=128 is the 16-way shard of a 2048-element parameter. Sequence 864 is that all-gather, not the packed weight-transfer broadcast. With 90a6b5593, all 16 ranks finished 310/310, in 0.44s on rank 0 and about 6.6s on the others, with no watchdog. Rank 0 still gathered on the send thread.

Scope. Validated with this commit: that 16+16 train-begin sync, the 4+4 train-begin sync, and unit tests that fail on the parent commit. The 2-trainer evidence used a monkeypatched cast, not this commit. Not run: a Qwen3-32B rerun (the original report also saw the stall there; a different bundle synced 32B in about 2s), a sync after an optimizer step, repeated syncs, the PEFT merged path, a checksum of the weights vLLM stored, the bounded-join and validation-error paths on GPU, vLLM >= 0.28, and the distillation trainer on GPU.

The pull request now says it fixes #7103, on the strength of the 16-trainer Qwen3-1.7B run.

[albertvillanova]

The mechanism question is settled. Thanks for running the A/B instead of defending the first explanation. Withdrawing lockstep and showing the manifest against the stream (3,441,149,952 bytes declared, 6,882,299,904 carried, on Qwen3-1.7B), with the 16+16 and 4+4 reproductions and a unit test that fails on the parent, is what I was asking for.

I checked the three things the duplication rule covers:

• _cast_stream_to_manifest is byte-identical in both weight_transfer.py copies.
• The send_weights regions are byte-identical.
• Dropping and self.weight_transfer is safe. _sync_weight_merged runs only when _lora_sync is falsy, and the only main-process self.weight_transfer = None is the _lora_sync branch, so all three guards were dead.

One thing left. The two new test helpers are not identical. In test_async_grpo_trainer.py, original and streamed come from _iter_vllm_named_params(model); in test_async_distillation_trainer.py they come from model.named_parameters() filtered on requires_grad. Both files already import _iter_vllm_named_params, and it is what the real _streaming_iter walks, so the distillation copy drifts from its twin and also exercises a slightly different parameter set than production. Using the helper in both makes the two additions byte-identical.

Nothing else blocking. Two notes, no action needed:

• Raising after thread.join(timeout=_CONTROL_TIMEOUT) can leave the daemon mid-broadcast with the NCCL group in an undefined state, but the process is failing at that point and a bounded failure beats the hang.
• Your scope list is clear that the PEFT merged path was not run on GPU, and that is the path the guard removals sit in. They are dead-code removals, so I am not asking for a run.

[sfc-gh-truwase]

@albertvillanova thanks — the distillation manifest test now walks original / streamed through _iter_vllm_named_params(model), same as the GRPO twin and production _streaming_iter. Landed in 2014272; from _tensors_handed_to_producer through TestManifestWeightStream the two additions match.

[qgallouedec]

Thanks @sfc-gh-truwase. I pushed 59dca9d on top, which reduces the fix to the dtype cast.

Reproduced on main with a single trainer rank (Qwen3-0.6B, FSDP2 with mixed_precision: bf16, vLLM 0.30, one GPU each). With dtype="bfloat16" the manifest declares 1,192,099,840 bytes and the stream carries 2,384,199,680. /update_weights returns 200, then the trainer hangs in thread.join() with the send thread parked in packed_nccl_broadcast_producer. The same run with dtype="float32", the default, syncs in 0.1 s, which is why @AmineDiro never hit it. With the cast, all three syncs take 0.1 s and step-1 KL is 0.00035, so vLLM gets the right weights 👌

I dropped the name/shape/count checks and the bounded join. The manifest and the stream both walk _iter_vllm_named_params on the same model, so only the dtype can disagree, and once it is cast the server can no longer return while the sender is blocked. The rejection tests go as well.