`WorkerWrap.load_weights` silently corrupts unquantized MoE rollouts after the first weight sync since `vllm>=0.20.0`

[jamesbraza]

WorkerWrap.load_weights calls vLLM's raw model.load_weights(...) directly. Per the upstream vllm-project/vllm#42821 (since vllm==0.20.0 was pulled in via #1628), that entrypoint is broken for unquantized MoE on FlashInfer backends (e.g. Qwen/Qwen3.6-35B-A3B).

SkyRL should move to self.model_runner.reload_weights(weights_iterator=...) (link), which is idempotent across repeated weight syncs.

+ from vllm.config import set_current_vllm_config

...

 def load_weights(self, request: bytes) -> None:
     ...
     weight_list = []
     for name, tensor in self._weight_receiver.receive_weights(request):
         weight_list.append((name, tensor))

-    self.model_runner.model.load_weights(weights=weight_list)
+    with set_current_vllm_config(self.vllm_config):
+        self.model_runner.reload_weights(weights_iterator=iter(weight_list))

     for weight in weight_list:
         del weight

This will also match vLLM's own reload_weights RPC.

[jamesbraza]

I did a bunch of runs with my proposed fix last weekend. The reload_weights swap-out is correct, but only sufficient when the sender ships the whole model in one WorkerWrap.load_weights call (e.g. chunks == 1). There is a second bug whenever chunks ≥ 2.

CudaIpcWeightTransferSender._send_chunks_legacy calls update_named_weights (and therein WorkerWrap.load_weights) once per WeightChunk. With FSDPWeightExtractor's default group_by_module=False, that's one chunk per parameter (equating to thousands per sync for nontrivial models).

1. Each per-chunk reload_weights runs finalize_layerwise_processing which iterates model.modules() and visits every submodule (e.g. norm layers, attention/MoE blocks).
2. A submodule is "represented in a chunk" iff one of its parameters appears in that chunk's WeightChunk.names; since each chunk carries only a small subset of the model's parameters, the vast majority of submodules are not represented in any given chunk and fall back to vLLM _place_kernel_tensors, which restores the submodule to the cached kernel tensors captured at initialize_layerwise_reload time.

This effectively rolls back any weights that earlier chunks in the same sync had loaded for that submodule, leaving the model with only the last chunk's submodules actually updated and every other submodule holding a stale kernel snapshot from earlier in the sync.

A possible complete fix would coalesce chunks on the receiver and call reload_weights once per sync (not once per WeightChunk). This is what I've been prototyping and it seems to work.

Cc @SumanthRH since he has a PR open.