nemo_rl.models.generation.vllm.vllm_sparse_delta#

Apply canonical sparse updates through vLLM’s native weight loaders.

Module Contents#

Classes#

_SparseWeightLoadMode

Turn native loader copies into sparse XOR or overwrite.

VllmSparseDeltaApplier

Own one dense GPU scratch buffer and delegate all placement to vLLM.

Functions#

_storage_key

_view_key

Data#

_TensorViewKey

_LoaderWeight

_LoaderObservation

API#

nemo_rl.models.generation.vllm.vllm_sparse_delta._TensorViewKey#

None

nemo_rl.models.generation.vllm.vllm_sparse_delta._LoaderWeight#

None

nemo_rl.models.generation.vllm.vllm_sparse_delta._LoaderObservation#

None

nemo_rl.models.generation.vllm.vllm_sparse_delta._storage_key(tensor: torch.Tensor) → int#
nemo_rl.models.generation.vllm.vllm_sparse_delta._view_key(
tensor: torch.Tensor,
) → nemo_rl.models.generation.vllm.vllm_sparse_delta._TensorViewKey#
class nemo_rl.models.generation.vllm.vllm_sparse_delta._SparseWeightLoadMode(
targets: set[int],
verification: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
)#

Bases: torch.utils._python_dispatch.TorchDispatchMode

Turn native loader copies into sparse XOR or overwrite.

Initialization

start(
name: str,
source: torch.Tensor,
operation: nemo_rl.utils.weight_transfer_sparse_codec.SparseOperation,
sample_limit: int,
exact_sentinel: int | None,
) → None#
_observe_xor_copy(
destination: torch.Tensor,
source: torch.Tensor,
) → bool#
_remember_changed(
destination: torch.Tensor,
changed: torch.Tensor,
) → None#
__torch_dispatch__(
func: Any,
_types: Any,
args: tuple[Any, ...] = (),
kwargs: dict[str, Any] | None = None,
) → Any#
finish() → None#

Record bounded target samples after the loader finishes transforms.

class nemo_rl.models.generation.vllm.vllm_sparse_delta.VllmSparseDeltaApplier(model_runner: Any, device: torch.device)#

Own one dense GPU scratch buffer and delegate all placement to vLLM.

Initialization

discover_native_skips(
state_dict_info: collections.abc.Mapping[str, tuple[tuple[int, ...], torch.dtype]],
) → set[str]#

Reserve scratch and classify rank-local skips and overwrite weights.

_source_tensor(item: dict[str, Any]) → torch.Tensor#
static _scatter_values(
source: torch.Tensor,
item: dict[str, Any],
locations: torch.Tensor,
values: torch.Tensor,
) → None#
_prepare_loader_weight(
item: dict[str, Any],
locations: torch.Tensor,
values: torch.Tensor,
) → nemo_rl.models.generation.vllm.vllm_sparse_delta._LoaderWeight#
_load_weights(
weights: collections.abc.Iterable[nemo_rl.models.generation.vllm.vllm_sparse_delta._LoaderWeight],
verification: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
) → tuple[Any, list[nemo_rl.models.generation.vllm.vllm_sparse_delta._LoaderObservation]]#
static _validate_loader_report(
loaded: Any,
observations: list[nemo_rl.models.generation.vllm.vllm_sparse_delta._LoaderObservation],
*,
allow_unknown_skips: bool,
) → None#
_apply_decoded_items(
items: collections.abc.Iterable[tuple[dict[str, Any], torch.Tensor, torch.Tensor]],
) → None#
_iter_sparse_payload(
payload: nemo_rl.utils.weight_transfer_sparse_codec.TensorPayload,
) → collections.abc.Iterator[nemo_rl.utils.weight_transfer_sparse_codec.SparseItem]#
update_weights_from_decoded_sparse_payload(
*payloads: bytes | str,
) → dict[str, Any]#
_load_sparse_payloads(
sources: tuple[str | io.BytesIO, ...],
) → dict[str, Any]#
synchronize_device() → None#
finish_sparse_delta_refit() → dict[str, Any]#

Synchronize and compare bounded samples of target entries just changed.