nemo_rl.utils.weight_transfer_sparse_codec#

Module Contents#

Classes#

_TensorPayloadBuilder

DeltaCompressionTracker

Source-side CPU or mmap baseline for sparse-delta refit.

Functions#

integer_dtype_for_element_size

dtype_from_name

sparse_operation

integer_view

encode_sparse_infos

merge_sparse_payloads

Combine encoded chunks without materializing dense source tensors.

sparse_locations_for_item

_encode_explicit_locations

Data#

NamedTensor

TensorBatch

SparseOperation

SparseInfo

TensorPayload

PreparedTensorPayload

SparseItem

_INTEGER_DTYPE_BY_SIZE

API#

nemo_rl.utils.weight_transfer_sparse_codec.NamedTensor#

None

nemo_rl.utils.weight_transfer_sparse_codec.TensorBatch#

None

nemo_rl.utils.weight_transfer_sparse_codec.SparseOperation#

None

nemo_rl.utils.weight_transfer_sparse_codec.SparseInfo#

None

nemo_rl.utils.weight_transfer_sparse_codec.TensorPayload#

None

nemo_rl.utils.weight_transfer_sparse_codec.PreparedTensorPayload#

None

nemo_rl.utils.weight_transfer_sparse_codec.SparseItem#

None

nemo_rl.utils.weight_transfer_sparse_codec._INTEGER_DTYPE_BY_SIZE#

None

class nemo_rl.utils.weight_transfer_sparse_codec._TensorPayloadBuilder#

Initialization

add_locations(locations: torch.Tensor) → tuple[int, int]#
add_values(values: torch.Tensor) → tuple[int, int, int]#
finish() → nemo_rl.utils.weight_transfer_sparse_codec.TensorPayload#
nemo_rl.utils.weight_transfer_sparse_codec.integer_dtype_for_element_size(element_size: int) → torch.dtype#
nemo_rl.utils.weight_transfer_sparse_codec.dtype_from_name(name: str) → torch.dtype#
nemo_rl.utils.weight_transfer_sparse_codec.sparse_operation(
value: object,
) → nemo_rl.utils.weight_transfer_sparse_codec.SparseOperation#
nemo_rl.utils.weight_transfer_sparse_codec.integer_view(tensor: torch.Tensor) → torch.Tensor#
nemo_rl.utils.weight_transfer_sparse_codec.encode_sparse_infos(
infos: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.SparseInfo],
) → nemo_rl.utils.weight_transfer_sparse_codec.TensorPayload#
nemo_rl.utils.weight_transfer_sparse_codec.merge_sparse_payloads(
payloads: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.TensorPayload],
) → nemo_rl.utils.weight_transfer_sparse_codec.TensorPayload#

Combine encoded chunks without materializing dense source tensors.

nemo_rl.utils.weight_transfer_sparse_codec.sparse_locations_for_item(
item: dict[str, Any],
packed_locations: torch.Tensor,
*,
device: torch.device | int | str,
dtype: torch.dtype = torch.int64,
) → torch.Tensor#
nemo_rl.utils.weight_transfer_sparse_codec._encode_explicit_locations(locations: torch.Tensor) → torch.Tensor#
class nemo_rl.utils.weight_transfer_sparse_codec.DeltaCompressionTracker(
config: nemo_rl.models.generation.vllm.config.VllmSparseRefitConfig,
)#

Source-side CPU or mmap baseline for sparse-delta refit.

Initialization

prepare_sparse_delta_payload(
tensors: nemo_rl.utils.weight_transfer_sparse_codec.TensorBatch,
) → nemo_rl.utils.weight_transfer_sparse_codec.PreparedTensorPayload#
_add_verification_samples(
metadata: list[dict[str, Any]],
) → None#
on_sync_succeeded() → None#
on_sync_failed() → None#
snapshot_baseline(
tensors: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.NamedTensor],
) → None#
_wait_for_baseline_commits() → None#
_commit_baseline_updates(
updates: collections.abc.Iterable[tuple[str, tuple[torch.Tensor, torch.Tensor]]],
) → None#
_baseline(
name: str,
shape: tuple[int, ...],
dtype: torch.dtype,
) → torch.Tensor#