feat: support non-colocated (broadcast) SGLang refit for the FSDP2/DTensor policy #3745

Summary

Non-colocated (weight_transfer_mode: broadcast) SGLang refit, added in #3612, supports only the Megatron policy backend. An FSDP2/DTensor policy with non-colocated SGLang generation is rejected:

# nemo_rl/weight_sync/sglang_weight_synchronizer.py
if not self._use_megatron:
    # Only Megatron implements the broadcast path; it depends on
    # AutoBridge restoring full HF tensors on trainer rank 0.
    raise NotImplementedError(
        "SGLang weight_transfer_mode='broadcast' is currently only "
        "supported for the Megatron policy backend."
    )

This is a feature request to implement the broadcast path for the FSDP2/DTensor policy worker, closing the one empty cell in the refit support matrix.

Current support matrix

vLLM covers all four cells; SGLang covers three.

Why this looks tractable

The comment attributes the restriction to AutoBridge, but that describes how the Megatron path is built rather than a capability the FSDP2 worker lacks. DTensorPolicyWorkerV2 already has a generator that yields full HF-named tensors, dtensor_params_generator(self.model, self.dtype), and it is already used both ways this feature needs:

• bucketed for SGLang colocated — iter_named_tensor_buckets(dtensor_params_generator(...)) in update_weights_to_sglang_colocated
• broadcast from rank 0 for vLLM non-colocated — packed_broadcast_producer(iterator=dtensor_params_generator(...), src=0) in broadcast_weights_for_collective

So both halves exist; they have just never been combined. Sketch, mirroring MegatronPolicyWorker.update_weights_to_sglang_distributed:

@torch.no_grad()
def update_weights_to_sglang_distributed(
    self, *, rollout_engines, rollout_engine_lock, buffer_size_bytes,
    target_precision="bf16", sglang_quantization_cfg=None,
) -> None:
    bucket_iter = iter_named_tensor_buckets(
        dtensor_params_generator(self.model, self.dtype),
        buffer_size_bytes=buffer_size_bytes,
    )
    if self.rank != 0:
        # Unsharding is collective: every rank must walk the iterator,
        # but only rank 0 broadcasts.
        for _ in bucket_iter:
            pass
        return
    broadcast_hf_buckets_via_distributed_impl(bucket_iterator=bucket_iter, ...)

Plus a connect_sglang_rollout_engines_distributed on the FSDP2 worker, and dropping the _use_megatron guard once both land.

Things that need care

• Every rank must drain the iterator. Unsharding is a collective, so a rank that returns early hangs the ones that don't. This exact hang was fixed on the Megatron path during feat(sglang): megatron backend weight refit for sglang rollouts #3612 review (nemo_rl/models/policy/utils.py, the placeholder-rank branch).
• bf16 only. The FSDP2 SGLang path currently rejects any other target_precision, so the first version would inherit that limit.
• Contiguity. dist.broadcast needs contiguous tensors; the Megatron path added an explicit .contiguous() for this reason.

Guard placement (smaller, separable fix)

Independent of implementing the feature: the current check runs inside sync_weights, so it fires at the first refit rather than at setup. A non-colocated FSDP2 v2 config passes every setup check, allocates the cluster, trains step 0 and completes a full rollout before failing.

The existing setup-time guard in nemo_rl/models/policy/lm_policy.py only rejects DTensor V1 + SGLang, and says nothing about colocation. Moving the _use_megatron check into create_weight_synchronizer — beside the checkpoint-engine guards that already reject unsupported combinations eagerly — would fail at t=0 instead.

Related

• feat(sglang): megatron backend weight refit for sglang rollouts #3612 — adds the broadcast transport (Megatron only)
• feat: extend checkpoint-engine refit to SGLang (and Megatron) generation backends #3288 — the same shape of gap for checkpoint-engine refit