nemo_rl.weight_sync.sglang_weight_synchronizer#

Weight synchronizers for the SGLang generation backend.

These run in the driver process, which is synced without a training-backend extra, so nothing here may import megatron.bridge or nemo_automodel: the refit only drives the policy and policy_generation facades.

The refit lifecycle — connect, pause, conditional KV invalidation, a begin/end weight-update session around the bucket transfer, then continue — is shared; the subclasses supply the transport-specific connect and transfer and own the GPU phase transitions around them.

Colocated:

  1. policy.offload_before_refit() – free GPU for staging

  2. generation.prepare_for_generation(tags=[“weights”]) – allocate buffers

  3. _refit() – Ray CUDA-IPC transfer

  4. policy.offload_after_refit() – restore optimizer state

  5. generation.prepare_for_generation(tags=[“kv_cache”]) – rebuild KV cache

Disaggregated:

  1. generation.prepare_for_generation(tags=[“weights”])

  2. _refit() – NCCL broadcast

  3. generation.prepare_for_generation(tags=[“kv_cache”])

The policy offload steps are skipped when disaggregated: the trainer keeps its GPUs to itself, so there is nothing to make room for.

prepare_for_generation runs on both paths. It is gated internally on sglang_server_config.needs_offload, which is an independent knob — with needs_offload: true (what every shipped config sets) these calls issue real resume_memory_occupation RPCs even when disaggregated, and the engines need them because finish_generation released that memory. They are only a no-op when needs_offload is false.

Module Contents#

Classes#

_SGLangWeightSynchronizer

Shared plumbing for the SGLang synchronizers.

SGLangColocatedWeightSynchronizer

Policy and SGLang engines share GPUs; weights move over Ray CUDA IPC.

SGLangDisaggregatedWeightSynchronizer

SGLang engines run on their own GPUs; weights move over NCCL broadcast.

API#

class nemo_rl.weight_sync.sglang_weight_synchronizer._SGLangWeightSynchronizer(
policy: Any,
generation: Any,
refit_buffer_size_gb: Optional[float] = None,
)#

Bases: nemo_rl.weight_sync.interfaces.WeightSynchronizer

Shared plumbing for the SGLang synchronizers.

Parameters:

  • policy – Policy object implementing ColocatablePolicyInterface.

  • generation – SGLangGeneration instance.

  • refit_buffer_size_gb – Fixed bucket size in GB for the weight transfer. If None, it is computed dynamically from free GPU memory.

Initialization

property is_stale: bool#
init_communicator() → None#
shutdown() → None#
_quantization_cfg() → dict#
abstractmethod _connect(
*,
rollout_engines: list,
engine_gpu_counts,
engine_gpu_offsets,
)#

Bring up the trainer-side transport for a new engine layout.

abstractmethod _send_buckets(
*,
rollout_engines: list,
buffer_size_bytes: int,
target_precision: str,
sglang_quantization_cfg: dict,
) → list#

Dispatch the transfer to the policy workers, returning Ray futures.

_reject_kv_scales(
kv_scales: Optional[dict[str, float]],
) → None#
_refit(buffer_size_bytes: int) → None#
_timed_refit(
timer: Optional[nemo_rl.utils.timer.Timer],
) → None#
_compute_buffer_size() → int#
class nemo_rl.weight_sync.sglang_weight_synchronizer.SGLangColocatedWeightSynchronizer(
policy: Any,
generation: Any,
refit_buffer_size_gb: Optional[float] = None,
)#

Bases: nemo_rl.weight_sync.sglang_weight_synchronizer._SGLangWeightSynchronizer

Policy and SGLang engines share GPUs; weights move over Ray CUDA IPC.

The trainer offloads before staging weights and re-offloads afterwards so the engines can take the memory back for their KV cache.

Initialization

_connect(*, rollout_engines, engine_gpu_counts, engine_gpu_offsets)#
_send_buckets(
*,
rollout_engines: list,
buffer_size_bytes: int,
target_precision: str,
sglang_quantization_cfg: dict,
) → list#
sync_weights(
*,
timer: Optional[nemo_rl.utils.timer.Timer] = None,
kv_scales: Optional[dict[str, float]] = None,
) → Optional[dict[str, float]]#
class nemo_rl.weight_sync.sglang_weight_synchronizer.SGLangDisaggregatedWeightSynchronizer(
policy: Any,
generation: Any,
refit_buffer_size_gb: Optional[float] = None,
)#

Bases: nemo_rl.weight_sync.sglang_weight_synchronizer._SGLangWeightSynchronizer

SGLang engines run on their own GPUs; weights move over NCCL broadcast.

No policy offload: the trainer is not competing with the engines for memory, and prepare_for_training onloads unconditionally anyway.

Initialization

_connect(*, rollout_engines, engine_gpu_counts, engine_gpu_offsets)#
_send_buckets(
*,
rollout_engines: list,
buffer_size_bytes: int,
target_precision: str,
sglang_quantization_cfg: dict,
) → list#
sync_weights(
*,
timer: Optional[nemo_rl.utils.timer.Timer] = None,
kv_scales: Optional[dict[str, float]] = None,
) → Optional[dict[str, float]]#