SGLangWeightSender

class torchrl.weight_update.llm.SGLangWeightSender(scheme: SGLangWeightSyncScheme)

Sends weights to SGLang workers using NCCL broadcast.

Parameters:

scheme – The SGLangWeightSyncScheme configuration.

flush_cache() → bool

Flush the SGLang server’s radix cache after weight update.

Returns:

True if cache was flushed successfully.

Return type:

bool

init_all_workers_group(model_metadata: dict[str, tuple[dtype, Size]]) → None

Initialize the NCCL communication group.

Parameters:

model_metadata – Dict mapping param names to (dtype, shape) tuples.

register_model(model: Any) → None

Register the model for weight extraction.

Parameters:

model – The PyTorch model to sync weights from.

update_weights(weights: dict[str, Tensor] | None = None) → None

Broadcast weights to SGLang workers.

Parameters:

weights – Optional dict of weights. If None, extracts from registered model.