nemo_rl.utils.weight_transfer_stream#

Shared sparse payload pipeline and S3 transport for remote vLLM refit.

Module Contents#

Classes#

SparseRefitTransport

Transport-specific callbacks used by the shared streaming pipeline.

_SparsePayloadBucket

_S3ObjectStore

Functions#

_s3_client

sparse_payload_checksum

decode_sparse_payload

iter_sparse_weight_chunks

_get_manifest_s3_store

sparse_export_chunk_size

_executor

init_sparse_delta_baseline_from_iterator

stream_sparse_delta_payloads

stream_sparse_delta_payloads_via_s3_manifest

download_s3_refit_payload

zstd_compress

Data#

_STREAM_LOCAL

_S3_PART_SIZE

_S3_MEMORY_LIMIT

API#

nemo_rl.utils.weight_transfer_stream._STREAM_LOCAL#

‘local(…)’

nemo_rl.utils.weight_transfer_stream._S3_PART_SIZE#

None

nemo_rl.utils.weight_transfer_stream._S3_MEMORY_LIMIT#

None

class nemo_rl.utils.weight_transfer_stream.SparseRefitTransport#

Transport-specific callbacks used by the shared streaming pipeline.

name: nemo_rl.models.generation.vllm.config.VllmRefitTransportName#

None

transfer_workers: int#

None

send: Callable[[bytes, int, int], dict[str, Any]]#

None

cleanup: Callable[[], None]#

None

class nemo_rl.utils.weight_transfer_stream._SparsePayloadBucket#
payloads: list[nemo_rl.utils.weight_transfer_sparse_codec.TensorPayload]#

None

dense_bytes: int#

0

encode_s: float#

0.0

next_index: int#

0

nemo_rl.utils.weight_transfer_stream._s3_client(region: str) → Any#
class nemo_rl.utils.weight_transfer_stream._S3ObjectStore(*, bucket: str, region: str)#

Initialization

put(key: str, body: bytes) → None#
get(key: str) → bytearray#
delete(key: str) → None#
_request(
method: str,
key: str,
body: bytes | None = None,
) → Any#
nemo_rl.utils.weight_transfer_stream.sparse_payload_checksum(body: bytes | bytearray) → str#
nemo_rl.utils.weight_transfer_stream.decode_sparse_payload(body: bytes | bytearray, checksum: str) → bytes#
nemo_rl.utils.weight_transfer_stream.iter_sparse_weight_chunks(
tensors: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.NamedTensor],
target_bytes: int,
) → collections.abc.Iterator[tuple[nemo_rl.utils.weight_transfer_sparse_codec.TensorBatch, float]]#
nemo_rl.utils.weight_transfer_stream._get_manifest_s3_store(
bucket: str,
region: str,
) → nemo_rl.utils.weight_transfer_stream._S3ObjectStore#
nemo_rl.utils.weight_transfer_stream.sparse_export_chunk_size(
delta_tracker: nemo_rl.utils.weight_transfer_sparse_codec.DeltaCompressionTracker,
transport: nemo_rl.models.generation.vllm.config.VllmRefitTransportName,
) → int#
nemo_rl.utils.weight_transfer_stream._executor(
key: str,
workers: int,
) → concurrent.futures.ThreadPoolExecutor#
nemo_rl.utils.weight_transfer_stream.init_sparse_delta_baseline_from_iterator(
iterator: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.NamedTensor],
*,
delta_tracker: nemo_rl.utils.weight_transfer_sparse_codec.DeltaCompressionTracker,
shard_rank: int,
shard_count: int,
transport: nemo_rl.models.generation.vllm.config.VllmRefitTransportName,
) → None#
nemo_rl.utils.weight_transfer_stream.stream_sparse_delta_payloads(
iterator: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.NamedTensor],
*,
delta_tracker: nemo_rl.utils.weight_transfer_sparse_codec.DeltaCompressionTracker,
transport: nemo_rl.utils.weight_transfer_stream.SparseRefitTransport,
shard_rank: int,
shard_count: int,
) → dict[str, int]#
nemo_rl.utils.weight_transfer_stream.stream_sparse_delta_payloads_via_s3_manifest(
iterator: collections.abc.Iterable[nemo_rl.utils.weight_transfer_sparse_codec.NamedTensor],
*,
delta_tracker: nemo_rl.utils.weight_transfer_sparse_codec.DeltaCompressionTracker,
refit_targets: collections.abc.Sequence[str],
transfer_id: str,
api_key_env_var: str | None,
timeout_s: float,
shard_rank: int,
shard_count: int,
) → dict[str, int]#
nemo_rl.utils.weight_transfer_stream.download_s3_refit_payload(
manifest: collections.abc.Mapping[str, Any],
) → bytes#
nemo_rl.utils.weight_transfer_stream.zstd_compress(raw: bytes, threads: int) → bytes#