NVIDIA-NeMo / NVIDIA-NeMo/RL

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

Open
#3,745 0 comments 0 reactions 1 assignee Claimed by @youngeunkwon0405 View on GitHub
Feature
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

## 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:

```python
# 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

| | FSDP2 / DTensor v2 | Megatron |
|---|---|---|
| vLLM colocated | ✅ `stream_weights_via_ipc_zmq` | ✅ |
| vLLM non-colocated | ✅ `broadcast_weights_for_collective` | ✅ |
| SGLang colocated | ✅ `update_weights_to_sglang_colocated` | ✅ |
| **SGLang non-colocated** | ❌ **this issue** | ✅ `update_weights_to_sglang_distributed` |

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`:

```python
@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 #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

- #3612 — adds the broadcast transport (Megatron only)
- #3288 — the same shape of gap for checkpoint-engine refit

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.