diff --git a/docs/design-docs/sparse-delta-refit.md b/docs/design-docs/sparse-delta-refit.md
new file mode 100644
index 00000000000..830d7f1bcec
--- /dev/null
+++ b/docs/design-docs/sparse-delta-refit.md
@@ -0,0 +1,660 @@
+# Remote Sparse-Delta vLLM Refit
+
+Remote sparse refit updates non-colocated vLLM workers without transferring a
+full checkpoint after every optimizer step. Megatron Bridge exports canonical
+Hugging Face (HF) tensors, and the policy workers compare their assigned export
+chunks against one distributed canonical CPU baseline. S3 and ZeroMQ share the
+codec, streaming pipeline, receiver, native-loader apply path, and commit
+protocol.
+
+The feature is opt-in and does not change existing NCCL, CUDA IPC, or packed
+refit behavior.
+
+## Supported scope
+
+Remote sparse refit requires:
+
+- a non-colocated Megatron policy and vLLM generation backend;
+- the same initial HF checkpoint on both clusters;
+- BF16 or FP16 unquantized rollout weights;
+- `kv_cache_dtype: auto`; and
+- `refit_transport: vllm_s3_sparse` or `vllm_zmq_sparse`.
+
+Validation rejects `quant_cfg`, `real_quant`, colocated or non-Megatron
+deployments, FP8 rollout weights, and FP8 KV-cache scales. The codec and
+overwrite apply path retain FP8 bit patterns, but they do not generate block
+scales or KV-cache scales, so that alone is not end-to-end FP8 support.
+Synchronous and asynchronous vLLM engines are supported, but the weight-version
+transition remains synchronous: generation pauses until every payload is
+applied and the global flush completes.
+
+The coordinator is currently integrated only with GRPO. PPO and distillation
+reject `refit_transport` during setup; extending them is tracked in
+[#3275](https://github.com/NVIDIA-NeMo/RL/issues/3275).
+
+## Architecture
+
+```mermaid
+%%{init: {"flowchart": {"curve": "linear", "nodeSpacing": 24, "rankSpacing": 30, "diagramPadding": 8}}}%%
+flowchart TB
+ subgraph SYSTEM[" "]
+ direction TB
+ SYNC["VllmRemoteSparseWeightSynchronizer
driver: orchestrates; no payload bytes"]
+
+ subgraph ROW[" "]
+ direction LR
+ subgraph TRAIN["Megatron policy cluster"]
+ direction TB
+ subgraph PACTORS["MegatronPolicyWorker x N - Ray actors"]
+ direction TB
+ REMOTE["MegatronRemoteSparseRefit"]
+ EXPORT["Megatron Bridge HF export
all policy ranks participate
deterministic owner: chunk_index modulo producers"]
+ TRACKER["DeltaCompressionTracker + bounded pipeline
distributed CPU or mmap baseline
compare -> XOR/overwrite -> encode -> zstd"]
+ REMOTE --> EXPORT --> TRACKER
+ end
+ end
+
+ subgraph VALUE["Cross-cluster value plane"]
+ direction TB
+ S3["S3 object
compressed payload bytes
one upload per owned payload"]
+ ZTREE["ZeroMQ binary relay tree
one cross-cluster root send per payload"]
+ end
+
+ subgraph GEN["vLLM generation cluster"]
+ direction TB
+ subgraph GACTORS["VllmGenerationWorker x M - Ray actors; one receiver per node"]
+ direction TB
+ RECEIVER["VllmSparseRefitReceiver"]
+ API["FastAPI control plane
/prepare /s3-manifest /zmq-flush /flush"]
+ ZMQ["ZmqSparseRefitServer
relay node + staged-ACK tree"]
+ QUEUE["Deduplicating bounded FIFO
batch staging + single apply stream"]
+ RECEIVER --> API -->|downloaded S3 body| QUEUE
+ RECEIVER --> ZMQ -->|in-process callback| QUEUE
+ end
+
+ subgraph VPROC["vLLM worker process x TP/PP/EP - reached by collective_rpc"]
+ APPLY["VllmInternalWorkerExtension -> VllmSparseDeltaApplier
decode -> reusable GPU scratch -> native load_weights()"]
+ end
+ QUEUE -->|path when ranks share a node; bytes otherwise| APPLY
+ end
+
+ TRAIN ~~~ VALUE ~~~ GEN
+ end
+
+ SYNC -. "Ray RPC: stream and finish(success)" .-> REMOTE
+ SYNC -. "HTTP: prepare, relay flush, global flush" .-> API
+ TRACKER -->|PUT compressed bytes| S3
+ TRACKER -. "POST manifest pointer" .-> API
+ S3 -->|GET compressed bytes per receiver node| API
+ TRACKER -->|DEALER to ROUTER: DATA bytes| ZTREE
+ ZTREE -->|relay DATA bytes| ZMQ
+ end
+
+ classDef coordinator fill:#f1edff,stroke:#7048e8,color:#3f248f,stroke-width:2px
+ classDef policy fill:#f2f8f2,stroke:#2f7d32,color:#245b27,stroke-width:1.5px
+ classDef s3 fill:#fff8e8,stroke:#9a6500,color:#704900,stroke-width:1.5px
+ classDef zmq fill:#edf9fa,stroke:#267783,color:#1e5962,stroke-width:1.5px
+ classDef receiver fill:#eef4ff,stroke:#2864dc,color:#17418f,stroke-width:1.5px
+ classDef queue fill:#fff1ef,stroke:#c43b31,color:#8d2721,stroke-width:1.5px
+ classDef apply fill:#f2f8f2,stroke:#2f7d32,color:#245b27,stroke-width:1.5px
+ class SYNC coordinator
+ class REMOTE,EXPORT,TRACKER policy
+ class S3 s3
+ class ZTREE,ZMQ zmq
+ class RECEIVER,API receiver
+ class QUEUE queue
+ class APPLY apply
+ style SYSTEM fill:transparent,stroke:transparent
+ style ROW fill:transparent,stroke:transparent
+ style TRAIN fill:#f8fbf8,stroke:#2f7d32,stroke-dasharray:5 5
+ style GEN fill:#f6f8ff,stroke:#2864dc,stroke-dasharray:5 5
+```
+
+*Figure 1. Dashed links are control messages or S3 pointers; solid links carry
+payload bytes or invoke the node-local apply path. Amber components are S3,
+teal components are ZeroMQ, and both transports share the sender pipeline,
+receiver queue, native-loader apply path, verification, and baseline commit.*
+
+| Path | Cross-cluster transfer | Receiver handoff |
+|---|---|---|
+| S3 | One object upload per owned payload; each receiver gets a small manifest pointer and downloads the object | FastAPI callback -> deduplicating FIFO -> `collective_rpc` |
+| ZeroMQ | One producer-to-root DATA send followed by binary-tree relay fanout | In-process relay callback -> the same FIFO -> `collective_rpc` |
+| Within one node | No S3 or ZeroMQ hop | Staged path when ranks share storage; serialized bytes otherwise; no CUDA IPC |
+
+| Responsibility | Implementation |
+|---|---|
+| Coordinate one transfer and commit | [`vllm_remote_sparse_weight_synchronizer.py`](../../nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py) |
+| Export canonical tensors and own the source tracker | [`megatron_remote_sparse_refit.py`](../../nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py) |
+| Track baselines and encode deltas | [`weight_transfer_sparse_codec.py`](../../nemo_rl/utils/weight_transfer_sparse_codec.py) |
+| Run the shared pipeline and S3 transport | [`weight_transfer_stream.py`](../../nemo_rl/utils/weight_transfer_stream.py) |
+| Share HTTP control-plane utilities | [`weight_transfer_http.py`](../../nemo_rl/utils/weight_transfer_http.py) |
+| Run the ZeroMQ transport and relay | [`weight_transfer_zmq.py`](../../nemo_rl/utils/weight_transfer_zmq.py) |
+| Queue receiver work and expose endpoints | [`vllm_sparse_refit.py`](../../nemo_rl/models/generation/vllm/vllm_sparse_refit.py) |
+| Apply canonical updates through native loaders | [`vllm_sparse_delta.py`](../../nemo_rl/models/generation/vllm/vllm_sparse_delta.py) |
+
+The baseline advances only after the receiver version is globally committed:
+
+```mermaid
+%%{init: {"flowchart": {"curve": "linear", "nodeSpacing": 24, "rankSpacing": 30, "diagramPadding": 8}}}%%
+flowchart LR
+ subgraph VERSION0["Initialization"]
+ direction TB
+ INIT["Same checkpoint version W0
loaded independently on both clusters"]
+ T0["Trainer GPU
W0"]
+ B0["CPU or mmap baseline
B0 = W0"]
+ V0["vLLM GPU
W0; no baseline copy"]
+ INIT -. "local load" .-> T0
+ INIT -. "local load" .-> V0
+ T0 -->|canonical export to local baseline| B0
+ end
+
+ subgraph VERSION1["Refit after optimizer step 1"]
+ direction TB
+ T1["Trainer GPU
W1"]
+ D1["delta1
changed bits of W1 vs B0
XOR; overwrite where required"]
+ V1["vLLM GPU
apply delta1 -> W1"]
+ A1["All receiver ACKs
relay flush + apply flush + verify"]
+ B1["Baseline commit
B1 = exact W1 source bits"]
+ T1 --> D1 --> V1 --> A1 --> B1
+ end
+
+ subgraph VERSION2["Refit after optimizer step 2"]
+ direction TB
+ T2["Trainer GPU
W2"]
+ D2["delta2
changed bits of W2 vs B1
XOR; overwrite where required"]
+ V2["vLLM GPU
apply delta2 -> W2"]
+ A2["All receiver ACKs
relay flush + apply flush + verify"]
+ B2["Baseline commit
B2 = exact W2 source bits"]
+ T2 --> D2 --> V2 --> A2 --> B2
+ end
+
+ NEXT["..."]
+
+ T0 -->|optimizer step 1| T1 -->|optimizer step 2| T2 --> NEXT
+ B0 -->|comparison reference| D1
+ B1 -->|comparison reference| D2
+ V0 -->|serve rollouts until commit| V1 -->|serve rollouts until commit| V2 --> NEXT
+ B0 -. "advance only after A1" .-> B1
+ B1 -. "advance only after A2" .-> B2
+
+ classDef init fill:#f1edff,stroke:#7048e8,color:#3f248f,stroke-width:1.5px
+ classDef weight fill:#f2f8f2,stroke:#2f7d32,color:#245b27,stroke-width:1.5px
+ classDef delta fill:#fff8e8,stroke:#9a6500,color:#704900,stroke-width:1.5px
+ classDef ack fill:#eef4ff,stroke:#2864dc,color:#17418f,stroke-width:1.5px
+ class INIT init
+ class T0,T1,T2,B0,B1,B2,V0,V1,V2 weight
+ class D1,D2 delta
+ class A1,A2 ack
+ style VERSION0 fill:#f7f8fa,stroke:#c8ced8
+ style VERSION1 fill:#f7f8fa,stroke:#c8ced8
+ style VERSION2 fill:#f7f8fa,stroke:#c8ced8
+```
+
+*Figure 2. One refit follows each configured training cadence, but only changed
+canonical bytes cross the refit value plane. No full checkpoint is transferred
+by this protocol at initialization or periodically. Every policy rank still
+participates in the intra-cluster Megatron Bridge export.*
+
+## Protocol
+
+### Baseline and ownership
+
+The policy initialization task starts baseline construction as soon as its
+workers are ready, while the independent vLLM model load continues.
+`VllmRemoteSparseWeightSynchronizer.init_communicator()` discovers the receiver
+endpoints and then joins the prelaunched baseline before setup returns. The
+first rollout therefore does not enter a redundant weight sync or race an
+unfinished snapshot with policy training.
+
+Every policy rank participates in Megatron Bridge's normal
+`export_hf_weights()` conversion. The bounded output chunks are assigned by
+`chunk_index % shard_count`, so only one producer snapshots, compares, encodes,
+and sends each canonical chunk. Across all producers, persistent baseline
+storage is approximately one full HF checkpoint, not one copy per producer.
+Skipped chunks are still produced transiently because Bridge collectives and
+conversion must run on every participating rank.
+
+Ownership is deterministic for a fixed canonical export order, tensor shapes,
+export chunk limit, and producer count. It is recomputed rather than stored in a
+manifest, so changing any of those inputs can move a chunk to another producer;
+within one transfer, every chunk still has exactly one owner.
+
+This deliberately keeps one representation and one tracker. There is no MCore
+policy-local baseline, conversion-task detector, affine projection, stable-name
+residual partition, or model-specific mapping logic in NeMo RL. QKV, MoE,
+Mamba, padded or tied weights, grouped exports, adapters, and custom Bridge
+postprocessing all follow Bridge's canonical export semantics.
+
+Baselines use file-backed `torch.from_file` tensors by default;
+`refit_cfg.baseline.in_memory: true` keeps them in RAM. File backing reduces
+anonymous resident-memory pressure but does not reduce logical baseline bytes.
+
+Baseline initialization also returns each canonical tensor's name, shape, and
+dtype. The synchronizer merges that metadata and asks every vLLM worker to
+reserve one reusable GPU byte buffer large enough for the largest canonical
+tensor. This does not export HF values or mutate vLLM weights, and it removes
+scratch allocation from the first timed refit.
+
+On a fresh run, vLLM already holds the shared checkpoint, so the redundant
+initial full sync is skipped. On resume, both clusters must still start from
+the same HF weight version; sparse refit does not reconstruct a rollout
+baseline from an arbitrary training checkpoint.
+
+### Compare and encode
+
+Each producer compares only its assigned canonical export chunks bytewise
+through an integer view with the same element width. The payload contains only
+changed locations and values. Let `H` be the full canonical checkpoint bytes,
+`W` the producer count, and `s` the changed element fraction. Aggregate
+persistent baseline, comparison, and wire volumes are approximately:
+
+```text
+per-producer baseline and compare ~= H / W
+aggregate baseline and compare = H
+aggregate wire = codec_metadata + changed_indices + sH
+```
+
+The unavoidable leading cost is still `Bridge_export(H)`: every rank must
+participate in TP/EP gathers, PP communication, and conversion before chunk
+ownership can discard unassigned output. Sparse refit therefore makes wire and
+receiver work proportional to `s`, but not Bridge export communication. The
+reported changed percentage is computed once from the assigned canonical chunks
+and aggregated across producers.
+
+`DeltaCompressionTracker` finds changed flat locations through the equal-width
+integer view. During receiver prewarm, the generic apply context observes each
+native loader's tensor copies. A name remains XOR-compatible only when every
+source is a same-dtype view of the canonical scratch storage and destination
+copies do not overlap. The receiver returns the union of incompatible names;
+the producer emits XOR for all other names and absolute overwrite values for
+that opaque set. This runtime classification contains no model-family or tensor
+layout rules. The pending source baseline always records the exact new source
+bits regardless of the wire operation.
+
+The producer pulls bounded export chunks and compares them in parallel. A
+separate bounded stage coalesces encoded chunks up to
+`sparse_bucket_size_bytes`, serializes them, and applies zstd level 1 before the
+transport executor. The measured S3 default keeps 64 MiB compare chunks smaller
+than the recommended 512 MiB wire bucket, preserving D2H/scan parallelism while
+reducing object and manifest count. ZeroMQ retains 256 MiB export chunks and
+uses the same 512 MiB wire-bucket default, producing 482 payloads across the
+current 32-producer topology while tree delivery overlaps later export chunks.
+
+Payload N can transfer while later chunks are compared and encoded. Source
+baselines do not commit until the complete transfer and all transport and
+receiver completion barriers succeed. Pipeline errors cancel outstanding local
+futures and propagate to the synchronizer.
+
+### Transport and apply
+
+| Property | S3 | ZeroMQ |
+|---|---|---|
+| Value plane | AWS CRT multipart `PUT_OBJECT` | DEALER to ROUTER relay |
+| Producer completion | HTTP object manifest accepted by every receiver | Relay registers payload and returns a staged ACK |
+| Retry identity | Object key + checksum | Transfer + producer + payload IDs + checksum |
+| Completion barrier | Receiver flush | Relay-tree flush, then receiver flush |
+| Cleanup | Delete after receiver replies; retry leftovers at stream end | Close producer sockets after sends; seal the transfer at flush |
+
+S3 uses 64 MiB multipart parts, a 2 GiB CRT client memory limit, and a 10 Gbps
+throughput target. ZeroMQ assigns each producer to one inference-cluster relay.
+That root applies the payload locally and forwards it through a balanced binary
+relay tree, so each payload crosses the inter-cluster boundary once rather than
+once per generation replica. The relay calls the same receiver decode/apply
+queue as S3 directly, without a loopback HTTP data hop.
+
+Each ZeroMQ relay validates and deduplicates `(transfer_id, producer_id,
+payload_id)`, submits its local apply and at most two child forwards to bounded
+executors, and acknowledges once that work is registered. The producer can then
+export and send later payloads while earlier tree delivery continues. After
+every producer finishes, the synchronizer calls `/nemo-rl/refit/zmq-flush` on
+all relays and checks that their staged payload count matches the
+producer-reported payload total. Only after all tree and apply futures succeed
+does it call the normal receiver flush. Tree delivery therefore overlaps the
+producer stream, but its remaining tail is still on the final critical path.
+
+The generic streamer accepts a `SparseRefitTransport` with `send()` and
+`cleanup()` methods. It owns export, compare, encode, serialization,
+backpressure, timing, and transfer concurrency. S3 and ZeroMQ only construct
+their transport state and call that streamer. Cleanup runs on each transfer
+worker after all sends finish, so failed S3 deletion retries and ZeroMQ socket
+closure occur on the same worker that created the resource.
+
+The receiver deduplicates payload identities and applies bounded batches on one
+FIFO worker thread. Each generation replica downloads a transport payload once.
+When its vLLM ranks share a node, the compact serialized payload is written
+directly under `/dev/shm` as soon as it arrives, without waiting for the batch
+to fill. Staging futures feed the serial collective apply worker, so download,
+decompression, staging, and earlier GPU applies can overlap. Queue depth limits
+submitted work to 32 batches by default; with batches of eight, that is roughly
+256 payloads plus the current partial batch.
+
+Locations use `int32` unless a single canonical tensor exceeds the signed
+32-bit index range; values remain grouped by dtype. For shared-node workers the
+collective RPC passes only staged file paths, and each rank uses
+`torch.load(..., mmap=True)`. The mmap is not a second baseline: it lets ranks
+share the compact payload's page cache instead of materializing independent CPU
+copies. Each worker decodes locations lazily while feeding the native loader.
+When ranks do not share a node, the receiver sends the same serialized bytes
+through one collective RPC and each worker decodes its copy.
+
+There is deliberately no TP/EP source plan. Every vLLM worker sees canonical
+sparse entries, scatters them into its reusable dense source buffer, and calls
+the model's native `load_weights()`. Native loaders own QKV, MoE, Mamba, TP, and
+EP placement; NeMo RL stores no route model or family-specific placement
+formulas and does not patch vLLM. This trades canonical-tensor scratch
+initialization and duplicated sparse H2D across ranks for independence from vLLM
+internals. Measure that cost on the target TP/EP topology; H2D no longer scales
+only with the worker-local sparse subset.
+
+During the untimed metadata prewarm, one no-op native-loader pass records names
+that issue no model-storage copy on a fixed rank and identifies loaders for
+which XOR cannot preserve copy semantics. Later refits skip rank-local names
+that the loader explicitly omitted and use overwrite for the incompatible
+union. The cache contains names only; it stores no placement offsets, tensor
+routes, or model-family rules.
+
+For ZeroMQ, the relay flush first drains every staged tree delivery. The final
+`/nemo-rl/refit/flush` then drains every receiver batch, synchronizes CUDA, and
+checks optional delta samples. Only after both barriers succeed does the source
+commit exact pending baseline bits in background CPU threads.
+
+> **Failure boundary:** source baseline commit is transactional, but receiver
+> writes are in place and are not rolled back. If a transfer fails after a
+> receiver accepts any payload, reload that receiver from a known-good weight
+> version before retrying. This is mandatory for XOR, because replaying an
+> already-applied XOR reverts those bits. Replaying overwrite is safe.
+> The synchronizer records this state as poisoned and rejects subsequent syncs
+> with recovery instructions. In-place recovery is tracked in
+> [#3274](https://github.com/NVIDIA-NeMo/RL/issues/3274).
+
+## Payload and native apply
+
+Each serialized payload is:
+
+```text
+(packed_location_bytes, packed_value_groups, tensor_metadata)
+```
+
+```mermaid
+%%{init: {"flowchart": {"curve": "linear", "nodeSpacing": 24, "rankSpacing": 30, "diagramPadding": 8}}}%%
+flowchart LR
+ CODEC["Shared encoder output
(locations, value groups, tensor metadata)"]
+ SERIAL["torch serialization + zstd
checksum over compressed body"]
+ CODEC --> SERIAL
+
+ subgraph S3F["S3 transport"]
+ direction TB
+ SO["S3 object body
compressed payload bytes"]
+ SM["HTTP manifest pointer
bucket, region, key, checksum,
transfer/producer/payload IDs, samples"]
+ SM -. "locates object" .-> SO
+ end
+
+ subgraph ZF["ZeroMQ DATA multipart"]
+ direction TB
+ ZK["Frame 1: DATA"]
+ ZM["Frame 2: JSON metadata
transfer/producer/payload IDs,
checksum, samples, optional API key"]
+ ZB["Frame 3: zstd-compressed torch payload"]
+ ZK --> ZM --> ZB
+ end
+
+ SERIAL -->|PUT bytes once| SO
+ SERIAL -->|DATA body| ZB
+ SERIAL -. "POST pointer" .-> SM
+ SO --> COMMON["Common receiver
checksum -> zstd -> torch deserialize"]
+ SM -. "identity + verification budget" .-> COMMON
+ ZB --> COMMON
+ ZM -. "identity + authentication" .-> COMMON
+ COMMON --> PT["Decoded payload tuple"]
+ PT --> PI["packed location bytes
range or uint16/32/64 deltas"]
+ PT --> PV["value groups by dtype
XOR or overwrite bits"]
+ PT --> PM["per-tensor metadata
HF name, shape, offsets, operation"]
+
+ classDef shared fill:#f1edff,stroke:#7048e8,color:#3f248f,stroke-width:1.5px
+ classDef s3 fill:#fff8e8,stroke:#9a6500,color:#704900,stroke-width:1.5px
+ classDef zmq fill:#edf9fa,stroke:#267783,color:#1e5962,stroke-width:1.5px
+ classDef decoded fill:#f2f8f2,stroke:#2f7d32,color:#245b27,stroke-width:1.5px
+ class CODEC,SERIAL shared
+ class SO,SM s3
+ class ZK,ZM,ZB zmq
+ class COMMON,PT,PI,PV,PM decoded
+```
+
+*Figure 3. S3 separates the object body from its control-plane manifest;
+ZeroMQ carries equivalent identity and body data as multipart frames. Both
+decode into the same codec tuple.*
+
+Contiguous locations use a range encoding. Other sorted locations are
+delta-encoded into the smallest lossless unsigned width among 16, 32, and 64
+bits. Metadata carries the HF name and shape, dtype, value offsets, location
+encoding, XOR or overwrite operation, and optional verification sample budget.
+
+HF coordinates are the canonical wire format because Megatron Bridge defines
+the training-to-HF mapping while vLLM owns the packed and sharded destination.
+For each item, the receiver resets its resident largest-tensor scratch buffer,
+scatters sparse values, and calls the model's native `load_weights()`.
+
+A storage-scoped PyTorch dispatch mode changes only copies into model parameter
+or buffer storage; it never encodes QKV, MoE, Mamba, TP, or EP geometry. For
+XOR, unchanged scratch bits are zero and target copies become bitwise XOR. The
+source must remain a view of scratch storage, dtypes must match, and overlapping
+destination copies fail closed. For overwrite, unchanged entries are NaN
+sentinels. The dispatch mode propagates the first sparse mask through subsequent
+native copies and writes only selected destination entries. It keeps no target
+backup because a partial batch already requires receiver reload. This supports
+native pointwise transforms and dtype casts without model-specific formulas.
+One-byte FP8 overwrite uses an exact bit sentinel and therefore requires a
+non-transforming loader; end-to-end quantized rollout refit remains out of
+scope.
+
+Native-loader return values distinguish an explicit skip from an unsupported
+apply. An empty loaded set is accepted, matching vLLM's handling of pipeline or
+expert ownership and checkpoint-only parameters such as inactive MTP weights.
+Loader exceptions propagate. A loader that reports a weight loaded without a
+supported target copy fails closed. There is no layout fallback, cached loader
+trace, or route model. NeMo RL's only integration point is the native model's
+public `load_weights()` behavior.
+
+## Configuration
+
+Configure the feature under `policy.generation`:
+
+```yaml
+policy:
+ generation:
+ backend: vllm
+ refit_transport: vllm_s3_sparse # or vllm_zmq_sparse
+ refit_cfg:
+ delta_compression:
+ encoding: xor # overwrite is selected automatically for opaque loaders
+ storage:
+ s3_bucket: my-refit-bucket # required only for vllm_s3_sparse
+ s3_region: us-east-1
+ baseline:
+ in_memory: false
+ verify_samples_per_payload: 0
+ colocated:
+ enabled: false
+ vllm_cfg:
+ async_engine: false # true is also supported
+ precision: bfloat16
+ kv_cache_dtype: auto
+ http_refit_api_key_env_var: NRL_REFIT_API_KEY
+ http_refit_server_port: 8081
+ zmq_refit_server_port: null
+```
+
+`refit_cfg` is optional; its Pydantic models resolve and log all defaults. S3
+fails during setup unless `refit_cfg.storage.s3_bucket` is nonempty. Its region
+and key prefix default to `us-east-1` and `nemo-rl-refit`. ZeroMQ requires
+routable TCP access to the relay port. The HTTP and ZeroMQ servers are
+plaintext, so use a trusted or encrypted network. When
+`http_refit_api_key_env_var` is set, the named variable must contain the same
+nonempty token on producers and receivers. Binding either server to all
+interfaces without a key emits a warning.
+
+| Control | Default |
+|---|---:|
+| `refit_cfg.delta_compression.encoding` | `xor` |
+| `refit_cfg.storage.s3_bucket` | none; required for S3 |
+| `refit_cfg.storage.s3_region` | `us-east-1` |
+| `refit_cfg.baseline.in_memory` | `false` |
+| `refit_cfg.verify_samples_per_payload` | `0` |
+
+Export chunks are capped by `sparse_bucket_size_bytes` and the packed tensor
+limit. The S3 defaults were selected by balanced 120B sweeps. Increase one
+concurrency control at a time; excess parallelism moves the bottleneck into
+host memory, Bridge export, relay-tree forwarding, or receiver apply.
+
+Run the checked-in ZeroMQ recipe with:
+
+```bash
+uv run python examples/run_grpo.py \
+ --config examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.yaml
+```
+
+For a diagnostic run, enable bounded transmitted-delta sampling through the
+logged config:
+
+```bash
+uv run python examples/run_grpo.py \
+ --config examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.yaml \
+ policy.generation.refit_cfg.verify_samples_per_payload=32
+```
+
+## Metrics and profiling
+
+| Signal | Meaning |
+|---|---|
+| `REFIT_BASELINE_INIT` | Baseline export and snapshot time |
+| `REFIT_RECEIVER_PREWARM` | GPU scratch reservation and rank-local native skip discovery |
+| `REFIT_{S3,ZMQ}_TIMING` | Producer wall time, stage service times, payloads, bytes, and density |
+| `REFIT_{S3,ZMQ}_DELTA_CHANGE` | Global changed and total element counts |
+| `REFIT_RECEIVER_TIMING` | Receiver staging span/wait, batches, apply, and verification |
+| `REFIT_{S3,ZMQ}_DELTA_VERIFY` | Sampled transmitted-delta accuracy |
+| `REFIT_ZMQ_RELAY_FLUSH` | Relay-flush wall time and aggregate fanout service time |
+| `REFIT_{S3,ZMQ}_GLOBAL_COMMIT` | Successful global flush |
+
+`total_s` is producer wall time. Stage fields such as `encode_s`, `s3_put_s`,
+and `zmq_send_s` are sums across concurrent tasks and can exceed `total_s`; do
+not add them as serial phases. Receiver responses also expose node
+decode/staging, worker deserialization, and native-loader apply time. These are
+concurrent sums as well, so compare them with receiver wall time rather than
+adding them. `refit/transfer/relay_flush_s` is the coordinator's ZeroMQ flush
+wall time; `fanout_service_s` is an aggregate concurrent service sum.
+
+Chunk counts include every canonical export chunk; changed and total elements
+include only the chunks assigned to that producer. Synchronizer metrics appear
+under `refit/delta/*`, `refit/delta_verify/*`, and `refit/transfer/*` in W&B and
+other configured loggers. End-to-end latency is
+`timing/train/prepare_for_generation/transfer_and_update_weights`.
+
+Nsight ranges cover baseline creation, policy streaming, and vLLM sparse apply.
+Relevant thread names start with `nrl-refit-`, `nrl-zmq-`, or
+`nrl-vllm-sparse-refit`.
+
+## Development and validation
+
+Keep transport changes behind the shared `stream_sparse_delta_payloads()`
+pipeline. A `SparseRefitTransport` provides `name`, `transfer_workers`,
+`send(body, payload_id, verification_candidates)`, and worker-local `cleanup()`
+only; it must not
+duplicate export, baseline tracking, encoding, backpressure, receiver queue, or
+apply logic. Retries must preserve payload identity and bytes, fan out to every
+required replica, and require a successful global flush before source commit.
+Never retry XOR after an uncertain or partial receiver apply.
+
+Do not add model-specific placement math or a persistent placement cache. New
+layouts must work through their native vLLM weight loader and the generic
+storage-scoped operation context. Tests should cover split, merged, transposed,
+overlapping, transformed, skipped, and dtype-cast copy behavior without naming
+model families. FP8 bit overwrite, contiguous ranges, and explicit locations
+also require coverage. Unknown names, transformed XOR, and overlapping XOR
+copies must fail closed.
+
+Codec changes must update encoder and decoder together, preserve 64-bit-safe
+locations, and commit exact source bits only after global success. Receiver
+changes must preserve FIFO application, bounded memory, error propagation,
+flush, CUDA synchronization, and clean shutdown.
+
+Do not reintroduce an MCore-local baseline, Bridge mapping duplication, or
+model-family projection formulas. Tests must cover complete canonical export,
+deterministic chunk ownership, transactional baseline updates, transport
+cleanup on success and failure, and unchanged producer/receiver overlap. Do not
+modify Megatron Bridge for a transport-specific hook.
+
+On the target topology, verify the exact commit, image digest, and checkpoint
+revision; validate fresh starts and same-version resumes; compare two balanced
+repetitions with an equivalent NCCL or full control; and require the requested
+changed density, one global commit, no traceback, and zero sampled mismatches.
+After failure injection, confirm that the source baseline does not commit and
+reload the receiver before retrying.
+
+## Refit bandwidth calculator
+
+[`refit_bandwidth_calculator.py`](../../tools/refit_bandwidth_calculator.py) is a
+projection of the latest zstd S3 and ZeroMQ measurements against a measured
+H100 NCCL envelope. It is not a general fabric or topology simulator.
+
+The sparse side uses the current 247.2 GB canonical-HF XOR results with a
+512 MiB sparse bucket. Both S3 and ZeroMQ have matched measured 3% and 5%
+anchors. End-to-end time and wire bytes scale linearly by indexed model bytes.
+Other model sizes and densities outside the measured points are explicit
+projections. Because the sparse values are end-to-end measurements rather than
+a link model, `--candidate-ethernet-gbps` does not rescale S3 or ZeroMQ. Text
+output reports the 512 MiB calibration, and JSON includes
+`sparse_bucket_size_bytes: 536870912`.
+
+| Transport | 3% | 5% |
+|---|---:|---:|
+| S3 | 20.234 s measured | 25.790 s measured |
+| ZeroMQ tree fanout | 24.095 s measured | 33.776 s measured |
+
+The measured anchors already include full Bridge export on every participating
+rank, deterministic chunk ownership, one aggregate checkpoint-sized canonical
+comparison across producers, transport work, and all completion barriers.
+Do not multiply comparison bytes by the producer count. The ZeroMQ anchor also
+includes its staged-ACK tree flush; the calculator does not model relay
+forwarding as an independent bandwidth term.
+
+The NCCL side interpolates `_NCCL_ANCHORS` in log model-size space. Those
+anchors were measured at 400 Gbps per rank and are projected onto the requested
+Ethernet rate as:
+
+```text
+T_ethernet = T_H100_IB * 400 / candidate_ethernet_gbps
+```
+
+`--candidate-ethernet-gbps` is raw bandwidth per rank, not aggregate node or
+cluster bandwidth. This makes NCCL and candidate Ethernet refer to the same
+per-rank link while leaving the independently measured sparse path unchanged.
+
+```bash
+uv run python tools/refit_bandwidth_calculator.py \
+ --model-size-gb 247.2 \
+ --changed-pct 3 \
+ --candidate-ethernet-gbps 25
+```
+
+The output reports the reference and projected NCCL envelopes, sparse latency,
+estimated wire bytes, and per-rank Ethernet crossover. Below the lower
+crossover sparse refit beats the complete NCCL envelope; above the upper
+crossover NCCL wins; between them the measured range has no single winner.
+`--json` emits the same fields for scripts.
+
+The production transport applies zstd level 1 to every payload, so the
+calculator intentionally has no synthetic raw-compression arm. Treat every
+model size other than 247.2 GB, density outside 3%-5% for either transport, and
+any different topology or parallel mapping as a projection rather than a
+performance claim.
+
+## Failure guide
+
+| Symptom | Action |
+|---|---|
+| Baseline is missing a tensor | Check baseline completion, checkpoint equality, and Bridge name mappings. |
+| No refit endpoint is found | Check worker startup, fixed ports, routing, and network policy. |
+| Every worker reports a tensor unloaded | Verify the canonical HF name and native loader; do not add placement formulas. |
+| A payload ID is reused with different bytes | Start a new transfer or resend the original payload unchanged. |
+| Changed percentage rises unexpectedly | Correlate `DELTA_CHANGE` with `GLOBAL_COMMIT` and baseline commit completion. |
+| Apply queue stalls | Inspect receiver timing; reduce producer or relay concurrency. |
+| A transfer fails after payload acceptance | Reload the receiver from a known-good checkpoint before retrying. |
diff --git a/docs/index.md b/docs/index.md
index f57dc285104..4e86e4e47b9 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -322,6 +322,7 @@ design-docs/uv.md
design-docs/dependency-management.md
design-docs/chat-datasets.md
design-docs/generation.md
+design-docs/sparse-delta-refit.md
design-docs/checkpointing.md
design-docs/loss-functions.md
design-docs/fsdp2-parallel-plan.md
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
index e59d1675464..44ae519a3f9 100644
--- a/examples/configs/grpo_math_1B.yaml
+++ b/examples/configs/grpo_math_1B.yaml
@@ -339,6 +339,8 @@ policy:
top_k: null
stop_token_ids: null
stop_strings: null
+ refit_transport: null # Set to "vllm_s3_sparse" or "vllm_zmq_sparse" for remote sparse-delta refit.
+ refit_cfg: null # Optional tuning and storage settings for remote sparse-delta refit.
mcore_generation_config:
async_engine: false
max_model_len: ${policy.max_total_sequence_length} # Engine-side max sequence length.
@@ -376,6 +378,9 @@ policy:
num_first_layers_in_bf16: 0
enable_vllm_metrics_logger: true # Set to true to enable vLLM internal metrics logger, turn off for better performance
vllm_metrics_logger_interval: 0.5 # Interval in seconds to collect vLLM logger metrics
+ http_refit_api_key_env_var: null # Optional env var containing the internal refit API key.
+ http_refit_server_port: null # Optional fixed port for Kubernetes targetPorts.
+ zmq_refit_server_port: null # Optional fixed ZeroMQ relay port for Kubernetes targetPorts.
vllm_kwargs: {}
colocated:
# true: generation shares training GPUs
diff --git a/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.yaml b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.yaml
new file mode 100644
index 00000000000..18ff88ca112
--- /dev/null
+++ b/examples/configs/recipes/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.yaml
@@ -0,0 +1,41 @@
+defaults: ./performance/grpo-qwen3-30ba3b-4n8g.yaml
+
+grpo:
+ num_prompts_per_step: 16
+ num_generations_per_prompt: 8
+ max_num_steps: 50
+ val_period: 1000
+
+checkpointing:
+ checkpoint_dir: results/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated
+
+policy:
+ train_global_batch_size: 128
+ generation_batch_size: 16
+ max_total_sequence_length: 2048
+ sequence_packing:
+ train_mb_tokens: 2048
+ logprob_mb_tokens: 4096
+ megatron_cfg:
+ activation_checkpointing: true
+ env_vars:
+ PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
+ generation:
+ refit_transport: vllm_zmq_sparse
+ refit_cfg:
+ delta_compression:
+ encoding: xor
+ verify_samples_per_payload: 0
+ baseline:
+ in_memory: false
+ colocated:
+ enabled: false
+ resources:
+ gpus_per_node: 8
+ num_nodes: 2
+
+logger:
+ log_dir: logs/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated
+ wandb:
+ project: nemo-rl-refit
+ name: grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated
diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py
index 786ffb4548c..7818a8dd418 100644
--- a/nemo_rl/algorithms/distillation.py
+++ b/nemo_rl/algorithms/distillation.py
@@ -217,6 +217,14 @@ def setup(
assert generation_config is not None, (
"A generation config in the PolicyConfig is required for distillation"
)
+ if (
+ generation_config["backend"] == "vllm"
+ and cast(VllmConfig, generation_config).get("refit_transport") is not None
+ ):
+ raise ValueError(
+ "Remote sparse refit is currently supported only by GRPO; distillation "
+ "support is tracked in https://github.com/NVIDIA-NeMo/RL/issues/3275."
+ )
# Disallow SP + packing for dtensor path
for cfg, who in ((policy_config, "student"), (teacher_config, "teacher")):
diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py
index 6de9fd7a6e8..2e3a472809f 100644
--- a/nemo_rl/algorithms/grpo.py
+++ b/nemo_rl/algorithms/grpo.py
@@ -101,6 +101,7 @@
from nemo_rl.models.generation.sglang.config import SGLangConfig
from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration
+from nemo_rl.models.generation.vllm.config import normalize_vllm_refit_config
from nemo_rl.models.megatron.router_replay import (
configure_vllm_for_router_replay,
router_replay_enabled,
@@ -354,6 +355,8 @@ def setup(
assert generation_config is not None, (
"A generation config in the PolicyConfig is required for GRPO"
)
+ if generation_config["backend"] == "vllm":
+ normalize_vllm_refit_config(cast(VllmConfig, generation_config))
# Set seed for all random number generators
set_seed(grpo_config["seed"])
@@ -907,6 +910,9 @@ def _spinup_nemo_gym(base_urls, model_name):
# vllm model loading prefers clean environment, initialize policy_generation before policy in colocated mode
backend = generation_config["backend"]
generation_config["model_name"] = policy_config["model_name"] # Needed for vLLM
+ remote_transport = None
+ remote_synchronizer_cls = None
+ remote_baseline_init_refs: list[Any] = []
# Dictionary to store worker initialization timing stats for logging
worker_init_timing_metrics = {}
@@ -967,6 +973,11 @@ def init_policy():
init_optimizer=True,
init_reference_model=init_reference_model,
)
+ if remote_transport is not None:
+ assert remote_synchronizer_cls is not None
+ remote_baseline_init_refs.extend(
+ remote_synchronizer_cls.start_baseline(p, remote_transport)
+ )
return p, time.perf_counter() - t0
def init_vllm():
@@ -1095,6 +1106,21 @@ def initialize_generation_with_policy(
elif backend == "vllm":
# vLLM generation: setup config, then initialize with policy
generation_config = cast(VllmConfig, generation_config)
+ if generation_config.get("refit_transport") is not None:
+ # Keep optional remote transport dependencies off the default path.
+ from nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer import (
+ VllmRemoteSparseWeightSynchronizer,
+ validate_vllm_remote_sparse_refit,
+ )
+
+ remote_transport = validate_vllm_remote_sparse_refit(
+ generation_config,
+ colocated=colocated_inference,
+ megatron_enabled=policy_config["megatron_cfg"]["enabled"],
+ )
+ assert remote_transport is not None
+ remote_synchronizer_cls = VllmRemoteSparseWeightSynchronizer
+
if generation_config["vllm_cfg"]["precision"] == "fp8":
assert loss_config.use_importance_sampling_correction, (
"Importance sampling must be enabled for vLLM FP8 generation for good convergence!"
@@ -1232,7 +1258,7 @@ def init_vllm_then_policy():
policy.print_node_ip_and_gpu_id()
# if it is not colocated inference, initialize collective communication for update weights
- if not colocated_inference:
+ if not colocated_inference and remote_transport is None:
t0 = time.perf_counter()
ip, port = train_cluster.get_master_address_and_port()
print(f"Using ip: {ip}, port: {port} for collective communication", flush=True)
@@ -1271,9 +1297,30 @@ def init_vllm_then_policy():
ray.get(futures_train + futures_inference)
worker_init_timing_metrics["collective_init_time_s"] = time.perf_counter() - t0
- state_dict_info = policy.prepare_refit_info()
- if policy_generation is not None:
- policy_generation.prepare_refit_info(state_dict_info)
+ if remote_transport is not None:
+ t0 = time.perf_counter()
+ assert isinstance(policy_generation, VllmGeneration)
+ assert remote_synchronizer_cls is not None
+ refit_config = generation_config["refit_cfg"]
+ assert refit_config is not None
+ policy_generation.weight_synchronizer = remote_synchronizer_cls(
+ policy,
+ policy_generation,
+ transport=remote_transport,
+ api_key_env_var=generation_config["vllm_cfg"].get(
+ "http_refit_api_key_env_var"
+ ),
+ request_timeout_s=refit_config.request_timeout_s,
+ baseline_init_refs=remote_baseline_init_refs,
+ )
+ policy_generation.weight_synchronizer.init_communicator()
+ worker_init_timing_metrics[f"vllm_{remote_transport}_sparse_init_time_s"] = (
+ time.perf_counter() - t0
+ )
+ else:
+ state_dict_info = policy.prepare_refit_info()
+ if policy_generation is not None:
+ policy_generation.prepare_refit_info(state_dict_info)
# Spin up non-colocated OPD teacher worker groups AFTER policy / vLLM are
# ready. Parallelizing with policy init races on Megatron-Bridge's HF->mcore
@@ -2033,7 +2080,7 @@ def refit_policy_generation(
_refit_buffer_size_gb: Optional[float] = None,
timer: Optional[Timer] = None,
kv_scales: Optional[dict[str, float]] = None,
-) -> None:
+) -> dict[str, float]:
"""Refit the policy generation interface with the latest policy weights.
Args:
@@ -2043,7 +2090,14 @@ def refit_policy_generation(
the buffer size is computed from remaining memory.
timer: Optional Timer used to time the prepare/transfer/update phase
kv_scales: Optional dictionary of KV cache scales for FP8 quantization.
+
+ Returns:
+ Scalar metrics reported by the selected weight synchronizer.
"""
+ synchronizer = getattr(policy_generation, "weight_synchronizer", None)
+ if synchronizer is not None:
+ return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}
+
# Megatron generation backend needs explicit suspend/resume around refits.
if isinstance(policy_generation, MegatronGeneration):
policy_generation.suspend_for_refit()
@@ -2144,6 +2198,16 @@ def refit_policy_generation(
if isinstance(policy_generation, MegatronGeneration):
policy_generation.resume_after_refit()
+ return {}
+
+
+def _initial_policy_generation_stale(
+ policy_generation: GenerationInterface, completed_steps: int
+) -> bool:
+ """Skip a fresh run's redundant sync when the synchronizer is already current."""
+ synchronizer = getattr(policy_generation, "weight_synchronizer", None)
+ return completed_steps > 0 or synchronizer is None or synchronizer.is_stale
+
def _log_mixed_rewards_and_advantages_information(
logger: Logger,
@@ -2333,7 +2397,6 @@ def grpo_train(
isinstance(policy_generation, MegatronGeneration)
and master_config.policy["generation"]["colocated"]["enabled"]
)
- POLICY_GENERATION_STALE = True # tracks if generation needs a refit before running
assert policy_generation is not None
# Check if we need to sync KV cache scales
@@ -2343,6 +2406,9 @@ def grpo_train(
# common config/state times
current_step = grpo_save_state["current_step"] # current step within an epoch
total_steps = grpo_save_state["total_steps"] # total steps across all epochs
+ POLICY_GENERATION_STALE = _initial_policy_generation_stale(
+ policy_generation, total_steps
+ )
max_num_steps = master_config.grpo[
"max_num_steps"
] # max number of steps to train for
@@ -2413,6 +2479,7 @@ def grpo_train(
# Run grpo/dapo training loop (single-turn)
for batch in wrapped_dataloader:
+ refit_metrics: dict[str, float] = {}
# A central place to store logging data that won't be deleted until the loop ends
metrics_logging_data = dict()
metrics = dict()
@@ -2489,7 +2556,7 @@ def grpo_train(
calibration_data, include_q=True
)["layers"]
- refit_policy_generation(
+ refit_metrics = refit_policy_generation(
policy,
policy_generation,
colocated_inference,
@@ -2944,7 +3011,7 @@ def grpo_train(
):
memory_tracker.snapshot_start_of_stage("Validation", dir())
if NEED_REFIT and POLICY_GENERATION_STALE:
- refit_policy_generation(
+ refit_metrics = refit_policy_generation(
policy,
policy_generation,
colocated_inference,
@@ -3296,6 +3363,8 @@ def grpo_train(
train_results, metrics, timing_metrics, master_config
)
+ if refit_metrics:
+ logger.log_metrics(refit_metrics, total_steps + 1, prefix="refit")
logger.log_metrics(metrics, total_steps + 1, prefix="train")
logger.log_metrics(
performance_metrics, total_steps + 1, prefix="performance"
@@ -3627,11 +3696,11 @@ def async_grpo_train(
isinstance(policy_generation, MegatronGeneration)
and master_config.policy["generation"]["colocated"]["enabled"]
)
- POLICY_GENERATION_STALE = True
assert policy_generation is not None
# Training state
step = grpo_save_state["current_step"]
+ POLICY_GENERATION_STALE = _initial_policy_generation_stale(policy_generation, step)
weight_version = step # Tracks refitted weight versions
consumed_samples = grpo_save_state["consumed_samples"]
total_valid_tokens = grpo_save_state.get(
@@ -3895,6 +3964,7 @@ def async_grpo_train(
# Main training loop
try:
while step < master_config.grpo["max_num_steps"]:
+ refit_metrics: dict[str, float] = {}
print(
f"\n{'=' * 25} Step {step + 1}/{master_config.grpo['max_num_steps']} {'=' * 25}"
)
@@ -4287,7 +4357,7 @@ def async_grpo_train(
# Only the actual refit/weight transfer should be counted as weight_sync
print("🔄 Performing policy generation refit...")
with timer.time("weight_sync"):
- refit_policy_generation(
+ refit_metrics = refit_policy_generation(
policy,
policy_generation,
colocated_inference,
@@ -4318,7 +4388,7 @@ def async_grpo_train(
trajectory_collector.pause.remote()
if NEED_REFIT and POLICY_GENERATION_STALE:
- refit_policy_generation(
+ refit_metrics = refit_policy_generation(
policy, policy_generation, colocated_inference
)
POLICY_GENERATION_STALE = False
@@ -4645,6 +4715,8 @@ def async_grpo_train(
merged_efficiency, total_wall_time, step + 1
)
+ if refit_metrics:
+ logger.log_metrics(refit_metrics, step + 1, prefix="refit")
logger.log_metrics(performance_metrics, step + 1, prefix="performance")
logger.log_metrics(metrics, step + 1, prefix="train")
logger.log_metrics(efficiency_loggable, step + 1, prefix="")
diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py
index 80340bbbf30..5bef74b5cef 100644
--- a/nemo_rl/algorithms/ppo.py
+++ b/nemo_rl/algorithms/ppo.py
@@ -228,6 +228,14 @@ def setup(
assert generation_config is not None, (
"A generation config in the PolicyConfig is required for PPO"
)
+ if (
+ generation_config["backend"] == "vllm"
+ and cast(VllmConfig, generation_config).get("refit_transport") is not None
+ ):
+ raise ValueError(
+ "Remote sparse refit is currently supported only by GRPO; PPO support "
+ "is tracked in https://github.com/NVIDIA-NeMo/RL/issues/3275."
+ )
if "megatron_cfg" in policy_config and policy_config["megatron_cfg"]["enabled"]:
policy_megatron_config = cast(MegatronConfig, policy_config["megatron_cfg"])
diff --git a/nemo_rl/models/generation/__init__.py b/nemo_rl/models/generation/__init__.py
index 76756fc7176..cc471b1b27b 100644
--- a/nemo_rl/models/generation/__init__.py
+++ b/nemo_rl/models/generation/__init__.py
@@ -45,7 +45,9 @@ def configure_generation_config(
if config["backend"] == "vllm":
config = cast(VllmConfig, config)
# set load_format
- config["vllm_cfg"]["load_format"] = "auto" if is_eval else "dummy"
+ config["vllm_cfg"]["load_format"] = (
+ "auto" if is_eval or config.get("refit_transport") else "dummy"
+ )
speculative_config = config.get("vllm_kwargs", {}).get("speculative_config")
if speculative_config and not is_eval and not has_refit_draft_weights:
# Speculative decoding needs real draft weights at startup, since the
diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py
index 44eb5d5d89c..6b831cce923 100644
--- a/nemo_rl/models/generation/vllm/config.py
+++ b/nemo_rl/models/generation/vllm/config.py
@@ -14,8 +14,12 @@
from typing import Any, Literal, NotRequired, TypedDict
+from pydantic import BaseModel, Field, NonNegativeInt, PositiveFloat, PositiveInt
+
from nemo_rl.models.generation.interfaces import GenerationConfig
+VllmRefitTransportName = Literal["s3", "zmq"]
+
class VllmSpecificArgs(TypedDict):
tensor_parallel_size: int
@@ -39,6 +43,12 @@ class VllmSpecificArgs(TypedDict):
# Exposing vLLM as a server is useful in instances where the multi-turn rollout is performed with utilities outside of NeMo RL, but the user still wants to take advantage of the refit logic in NeMo RL that keeps the policy and generation up to date.
# Currently it will expose the /tokenize and /v1/chat/completions endpoints. Later on we may expose /v1/completions or /v1/responses.
expose_http_server: NotRequired[bool]
+ # Environment variable containing the internal refit API key.
+ http_refit_api_key_env_var: NotRequired[str | None]
+ # Fixed internal refit endpoint port for stable Kubernetes targetPorts.
+ http_refit_server_port: NotRequired[int | None]
+ # Fixed ZeroMQ relay port for stable Kubernetes targetPorts.
+ zmq_refit_server_port: NotRequired[int | None]
# These kwargs are passed to the vllm.LLM HTTP server Chat Completions endpoint config. Typically this will include things like tool parser, chat template, etc
http_server_serving_chat_kwargs: NotRequired[dict[str, Any]]
# Miscellaneous top level vLLM HTTP server arguments.
@@ -52,9 +62,61 @@ class VllmSpecificArgs(TypedDict):
reasoning_parser_plugin: NotRequired[str]
+class VllmDeltaCompressionConfig(BaseModel, extra="allow"):
+ encoding: Literal["xor", "overwrite"] = "xor"
+ sparse_bucket_size_bytes: PositiveInt = 512 * 1024**2
+ export_chunk_bytes: dict[str, PositiveInt] = Field(
+ default_factory=lambda: {"s3": 64 * 1024**2, "zmq": 256 * 1024**2}
+ )
+ zstd_threads: dict[str, NonNegativeInt] = Field(
+ default_factory=lambda: {"s3": 0, "zmq": 0}
+ )
+
+
+class VllmRefitStorageConfig(BaseModel, extra="allow"):
+ s3_bucket: str | None = None
+ s3_region: str = "us-east-1"
+ s3_prefix: str = "nemo-rl-refit"
+ staging_dir: str = "/dev/shm"
+
+
+class VllmRefitBaselineConfig(BaseModel, extra="allow"):
+ in_memory: bool = False
+ mmap_dir: str | None = None
+
+
+class VllmRefitTuningConfig(BaseModel, extra="allow"):
+ encode_workers: dict[str, PositiveInt] = Field(
+ default_factory=lambda: {"s3": 8, "zmq": 8}
+ )
+ transfer_workers: dict[str, PositiveInt] = Field(
+ default_factory=lambda: {"s3": 32, "zmq": 4}
+ )
+ zmq_retries: NonNegativeInt = 3
+ zmq_relay_payload_workers: PositiveInt = 16
+ zmq_relay_forward_workers: PositiveInt = 8
+ apply_queue_depth: PositiveInt = 32
+ apply_batch_size: PositiveInt = 8
+ partition_workers: PositiveInt = 8
+
+
+class VllmRefitConfig(BaseModel, extra="allow"):
+ delta_compression: VllmDeltaCompressionConfig = Field(
+ default_factory=VllmDeltaCompressionConfig
+ )
+ storage: VllmRefitStorageConfig = Field(default_factory=VllmRefitStorageConfig)
+ baseline: VllmRefitBaselineConfig = Field(default_factory=VllmRefitBaselineConfig)
+ tuning: VllmRefitTuningConfig = Field(default_factory=VllmRefitTuningConfig)
+ verify_samples_per_payload: NonNegativeInt = 0
+ request_timeout_s: PositiveFloat = 600.0
+
+
class VllmConfig(GenerationConfig):
vllm_cfg: VllmSpecificArgs
vllm_kwargs: NotRequired[dict[str, Any]]
+ # Null uses NCCL; remote sparse refit supports S3 or ZeroMQ value planes.
+ refit_transport: NotRequired[Literal["vllm_s3_sparse", "vllm_zmq_sparse"] | None]
+ refit_cfg: NotRequired[VllmRefitConfig | None]
# quantization config
quant_cfg: NotRequired[str | None]
@@ -63,3 +125,12 @@ class VllmConfig(GenerationConfig):
# modules. This is intended for ModelOpt NVFP4 rollout experiments.
real_quant: NotRequired[bool]
real_quant_ignore: NotRequired[list[str]]
+
+
+def normalize_vllm_refit_config(config: VllmConfig) -> VllmRefitConfig | None:
+ """Resolve sparse-refit defaults into the generation config."""
+ if config.get("refit_transport") is None:
+ return None
+ refit_config = VllmRefitConfig.model_validate(config.get("refit_cfg") or {})
+ config["refit_cfg"] = refit_config
+ return refit_config
diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py
index f6b2e5adb94..8344b77330c 100644
--- a/nemo_rl/models/generation/vllm/vllm_backend.py
+++ b/nemo_rl/models/generation/vllm/vllm_backend.py
@@ -13,6 +13,7 @@
# limitations under the License.
import gc
import re
+import socket
import traceback
from typing import Any
@@ -92,6 +93,8 @@ def _read_mtp_layer_weights_from_checkpoint(
class VllmInternalWorkerExtension:
+ _sparse_delta_applier: Any = None
+
def bind_numa(self) -> bool:
"""Pin this TP worker to its GPU's NUMA-local CPUs/memory.
@@ -137,6 +140,10 @@ def report_device_id(self) -> str:
return get_device_uuid(self.device.index)
+ def report_node_hostname(self) -> str:
+ """Return the host shared by worker processes on this node."""
+ return socket.gethostname()
+
def get_zmq_address(self):
"""Get the ZMQ address for the current device."""
return f"ipc:///tmp/{self.report_device_id()}.sock"
@@ -166,6 +173,13 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None:
"""
self.state_dict_info = state_dict_info # pyrefly: ignore[implicitly-defined-attribute] This class does not define __init__ so assignments like this should be ignored
+ def prepare_sparse_delta_refit_info(
+ self, state_dict_info: dict[str, tuple[tuple[int, ...], torch.dtype]]
+ ) -> list[str]:
+ """Reserve scratch space and report weights that require overwrite."""
+ applier = self._get_sparse_delta_applier()
+ return sorted(applier.discover_native_skips(state_dict_info))
+
def _maybe_process_fp8_kv_cache(self) -> None:
"""Process weights after loading for FP8 KV cache (static scales)."""
use_fp8_kv_cache = False
@@ -353,6 +367,19 @@ def _load_weights(self, weights):
self._load_draft_weights(draft_weights)
+ def _get_sparse_delta_applier(self) -> Any:
+ if self._sparse_delta_applier is None:
+ # Avoid importing sparse-refit code for existing refit transports.
+ from nemo_rl.models.generation.vllm.vllm_sparse_delta import (
+ VllmSparseDeltaApplier,
+ )
+
+ self._sparse_delta_applier = VllmSparseDeltaApplier(
+ self.model_runner,
+ self.device,
+ )
+ return self._sparse_delta_applier
+
@wrap_with_nvtx_name("vllm_internal_worker_extension/update_weights_via_ipc_zmq")
def update_weights_via_ipc_zmq(self) -> bool:
"""Receive and update model weights via ZMQ IPC socket.
@@ -482,6 +509,18 @@ def update_weights_from_collective(self) -> bool:
torch.cuda.empty_cache()
return True
+ def update_weights_from_decoded_sparse_payload(
+ self, *payloads: bytes | str
+ ) -> dict[str, Any]:
+ applier = self._get_sparse_delta_applier()
+ return applier.update_weights_from_decoded_sparse_payload(*payloads)
+
+ def synchronize_device(self) -> None:
+ self._get_sparse_delta_applier().synchronize_device()
+
+ def finish_sparse_delta_refit(self) -> dict[str, Any]:
+ return self._get_sparse_delta_applier().finish_sparse_delta_refit()
+
def cleanup(self) -> None:
"""Shutdown and cleanup resources."""
# Close ZMQ socket and context if they exist
diff --git a/nemo_rl/models/generation/vllm/vllm_generation.py b/nemo_rl/models/generation/vllm/vllm_generation.py
index d68bf512bcc..e889fe9ed2d 100644
--- a/nemo_rl/models/generation/vllm/vllm_generation.py
+++ b/nemo_rl/models/generation/vllm/vllm_generation.py
@@ -43,6 +43,7 @@
compute_spec_decode_metrics,
resolve_generation_worker_cls,
)
+from nemo_rl.weight_sync.interfaces import WeightSynchronizer
logger = logging.getLogger(__name__)
@@ -102,6 +103,7 @@ def __init__(
# Store config
self.cfg = config
self._defer_model_load = defer_model_load
+ self.weight_synchronizer: WeightSynchronizer | None = None
self.tp_size = self.cfg["vllm_cfg"]["tensor_parallel_size"]
self.pp_size = self.cfg["vllm_cfg"]["pipeline_parallel_size"]
self.ep_size = self.cfg["vllm_cfg"]["expert_parallel_size"]
@@ -903,6 +905,8 @@ def finish_generation(self, *args: Any, **kwargs: Any) -> bool:
def shutdown(self) -> bool:
"""Shut down all vLLM workers and clean up resources."""
try:
+ if self.weight_synchronizer is not None:
+ self.weight_synchronizer.shutdown()
# Use the worker group's shutdown method with the worker's cleanup method
return self.worker_group.shutdown(cleanup_method="shutdown")
except Exception as e:
diff --git a/nemo_rl/models/generation/vllm/vllm_sparse_delta.py b/nemo_rl/models/generation/vllm/vllm_sparse_delta.py
new file mode 100644
index 00000000000..1d77e7f388c
--- /dev/null
+++ b/nemo_rl/models/generation/vllm/vllm_sparse_delta.py
@@ -0,0 +1,525 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Apply canonical sparse updates through vLLM's native weight loaders."""
+
+import io
+import time
+from collections.abc import Iterable, Iterator, Mapping
+from math import prod
+from typing import Any, cast
+
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+from nemo_rl.utils import weight_transfer_sparse_codec as sparse_codec
+from nemo_rl.utils.nsys import wrap_with_nvtx_name
+
+_TensorViewKey = tuple[int, int, tuple[int, ...], tuple[int, ...]]
+_LoaderWeight = tuple[str, torch.Tensor, sparse_codec.SparseOperation, int, int | None]
+_LoaderObservation = tuple[str, int, bool]
+
+
+def _storage_key(tensor: torch.Tensor) -> int:
+ return tensor.untyped_storage()._cdata
+
+
+def _view_key(tensor: torch.Tensor) -> _TensorViewKey:
+ return (
+ _storage_key(tensor),
+ int(tensor.storage_offset()),
+ tuple(map(int, tensor.shape)),
+ tuple(map(int, tensor.stride())),
+ )
+
+
+class _SparseWeightLoadMode(TorchDispatchMode):
+ """Turn native loader copies into sparse XOR or overwrite."""
+
+ def __init__(
+ self,
+ targets: set[int],
+ verification: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
+ ) -> None:
+ super().__init__()
+ self._targets = targets
+ self._verification = verification
+ self._source_storage = 0
+ self._source_name = ""
+ self._operation: sparse_codec.SparseOperation = "overwrite"
+ self._sample_limit = 0
+ self._exact_sentinel: int | None = None
+ self._active_masks: dict[_TensorViewKey, torch.Tensor] = {}
+ self._verification_masks: dict[
+ _TensorViewKey, tuple[torch.Tensor, torch.Tensor]
+ ] = {}
+ self._xor_spans: dict[int, list[tuple[int, int]]] = {}
+ self.copies = 0
+ self.xor_compatible = True
+
+ def start(
+ self,
+ name: str,
+ source: torch.Tensor,
+ operation: sparse_codec.SparseOperation,
+ sample_limit: int,
+ exact_sentinel: int | None,
+ ) -> None:
+ self._source_name = name
+ self._source_storage = _storage_key(source)
+ self._operation = operation
+ self._sample_limit = sample_limit
+ self._exact_sentinel = exact_sentinel
+ self._active_masks.clear()
+ self._verification_masks.clear()
+ self._xor_spans.clear()
+ self.copies = 0
+ self.xor_compatible = True
+
+ def _observe_xor_copy(
+ self, destination: torch.Tensor, source: torch.Tensor
+ ) -> bool:
+ if (
+ _storage_key(source) != self._source_storage
+ or source.dtype != destination.dtype
+ ):
+ self.xor_compatible = False
+ return False
+ origin = int(destination.storage_offset())
+ extents = [
+ (int(size) - 1) * int(stride)
+ for size, stride in zip(
+ destination.shape, destination.stride(), strict=True
+ )
+ ]
+ span = (
+ origin + sum(min(0, extent) for extent in extents),
+ origin + sum(max(0, extent) for extent in extents),
+ )
+ spans = self._xor_spans.setdefault(_storage_key(destination), [])
+ if any(span[0] <= other[1] and other[0] <= span[1] for other in spans):
+ self.xor_compatible = False
+ return False
+ spans.append(span)
+ return True
+
+ def _remember_changed(
+ self, destination: torch.Tensor, changed: torch.Tensor
+ ) -> None:
+ if self._sample_limit <= 0:
+ return
+ view_key = _view_key(destination)
+ previous = self._verification_masks.get(view_key)
+ if previous is not None:
+ changed = previous[1] | changed
+ self._verification_masks[view_key] = (destination, changed)
+
+ def __torch_dispatch__(
+ self,
+ func: Any,
+ _types: Any,
+ args: tuple[Any, ...] = (),
+ kwargs: dict[str, Any] | None = None,
+ ) -> Any:
+ if func is not torch.ops.aten.copy_.default:
+ return func(*args, **(kwargs or {}))
+
+ destination, source = cast(tuple[torch.Tensor, torch.Tensor], args[:2])
+ if _storage_key(destination) not in self._targets:
+ return func(*args, **(kwargs or {}))
+
+ self.copies += 1
+ xor_compatible = self._observe_xor_copy(destination, source)
+ if self._operation == "overwrite":
+ if not source.dtype.is_floating_point:
+ raise RuntimeError("Sparse overwrite requires a floating-point loader.")
+ source = source.expand_as(destination)
+ view_key = _view_key(destination)
+ if self._exact_sentinel is not None:
+ if (
+ _storage_key(source) != self._source_storage
+ or source.dtype != destination.dtype
+ ):
+ raise RuntimeError(
+ "Exact FP8 overwrite cannot pass through a transforming loader."
+ )
+ source_bits = sparse_codec.integer_view(source)
+ changed = source_bits.ne(self._exact_sentinel)
+ destination_bits = sparse_codec.integer_view(destination)
+ destination_bits.masked_scatter_(
+ changed, source_bits.masked_select(changed)
+ )
+ self._remember_changed(destination, changed)
+ return destination
+ changed = self._active_masks.get(view_key)
+ if changed is None or _storage_key(source) == self._source_storage:
+ changed = ~torch.isnan(source)
+ self._active_masks[view_key] = changed
+ destination.masked_scatter_(
+ changed, source.masked_select(changed).to(destination.dtype)
+ )
+ self._remember_changed(destination, changed)
+ return destination
+
+ if not xor_compatible:
+ raise RuntimeError(
+ f"XOR for {self._source_name!r} cannot pass through this native "
+ "loader without changing semantics."
+ )
+ source = source.expand_as(destination)
+ destination_bits = sparse_codec.integer_view(destination)
+ source_bits = sparse_codec.integer_view(source)
+ changed = source_bits.ne(0)
+ values = destination_bits.masked_select(changed).bitwise_xor(
+ source_bits.masked_select(changed)
+ )
+ destination_bits.masked_scatter_(changed, values)
+ self._remember_changed(destination, changed)
+ return destination
+
+ @torch.no_grad()
+ def finish(self) -> None:
+ """Record bounded target samples after the loader finishes transforms."""
+ for target, changed in self._verification_masks.values():
+ if self._sample_limit <= 0:
+ break
+ if not target.is_contiguous():
+ continue
+ locations = changed.reshape(-1).nonzero().reshape(-1)[: self._sample_limit]
+ if locations.numel():
+ target_bits = sparse_codec.integer_view(target).reshape(-1)
+ self._verification.append(
+ (target, locations, target_bits.index_select(0, locations).clone())
+ )
+ self._sample_limit -= locations.numel()
+ self._verification_masks.clear()
+
+
+class VllmSparseDeltaApplier:
+ """Own one dense GPU scratch buffer and delegate all placement to vLLM."""
+
+ def __init__(self, model_runner: Any, device: torch.device) -> None:
+ self.model_runner = model_runner
+ self._cuda_device_index = device.index
+ model = model_runner.model
+ self._target_storages = {
+ _storage_key(tensor) for tensor in (*model.parameters(), *model.buffers())
+ }
+ self._scratch = torch.empty(0, dtype=torch.uint8, device=device)
+ self._skipped_names: set[str] = set()
+ self._verification: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
+
+ def discover_native_skips(
+ self, state_dict_info: Mapping[str, tuple[tuple[int, ...], torch.dtype]]
+ ) -> set[str]:
+ """Reserve scratch and classify rank-local skips and overwrite weights."""
+ required = max(
+ (prod(shape) * dtype.itemsize for shape, dtype in state_dict_info.values()),
+ default=0,
+ )
+ if required > self._scratch.numel():
+ self._scratch = torch.empty(
+ required, dtype=torch.uint8, device=self._scratch.device
+ )
+ pending = [
+ (name, shape, dtype)
+ for name, (shape, dtype) in state_dict_info.items()
+ if dtype.is_floating_point
+ ]
+ if not pending:
+ return set()
+
+ def weights() -> Iterator[_LoaderWeight]:
+ for name, shape, dtype in pending:
+ source = self._source_tensor(
+ {"shape": shape, "dtype": str(dtype).removeprefix("torch.")}
+ )
+ source.fill_(float("nan"))
+ exact_sentinel = (
+ int(sparse_codec.integer_view(source).reshape(-1)[0].item())
+ if source.element_size() == 1
+ else None
+ )
+ yield name, source, "overwrite", 0, exact_sentinel
+
+ loaded, observations = self._load_weights(weights(), [])
+ if len(observations) != len(pending):
+ raise RuntimeError(
+ "Native loader did not consume all sparse weight metadata."
+ )
+ self._validate_loader_report(loaded, observations, allow_unknown_skips=True)
+ if loaded is not None:
+ self._skipped_names.update(
+ name for name, copies, _ in observations if copies == 0
+ )
+ return {
+ name
+ for name, copies, xor_compatible in observations
+ if copies and not xor_compatible
+ }
+
+ def _source_tensor(self, item: dict[str, Any]) -> torch.Tensor:
+ shape = tuple(int(dim) for dim in item["shape"])
+ dtype = sparse_codec.dtype_from_name(str(item["dtype"]))
+ byte_count = prod(shape) * dtype.itemsize
+ if byte_count > self._scratch.numel():
+ self._scratch = torch.empty(
+ byte_count, dtype=torch.uint8, device=self._scratch.device
+ )
+ return self._scratch[:byte_count].view(dtype).view(shape)
+
+ @staticmethod
+ def _scatter_values(
+ source: torch.Tensor,
+ item: dict[str, Any],
+ locations: torch.Tensor,
+ values: torch.Tensor,
+ ) -> None:
+ source_bits = sparse_codec.integer_view(source).reshape(-1)
+ expected_dtype = sparse_codec.integer_dtype_for_element_size(
+ source.element_size()
+ )
+ if values.dtype != expected_dtype:
+ raise RuntimeError(
+ f"Sparse values have the wrong dtype for {item['name']!r}."
+ )
+ values = values.to(device=source.device, non_blocking=True)
+ if item["index_encoding"] == "range":
+ source_bits.narrow(0, int(item["range_start"]), values.numel()).copy_(
+ values
+ )
+ return
+ source_bits.index_copy_(
+ 0,
+ locations.to(device=source.device, dtype=torch.int64, non_blocking=True),
+ values,
+ )
+
+ def _prepare_loader_weight(
+ self,
+ item: dict[str, Any],
+ locations: torch.Tensor,
+ values: torch.Tensor,
+ ) -> _LoaderWeight:
+ operation = sparse_codec.sparse_operation(item["operation"])
+ source = self._source_tensor(item)
+ exact_sentinel = None
+ if operation == "xor":
+ source.zero_()
+ elif not source.dtype.is_floating_point:
+ raise RuntimeError("Sparse overwrite requires a floating-point source.")
+ else:
+ source.fill_(float("nan"))
+ if source.element_size() == 1:
+ source_bits = sparse_codec.integer_view(source)
+ exact_sentinel = int(source_bits.reshape(-1)[0].item())
+ if bool(values.eq(exact_sentinel).any()):
+ exact_sentinel ^= 0x80
+ if bool(values.eq(exact_sentinel).any()):
+ raise RuntimeError(
+ "FP8 sparse overwrite exhausted its sentinel values."
+ )
+ source_bits.fill_(exact_sentinel)
+ self._scatter_values(source, item, locations, values)
+
+ return (
+ str(item["name"]),
+ source,
+ operation,
+ int(item.get("verification_samples", 0)),
+ exact_sentinel,
+ )
+
+ def _load_weights(
+ self,
+ weights: Iterable[_LoaderWeight],
+ verification: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
+ ) -> tuple[Any, list[_LoaderObservation]]:
+ mode = _SparseWeightLoadMode(self._target_storages, verification)
+ yielded_names: list[str] = []
+ observations: list[_LoaderObservation] = []
+
+ def observed_weights() -> Iterator[tuple[str, torch.Tensor]]:
+ active = False
+ for name, source, operation, sample_limit, exact_sentinel in weights:
+ if active:
+ mode.finish()
+ observations.append(
+ (yielded_names[-1], mode.copies, mode.xor_compatible)
+ )
+ mode.start(name, source, operation, sample_limit, exact_sentinel)
+ active = True
+ yielded_names.append(name)
+ yield name, source
+ if active:
+ mode.finish()
+ observations.append(
+ (yielded_names[-1], mode.copies, mode.xor_compatible)
+ )
+
+ with torch.no_grad(), mode:
+ loaded = self.model_runner.model.load_weights(observed_weights())
+ if len(observations) != len(yielded_names):
+ raise RuntimeError("Native loader did not consume all sparse weights.")
+ return loaded, observations
+
+ @staticmethod
+ def _validate_loader_report(
+ loaded: Any,
+ observations: list[_LoaderObservation],
+ *,
+ allow_unknown_skips: bool,
+ ) -> None:
+ copied = sum(copies > 0 for _, copies, _ in observations)
+ if loaded is None:
+ if not allow_unknown_skips and copied != len(observations):
+ raise RuntimeError(
+ "Native loader did not report whether uncopied sparse weights "
+ "were skipped."
+ )
+ elif len(loaded) > copied:
+ raise RuntimeError(
+ "Native loader reported a loaded sparse weight without a supported "
+ "target copy."
+ )
+
+ def _apply_decoded_items(
+ self,
+ items: Iterable[tuple[dict[str, Any], torch.Tensor, torch.Tensor]],
+ ) -> None:
+ def weights() -> Iterator[_LoaderWeight]:
+ for item, locations, values in items:
+ yield self._prepare_loader_weight(item, locations, values)
+
+ loaded, observations = self._load_weights(weights(), self._verification)
+ self._validate_loader_report(loaded, observations, allow_unknown_skips=False)
+
+ def _iter_sparse_payload(
+ self, payload: sparse_codec.TensorPayload
+ ) -> Iterator[sparse_codec.SparseItem]:
+ packed_locations, value_groups, metadata = payload
+ for item in metadata:
+ if str(item["name"]) in self._skipped_names:
+ continue
+ value_start = int(item["value_start"])
+ value_end = int(item["value_end"])
+ location_dtype = (
+ torch.int32
+ if prod(item["shape"]) <= torch.iinfo(torch.int32).max
+ else torch.int64
+ )
+ yield (
+ item,
+ sparse_codec.sparse_locations_for_item(
+ item,
+ packed_locations,
+ device="cpu",
+ dtype=location_dtype,
+ ),
+ value_groups[int(item["value_group"])][value_start:value_end],
+ )
+
+ @wrap_with_nvtx_name(
+ "vllm_internal_worker_extension/update_weights_from_decoded_sparse_payload"
+ )
+ def update_weights_from_decoded_sparse_payload(
+ self, *payloads: bytes | str
+ ) -> dict[str, Any]:
+ return self._load_sparse_payloads(
+ tuple(
+ io.BytesIO(payload) if isinstance(payload, bytes) else payload
+ for payload in payloads
+ )
+ )
+
+ def _load_sparse_payloads(
+ self, sources: tuple[str | io.BytesIO, ...]
+ ) -> dict[str, Any]:
+ started = time.perf_counter()
+ deserialize_s = [0.0]
+
+ def sparse_items() -> Iterator[sparse_codec.SparseItem]:
+ for source in sources:
+ item_started = time.perf_counter()
+ payload = cast(
+ sparse_codec.TensorPayload,
+ torch.load(
+ source,
+ map_location="cpu",
+ weights_only=True,
+ mmap=isinstance(source, str),
+ ),
+ )
+ deserialize_s[0] += time.perf_counter() - item_started
+ yield from self._iter_sparse_payload(payload)
+
+ item_started = time.perf_counter()
+ self._apply_decoded_items(sparse_items())
+ sparse_apply_s = time.perf_counter() - item_started
+ return {
+ "ok": True,
+ "receiver_deserialize_s": deserialize_s[0],
+ "receiver_sparse_apply_s": sparse_apply_s,
+ "receiver_total_s": time.perf_counter() - started,
+ }
+
+ def synchronize_device(self) -> None:
+ if torch.cuda.is_available():
+ torch.cuda.synchronize(self._cuda_device_index)
+
+ def finish_sparse_delta_refit(self) -> dict[str, Any]:
+ """Synchronize and compare bounded samples of target entries just changed."""
+ self.synchronize_device()
+ verification, self._verification = self._verification, []
+ stats = (
+ torch.zeros(4, device=verification[0][0].device) if verification else None
+ )
+ samples = 0
+ with torch.no_grad():
+ for target, locations, expected_bits in verification:
+ actual_bits = (
+ sparse_codec.integer_view(target)
+ .reshape(-1)
+ .index_select(0, locations)
+ )
+ bit_mismatches = actual_bits.ne(expected_bits)
+ actual = actual_bits.view(target.dtype).float()
+ expected = expected_bits.view(target.dtype).float()
+ difference = torch.nan_to_num(
+ torch.where(
+ bit_mismatches,
+ (actual - expected).abs(),
+ torch.zeros_like(actual),
+ ),
+ nan=float("inf"),
+ )
+ assert stats is not None
+ stats[0] += difference.sum()
+ stats[1] = torch.maximum(stats[1], difference.max())
+ stats[2] += bit_mismatches.sum()
+ stats[3] += (
+ bit_mismatches
+ & ~torch.isclose(actual, expected, rtol=1e-6, atol=1e-8)
+ ).sum()
+ samples += actual.numel()
+ values = [0.0] * 4 if stats is None else stats.cpu().tolist()
+ return {
+ "ok": True,
+ "verification_candidates": samples,
+ "verification_samples": samples,
+ "verification_exact_mismatches": int(values[2]),
+ "verification_mismatches": int(values[3]),
+ "verification_abs_sum": float(values[0]),
+ "verification_max_abs": float(values[1]),
+ }
diff --git a/nemo_rl/models/generation/vllm/vllm_sparse_refit.py b/nemo_rl/models/generation/vllm/vllm_sparse_refit.py
new file mode 100644
index 00000000000..b298d0674d1
--- /dev/null
+++ b/nemo_rl/models/generation/vllm/vllm_sparse_refit.py
@@ -0,0 +1,558 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Remote sparse-refit receiver lifecycle for vLLM generation workers."""
+
+import asyncio
+import hmac
+import logging
+import os
+import tempfile
+import threading
+import time
+from collections.abc import Mapping
+from concurrent.futures import Future, ThreadPoolExecutor
+from functools import cache
+from typing import Any, Literal, NamedTuple, cast
+
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+
+from nemo_rl.distributed.virtual_cluster import (
+ DEFAULT_GENERATION_PORT_RANGE_HIGH,
+ DEFAULT_GENERATION_PORT_RANGE_LOW,
+ _get_free_port_local,
+ _get_node_ip_local,
+)
+from nemo_rl.models.generation.vllm.config import VllmRefitConfig
+from nemo_rl.utils import weight_transfer_sparse_codec as sparse_codec
+from nemo_rl.utils.weight_transfer_http import (
+ G_VLLM_REFIT_API_KEY_HEADER,
+ G_VLLM_REFIT_FLUSH_PATH,
+ G_VLLM_REFIT_PREPARE_PATH,
+ G_VLLM_REFIT_S3_MANIFEST_PATH,
+ G_VLLM_REFIT_ZMQ_FLUSH_PATH,
+ merge_vllm_refit_metrics,
+ vllm_refit_api_key,
+)
+from nemo_rl.utils.weight_transfer_stream import (
+ decode_sparse_payload,
+ download_s3_refit_payload,
+)
+from nemo_rl.utils.weight_transfer_zmq import (
+ ZmqSparseRefitServer,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@cache
+def _warn_unauthenticated_refit_server(transport: str) -> None:
+ logger.warning(
+ "%s sparse-refit server is binding 0.0.0.0 without an API key; "
+ "weight-write endpoints are reachable from the network.",
+ transport,
+ )
+
+
+class _StagedSparsePayload(NamedTuple):
+ path: str
+ started_at: float
+ finished_at: float
+ save_s: float
+
+
+def _stage_sparse_payload(
+ serialized: bytes,
+ staging_dir: str,
+) -> _StagedSparsePayload:
+ started_at = time.perf_counter()
+ descriptor, path = tempfile.mkstemp(
+ prefix="nemo_rl_refit_", suffix=".pt", dir=staging_dir
+ )
+ started = time.perf_counter()
+ try:
+ with os.fdopen(descriptor, "wb") as handle:
+ handle.write(serialized)
+ save_s = time.perf_counter() - started
+ except Exception:
+ os.unlink(path)
+ raise
+ finished_at = time.perf_counter()
+ return _StagedSparsePayload(
+ path,
+ started_at,
+ finished_at,
+ save_s,
+ )
+
+
+class VllmSparseRefitReceiver:
+ """Own the optional transport server, apply queue, and relay resources."""
+
+ def __init__(self, worker: Any) -> None:
+ self._worker = worker
+ self._refit_config = VllmRefitConfig.model_validate(
+ worker.cfg.get("refit_cfg") or {}
+ )
+ tuning = self._refit_config.tuning
+ self._refit_apply_queue_condition = threading.Condition()
+ self._refit_apply_executor = ThreadPoolExecutor(
+ max_workers=1,
+ thread_name_prefix="nrl-vllm-sparse-refit",
+ )
+ self._refit_apply_futures: list[Future[dict[str, Any]]] = []
+ self._refit_apply_pending_payloads: list[
+ bytes | Future[_StagedSparsePayload]
+ ] = []
+ self._refit_seen_payloads: dict[tuple[str, int, int], str] = {}
+ self._refit_workers_share_node = False
+ self._refit_apply_queue_depth = tuning.apply_queue_depth
+ self._refit_apply_batch_size = tuning.apply_batch_size
+ self._refit_partition_executor = ThreadPoolExecutor(
+ max_workers=tuning.partition_workers,
+ thread_name_prefix="nrl-vllm-sparse-partition",
+ )
+ self._refit_verification_candidates = 0
+ self._refit_batch_staging_dir = self._refit_config.storage.staging_dir
+ self._refit_http_server: tuple[Any, threading.Thread, str] | None = None
+ self._zmq_refit_server: tuple[ZmqSparseRefitServer, str] | None = None
+ self._refit_async_loop: asyncio.AbstractEventLoop | None = None
+
+ def set_worker_hostnames(self, hostnames: list[str]) -> None:
+ self._refit_workers_share_node = len(set(hostnames)) == 1
+
+ def start_sync_server(self) -> None:
+ llm = self._worker.llm
+ if llm is None:
+ raise RuntimeError("vLLM is not initialized on this worker.")
+ self.set_worker_hostnames(llm.collective_rpc("report_node_hostname", args=()))
+ self._setup_vllm_refit_server()
+
+ def shutdown(self) -> None:
+ self.stop_zmq_sparse_refit_relay()
+ if self._refit_http_server is not None:
+ self._refit_http_server[0].should_exit = True
+
+ self._flush_queued_sparse_payloads()
+ self._refit_apply_executor.shutdown(wait=True)
+ self._refit_partition_executor.shutdown(wait=True)
+
+ if self._refit_http_server is not None:
+ self._refit_http_server[1].join(timeout=5.0)
+ self._refit_http_server = None
+
+ def _enqueue_sparse_payload_apply(
+ self,
+ payload: bytes,
+ payload_key: tuple[str, int, int],
+ checksum: str,
+ verification_candidates: int = 0,
+ ) -> dict[str, Any]:
+ completed: list[Future[dict[str, Any]]] = []
+ with self._refit_apply_queue_condition:
+ seen_checksum = self._refit_seen_payloads.get(payload_key)
+ if seen_checksum is not None:
+ if seen_checksum != checksum:
+ raise ValueError(
+ "A sparse refit payload ID was reused with different data."
+ )
+ return {"ok": True, "payloads": 0, "duplicate": True}
+ while (
+ len(self._refit_apply_futures) >= self._refit_apply_queue_depth
+ and not self._refit_apply_futures[0].done()
+ ):
+ self._refit_apply_queue_condition.wait()
+ while self._refit_apply_futures and self._refit_apply_futures[0].done():
+ completed.append(self._refit_apply_futures.pop(0))
+ response = self._collect_refit_apply_results(completed)
+ self._refit_seen_payloads[payload_key] = checksum
+ self._refit_verification_candidates += verification_candidates
+ pending: bytes | Future[_StagedSparsePayload] = payload
+ if self._refit_workers_share_node:
+ pending = self._refit_partition_executor.submit(
+ _stage_sparse_payload,
+ payload,
+ self._refit_batch_staging_dir,
+ )
+ self._refit_apply_pending_payloads.append(pending)
+ if len(self._refit_apply_pending_payloads) == self._refit_apply_batch_size:
+ self._submit_pending_sparse_payloads()
+ return response
+
+ def _submit_pending_sparse_payloads(self) -> None:
+ payloads = tuple(self._refit_apply_pending_payloads)
+ self._refit_apply_pending_payloads.clear()
+ apply = (
+ self.update_weights_from_staged_sparse_payloads
+ if self._refit_workers_share_node
+ else self.update_weights_from_serialized_sparse_payloads
+ )
+ future = self._refit_apply_executor.submit(cast(Any, apply), payloads)
+ self._refit_apply_futures.append(future)
+ future.add_done_callback(self._notify_refit_apply_waiters)
+
+ def _notify_refit_apply_waiters(self, _future: Future[dict[str, Any]]) -> None:
+ with self._refit_apply_queue_condition:
+ self._refit_apply_queue_condition.notify_all()
+
+ def _collect_refit_apply_results(
+ self,
+ futures: list[Future[dict[str, Any]]],
+ ) -> dict[str, Any]:
+ results = [future.result() for future in futures]
+ return {
+ "ok": True,
+ "payloads": sum(int(result.get("payloads", 0)) for result in results),
+ **merge_vllm_refit_metrics({}, results, maximum=False),
+ }
+
+ @staticmethod
+ def _refit_collective_response(worker_results: Any) -> dict[str, Any]:
+ results = cast(list[dict[str, Any]], worker_results)
+ return {
+ "ok": True,
+ **merge_vllm_refit_metrics(
+ {}, results, maximum=True, candidate_maximum=True
+ ),
+ }
+
+ def _refit_collective_rpc(
+ self,
+ method: str,
+ args: tuple[Any, ...],
+ ) -> Any:
+ llm = self._worker.llm
+ if llm is None:
+ raise RuntimeError("vLLM is not initialized on this worker.")
+ if not self._worker.cfg["vllm_cfg"]["async_engine"]:
+ return llm.collective_rpc(method, args=args)
+ if self._refit_async_loop is None:
+ raise RuntimeError("The async vLLM refit server loop is not initialized.")
+ return asyncio.run_coroutine_threadsafe(
+ llm.collective_rpc(method, args=args),
+ self._refit_async_loop,
+ ).result()
+
+ def update_weights_from_serialized_sparse_payloads(
+ self,
+ serialized_payloads: tuple[bytes, ...],
+ ) -> dict[str, Any]:
+ """Apply a FIFO batch of sparse deltas through one collective RPC."""
+ response = self._refit_collective_response(
+ self._refit_collective_rpc(
+ "update_weights_from_decoded_sparse_payload",
+ serialized_payloads,
+ )
+ )
+ response["payloads"] = len(serialized_payloads)
+ return response
+
+ def update_weights_from_staged_sparse_payloads(
+ self,
+ staged_payloads: tuple[Future[_StagedSparsePayload], ...],
+ ) -> dict[str, Any]:
+ started = time.perf_counter()
+ staged: list[_StagedSparsePayload] = []
+ try:
+ stage_error = None
+ for future in staged_payloads:
+ try:
+ staged.append(future.result())
+ except Exception as exc:
+ stage_error = stage_error or exc
+ if stage_error is not None:
+ raise stage_error
+ stage_wait_s = time.perf_counter() - started
+ try:
+ response = self._refit_collective_response(
+ self._refit_collective_rpc(
+ "update_weights_from_decoded_sparse_payload",
+ tuple(payload.path for payload in staged),
+ )
+ )
+ except Exception:
+ # Drain peers before removing shared batch files.
+ self._refit_collective_rpc("synchronize_device", ())
+ raise
+ finally:
+ for payload in staged:
+ os.unlink(payload.path)
+ worker_total_s = float(response.get("receiver_total_s", 0.0))
+ response.update(
+ receiver_node_deserialize_s=0.0,
+ receiver_stage_s=(
+ max(payload.finished_at for payload in staged)
+ - min(payload.started_at for payload in staged)
+ ),
+ receiver_stage_save_s=max(
+ (payload.save_s for payload in staged), default=0.0
+ ),
+ receiver_stage_wait_s=stage_wait_s,
+ )
+ response["receiver_worker_total_s"] = worker_total_s
+ response["receiver_total_s"] = time.perf_counter() - started
+ response["payloads"] = len(staged_payloads)
+ return response
+
+ def _flush_queued_sparse_payloads(self) -> dict[str, Any]:
+ started = time.perf_counter()
+ with self._refit_apply_queue_condition:
+ if self._refit_apply_pending_payloads:
+ self._submit_pending_sparse_payloads()
+ futures = list(self._refit_apply_futures)
+ self._refit_apply_futures.clear()
+ self._refit_apply_queue_condition.notify_all()
+ payload_count = len(self._refit_seen_payloads)
+ batch_count = (
+ payload_count + self._refit_apply_batch_size - 1
+ ) // self._refit_apply_batch_size
+ response = self._collect_refit_apply_results(futures)
+ if futures:
+ verification = self._refit_collective_response(
+ self._refit_collective_rpc("finish_sparse_delta_refit", ())
+ )
+ if self._refit_verification_candidates:
+ verification["verification_candidates"] = (
+ self._refit_verification_candidates
+ )
+ response.update(verification)
+ with self._refit_apply_queue_condition:
+ self._refit_seen_payloads.clear()
+ self._refit_verification_candidates = 0
+ response.update(
+ payloads=payload_count,
+ batches=batch_count,
+ seconds=time.perf_counter() - started,
+ )
+ if futures:
+ print(
+ "REFIT_RECEIVER_TIMING "
+ f"payloads={payload_count} batches={batch_count} "
+ f"total_s={response['seconds']:.3f} "
+ f"payload_total_s={response.get('receiver_total_s', 0.0):.3f} "
+ f"delta_verify_candidates="
+ f"{response.get('verification_candidates', 0)} "
+ f"delta_verify_samples={response.get('verification_samples', 0)} "
+ f"delta_verify_exact_mismatches="
+ f"{response.get('verification_exact_mismatches', 0)} "
+ f"delta_verify_mismatches="
+ f"{response.get('verification_mismatches', 0)} "
+ f"delta_verify_max_abs="
+ f"{response.get('verification_max_abs', 0.0):.8g}",
+ flush=True,
+ )
+ return response
+
+ def _prepare_sparse_refit_info(self, request: dict[str, Any]) -> dict[str, Any]:
+ started = time.perf_counter()
+ state_dict_info = {
+ name: (tuple(shape), sparse_codec.dtype_from_name(dtype))
+ for name, (shape, dtype) in request["tensors"].items()
+ }
+ worker_results = self._refit_collective_rpc(
+ "prepare_sparse_delta_refit_info", (state_dict_info,)
+ )
+ overwrite_names = sorted({name for names in worker_results for name in names})
+ seconds = time.perf_counter() - started
+ print(
+ f"REFIT_RECEIVER_PREWARM tensors={len(state_dict_info)} "
+ f"overwrite_tensors={len(overwrite_names)} seconds={seconds:.3f}",
+ flush=True,
+ )
+ return {
+ "ok": True,
+ "tensors": len(state_dict_info),
+ "overwrite_names": overwrite_names,
+ "seconds": seconds,
+ }
+
+ async def _apply_s3_manifest_payload(
+ self,
+ manifest: dict[str, Any],
+ ) -> dict[str, Any]:
+ started = time.perf_counter()
+ body = await asyncio.to_thread(download_s3_refit_payload, manifest)
+ download_s = time.perf_counter() - started
+ key = str(manifest["key"])
+ checksum = str(manifest["checksum"])
+ result = await asyncio.to_thread(
+ self._enqueue_sparse_payload_apply,
+ body,
+ (key, -1, -1),
+ checksum,
+ int(manifest["verification_candidates"]),
+ )
+ result["receiver_s3_download_s"] = download_s
+ return result
+
+ def _apply_zmq_payload(
+ self, compressed: bytes, metadata: Mapping[str, Any]
+ ) -> dict[str, Any]:
+ started = time.perf_counter()
+ checksum = str(metadata["checksum"])
+ payload = decode_sparse_payload(compressed, checksum)
+ decode_s = time.perf_counter() - started
+ result = self._enqueue_sparse_payload_apply(
+ payload,
+ (
+ str(metadata["transfer_id"]),
+ int(metadata["producer_id"]),
+ int(metadata["payload_id"]),
+ ),
+ checksum,
+ int(metadata["verification_candidates"]),
+ )
+ result["receiver_zmq_decode_s"] = decode_s
+ return result
+
+ def setup_api_server(self, app: Any) -> None:
+ cfg = self._worker.cfg
+ token = vllm_refit_api_key(cfg["vllm_cfg"].get("http_refit_api_key_env_var"))
+
+ async def respond(
+ raw_request: Request,
+ action: Literal["prepare", "s3", "flush", "zmq_flush"],
+ ) -> JSONResponse:
+ if cfg["vllm_cfg"]["async_engine"]:
+ self._refit_async_loop = asyncio.get_running_loop()
+ supplied_token = raw_request.headers.get(G_VLLM_REFIT_API_KEY_HEADER)
+ if token is not None and (
+ supplied_token is None or not hmac.compare_digest(token, supplied_token)
+ ):
+ return JSONResponse(
+ content={"ok": False, "error": "unauthorized"}, status_code=403
+ )
+ try:
+ if action == "prepare":
+ result = await asyncio.to_thread(
+ self._prepare_sparse_refit_info,
+ await raw_request.json(),
+ )
+ elif action == "s3":
+ result = await self._apply_s3_manifest_payload(
+ await raw_request.json()
+ )
+ elif action == "zmq_flush":
+ body = await raw_request.json()
+ result = await asyncio.to_thread(
+ self.flush_zmq_sparse_refit_relay,
+ str(body["transfer_id"]),
+ int(body.get("expected_payloads", 0)),
+ )
+ else:
+ result = await asyncio.to_thread(self._flush_queued_sparse_payloads)
+ except Exception as exc:
+ result = {"ok": False, "error": str(exc)}
+ return JSONResponse(
+ content=result,
+ status_code=200 if result.get("ok") is True else 500,
+ )
+
+ def endpoint(
+ action: Literal["prepare", "s3", "flush", "zmq_flush"],
+ ):
+ async def handle(raw_request: Request) -> JSONResponse:
+ return await respond(raw_request, action)
+
+ return handle
+
+ for path, action in (
+ (G_VLLM_REFIT_S3_MANIFEST_PATH, "s3"),
+ (G_VLLM_REFIT_PREPARE_PATH, "prepare"),
+ (G_VLLM_REFIT_FLUSH_PATH, "flush"),
+ (G_VLLM_REFIT_ZMQ_FLUSH_PATH, "zmq_flush"),
+ ):
+ app.add_api_route(path, endpoint(action), methods=["POST"])
+
+ def report_refit_server_base_url(self) -> str | None:
+ if self._refit_http_server is not None:
+ return self._refit_http_server[2]
+ base_url = getattr(self._worker, "base_url", None)
+ return base_url.removesuffix("/v1") if base_url else None
+
+ def start_zmq_sparse_refit_relay(self) -> str:
+ if self._zmq_refit_server is not None:
+ return self._zmq_refit_server[1]
+ cfg = self._worker.cfg
+ port = cfg["vllm_cfg"].get("zmq_refit_server_port") or _get_free_port_local(
+ cfg.get("port_range_low", DEFAULT_GENERATION_PORT_RANGE_LOW),
+ cfg.get("port_range_high", DEFAULT_GENERATION_PORT_RANGE_HIGH),
+ )
+ server = ZmqSparseRefitServer(
+ self._apply_zmq_payload,
+ bind_address=f"tcp://0.0.0.0:{port}",
+ api_key_env_var=cfg["vllm_cfg"].get("http_refit_api_key_env_var"),
+ timeout_s=self._refit_config.request_timeout_s,
+ tuning=self._refit_config.tuning,
+ )
+ if (
+ vllm_refit_api_key(cfg["vllm_cfg"].get("http_refit_api_key_env_var"))
+ is None
+ ):
+ _warn_unauthenticated_refit_server("ZeroMQ")
+ server.start()
+ address = f"tcp://{_get_node_ip_local()}:{port}"
+ self._zmq_refit_server = (server, address)
+ print(f"Starting vLLM ZeroMQ refit relay on {address}", flush=True)
+ return address
+
+ def configure_zmq_sparse_refit_relay(self, relay_addresses: list[str]) -> None:
+ if self._zmq_refit_server is None:
+ raise RuntimeError("ZeroMQ sparse refit relay is not running.")
+ server, own_address = self._zmq_refit_server
+ server.configure_tree(
+ relay_addresses,
+ own_address=own_address,
+ )
+
+ def stop_zmq_sparse_refit_relay(self) -> None:
+ if self._zmq_refit_server is not None:
+ self._zmq_refit_server[0].close()
+ self._zmq_refit_server = None
+
+ def flush_zmq_sparse_refit_relay(
+ self, transfer_id: str, expected_payloads: int = 0
+ ) -> dict[str, Any]:
+ if self._zmq_refit_server is None:
+ raise RuntimeError("ZeroMQ sparse refit relay is not running.")
+ return self._zmq_refit_server[0].flush(transfer_id, expected_payloads)
+
+ def _setup_vllm_refit_server(self) -> None:
+ app = FastAPI()
+ self.setup_api_server(app)
+ cfg = self._worker.cfg
+ port = cfg["vllm_cfg"].get("http_refit_server_port") or _get_free_port_local(
+ cfg.get("port_range_low", DEFAULT_GENERATION_PORT_RANGE_LOW),
+ cfg.get("port_range_high", DEFAULT_GENERATION_PORT_RANGE_HIGH),
+ )
+ if (
+ vllm_refit_api_key(cfg["vllm_cfg"].get("http_refit_api_key_env_var"))
+ is None
+ ):
+ _warn_unauthenticated_refit_server("HTTP")
+ server = uvicorn.Server(
+ uvicorn.Config(
+ app,
+ host="0.0.0.0",
+ port=port,
+ timeout_keep_alive=120,
+ )
+ )
+ thread = threading.Thread(target=server.run, daemon=True)
+ thread.start()
+ base_url = f"http://{_get_node_ip_local()}:{port}"
+ self._refit_http_server = (server, thread, base_url)
+ print(f"Starting vLLM refit server on {base_url}", flush=True)
diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py
index 8329da0b93c..b3a8aa2b63d 100644
--- a/nemo_rl/models/generation/vllm/vllm_worker.py
+++ b/nemo_rl/models/generation/vllm/vllm_worker.py
@@ -247,6 +247,14 @@ def __init__(
self._init_config(
config, bundle_indices, fraction_of_gpus, seed, extra_env_vars
)
+ self._sparse_refit_receiver: Any = None
+ if self.is_model_owner and self.cfg.get("refit_transport") is not None:
+ # Avoid receiver dependencies and threads for existing refit transports.
+ from nemo_rl.models.generation.vllm.vllm_sparse_refit import (
+ VllmSparseRefitReceiver,
+ )
+
+ self._sparse_refit_receiver = VllmSparseRefitReceiver(self)
if not self.is_model_owner:
return
@@ -650,6 +658,27 @@ def _get_raw_spec_counters(self) -> dict[str, float | list[float]]:
metrics[metric.name] = metric.value
return metrics
+ def report_refit_server_base_url(self) -> str | None:
+ receiver = self._sparse_refit_receiver
+ return receiver.report_refit_server_base_url() if receiver is not None else None
+
+ def start_zmq_sparse_refit_relay(self) -> str:
+ receiver = self._sparse_refit_receiver
+ if receiver is None:
+ raise RuntimeError("Remote sparse refit is not enabled for this worker.")
+ return receiver.start_zmq_sparse_refit_relay()
+
+ def configure_zmq_sparse_refit_relay(self, relay_addresses: list[str]) -> None:
+ receiver = self._sparse_refit_receiver
+ if receiver is None:
+ raise RuntimeError("Remote sparse refit is not enabled for this worker.")
+ receiver.configure_zmq_sparse_refit_relay(relay_addresses)
+
+ def stop_zmq_sparse_refit_relay(self) -> None:
+ receiver = self._sparse_refit_receiver
+ if receiver is not None:
+ receiver.stop_zmq_sparse_refit_relay()
+
class VllmGenerationWorkerImpl(BaseVllmGenerationWorker):
def _create_engine(self, llm_kwargs: dict[str, Any]) -> None:
@@ -665,6 +694,8 @@ def post_init(self):
self.llm.collective_rpc(
"load_mtp_weights_from_disk", args=(self.model_name,)
)
+ if self._sparse_refit_receiver is not None:
+ self._sparse_refit_receiver.start_sync_server()
def init_collective(
self,
@@ -1090,6 +1121,9 @@ def wake_up(self, **kwargs):
def shutdown(self) -> bool:
"""Clean up vLLM resources."""
try:
+ if self._sparse_refit_receiver is not None:
+ self._sparse_refit_receiver.shutdown()
+
if self.llm is not None:
# Clean up extension resources (e.g., ZMQ sockets)
self.llm.collective_rpc("cleanup", args=tuple())
diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py
index eae64f3f079..2f50e4a150e 100644
--- a/nemo_rl/models/generation/vllm/vllm_worker_async.py
+++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py
@@ -429,6 +429,9 @@ async def post_init_async(self):
await self.llm.collective_rpc(
"load_mtp_weights_from_disk", args=(self.model_name,)
)
+ if self._sparse_refit_receiver is not None:
+ hostnames = await self.llm.collective_rpc("report_node_hostname", args=())
+ self._sparse_refit_receiver.set_worker_hostnames(hostnames)
async def get_reserved_url(self) -> Optional[str]:
"""Return the URL from the reserved socket, available before model loading."""
@@ -911,6 +914,8 @@ def _setup_vllm_server(self) -> "tuple[threading.Thread, str, uvicorn.Server]":
app = FastAPI()
app = self._setup_vllm_openai_api_server(app)
+ if self._sparse_refit_receiver is not None:
+ self._sparse_refit_receiver.setup_api_server(app)
########################################
# Server spinup
@@ -1535,6 +1540,9 @@ async def wake_up_async(self, **kwargs):
async def shutdown(self) -> bool:
"""Clean up vLLM resources."""
try:
+ if self._sparse_refit_receiver is not None:
+ await asyncio.to_thread(self._sparse_refit_receiver.shutdown)
+
if self.llm is not None:
# Clean up extension resources (e.g., ZMQ sockets)
await self.llm.collective_rpc("cleanup", args=tuple())
diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py
index 6eea8bc4f8e..dea0c079a14 100644
--- a/nemo_rl/models/policy/workers/megatron_policy_worker.py
+++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py
@@ -153,6 +153,7 @@ class MegatronPolicyWorkerImpl(
# begin/abort; None when no step is open. Declared at class level so
# ``self._train_step_state = None`` after finish/abort type-checks.
_train_step_state: Optional[dict[str, Any]] = None
+ _remote_sparse_refit: Any = None
def __repr__(self):
"""Customizes the actor's prefix in the Ray logs.
@@ -1795,6 +1796,60 @@ def _get_model_config(self):
return model.config
return None
+ @torch.no_grad()
+ @wrap_with_nvtx_name("megatron_policy_worker/init_remote_sparse_delta_baseline")
+ def init_remote_sparse_delta_baseline(
+ self,
+ *,
+ shard_rank: int,
+ shard_count: int,
+ transport: str,
+ ) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
+ return self._require_remote_sparse_refit().initialize_baseline(
+ shard_rank=shard_rank,
+ shard_count=shard_count,
+ transport=transport,
+ )
+
+ @torch.no_grad()
+ @wrap_with_nvtx_name("megatron_policy_worker/stream_remote_sparse_weights")
+ def stream_remote_sparse_weights(
+ self,
+ transport: str,
+ targets: list[str],
+ *,
+ transfer_id: str,
+ api_key_env_var: Optional[str],
+ timeout_s: float,
+ shard_rank: int,
+ shard_count: int,
+ overwrite_names: list[str],
+ ) -> dict[str, int]:
+ return self._require_remote_sparse_refit().stream(
+ transport,
+ targets,
+ transfer_id=transfer_id,
+ api_key_env_var=api_key_env_var,
+ timeout_s=timeout_s,
+ shard_rank=shard_rank,
+ shard_count=shard_count,
+ overwrite_names=overwrite_names,
+ )
+
+ def _require_remote_sparse_refit(self) -> Any:
+ if self._remote_sparse_refit is None:
+ from nemo_rl.models.policy.workers.megatron_remote_sparse_refit import (
+ MegatronRemoteSparseRefit,
+ )
+
+ refit_config = self.cfg["generation"]["refit_cfg"]
+ assert refit_config is not None
+ self._remote_sparse_refit = MegatronRemoteSparseRefit(self, refit_config)
+ return self._remote_sparse_refit
+
+ def finish_remote_sparse_delta_sync(self, *, succeeded: bool) -> None:
+ self._require_remote_sparse_refit().finish(succeeded)
+
def _calculate_refit_param_info(self) -> list[tuple[str, int]]:
"""Calculate parameter information for refit.
@@ -1866,7 +1921,7 @@ def _iter_params_with_optional_kv_scales(
base_iter = self.megatron_bridge.export_hf_weights(
[self.model],
show_progress=False,
- conversion_tasks=self.refit_conversion_tasks, # used for metadata caching
+ conversion_tasks=self.refit_conversion_tasks,
)
# Yield the original parameters first.
diff --git a/nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py b/nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py
new file mode 100644
index 00000000000..dced0511ee0
--- /dev/null
+++ b/nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py
@@ -0,0 +1,92 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Canonical Hugging Face sparse-refit state for a Megatron policy worker."""
+
+from typing import Any
+
+import torch
+
+from nemo_rl.models.generation.vllm.config import (
+ VllmRefitConfig,
+ VllmRefitTransportName,
+)
+from nemo_rl.utils.weight_transfer_sparse_codec import DeltaCompressionTracker
+from nemo_rl.utils.weight_transfer_stream import (
+ init_sparse_delta_baseline_from_iterator,
+ stream_sparse_delta_payloads_via_s3_manifest,
+)
+from nemo_rl.utils.weight_transfer_zmq import stream_sparse_delta_payloads_via_zmq
+
+
+class MegatronRemoteSparseRefit:
+ def __init__(self, worker: Any, refit_config: VllmRefitConfig) -> None:
+ self._worker = worker
+ self._tracker = DeltaCompressionTracker(refit_config)
+
+ def initialize_baseline(
+ self,
+ *,
+ shard_rank: int,
+ shard_count: int,
+ transport: VllmRefitTransportName,
+ ) -> dict[str, tuple[tuple[int, ...], torch.dtype]]:
+ init_sparse_delta_baseline_from_iterator(
+ self._worker._iter_params_with_optional_kv_scales(),
+ delta_tracker=self._tracker,
+ shard_rank=shard_rank,
+ shard_count=shard_count,
+ transport=transport,
+ )
+ return {
+ name: (tuple(tensor.shape), tensor.dtype)
+ for name, tensor in self._tracker.baseline.items()
+ }
+
+ def stream(
+ self,
+ transport: VllmRefitTransportName,
+ targets: list[str],
+ *,
+ transfer_id: str,
+ api_key_env_var: str | None,
+ timeout_s: float,
+ shard_rank: int,
+ shard_count: int,
+ overwrite_names: list[str],
+ ) -> dict[str, int]:
+ streamer = {
+ "s3": stream_sparse_delta_payloads_via_s3_manifest,
+ "zmq": stream_sparse_delta_payloads_via_zmq,
+ }[transport]
+ self._tracker.overwrite_names = frozenset(overwrite_names)
+ result = streamer(
+ self._worker._iter_params_with_optional_kv_scales(),
+ delta_tracker=self._tracker,
+ transfer_id=transfer_id,
+ refit_targets=targets,
+ api_key_env_var=api_key_env_var,
+ timeout_s=timeout_s,
+ shard_rank=shard_rank,
+ shard_count=shard_count,
+ )
+ if torch.cuda.is_available():
+ torch.cuda.synchronize()
+ return result
+
+ def finish(self, succeeded: bool) -> None:
+ if succeeded:
+ self._tracker.on_sync_succeeded()
+ else:
+ self._tracker.on_sync_failed()
diff --git a/nemo_rl/utils/weight_transfer_http.py b/nemo_rl/utils/weight_transfer_http.py
new file mode 100644
index 00000000000..685bc9349d9
--- /dev/null
+++ b/nemo_rl/utils/weight_transfer_http.py
@@ -0,0 +1,135 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""HTTP control-plane utilities shared by remote sparse-refit transports."""
+
+import os
+import threading
+from collections.abc import Iterable, Mapping, Sequence
+from concurrent.futures import ThreadPoolExecutor
+from functools import cache
+from typing import Any
+
+import requests
+from urllib3.util.retry import Retry
+
+G_VLLM_REFIT_S3_MANIFEST_PATH = "/nemo-rl/refit/s3-manifest"
+G_VLLM_REFIT_PREPARE_PATH = "/nemo-rl/refit/prepare"
+G_VLLM_REFIT_FLUSH_PATH = "/nemo-rl/refit/flush"
+G_VLLM_REFIT_ZMQ_FLUSH_PATH = "/nemo-rl/refit/zmq-flush"
+G_VLLM_REFIT_API_KEY_HEADER = "x-nemo-rl-refit-key"
+_HTTP_LOCAL = threading.local()
+_HTTP_ADAPTER = requests.adapters.HTTPAdapter(
+ pool_connections=64,
+ pool_maxsize=64,
+ max_retries=Retry(
+ total=3,
+ backoff_factor=0.25,
+ status_forcelist=(502, 503, 504),
+ allowed_methods={"POST"},
+ ),
+)
+
+
+def vllm_refit_endpoints(base_urls: Sequence[str], path: str) -> list[str]:
+ return list(
+ dict.fromkeys(
+ f"{url.strip().rstrip('/')}{path}" for url in base_urls if url.strip()
+ )
+ )
+
+
+def vllm_refit_api_key(api_key_env_var: str | None) -> str | None:
+ if not api_key_env_var:
+ return None
+ token = os.environ.get(api_key_env_var)
+ if not token:
+ raise RuntimeError(
+ "vLLM sparse refit API key env var "
+ f"{api_key_env_var!r} is configured but unset or empty."
+ )
+ return token
+
+
+def refit_http_session() -> requests.Session:
+ session = getattr(_HTTP_LOCAL, "session", None)
+ if session is None:
+ session = requests.Session()
+ session.mount("http://", _HTTP_ADAPTER)
+ session.mount("https://", _HTTP_ADAPTER)
+ _HTTP_LOCAL.session = session
+ return session
+
+
+@cache
+def _http_executor(workers: int) -> ThreadPoolExecutor:
+ return ThreadPoolExecutor(max_workers=workers, thread_name_prefix="nrl-refit-http")
+
+
+def post_vllm_refit_endpoints(
+ endpoint_urls: Sequence[str],
+ body: Mapping[str, Any],
+ *,
+ api_key: str | None,
+ timeout_s: float,
+) -> list[dict[str, Any]]:
+ request_headers = {}
+ if api_key:
+ request_headers[G_VLLM_REFIT_API_KEY_HEADER] = api_key
+
+ def post(url: str) -> dict[str, Any]:
+ response = refit_http_session().post(
+ url,
+ json=body,
+ headers=request_headers,
+ timeout=timeout_s,
+ )
+ try:
+ result: dict[str, Any] = response.json() if response.content else {}
+ except requests.exceptions.JSONDecodeError:
+ result = {}
+ if response.status_code >= 400 or result.get("ok") is not True:
+ raise RuntimeError(
+ f"vLLM refit failed for {url}: HTTP {response.status_code}: "
+ f"{response.text[:512]}"
+ )
+ return result
+
+ return list(_http_executor(len(endpoint_urls)).map(post, endpoint_urls))
+
+
+def merge_vllm_refit_metrics(
+ result: dict[str, Any],
+ metrics: Iterable[Mapping[str, Any]],
+ *,
+ maximum: bool,
+ candidate_maximum: bool | None = None,
+) -> dict[str, Any]:
+ for metric in metrics:
+ for key, value in metric.items():
+ if key.startswith("receiver_") and key.endswith("_s"):
+ number, use_maximum = float(value), maximum
+ elif candidate_maximum is not None and key.startswith("verification_"):
+ number = value
+ use_maximum = key == "verification_max_abs" or (
+ key == "verification_candidates" and candidate_maximum
+ )
+ else:
+ continue
+ if key in result:
+ number = (
+ max(result[key], number) if use_maximum else result[key] + number
+ )
+ result[key] = number
+ return result
diff --git a/nemo_rl/utils/weight_transfer_sparse_codec.py b/nemo_rl/utils/weight_transfer_sparse_codec.py
new file mode 100644
index 00000000000..8e8eae54735
--- /dev/null
+++ b/nemo_rl/utils/weight_transfer_sparse_codec.py
@@ -0,0 +1,359 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import tempfile
+import threading
+from collections.abc import Iterable
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Literal
+
+import numpy as np
+import torch
+
+from nemo_rl.models.generation.vllm.config import VllmRefitConfig
+
+NamedTensor = tuple[str, torch.Tensor]
+TensorBatch = list[NamedTensor]
+SparseOperation = Literal["xor", "overwrite"]
+SparseInfo = tuple[str, torch.Tensor, torch.Tensor, torch.Tensor, SparseOperation]
+TensorPayload = tuple[torch.Tensor, tuple[torch.Tensor, ...], list[dict[str, Any]]]
+PreparedTensorPayload = tuple[TensorPayload, int, int]
+SparseItem = tuple[dict[str, Any], torch.Tensor, torch.Tensor]
+
+_INTEGER_DTYPE_BY_SIZE = {
+ 1: torch.uint8,
+ 2: torch.int16,
+ 4: torch.int32,
+ 8: torch.int64,
+}
+
+
+class _TensorPayloadBuilder:
+ def __init__(self) -> None:
+ self.locations: list[torch.Tensor] = []
+ self.value_parts: list[list[torch.Tensor]] = []
+ self.value_group_by_dtype: dict[torch.dtype, int] = {}
+ self.value_offsets: list[int] = []
+ self.metadata: list[dict[str, Any]] = []
+ self.index_offset = 0
+
+ def add_locations(self, locations: torch.Tensor) -> tuple[int, int]:
+ start = self.index_offset
+ if locations.numel():
+ self.locations.append(locations)
+ self.index_offset += locations.numel()
+ return start, self.index_offset
+
+ def add_values(self, values: torch.Tensor) -> tuple[int, int, int]:
+ group = self.value_group_by_dtype.get(values.dtype)
+ if group is None:
+ group = len(self.value_parts)
+ self.value_group_by_dtype[values.dtype] = group
+ self.value_parts.append([])
+ self.value_offsets.append(0)
+ start = self.value_offsets[group]
+ self.value_parts[group].append(values)
+ self.value_offsets[group] += values.numel()
+ return group, start, self.value_offsets[group]
+
+ def finish(self) -> TensorPayload:
+ indices = (
+ torch.cat(self.locations)
+ if self.locations
+ else torch.empty(0, dtype=torch.uint8)
+ )
+ values = tuple(
+ torch.cat(parts) if len(parts) > 1 else parts[0]
+ for parts in self.value_parts
+ )
+ return indices, values, self.metadata
+
+
+def integer_dtype_for_element_size(element_size: int) -> torch.dtype:
+ try:
+ return _INTEGER_DTYPE_BY_SIZE[element_size]
+ except KeyError as error:
+ raise ValueError(f"Unsupported tensor element size {element_size}.") from error
+
+
+def dtype_from_name(name: str) -> torch.dtype:
+ dtype = getattr(torch, name, None)
+ if not isinstance(dtype, torch.dtype):
+ raise ValueError(f"Unsupported sparse-refit tensor dtype {name!r}.")
+ integer_dtype_for_element_size(dtype.itemsize)
+ return dtype
+
+
+def sparse_operation(value: object) -> SparseOperation:
+ if value == "xor" or value == "overwrite":
+ return value
+ raise ValueError(f"Unsupported sparse-refit operation {value!r}.")
+
+
+def integer_view(tensor: torch.Tensor) -> torch.Tensor:
+ return tensor.view(integer_dtype_for_element_size(tensor.element_size()))
+
+
+def encode_sparse_infos(
+ infos: Iterable[SparseInfo],
+) -> TensorPayload:
+ payload = _TensorPayloadBuilder()
+ for name, tensor, raw_locations, raw_values, operation in infos:
+ count = int(raw_values.numel())
+ if count == 1 or int(raw_locations[-1] - raw_locations[0] + 1) == count:
+ index_start = index_end = payload.index_offset
+ location_metadata = {
+ "index_encoding": "range",
+ "range_start": int(raw_locations[0]),
+ }
+ else:
+ location_tensor = _encode_explicit_locations(raw_locations)
+ location_metadata = {"index_encoding": "deltas"}
+ index_start, index_end = payload.add_locations(location_tensor)
+ value_group, value_start, value_end = payload.add_values(raw_values)
+ payload.metadata.append(
+ {
+ "name": name,
+ "shape": tuple(int(dim) for dim in tensor.shape),
+ "dtype": str(tensor.dtype).removeprefix("torch."),
+ "operation": operation,
+ "index_start": index_start,
+ "index_end": index_end,
+ "value_group": value_group,
+ "value_start": value_start,
+ "value_end": value_end,
+ **location_metadata,
+ }
+ )
+ return payload.finish()
+
+
+def merge_sparse_payloads(payloads: Iterable[TensorPayload]) -> TensorPayload:
+ """Combine encoded chunks without materializing dense source tensors."""
+ payload = _TensorPayloadBuilder()
+ for locations, value_groups, items in payloads:
+ group_remap = {}
+ group_starts = {}
+ for old_group, values in enumerate(value_groups):
+ new_group, start, _ = payload.add_values(values)
+ group_remap[old_group] = new_group
+ group_starts[old_group] = start
+ index_offset, _ = payload.add_locations(locations)
+ for item in items:
+ merged = dict(item)
+ merged["index_start"] = int(item["index_start"]) + index_offset
+ merged["index_end"] = int(item["index_end"]) + index_offset
+ old_group = int(item["value_group"])
+ merged["value_group"] = group_remap[old_group]
+ merged["value_start"] = int(item["value_start"]) + group_starts[old_group]
+ merged["value_end"] = int(item["value_end"]) + group_starts[old_group]
+ payload.metadata.append(merged)
+ return payload.finish()
+
+
+def sparse_locations_for_item(
+ item: dict[str, Any],
+ packed_locations: torch.Tensor,
+ *,
+ device: torch.device | int | str,
+ dtype: torch.dtype = torch.int64,
+) -> torch.Tensor:
+ if dtype not in (torch.int32, torch.int64):
+ raise ValueError(f"Unsupported sparse location dtype {dtype}.")
+ count = int(item["value_end"]) - int(item["value_start"])
+ if item["index_encoding"] == "range":
+ start = int(item["range_start"])
+ return torch.arange(start, start + count, dtype=dtype, device=device)
+
+ index_start, index_end = int(item["index_start"]), int(item["index_end"])
+ raw = packed_locations[index_start:index_end].detach().cpu().numpy()
+ delta_dtype = {2: np.uint16, 4: np.uint32, 8: np.uint64}[raw.size // count]
+ location_dtype = np.int32 if dtype == torch.int32 else np.int64
+ locations = raw.view(delta_dtype).astype(location_dtype, copy=False)
+ locations += 1
+ np.cumsum(locations, out=locations)
+ locations -= 1
+ return torch.from_numpy(locations).to(device=device)
+
+
+def _encode_explicit_locations(
+ locations: torch.Tensor,
+) -> torch.Tensor:
+ indices = locations.detach().cpu().numpy().astype(np.int64, copy=False)
+ deltas = np.diff(indices, prepend=-1) - 1
+ max_delta = int(deltas.max())
+ dtype = (
+ np.uint16
+ if max_delta <= 0xFFFF
+ else np.uint32
+ if max_delta <= 0xFFFFFFFF
+ else np.uint64
+ )
+ raw = deltas.astype(dtype, copy=False).tobytes()
+ return torch.from_numpy(np.frombuffer(raw, dtype=np.uint8).copy())
+
+
+class DeltaCompressionTracker:
+ """Source-side CPU or mmap baseline for sparse-delta refit."""
+
+ def __init__(
+ self,
+ config: VllmRefitConfig,
+ ) -> None:
+ self.refit_config = config
+ delta_config = config.delta_compression
+ self.sparse_bucket_size_bytes = delta_config.sparse_bucket_size_bytes
+ self.encoding = sparse_operation(delta_config.encoding)
+ self.overwrite_names: frozenset[str] = frozenset()
+ self.verification_samples = config.verify_samples_per_payload
+ self.baseline_in_memory = config.baseline.in_memory
+ self.baseline_mmap_dir = config.baseline.mmap_dir
+ self.baseline: dict[str, torch.Tensor] = {}
+ self._pending_updates: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}
+ self._pending_updates_lock = threading.Lock()
+ self._baseline_commits: tuple[Any, ...] = ()
+ self._baseline_commit_executor = ThreadPoolExecutor(
+ max_workers=4, thread_name_prefix="nrl-refit-baseline"
+ )
+
+ def prepare_sparse_delta_payload(
+ self, tensors: TensorBatch
+ ) -> PreparedTensorPayload:
+ self._wait_for_baseline_commits()
+ sparse_infos: list[SparseInfo] = []
+ pending_updates = {}
+ changed_elements = total_elements = 0
+ for name, tensor in tensors:
+ baseline = self.baseline.get(name)
+ if baseline is None:
+ raise RuntimeError(f"Sparse delta baseline is missing {name!r}.")
+ current = tensor.detach().cpu().contiguous()
+ baseline_bits = integer_view(baseline).view(-1)
+ current_bits = integer_view(current).view(-1)
+ locations = current_bits.ne(baseline_bits).nonzero().view(-1)
+ total_elements += current.numel()
+ changed_elements += locations.numel()
+ if locations.numel():
+ current_values = current_bits[locations]
+ pending_updates[name] = (locations, current_values)
+ operation: SparseOperation = (
+ "overwrite" if name in self.overwrite_names else self.encoding
+ )
+ values = (
+ current_values.bitwise_xor(baseline_bits[locations])
+ if operation == "xor"
+ else current_values
+ )
+ sparse_infos.append(
+ (
+ name,
+ current,
+ locations,
+ values,
+ operation,
+ )
+ )
+ with self._pending_updates_lock:
+ self._pending_updates.update(pending_updates)
+ payload = encode_sparse_infos(sparse_infos)
+ if self.verification_samples:
+ self._add_verification_samples(payload[2])
+ return payload, changed_elements, total_elements
+
+ def _add_verification_samples(
+ self,
+ metadata: list[dict[str, Any]],
+ ) -> None:
+ sizes = [int(item["value_end"]) - int(item["value_start"]) for item in metadata]
+ total = sum(sizes)
+ count = min(self.verification_samples, total)
+ sample_ranks = [
+ ((2 * index + 1) * total) // (2 * count) for index in range(count)
+ ]
+ sample_index = offset = 0
+ for item, size in zip(metadata, sizes, strict=True):
+ end = offset + size
+ samples = 0
+ while sample_index < count and sample_ranks[sample_index] < end:
+ samples += 1
+ sample_index += 1
+ if samples:
+ item["verification_samples"] = samples
+ offset = end
+
+ def on_sync_succeeded(self) -> None:
+ with self._pending_updates_lock:
+ pending_updates, self._pending_updates = self._pending_updates, {}
+ items = list(pending_updates.items())
+ workers = min(4, len(items))
+ self._baseline_commits = tuple(
+ self._baseline_commit_executor.submit(
+ self._commit_baseline_updates, items[worker::workers]
+ )
+ for worker in range(workers)
+ )
+
+ def on_sync_failed(self) -> None:
+ with self._pending_updates_lock:
+ self._pending_updates.clear()
+
+ def snapshot_baseline(self, tensors: Iterable[NamedTensor]) -> None:
+ self._wait_for_baseline_commits()
+ for name, tensor in tensors:
+ baseline = self._baseline(name, tuple(tensor.shape), tensor.dtype)
+ baseline.view(torch.uint8).view(-1).copy_(
+ tensor.detach().cpu().contiguous().view(torch.uint8).view(-1)
+ )
+
+ def _wait_for_baseline_commits(self) -> None:
+ for commit in self._baseline_commits:
+ commit.result()
+ self._baseline_commits = ()
+
+ def _commit_baseline_updates(
+ self,
+ updates: Iterable[tuple[str, tuple[torch.Tensor, torch.Tensor]]],
+ ) -> None:
+ for name, (locations, values) in updates:
+ target = integer_view(self.baseline[name]).view(-1)
+ count = locations.numel()
+ first = int(locations[0])
+ if int(locations[-1]) - first + 1 == count:
+ target[first : first + count].copy_(values)
+ continue
+ target.index_copy_(0, locations, values)
+
+ def _baseline(
+ self,
+ name: str,
+ shape: tuple[int, ...],
+ dtype: torch.dtype,
+ ) -> torch.Tensor:
+ if name in self.baseline:
+ return self.baseline[name]
+ numel = torch.Size(shape).numel()
+ nbytes = numel * dtype.itemsize
+ if self.baseline_in_memory:
+ storage = torch.empty(nbytes, dtype=torch.uint8)
+ else:
+ with tempfile.NamedTemporaryFile(
+ prefix="nrl-refit-baseline-", dir=self.baseline_mmap_dir
+ ) as handle:
+ handle.truncate(nbytes)
+ storage = torch.from_file(
+ handle.name, shared=True, size=nbytes, dtype=torch.uint8
+ )
+ baseline = storage.view(dtype).view(shape)
+ self.baseline[name] = baseline
+ return baseline
diff --git a/nemo_rl/utils/weight_transfer_stream.py b/nemo_rl/utils/weight_transfer_stream.py
new file mode 100644
index 00000000000..b39d0804df9
--- /dev/null
+++ b/nemo_rl/utils/weight_transfer_stream.py
@@ -0,0 +1,615 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared sparse payload pipeline and S3 transport for remote vLLM refit."""
+
+import hashlib
+import io
+import threading
+import time
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
+from contextlib import suppress
+from dataclasses import dataclass
+from functools import cache
+from typing import Any, Callable
+from urllib.parse import quote
+
+import torch
+import zstandard
+
+from nemo_rl.models.generation.vllm.config import VllmRefitTransportName
+from nemo_rl.utils.packed_tensor import get_target_packed_tensor_size
+from nemo_rl.utils.weight_transfer_http import (
+ G_VLLM_REFIT_S3_MANIFEST_PATH,
+ merge_vllm_refit_metrics,
+ post_vllm_refit_endpoints,
+ vllm_refit_api_key,
+ vllm_refit_endpoints,
+)
+from nemo_rl.utils.weight_transfer_sparse_codec import (
+ DeltaCompressionTracker,
+ NamedTensor,
+ TensorBatch,
+ TensorPayload,
+ merge_sparse_payloads,
+)
+
+_STREAM_LOCAL = threading.local()
+# A 64 MiB CRT part and 2 GiB client cap allow 32 in-flight parts, matching the
+# upload concurrency selected by the balanced 120B S3 sweeps while bounding RAM.
+_S3_PART_SIZE = 64 * 1024**2
+_S3_MEMORY_LIMIT = 2 * 1024**3
+
+
+@dataclass(frozen=True)
+class SparseRefitTransport:
+ """Transport-specific callbacks used by the shared streaming pipeline."""
+
+ name: VllmRefitTransportName
+ transfer_workers: int
+ send: Callable[[bytes, int, int], dict[str, Any]]
+ cleanup: Callable[[], None]
+
+
+@dataclass
+class _SparsePayloadBucket:
+ payloads: list[TensorPayload]
+ dense_bytes: int = 0
+ encode_s: float = 0.0
+ next_index: int = 0
+
+
+@cache
+def _s3_client(region: str) -> Any:
+ # Keep the AWS runtime unloaded for ZeroMQ-only jobs.
+ from awscrt.auth import AwsCredentialsProvider
+ from awscrt.io import ClientBootstrap, DefaultHostResolver, EventLoopGroup
+ from awscrt.s3 import S3Client, create_default_s3_signing_config
+
+ event_loop_group = EventLoopGroup()
+ bootstrap = ClientBootstrap(
+ event_loop_group,
+ DefaultHostResolver(event_loop_group),
+ )
+ return S3Client(
+ bootstrap=bootstrap,
+ region=region,
+ signing_config=create_default_s3_signing_config(
+ region=region,
+ credential_provider=AwsCredentialsProvider.new_default_chain(bootstrap),
+ ),
+ part_size=_S3_PART_SIZE,
+ multipart_upload_threshold=_S3_PART_SIZE,
+ throughput_target_gbps=10.0,
+ memory_limit=_S3_MEMORY_LIMIT,
+ )
+
+
+class _S3ObjectStore:
+ def __init__(self, *, bucket: str, region: str) -> None:
+ from awscrt.s3 import S3RequestType
+
+ self.bucket = bucket
+ self.region = region
+ self._client = _s3_client(region)
+ self._request_type = S3RequestType
+
+ def put(self, key: str, body: bytes) -> None:
+ self._client.make_request(
+ type=self._request_type.PUT_OBJECT,
+ request=self._request("PUT", key, body),
+ ).finished_future.result()
+
+ def get(self, key: str) -> bytearray:
+ from awscrt.http import HttpHeaders
+
+ body = bytearray()
+
+ def on_headers(
+ status_code: int,
+ headers: list[tuple[str, str]],
+ **_kwargs: Any,
+ ) -> None:
+ nonlocal body
+ if status_code != 200:
+ raise RuntimeError(f"S3 GET returned HTTP {status_code}.")
+ length = HttpHeaders(headers).get("content-length")
+ if length is None:
+ raise RuntimeError("S3 GET response omitted content-length.")
+ body = bytearray(int(length))
+
+ def on_body(chunk: bytes, offset: int, **_kwargs: Any) -> None:
+ body[offset : offset + len(chunk)] = chunk
+
+ self._client.make_request(
+ type=self._request_type.GET_OBJECT,
+ request=self._request("GET", key),
+ on_headers=on_headers,
+ on_body=on_body,
+ ).finished_future.result()
+ return body
+
+ def delete(self, key: str) -> None:
+ self._client.make_request(
+ type=self._request_type.DEFAULT,
+ request=self._request("DELETE", key),
+ operation_name="DeleteObject",
+ ).finished_future.result()
+
+ def _request(self, method: str, key: str, body: bytes | None = None) -> Any:
+ from awscrt.http import HttpHeaders, HttpRequest
+
+ headers = HttpHeaders(
+ [("host", f"{self.bucket}.s3.{self.region}.amazonaws.com")]
+ )
+ if body is not None:
+ headers.add("content-length", str(len(body)))
+ headers.add("content-type", "application/octet-stream")
+ elif method == "DELETE":
+ headers.add("content-length", "0")
+ return HttpRequest(
+ method,
+ f"/{quote(key, safe='/~')}",
+ headers,
+ io.BytesIO(body) if body is not None else None,
+ )
+
+
+def sparse_payload_checksum(body: bytes | bytearray) -> str:
+ return hashlib.blake2b(body, digest_size=16).hexdigest()
+
+
+def decode_sparse_payload(body: bytes | bytearray, checksum: str) -> bytes:
+ actual = sparse_payload_checksum(body)
+ if actual != checksum:
+ raise ValueError(
+ f"Sparse refit payload checksum mismatch: expected={checksum}, actual={actual}."
+ )
+ decompressor = getattr(_STREAM_LOCAL, "zstd_decompressor", None)
+ if decompressor is None:
+ decompressor = zstandard.ZstdDecompressor()
+ _STREAM_LOCAL.zstd_decompressor = decompressor
+ return decompressor.decompress(body)
+
+
+def iter_sparse_weight_chunks(
+ tensors: Iterable[NamedTensor], target_bytes: int
+) -> Iterator[tuple[TensorBatch, float]]:
+ iterator = iter(tensors)
+ pending = None
+ while True:
+ started = time.perf_counter()
+ chunk = [pending] if pending is not None else []
+ size = pending[1].numel() * pending[1].element_size() if pending else 0
+ pending = None
+ for item in iterator:
+ item_size = item[1].numel() * item[1].element_size()
+ if chunk and size + item_size > target_bytes:
+ pending = item
+ break
+ chunk.append(item)
+ size += item_size
+ if size >= target_bytes:
+ break
+ export_pull_s = time.perf_counter() - started
+ if not chunk:
+ return
+ yield chunk, export_pull_s
+
+
+@cache
+def _get_manifest_s3_store(bucket: str, region: str) -> _S3ObjectStore:
+ return _S3ObjectStore(bucket=bucket, region=region)
+
+
+def sparse_export_chunk_size(
+ delta_tracker: DeltaCompressionTracker,
+ transport: VllmRefitTransportName,
+) -> int:
+ requested = delta_tracker.refit_config.delta_compression.export_chunk_bytes[
+ transport
+ ]
+ if torch.cuda.is_available():
+ requested = min(requested, get_target_packed_tensor_size())
+ return min(requested, delta_tracker.sparse_bucket_size_bytes)
+
+
+@cache
+def _executor(key: str, workers: int) -> ThreadPoolExecutor:
+ return ThreadPoolExecutor(max_workers=workers, thread_name_prefix=f"nrl-{key}")
+
+
+def init_sparse_delta_baseline_from_iterator(
+ iterator: Iterable[NamedTensor],
+ *,
+ delta_tracker: DeltaCompressionTracker,
+ shard_rank: int,
+ shard_count: int,
+ transport: VllmRefitTransportName,
+) -> None:
+ start_s = time.perf_counter()
+ export_chunk_size = sparse_export_chunk_size(delta_tracker, transport)
+
+ chunk_count = 0
+ export_pull_s = snapshot_s = 0.0
+ for chunk_index, (chunk, pull_s) in enumerate(
+ iter_sparse_weight_chunks(iterator, export_chunk_size)
+ ):
+ chunk_count = chunk_index + 1
+ export_pull_s += pull_s
+ if chunk_index % shard_count != shard_rank:
+ continue
+ started = time.perf_counter()
+ delta_tracker.snapshot_baseline(chunk)
+ snapshot_s += time.perf_counter() - started
+ print(
+ "REFIT_BASELINE_INIT "
+ f"event=end chunks={chunk_count} export_pull_s={export_pull_s:.3f} "
+ f"snapshot_s={snapshot_s:.3f} "
+ f"seconds={time.perf_counter() - start_s:.3f}",
+ flush=True,
+ )
+
+
+def stream_sparse_delta_payloads(
+ iterator: Iterable[NamedTensor],
+ *,
+ delta_tracker: DeltaCompressionTracker,
+ transport: SparseRefitTransport,
+ shard_rank: int,
+ shard_count: int,
+) -> dict[str, int]:
+ prefix = transport.name.upper()
+ refit_config = delta_tracker.refit_config
+ encode_workers = refit_config.tuning.encode_workers[transport.name]
+ encode_executor = _executor(f"refit-{transport.name}-encode", encode_workers)
+ serialize_workers = min(4, encode_workers)
+ serialize_executor = _executor(
+ f"refit-{transport.name}-serialize", serialize_workers
+ )
+ transfer_executor = ThreadPoolExecutor(
+ max_workers=transport.transfer_workers,
+ thread_name_prefix=f"nrl-refit-{transport.name}-transfer",
+ )
+ export_chunk_size = sparse_export_chunk_size(delta_tracker, transport.name)
+
+ def encode_chunk(
+ chunk: TensorBatch,
+ ) -> tuple[TensorPayload | None, float, int, int, int]:
+ dense_bytes = sum(tensor.numel() * tensor.element_size() for _, tensor in chunk)
+ started = time.perf_counter()
+ payload, changed_elements, total_elements = (
+ delta_tracker.prepare_sparse_delta_payload(chunk)
+ )
+ encode_s = time.perf_counter() - started
+ return (
+ payload if payload[2] else None,
+ encode_s,
+ changed_elements,
+ total_elements,
+ dense_bytes,
+ )
+
+ def serialize_payloads(
+ payloads: tuple[TensorPayload, ...], encode_s: float
+ ) -> tuple[bytes, int, dict[str, float]]:
+ started = time.perf_counter()
+ buffer = io.BytesIO()
+ merged = merge_sparse_payloads(payloads)
+ torch.save(merged, buffer)
+ raw_body = buffer.getvalue()
+ serialize_s = time.perf_counter() - started
+ started = time.perf_counter()
+ body = zstd_compress(
+ raw_body,
+ refit_config.delta_compression.zstd_threads[transport.name],
+ )
+ compress_s = time.perf_counter() - started
+ return (
+ body,
+ sum(int(item.get("verification_samples", 0)) for item in merged[2]),
+ {
+ "encode_s": encode_s,
+ "serialize_s": serialize_s,
+ "compress_s": compress_s,
+ },
+ )
+
+ def transfer_payload(
+ encoded: tuple[bytes, int, dict[str, float]], payload_index: int
+ ) -> dict[str, Any]:
+ body, verification_candidates, encode_timing = encoded
+ result = transport.send(body, payload_index, verification_candidates)
+ result.update(
+ body_size=len(body),
+ **encode_timing,
+ )
+ return result
+
+ timing: dict[str, float] = {}
+ receiver_timing: dict[str, float] = {}
+ counts = {
+ "payloads": 0,
+ "wire_bytes": 0,
+ "changed_elements": 0,
+ "total_elements": 0,
+ }
+ chunk_count = 0
+ export_pull_s = 0.0
+ encode_inflight: dict[Any, None] = {}
+ serialize_inflight: dict[Any, int] = {}
+ transfer_inflight: dict[Any, None] = {}
+ transfer_submitted = False
+ max_encode_inflight = encode_workers * 2
+ max_serialize_inflight = serialize_workers * 2
+ max_transfer_inflight = transport.transfer_workers * 2
+ bucket = _SparsePayloadBucket([])
+
+ def resolved(inflight: dict[Any, Any], *, block: bool) -> Iterator[tuple[Any, Any]]:
+ if not inflight:
+ return
+ completed = (
+ wait(inflight, return_when=FIRST_COMPLETED)[0]
+ if block
+ else tuple(future for future in inflight if future.done())
+ )
+ for future in completed:
+ metadata = inflight.pop(future)
+ yield metadata, future.result()
+
+ def collect_transfers(*, block: bool) -> None:
+ for _, result in resolved(transfer_inflight, block=block):
+ counts["payloads"] += 1
+ counts["wire_bytes"] += int(result["body_size"])
+ for key, value in result.items():
+ if key.endswith("_s"):
+ timing[key] = timing.get(key, 0.0) + float(value)
+ merge_vllm_refit_metrics(
+ receiver_timing, [result["receiver"]], maximum=False
+ )
+
+ def collect_serialized(*, block: bool) -> None:
+ nonlocal transfer_submitted
+ for index, encoded in resolved(serialize_inflight, block=block):
+ while len(transfer_inflight) >= max_transfer_inflight:
+ collect_transfers(block=True)
+ transfer_submitted = True
+ transfer_inflight[
+ transfer_executor.submit(transfer_payload, encoded, index)
+ ] = None
+ collect_transfers(block=False)
+
+ def submit_bucket() -> None:
+ if not bucket.payloads:
+ return
+ while len(serialize_inflight) >= max_serialize_inflight:
+ collect_serialized(block=True)
+ serialize_inflight[
+ serialize_executor.submit(
+ serialize_payloads, tuple(bucket.payloads), bucket.encode_s
+ )
+ ] = bucket.next_index
+ bucket.next_index += 1
+ bucket.payloads.clear()
+ bucket.dense_bytes = 0
+ bucket.encode_s = 0.0
+ collect_serialized(block=False)
+
+ def consume_encoded(encoded: Any) -> None:
+ payload, encode_s, changed_elements, total_elements, dense_bytes = encoded
+ counts["changed_elements"] += changed_elements
+ counts["total_elements"] += total_elements
+ if payload is None:
+ return
+ if (
+ bucket.payloads
+ and bucket.dense_bytes + dense_bytes
+ > delta_tracker.sparse_bucket_size_bytes
+ ):
+ submit_bucket()
+ bucket.payloads.append(payload)
+ bucket.dense_bytes += dense_bytes
+ bucket.encode_s += encode_s
+ if bucket.dense_bytes >= delta_tracker.sparse_bucket_size_bytes:
+ submit_bucket()
+
+ def drain_encodes() -> None:
+ for _, encoded in resolved(encode_inflight, block=True):
+ consume_encoded(encoded)
+
+ stream_start = time.perf_counter()
+ try:
+ for chunk_index, (chunk, pull_s) in enumerate(
+ iter_sparse_weight_chunks(iterator, export_chunk_size)
+ ):
+ chunk_count = chunk_index + 1
+ export_pull_s += pull_s
+ if chunk_index % shard_count != shard_rank:
+ continue
+ while len(encode_inflight) >= max_encode_inflight:
+ drain_encodes()
+ encode_inflight[encode_executor.submit(encode_chunk, chunk)] = None
+
+ while encode_inflight:
+ drain_encodes()
+ submit_bucket()
+ while serialize_inflight:
+ collect_serialized(block=True)
+ while transfer_inflight:
+ collect_transfers(block=True)
+ except Exception:
+ for futures in (encode_inflight, serialize_inflight, transfer_inflight):
+ for future in futures:
+ future.cancel()
+ if futures:
+ wait(futures)
+ raise
+ finally:
+ try:
+ if transfer_submitted:
+ barrier = threading.Barrier(transport.transfer_workers)
+
+ def cleanup_transport(_index: int) -> None:
+ barrier.wait()
+ transport.cleanup()
+
+ list(
+ transfer_executor.map(
+ cleanup_transport, range(transport.transfer_workers)
+ )
+ )
+ finally:
+ transfer_executor.shutdown(wait=True, cancel_futures=True)
+
+ report = {
+ "total_s": time.perf_counter() - stream_start,
+ "export_pull_s": export_pull_s,
+ **timing,
+ "payloads": counts["payloads"],
+ "chunks": chunk_count,
+ "wire_mb": counts["wire_bytes"] / 1e6,
+ "encode_workers": encode_workers,
+ "export_chunk_mb": export_chunk_size / 1e6,
+ "shard_rank": shard_rank,
+ "shard_count": shard_count,
+ "changed_elements": counts["changed_elements"],
+ "total_elements": counts["total_elements"],
+ "changed_pct": 100.0
+ * counts["changed_elements"]
+ / max(counts["total_elements"], 1),
+ }
+ report.update(receiver_timing)
+ print(
+ f"REFIT_{prefix}_TIMING "
+ + " ".join(f"{key}={value}" for key, value in report.items()),
+ flush=True,
+ )
+ return {
+ "payloads": counts["payloads"],
+ "changed_elements": counts["changed_elements"],
+ "total_elements": counts["total_elements"],
+ }
+
+
+def stream_sparse_delta_payloads_via_s3_manifest(
+ iterator: Iterable[NamedTensor],
+ *,
+ delta_tracker: DeltaCompressionTracker,
+ refit_targets: Sequence[str],
+ transfer_id: str,
+ api_key_env_var: str | None,
+ timeout_s: float,
+ shard_rank: int,
+ shard_count: int,
+) -> dict[str, int]:
+ endpoint_urls = vllm_refit_endpoints(refit_targets, G_VLLM_REFIT_S3_MANIFEST_PATH)
+ if not endpoint_urls:
+ raise ValueError("At least one vLLM S3 refit URL is required.")
+ refit_config = delta_tracker.refit_config
+ bucket = (refit_config.storage.s3_bucket or "").strip()
+ if not bucket:
+ raise RuntimeError(
+ "policy.generation.refit_cfg.storage.s3_bucket must be set for S3 refit."
+ )
+ region = refit_config.storage.s3_region.strip()
+ if not region:
+ raise ValueError("refit_cfg.storage.s3_region must not be empty.")
+ store = _get_manifest_s3_store(bucket, region)
+ object_prefix = refit_config.storage.s3_prefix.strip("/")
+ run_prefix = (
+ f"{object_prefix}/{transfer_id}/{shard_rank:06d}"
+ if object_prefix
+ else f"{transfer_id}/{shard_rank:06d}"
+ )
+ api_key = vllm_refit_api_key(api_key_env_var)
+ keys = threading.local()
+
+ def send(
+ body: bytes, payload_id: int, verification_candidates: int
+ ) -> dict[str, Any]:
+ key = f"{run_prefix}/{payload_id:06d}.pt"
+ pending = getattr(keys, "values", None)
+ if pending is None:
+ pending = keys.values = []
+ pending.append(key)
+ try:
+ started = time.perf_counter()
+ store.put(key, body)
+ s3_put_s = time.perf_counter() - started
+ started = time.perf_counter()
+ responses = post_vllm_refit_endpoints(
+ endpoint_urls,
+ {
+ "bucket": store.bucket,
+ "region": store.region,
+ "key": key,
+ "checksum": sparse_payload_checksum(body),
+ "verification_candidates": verification_candidates,
+ },
+ api_key=api_key,
+ timeout_s=timeout_s,
+ )
+ result = {
+ "s3_put_s": s3_put_s,
+ "manifest_post_s": time.perf_counter() - started,
+ "receiver": merge_vllm_refit_metrics({}, responses, maximum=True),
+ }
+ finally:
+ try:
+ store.delete(key)
+ except Exception:
+ pass
+ else:
+ pending.remove(key)
+ return result
+
+ def cleanup() -> None:
+ for key in getattr(keys, "values", ()):
+ with suppress(Exception):
+ store.delete(key)
+ keys.values = []
+
+ return stream_sparse_delta_payloads(
+ iterator,
+ delta_tracker=delta_tracker,
+ transport=SparseRefitTransport(
+ name="s3",
+ transfer_workers=refit_config.tuning.transfer_workers["s3"],
+ send=send,
+ cleanup=cleanup,
+ ),
+ shard_rank=shard_rank,
+ shard_count=shard_count,
+ )
+
+
+def download_s3_refit_payload(
+ manifest: Mapping[str, Any],
+) -> bytes:
+ body = _get_manifest_s3_store(str(manifest["bucket"]), str(manifest["region"])).get(
+ str(manifest["key"])
+ )
+ return decode_sparse_payload(body, str(manifest["checksum"]))
+
+
+def zstd_compress(raw: bytes, threads: int) -> bytes:
+ compressor = getattr(_STREAM_LOCAL, "zstd_compressor", None)
+ if compressor is None:
+ compressor = zstandard.ZstdCompressor(
+ level=1,
+ threads=threads,
+ )
+ _STREAM_LOCAL.zstd_compressor = compressor
+ return compressor.compress(raw)
diff --git a/nemo_rl/utils/weight_transfer_zmq.py b/nemo_rl/utils/weight_transfer_zmq.py
new file mode 100644
index 00000000000..c7f13577e28
--- /dev/null
+++ b/nemo_rl/utils/weight_transfer_zmq.py
@@ -0,0 +1,550 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transactional ZeroMQ value plane for remote sparse vLLM refit."""
+
+import hmac
+import json
+import threading
+import time
+import uuid
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import suppress
+from dataclasses import dataclass, field
+from typing import Any
+
+import zmq
+
+from nemo_rl.models.generation.vllm.config import VllmRefitTuningConfig
+from nemo_rl.utils.weight_transfer_http import (
+ merge_vllm_refit_metrics,
+ vllm_refit_api_key,
+)
+from nemo_rl.utils.weight_transfer_sparse_codec import (
+ DeltaCompressionTracker,
+ NamedTensor,
+)
+from nemo_rl.utils.weight_transfer_stream import (
+ SparseRefitTransport,
+ sparse_payload_checksum,
+ stream_sparse_delta_payloads,
+)
+
+_PROTOCOL = "nemo-rl-sparse-zmq-v1"
+_DATA = b"DATA"
+_ACK = b"ACK"
+_NACK = b"NACK"
+
+
+@dataclass
+class _RelayTransfer:
+ checksums: dict[tuple[int, int], str] = field(default_factory=dict)
+ futures: list[Future[dict[str, Any]]] = field(default_factory=list)
+
+
+def _json_bytes(value: Mapping[str, Any]) -> bytes:
+ return json.dumps(value, separators=(",", ":"), sort_keys=True).encode()
+
+
+def _configure_socket(socket: zmq.Socket, high_water_mark: int) -> None:
+ socket.setsockopt(zmq.LINGER, 0)
+ socket.setsockopt(zmq.SNDHWM, high_water_mark)
+ socket.setsockopt(zmq.RCVHWM, high_water_mark)
+ socket.setsockopt(zmq.TCP_KEEPALIVE, 1)
+
+
+class ZmqSparseRefitClient:
+ """One-thread DEALER client with retry-safe payload identifiers."""
+
+ def __init__(
+ self,
+ address: str,
+ *,
+ timeout_s: float,
+ producer_id: int,
+ retries: int,
+ api_key: str | None = None,
+ ) -> None:
+ self._address = address
+ self._timeout_ms = max(1, int(timeout_s * 1000))
+ self._producer_id = producer_id
+ self._retries = retries
+ self._api_key = api_key
+ self._socket = zmq.Context.instance().socket(zmq.DEALER)
+ _configure_socket(self._socket, 2)
+ self._socket.setsockopt(
+ zmq.IDENTITY, f"nrl-{producer_id}-{uuid.uuid4().hex}".encode()
+ )
+ self._socket.setsockopt(zmq.IMMEDIATE, 1)
+ self._socket.setsockopt(zmq.SNDTIMEO, self._timeout_ms)
+ self._socket.connect(address)
+
+ def send_payload(
+ self,
+ *,
+ transfer_id: str,
+ payload_id: int,
+ checksum: str,
+ verification_candidates: int,
+ body: bytes,
+ relay_root: str | None = None,
+ producer_id: int | None = None,
+ ) -> dict[str, Any]:
+ producer_id = self._producer_id if producer_id is None else producer_id
+ metadata = {
+ "protocol": _PROTOCOL,
+ "transfer_id": transfer_id,
+ "producer_id": producer_id,
+ "payload_id": payload_id,
+ "checksum": checksum,
+ "verification_candidates": verification_candidates,
+ }
+ if relay_root is not None:
+ metadata["relay_root"] = relay_root
+ if self._api_key is not None:
+ metadata["api_key"] = self._api_key
+ metadata_frame = _json_bytes(metadata)
+ for attempt in range(self._retries + 1):
+ try:
+ self._socket.send_multipart(
+ [_DATA, metadata_frame, body],
+ copy=False,
+ )
+ except zmq.Again:
+ if attempt == self._retries:
+ break
+ continue
+
+ deadline = time.monotonic() + self._timeout_ms / 1000
+ while True:
+ remaining_ms = max(1, int((deadline - time.monotonic()) * 1000))
+ if deadline <= time.monotonic() or not self._socket.poll(
+ remaining_ms, zmq.POLLIN
+ ):
+ break
+ frames = self._socket.recv_multipart()
+ if len(frames) != 2:
+ continue
+ kind, raw_reply = frames
+ reply = json.loads(raw_reply)
+ if kind == _NACK and "transfer_id" not in reply:
+ raise RuntimeError(f"ZeroMQ sparse refit rejected payload: {reply}")
+ reply_key = (
+ reply.get("transfer_id"),
+ reply.get("producer_id"),
+ reply.get("payload_id"),
+ )
+ if reply_key != (transfer_id, producer_id, payload_id):
+ continue
+ if kind == _ACK and reply.get("ok") is True:
+ return reply
+ raise RuntimeError(f"ZeroMQ sparse refit rejected payload: {reply}")
+
+ if attempt < self._retries:
+ time.sleep(min(0.05 * 2**attempt, 0.5))
+
+ raise TimeoutError(
+ f"Timed out sending sparse refit payload {payload_id} to {self._address}."
+ )
+
+ def close(self) -> None:
+ self._socket.close()
+
+
+class ZmqSparseRefitServer:
+ """Bounded ROUTER relay that applies locally and fans out through a tree."""
+
+ def __init__(
+ self,
+ apply_payload: Callable[[bytes, Mapping[str, Any]], dict[str, Any]],
+ *,
+ bind_address: str,
+ api_key_env_var: str | None,
+ timeout_s: float,
+ tuning: VllmRefitTuningConfig,
+ ) -> None:
+ self._apply_payload = apply_payload
+ self._bind_address = bind_address
+ self._token = vllm_refit_api_key(api_key_env_var)
+ self._timeout_s = timeout_s
+ self._retries = tuning.zmq_retries
+ self._stop = threading.Event()
+ self._ready = threading.Event()
+ self._thread: threading.Thread | None = None
+ self._endpoint: str | None = None
+ self._error: Exception | None = None
+ self._payload_workers = tuning.zmq_relay_payload_workers
+ self._transfer_lock = threading.Lock()
+ self._transfer_condition = threading.Condition(self._transfer_lock)
+ self._transfers: dict[str, _RelayTransfer] = {}
+ self._flush_results: dict[str, Future[dict[str, Any]]] = {}
+ self._tree: tuple[tuple[str, ...], int] | None = None
+ self._forward_local = threading.local()
+ self._forward_clients: list[ZmqSparseRefitClient] = []
+ self._payload_executor = ThreadPoolExecutor(
+ max_workers=self._payload_workers,
+ thread_name_prefix="nrl-zmq-payload",
+ )
+ self._forward_executor = ThreadPoolExecutor(
+ max_workers=tuning.zmq_relay_forward_workers,
+ thread_name_prefix="nrl-zmq-forward",
+ )
+
+ def configure_tree(
+ self,
+ relay_addresses: Sequence[str],
+ *,
+ own_address: str,
+ ) -> None:
+ addresses = tuple(dict.fromkeys(relay_addresses))
+ self._tree = (addresses, addresses.index(own_address))
+
+ def start(self) -> str:
+ self._thread = threading.Thread(
+ target=self._run,
+ name="nrl-zmq-refit-relay",
+ daemon=True,
+ )
+ self._thread.start()
+ if not self._ready.wait(timeout=10.0):
+ raise RuntimeError("Timed out starting the ZeroMQ sparse refit relay.")
+ if self._error is not None:
+ raise RuntimeError(
+ "Failed to start the ZeroMQ sparse refit relay."
+ ) from self._error
+ assert self._endpoint is not None
+ return self._endpoint
+
+ def close(self) -> None:
+ self._stop.set()
+ if self._thread is not None:
+ self._thread.join(timeout=max(5.0, self._timeout_s))
+ if self._thread.is_alive():
+ raise RuntimeError("Timed out stopping the ZeroMQ sparse refit relay.")
+ self._thread = None
+
+ def flush(self, transfer_id: str, expected_payloads: int = 0) -> dict[str, Any]:
+ """Wait for every staged fanout belonging to one transfer."""
+ with self._transfer_condition:
+ completion = self._flush_results.get(transfer_id)
+ if completion is None:
+ ready = self._transfer_condition.wait_for(
+ lambda: (
+ len(
+ self._transfers.get(transfer_id, _RelayTransfer()).checksums
+ )
+ >= expected_payloads
+ ),
+ timeout=self._timeout_s,
+ )
+ if not ready:
+ raise TimeoutError(
+ f"Timed out waiting for {expected_payloads} ZeroMQ payloads "
+ f"for transfer {transfer_id}."
+ )
+ completion = self._flush_results.get(transfer_id)
+ if completion is None:
+ completion = Future()
+ self._flush_results[transfer_id] = completion
+ staged = self._transfers.pop(transfer_id, _RelayTransfer())
+ else:
+ staged = None
+ if staged is None:
+ return completion.result()
+
+ try:
+ started = time.perf_counter()
+ results: list[dict[str, Any]] = []
+ first_error: Exception | None = None
+ for future in staged.futures:
+ try:
+ results.append(future.result())
+ except Exception as exc:
+ if first_error is None:
+ first_error = exc
+ if first_error is not None:
+ raise RuntimeError(
+ f"ZeroMQ relay fanout failed for transfer {transfer_id}: "
+ f"{first_error}"
+ ) from first_error
+
+ merged = merge_vllm_refit_metrics({}, results, maximum=False)
+ merged.update(
+ ok=True,
+ payloads=len(staged.checksums),
+ receiver_relay_flush_s=time.perf_counter() - started,
+ )
+ except Exception as exc:
+ completion.set_exception(exc)
+ raise
+ completion.set_result(merged)
+ return merged
+
+ def _fanout(
+ self,
+ body: bytes,
+ metadata: Mapping[str, Any],
+ ) -> dict[str, Any]:
+ started = time.perf_counter()
+ result = self._apply_payload(body, metadata)
+ result["receiver_relay_fanout_s"] = time.perf_counter() - started
+ return result
+
+ def _forward(
+ self,
+ body: bytes,
+ metadata: Mapping[str, Any],
+ address: str,
+ relay_root: str,
+ ) -> dict[str, Any]:
+ clients = getattr(self._forward_local, "clients", None)
+ if clients is None:
+ clients = {}
+ self._forward_local.clients = clients
+ client = clients.get(address)
+ if client is None:
+ client = ZmqSparseRefitClient(
+ address,
+ timeout_s=self._timeout_s,
+ producer_id=0,
+ retries=self._retries,
+ api_key=self._token,
+ )
+ clients[address] = client
+ with self._transfer_lock:
+ self._forward_clients.append(client)
+ started = time.perf_counter()
+ client.send_payload(
+ transfer_id=str(metadata["transfer_id"]),
+ payload_id=int(metadata["payload_id"]),
+ checksum=str(metadata["checksum"]),
+ verification_candidates=int(metadata["verification_candidates"]),
+ body=body,
+ relay_root=relay_root,
+ producer_id=int(metadata["producer_id"]),
+ )
+ return {"receiver_relay_forward_s": time.perf_counter() - started}
+
+ @staticmethod
+ def _send_reply(
+ socket: zmq.Socket,
+ identity: bytes,
+ kind: bytes,
+ reply: Mapping[str, Any],
+ ) -> None:
+ with suppress(zmq.ZMQError):
+ socket.send_multipart(
+ [identity, kind, _json_bytes(reply)], flags=zmq.NOBLOCK
+ )
+
+ def _parse_data_message(
+ self,
+ frames: list[bytes],
+ ) -> tuple[bytes, tuple[str, int, int], bytes, dict[str, Any]]:
+ if len(frames) != 4:
+ raise ValueError(f"Expected 4 ZeroMQ frames, received {len(frames)}.")
+ identity, kind, raw_metadata, body = frames
+ if kind != _DATA:
+ raise ValueError(f"Unsupported ZeroMQ sparse refit message {kind!r}.")
+ metadata = json.loads(raw_metadata)
+ if metadata.get("protocol") != _PROTOCOL:
+ raise ValueError("Unsupported ZeroMQ sparse refit protocol.")
+ supplied_token = metadata.get("api_key")
+ if self._token is not None and (
+ not isinstance(supplied_token, str)
+ or not hmac.compare_digest(self._token, supplied_token)
+ ):
+ raise PermissionError("ZeroMQ sparse refit producer authentication failed.")
+ transfer_id = str(metadata["transfer_id"])
+ producer_id = int(metadata["producer_id"])
+ payload_id = int(metadata["payload_id"])
+ checksum = str(metadata["checksum"])
+ verification_candidates = int(metadata["verification_candidates"])
+ if (
+ not transfer_id
+ or producer_id < 0
+ or payload_id < 0
+ or not checksum
+ or verification_candidates < 0
+ ):
+ raise ValueError("Invalid ZeroMQ sparse refit payload identity.")
+ relay_root = metadata.get("relay_root")
+ if relay_root is not None and (
+ self._tree is None or relay_root not in self._tree[0]
+ ):
+ raise ValueError("Invalid ZeroMQ relay root.")
+ return identity, (transfer_id, producer_id, payload_id), body, metadata
+
+ def _run(self) -> None:
+ context = zmq.Context()
+ socket = context.socket(zmq.ROUTER)
+ try:
+ _configure_socket(socket, 16)
+ socket.setsockopt(zmq.ROUTER_MANDATORY, 1)
+ socket.bind(self._bind_address)
+ self._endpoint = socket.getsockopt_string(zmq.LAST_ENDPOINT)
+ self._ready.set()
+
+ while not self._stop.is_set():
+ if socket.poll(10, zmq.POLLIN):
+ frames = socket.recv_multipart()
+ identity = frames[0] if frames else b""
+ try:
+ identity, key, body, metadata = self._parse_data_message(frames)
+ transfer_id, producer_id, payload_id = key
+ payload_key = (producer_id, payload_id)
+ checksum = str(metadata["checksum"])
+ with self._transfer_lock:
+ if transfer_id in self._flush_results:
+ raise RuntimeError(
+ "ZeroMQ sparse refit transfer is already flushed."
+ )
+ staged = self._transfers.setdefault(
+ transfer_id, _RelayTransfer()
+ )
+ previous = staged.checksums.get(payload_key)
+ if previous is not None and previous != checksum:
+ raise ValueError(
+ "Conflicting ZeroMQ sparse refit payload checksum."
+ )
+ if previous is None:
+ staged.checksums[payload_key] = checksum
+ staged.futures.append(
+ self._payload_executor.submit(
+ self._fanout, body, metadata
+ )
+ )
+ if self._tree is not None:
+ addresses, own_index = self._tree
+ relay_root = str(
+ metadata.get("relay_root", addresses[own_index])
+ )
+ root_index = addresses.index(relay_root)
+ node_index = (own_index - root_index) % len(
+ addresses
+ )
+ for child_index in (
+ 2 * node_index + 1,
+ 2 * node_index + 2,
+ ):
+ if child_index < len(addresses):
+ child = addresses[
+ (root_index + child_index)
+ % len(addresses)
+ ]
+ staged.futures.append(
+ self._forward_executor.submit(
+ self._forward,
+ body,
+ metadata,
+ child,
+ relay_root,
+ )
+ )
+ self._transfer_condition.notify_all()
+ self._send_reply(
+ socket,
+ identity,
+ _ACK,
+ {
+ "ok": True,
+ "staged": True,
+ "transfer_id": transfer_id,
+ "producer_id": producer_id,
+ "payload_id": payload_id,
+ },
+ )
+ except Exception as exc:
+ self._send_reply(
+ socket,
+ identity,
+ _NACK,
+ {"ok": False, "error": str(exc)},
+ )
+ except Exception as exc:
+ self._error = exc
+ self._ready.set()
+ finally:
+ self._payload_executor.shutdown(wait=True, cancel_futures=True)
+ self._forward_executor.shutdown(wait=True, cancel_futures=True)
+ for client in self._forward_clients:
+ client.close()
+ socket.close()
+ context.term()
+
+
+def stream_sparse_delta_payloads_via_zmq(
+ iterator: Iterable[NamedTensor],
+ *,
+ delta_tracker: DeltaCompressionTracker,
+ refit_targets: Sequence[str],
+ transfer_id: str,
+ api_key_env_var: str | None,
+ timeout_s: float,
+ shard_rank: int,
+ shard_count: int,
+) -> dict[str, int]:
+ addresses = [address.strip() for address in refit_targets if address.strip()]
+ if not addresses:
+ raise ValueError("At least one ZeroMQ sparse refit address is required.")
+ address = addresses[shard_rank % len(addresses)]
+ tuning = delta_tracker.refit_config.tuning
+ api_key = vllm_refit_api_key(api_key_env_var)
+ local = threading.local()
+
+ def send(
+ body: bytes, payload_id: int, verification_candidates: int
+ ) -> dict[str, Any]:
+ client = getattr(local, "client", None)
+ if client is None:
+ client = ZmqSparseRefitClient(
+ address,
+ timeout_s=timeout_s,
+ producer_id=shard_rank,
+ retries=tuning.zmq_retries,
+ api_key=api_key,
+ )
+ local.client = client
+ started = time.perf_counter()
+ reply = client.send_payload(
+ transfer_id=transfer_id,
+ payload_id=payload_id,
+ checksum=sparse_payload_checksum(body),
+ verification_candidates=verification_candidates,
+ body=body,
+ )
+ return {
+ "zmq_send_s": time.perf_counter() - started,
+ "receiver": reply,
+ }
+
+ def cleanup() -> None:
+ client = getattr(local, "client", None)
+ if client is not None:
+ client.close()
+ del local.client
+
+ return stream_sparse_delta_payloads(
+ iterator,
+ delta_tracker=delta_tracker,
+ transport=SparseRefitTransport(
+ name="zmq",
+ transfer_workers=tuning.transfer_workers["zmq"],
+ send=send,
+ cleanup=cleanup,
+ ),
+ shard_rank=shard_rank,
+ shard_count=shard_count,
+ )
diff --git a/nemo_rl/weight_sync/interfaces.py b/nemo_rl/weight_sync/interfaces.py
index e60317be36d..f0e0817d63a 100644
--- a/nemo_rl/weight_sync/interfaces.py
+++ b/nemo_rl/weight_sync/interfaces.py
@@ -62,7 +62,7 @@ def sync_weights(
*,
timer: Optional[Timer] = None,
kv_scales: Optional[dict[str, float]] = None,
- ) -> None:
+ ) -> Optional[dict[str, float]]:
"""Transfer the latest policy weights to the generation backend.
This method encapsulates the full sync lifecycle:
@@ -88,6 +88,9 @@ def sync_weights(
which forwards them to ``policy.broadcast_weights_for_collective()``.
IPC and HTTP transports ignore this parameter.
+ Returns:
+ Optional transport-specific scalar metrics for the current sync.
+
Raises:
RuntimeError: If the weight transfer fails.
"""
diff --git a/nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py b/nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py
new file mode 100644
index 00000000000..eda7045d9cd
--- /dev/null
+++ b/nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py
@@ -0,0 +1,402 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared S3/ZeroMQ sparse synchronizer for remote non-colocated vLLM refit."""
+
+import time
+import uuid
+from collections import defaultdict
+from contextlib import nullcontext, suppress
+from typing import Any
+
+import ray
+
+from nemo_rl.models.generation.vllm.config import (
+ VllmConfig,
+ normalize_vllm_refit_config,
+)
+from nemo_rl.utils.timer import Timer
+from nemo_rl.utils.weight_transfer_http import (
+ G_VLLM_REFIT_FLUSH_PATH,
+ G_VLLM_REFIT_PREPARE_PATH,
+ G_VLLM_REFIT_ZMQ_FLUSH_PATH,
+ merge_vllm_refit_metrics,
+ post_vllm_refit_endpoints,
+ vllm_refit_api_key,
+ vllm_refit_endpoints,
+)
+from nemo_rl.weight_sync.interfaces import WeightSynchronizer
+
+_REMOTE_SPARSE_TRANSPORTS = {
+ "vllm_s3_sparse": "s3",
+ "vllm_zmq_sparse": "zmq",
+}
+
+
+def validate_vllm_remote_sparse_refit(
+ config: VllmConfig,
+ *,
+ colocated: bool,
+ megatron_enabled: bool,
+) -> str | None:
+ """Validate the optional config and return its internal transport name."""
+ transport = config.get("refit_transport")
+ if transport is None:
+ return None
+ if transport not in _REMOTE_SPARSE_TRANSPORTS:
+ raise ValueError(f"Unsupported vLLM refit transport {transport!r}.")
+ vllm_cfg = config["vllm_cfg"]
+ refit_config = normalize_vllm_refit_config(config)
+ assert refit_config is not None
+ if (
+ colocated
+ or not megatron_enabled
+ or vllm_cfg["precision"] == "fp8"
+ or vllm_cfg["kv_cache_dtype"].startswith("fp8")
+ or config.get("quant_cfg")
+ or config.get("real_quant")
+ ):
+ raise ValueError(
+ f"{transport} requires a non-colocated Megatron policy, BF16/FP16 "
+ "vLLM, and an unquantized rollout."
+ )
+ if transport == "vllm_s3_sparse" and not (
+ refit_config.storage.s3_bucket and refit_config.storage.s3_bucket.strip()
+ ):
+ raise ValueError(
+ "vllm_s3_sparse requires policy.generation.refit_cfg.storage.s3_bucket."
+ )
+ return _REMOTE_SPARSE_TRANSPORTS[transport]
+
+
+class VllmRemoteSparseWeightSynchronizer(WeightSynchronizer):
+ def __init__(
+ self,
+ policy: Any,
+ generation: Any,
+ *,
+ transport: str,
+ api_key_env_var: str | None = None,
+ request_timeout_s: float = 600.0,
+ baseline_init_refs: list[Any] | None = None,
+ ) -> None:
+ self._policy = policy
+ self._generation = generation
+ self._transport = transport
+ self._api_key_env_var = api_key_env_var
+ self._request_timeout_s = request_timeout_s
+ self._refit_urls: list[str] = []
+ self._targets: list[str] = []
+ self._overwrite_names: list[str] = []
+ self._baseline_init_refs = list(baseline_init_refs or ())
+ self._baseline_commit_refs: list[Any] = []
+ self._stale = True
+ self._poisoned = False
+
+ def sync_weights(
+ self,
+ *,
+ timer: Timer | None = None,
+ kv_scales: dict[str, float] | None = None,
+ ) -> dict[str, float]:
+ if self._poisoned:
+ raise RuntimeError(
+ "Sparse refit synchronizer was poisoned by a prior failed sync: "
+ "receivers may be partially applied while the trainer baseline is "
+ "uncommitted, so re-applying deltas would double-XOR and corrupt "
+ "weights. Reload the rollout workers from a known-good checkpoint "
+ "and re-initialize the communicator before syncing again. See "
+ "https://github.com/NVIDIA-NeMo/RL/issues/3274."
+ )
+ context = (
+ timer.time("prepare_for_generation/transfer_and_update_weights")
+ if timer
+ else nullcontext()
+ )
+ with context:
+ if self._baseline_commit_refs:
+ ray.get(self._baseline_commit_refs)
+ self._baseline_commit_refs.clear()
+ if not self._generation.invalidate_kv_cache():
+ raise RuntimeError(
+ f"vLLM KV cache invalidation failed before {self._transport} "
+ "weight update."
+ )
+ if self._baseline_init_refs:
+ ray.get(self._baseline_init_refs)
+ self._baseline_init_refs.clear()
+
+ succeeded = False
+ relay_flushed = False
+ relay_flush_s = 0.0
+ transfer_id = uuid.uuid4().hex
+ try:
+ results = ray.get(
+ self._run_policy_workers(
+ "stream_remote_sparse_weights",
+ transport=self._transport,
+ targets=self._targets,
+ transfer_id=transfer_id,
+ api_key_env_var=self._api_key_env_var,
+ timeout_s=self._request_timeout_s,
+ overwrite_names=self._overwrite_names,
+ )
+ )
+ payloads = sum(result["payloads"] for result in results)
+ changed = sum(result["changed_elements"] for result in results)
+ total = sum(result["total_elements"] for result in results)
+ changed_pct = 100.0 * changed / max(total, 1)
+ print(
+ f"REFIT_{self._transport.upper()}_DELTA_CHANGE "
+ f"changed_elements={changed} total_elements={total} "
+ f"changed_pct={changed_pct:.8g}",
+ flush=True,
+ )
+
+ verification: defaultdict[str, float] = defaultdict(float)
+ commit_s = 0.0
+ if payloads:
+ if self._transport == "zmq":
+ started = time.perf_counter()
+ relay_results = self._request_receivers(
+ G_VLLM_REFIT_ZMQ_FLUSH_PATH,
+ {
+ "transfer_id": transfer_id,
+ "expected_payloads": payloads,
+ },
+ )
+ relay_flush_s = time.perf_counter() - started
+ relay_flushed = True
+ if any(
+ int(result.get("payloads", 0)) != payloads
+ for result in relay_results
+ ):
+ raise RuntimeError(
+ f"ZeroMQ relays did not all stage {payloads} payloads."
+ )
+ print(
+ "REFIT_ZMQ_RELAY_FLUSH "
+ f"transfer_id={transfer_id} payloads={payloads} "
+ f"seconds={relay_flush_s:.3f} "
+ "fanout_service_s="
+ f"{sum(float(result.get('receiver_relay_fanout_s', 0.0)) for result in relay_results):.3f}",
+ flush=True,
+ )
+ started = time.perf_counter()
+ verification.update(
+ merge_vllm_refit_metrics(
+ {},
+ self._request_receivers(G_VLLM_REFIT_FLUSH_PATH, {}),
+ maximum=True,
+ candidate_maximum=False,
+ )
+ )
+ candidates = int(verification["verification_candidates"])
+ samples = int(verification["verification_samples"])
+ exact = int(verification["verification_exact_mismatches"])
+ mismatches = int(verification["verification_mismatches"])
+ abs_sum = float(verification["verification_abs_sum"])
+ max_abs = float(verification["verification_max_abs"])
+ if candidates or samples:
+ print(
+ f"REFIT_{self._transport.upper()}_DELTA_VERIFY "
+ f"candidates={candidates} samples={samples} "
+ f"exact_mismatches={exact} mismatches={mismatches} "
+ f"mean_abs={abs_sum / max(samples, 1):.8g} "
+ f"max_abs={max_abs:.8g}",
+ flush=True,
+ )
+ if mismatches:
+ raise RuntimeError(
+ f"Sparse refit sampled {mismatches} mismatched deltas "
+ f"out of {samples}."
+ )
+ commit_s = time.perf_counter() - started
+ print(
+ f"REFIT_{self._transport.upper()}_GLOBAL_COMMIT "
+ f"transfer_id={transfer_id} payloads={payloads} "
+ f"seconds={commit_s:.3f}",
+ flush=True,
+ )
+ succeeded = True
+ finally:
+ if not succeeded:
+ self._poisoned = True
+ if self._transport == "zmq" and not relay_flushed:
+ with suppress(Exception):
+ self._request_receivers(
+ G_VLLM_REFIT_ZMQ_FLUSH_PATH,
+ {"transfer_id": transfer_id},
+ timeout_s=min(self._request_timeout_s, 60.0),
+ )
+ with suppress(Exception):
+ self._request_receivers(
+ G_VLLM_REFIT_FLUSH_PATH,
+ {},
+ timeout_s=min(self._request_timeout_s, 60.0),
+ )
+ self._baseline_commit_refs = (
+ self._policy.worker_group.run_all_workers_single_data(
+ "finish_remote_sparse_delta_sync", succeeded=succeeded
+ )
+ )
+
+ self._stale = False
+ samples = int(verification["verification_samples"])
+ mismatches = int(verification["verification_mismatches"])
+ metrics = {
+ "delta/changed_elements": float(changed),
+ "delta/total_elements": float(total),
+ "delta/changed_pct": changed_pct,
+ "delta_verify/mismatch_pct": 100.0 * mismatches / max(samples, 1),
+ "delta_verify/mean_abs": float(verification["verification_abs_sum"])
+ / max(samples, 1),
+ "transfer/payloads": float(payloads),
+ "transfer/relay_flush_s": relay_flush_s,
+ "transfer/global_commit_s": commit_s,
+ }
+ metrics.update(
+ {
+ f"delta_verify/{key}": float(verification[f"verification_{key}"])
+ for key in (
+ "candidates",
+ "samples",
+ "exact_mismatches",
+ "mismatches",
+ "max_abs",
+ )
+ }
+ )
+ return metrics
+
+ @property
+ def is_stale(self) -> bool:
+ return self._stale
+
+ def mark_stale(self) -> None:
+ self._stale = True
+
+ def _run_policy_workers(self, method_name: str, **kwargs: Any) -> list[Any]:
+ workers = self._policy.worker_group
+ count = len(workers.workers)
+ return workers.run_all_workers_multiple_data(
+ method_name,
+ common_kwargs={**kwargs, "shard_count": count},
+ shard_rank=list(range(count)),
+ )
+
+ def _run_generation_workers(self, method_name: str, **kwargs: Any) -> list[Any]:
+ workers = self._generation.worker_group
+ return ray.get(
+ workers.run_all_workers_single_data(
+ method_name,
+ run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
+ **kwargs,
+ )
+ )
+
+ def _request_receivers(
+ self,
+ path: str,
+ body: dict[str, Any],
+ *,
+ timeout_s: float | None = None,
+ ) -> list[dict[str, Any]]:
+ return post_vllm_refit_endpoints(
+ vllm_refit_endpoints(self._refit_urls, path),
+ body,
+ api_key=vllm_refit_api_key(self._api_key_env_var),
+ timeout_s=self._request_timeout_s if timeout_s is None else timeout_s,
+ )
+
+ @staticmethod
+ def start_baseline(policy: Any, transport: str) -> list[Any]:
+ workers = policy.worker_group
+ count = len(workers.workers)
+ return workers.run_all_workers_multiple_data(
+ "init_remote_sparse_delta_baseline",
+ common_kwargs={"transport": transport, "shard_count": count},
+ shard_rank=list(range(count)),
+ )
+
+ @staticmethod
+ def _merge_refit_info(parts: list[dict[str, Any]]) -> dict[str, Any]:
+ merged = {}
+ for part in parts:
+ for name, info in part.items():
+ if name in merged and merged[name] != info:
+ raise ValueError(f"Conflicting sparse refit metadata for {name!r}.")
+ merged[name] = info
+ return merged
+
+ def init_communicator(self) -> None:
+ if not self._baseline_init_refs:
+ self._baseline_init_refs = self.start_baseline(
+ self._policy, self._transport
+ )
+ self._refit_urls = [
+ url
+ for url in self._run_generation_workers("report_refit_server_base_url")
+ if url
+ ]
+ self._targets = self._refit_urls
+ if self._transport == "zmq":
+ self._targets = [
+ address
+ for address in self._run_generation_workers(
+ "start_zmq_sparse_refit_relay"
+ )
+ if address
+ ]
+ if not self._refit_urls or not self._targets:
+ raise ValueError(
+ f"vLLM {self._transport} sparse refit endpoints are missing."
+ )
+ if self._transport == "zmq":
+ self._run_generation_workers(
+ "configure_zmq_sparse_refit_relay", relay_addresses=self._targets
+ )
+ state_dict_info = self._merge_refit_info(
+ ray.get(list(self._baseline_init_refs))
+ )
+ self._baseline_init_refs.clear()
+ responses = self._request_receivers(
+ G_VLLM_REFIT_PREPARE_PATH,
+ {
+ "tensors": {
+ name: [list(shape), str(dtype).removeprefix("torch.")]
+ for name, (shape, dtype) in state_dict_info.items()
+ }
+ },
+ )
+ self._overwrite_names = sorted(
+ {
+ name
+ for response in responses
+ for name in response.get("overwrite_names", ())
+ }
+ )
+ self._stale = False
+
+ def shutdown(self) -> None:
+ for ref in self._baseline_init_refs + self._baseline_commit_refs:
+ ray.cancel(ref, force=False)
+ if self._transport == "zmq":
+ self._run_generation_workers("stop_zmq_sparse_refit_relay")
+ self._baseline_init_refs.clear()
+ self._baseline_commit_refs.clear()
+ self._refit_urls.clear()
+ self._targets.clear()
+ self._overwrite_names.clear()
+ self._stale = True
diff --git a/pyproject.toml b/pyproject.toml
index 11dc7691676..63b1de062e3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -66,6 +66,8 @@ dependencies = [
"nccl4py; sys_platform != 'darwin'", # for non-colocated refit
"cuda-bindings; sys_platform != 'darwin'", # for non-colocated refit
"pybase64", # for sglang refit
+ "awscrt>=0.35.0", # for parallel S3 refit transport
+ "zstandard", # for sparse refit body compression
"nvidia-cudnn-cu13==9.20.0.48; sys_platform != 'darwin'", # for transformer-engine no build isolation
# tilelang — replacement Triton kernel mamba-ssm requires when
# Triton >= 3.4.0 on Hopper, see github.com/state-spaces/mamba#640.
diff --git a/pyrefly.toml b/pyrefly.toml
index 117ec1fdcd7..75315f5d862 100644
--- a/pyrefly.toml
+++ b/pyrefly.toml
@@ -18,6 +18,7 @@ replace-imports-with-any = [
"numpy.*",
"sphinx.*",
"docutils.*",
+ "zstandard.*",
]
project-includes = [
# TODO: enable these once we have 100 correctness
@@ -170,6 +171,8 @@ project-includes = [
"nemo_rl/models/generation/vllm/reasoning_parsers/nano_v3_reasoning_parser.py",
"nemo_rl/models/generation/vllm/utils.py",
"nemo_rl/models/generation/vllm/vllm_backend.py",
+ "nemo_rl/models/generation/vllm/vllm_sparse_delta.py",
+ "nemo_rl/models/generation/vllm/vllm_sparse_refit.py",
"nemo_rl/models/huggingface/__init__.py",
"nemo_rl/models/megatron/__init__.py",
"nemo_rl/models/megatron/draft/__init__.py",
@@ -177,6 +180,7 @@ project-includes = [
"nemo_rl/models/policy/interfaces.py",
"nemo_rl/models/policy/utils.py",
"nemo_rl/models/policy/workers/__init__.py",
+ "nemo_rl/models/policy/workers/megatron_remote_sparse_refit.py",
"nemo_rl/models/policy/workers/patches.py",
"nemo_rl/models/value/__init__.py",
"nemo_rl/models/value/config.py",
@@ -193,18 +197,24 @@ project-includes = [
"nemo_rl/utils/r3_trace.py",
"nemo_rl/utils/timer.py",
"nemo_rl/utils/venvs.py",
+ "nemo_rl/utils/weight_transfer_http.py",
+ "nemo_rl/utils/weight_transfer_sparse_codec.py",
+ "nemo_rl/utils/weight_transfer_stream.py",
+ "nemo_rl/utils/weight_transfer_zmq.py",
"nemo_rl/weight_sync/__init__.py",
"nemo_rl/weight_sync/collective_weight_synchronizer.py",
"nemo_rl/weight_sync/factory.py",
"nemo_rl/weight_sync/http_weight_synchronizer.py",
"nemo_rl/weight_sync/interfaces.py",
"nemo_rl/weight_sync/ipc_weight_synchronizer.py",
+ "nemo_rl/weight_sync/vllm_remote_sparse_weight_synchronizer.py",
"tools/model_diagnostics/1.max_model_len_respected.py",
"tools/model_diagnostics/2.long_generation_decode_vs_prefill.py",
"tools/model_diagnostics/3.check_and_reinit_hf_model_embeddings_untrained.py",
"tools/model_diagnostics/4.vllm_precision_compilation_test.py",
"tools/model_diagnostics/5.prefix_caching_nan.py",
"tools/model_diagnostics/6.vllm_routed_experts_completeness.py",
+ "tools/refit_bandwidth_calculator.py",
"tools/x_token/__init__.py",
"tools/x_token/reapply_exact_map.py",
"tools/x_token/sort_and_cut_projection_matrix.py",
diff --git a/tests/test_suites/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.sh b/tests/test_suites/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.sh
new file mode 100755
index 00000000000..2bea0a11ace
--- /dev/null
+++ b/tests/test_suites/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+source $SCRIPT_DIR/common.env
+
+# ===== BEGIN CONFIG =====
+NUM_NODES=4
+STEPS_PER_RUN=50
+MAX_STEPS=50
+NUM_RUNS=$(( (MAX_STEPS + STEPS_PER_RUN - 1) / STEPS_PER_RUN ))
+NUM_MINUTES=65
+# ===== END CONFIG =====
+
+exit_if_max_steps_reached
+
+cd $PROJECT_ROOT
+uv run examples/run_grpo.py \
+ --config $CONFIG_PATH \
+ grpo.max_num_steps=$MAX_STEPS \
+ logger.log_dir=$LOG_DIR \
+ logger.wandb_enabled=True \
+ logger.wandb.project=nemo-rl-refit \
+ logger.wandb.name=$EXP_NAME \
+ logger.monitor_gpus=True \
+ logger.tensorboard_enabled=True \
+ checkpointing.enabled=False \
+ $@ \
+ 2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+MAX_RECORDED_STEP=$(jq -r 'if has("train/loss") then (."train/loss" | keys | map(tonumber) | max // 0) else 0 end' $JSON_METRICS)
+if [[ $MAX_RECORDED_STEP -lt $MAX_STEPS ]]; then
+ echo "[ERROR] Expected train/loss through step $MAX_STEPS, found step $MAX_RECORDED_STEP"
+ exit 1
+fi
+
+uv run tests/check_metrics.py $JSON_METRICS \
+ 'median(data["train/token_mult_prob_error"]) < 1.03' \
+ "data[\"train/token_mult_prob_error\"][\"$MAX_STEPS\"] < 1.03" \
+ 'ratio_above(data["train/token_mult_prob_error"], 1.03) < 0.05' \
+ "data[\"train/reward\"][\"$MAX_STEPS\"] > 0.2" \
+ 'min(data["refit/transfer/payloads"]) > 0' \
+ 'min(data["refit/transfer/relay_flush_s"]) > 0' \
+ 'min(data["refit/delta/changed_pct"]) > 0' \
+ 'max(data["refit/delta/changed_pct"]) <= 5'
diff --git a/tests/test_suites/nightly.txt b/tests/test_suites/nightly.txt
index 81e8ebc90c4..a30f4f7c862 100644
--- a/tests/test_suites/nightly.txt
+++ b/tests/test_suites/nightly.txt
@@ -93,6 +93,7 @@ tests/test_suites/llm/grpo-qwen3-8b-base-dapo-2n8g-long-megatron-qa-nvfp4-w4a16.
# Non-colocated
tests/test_suites/llm/grpo-llama3.1-8b-instruct-2n8g-fsdp2tp1-noncolocated.sh
tests/test_suites/llm/grpo-llama3.2-1b-instruct-2n8g-megatron_generation-noncolocated.sh
+tests/test_suites/llm/grpo-qwen3-30ba3b-4n8g-megatron-zmq-deltaweight-noncolocated.sh
# Nemotron Super 49B
#https://github.com/NVIDIA-NeMo/RL/issues/1374
diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py
index df91ba56a0f..5e26f583218 100644
--- a/tests/unit/algorithms/test_grpo.py
+++ b/tests/unit/algorithms/test_grpo.py
@@ -33,6 +33,7 @@
_apply_mask_sample_filter,
_apply_message_level_advantage_penalties,
_default_grpo_save_state,
+ _initial_policy_generation_stale,
_raise_if_reward_penalties_enabled_without_nemo_gym,
_resolve_message_level_advantage_penalties,
_should_use_async_rollouts,
@@ -115,6 +116,17 @@ def test_missing_mask_sample_is_noop(self):
)
+def test_initial_policy_generation_stale() -> None:
+ generation = MagicMock()
+ generation.weight_synchronizer.is_stale = False
+
+ assert not _initial_policy_generation_stale(generation, completed_steps=0)
+ assert _initial_policy_generation_stale(generation, completed_steps=1)
+
+ generation.weight_synchronizer.is_stale = True
+ assert _initial_policy_generation_stale(generation, completed_steps=0)
+
+
@pytest.fixture
def mock_grpo_components():
# Create mock components
@@ -1863,7 +1875,9 @@ def fake_batched_message_log_to_flat_message(*_args, **_kwargs):
lambda *_args, **_kwargs: (torch.tensor([0.1]), torch.tensor([1.0])),
)
monkeypatch.setattr(
- grpo_mod, "refit_policy_generation", lambda *_args, **_kwargs: None
+ grpo_mod,
+ "refit_policy_generation",
+ lambda *_args, **_kwargs: {"delta/changed_pct": 4.0},
)
monkeypatch.setattr(
grpo_mod, "print_performance_metrics", lambda *_args, **_kwargs: {}
@@ -1913,6 +1927,11 @@ def fake_batched_message_log_to_flat_message(*_args, **_kwargs):
assert train_metrics["min_seq_mult_prob_error_after_mask"] == 1.0
assert train_metrics["num_masked_seqs_by_logprob_error"] == 2
assert train_metrics["masked_correct_pct"] == 0.5
+ assert any(
+ call.args[0] == {"delta/changed_pct": 4.0}
+ and call.kwargs.get("prefix") == "refit"
+ for call in mock_grpo_components["logger"].log_metrics.call_args_list
+ )
def test_grpo_train_shutdown_on_epoch_completion(mock_grpo_components, tmp_path):
diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py
index ec659bf7602..3017aa9941a 100644
--- a/tests/unit/models/generation/test_vllm_generation.py
+++ b/tests/unit/models/generation/test_vllm_generation.py
@@ -447,6 +447,18 @@ def test_configure_generation_config_uses_real_startup_weights_without_draft_ref
assert configured["vllm_cfg"]["load_format"] == "auto"
+@pytest.mark.parametrize("transport", ["vllm_s3_sparse", "vllm_zmq_sparse"])
+def test_configure_generation_config_uses_real_delta_baseline(transport: str):
+ vllm_config = deepcopy(basic_vllm_test_config)
+ vllm_config["refit_transport"] = transport
+
+ configured = configure_generation_config(
+ vllm_config, MagicMock(pad_token_id=0, eos_token_id=1)
+ )
+
+ assert configured["vllm_cfg"]["load_format"] == "auto"
+
+
def test_configure_generation_config_keeps_dummy_startup_weights_with_draft_refit():
"""Speculative training can keep dummy startup weights when draft refit is available."""
vllm_config = deepcopy(basic_vllm_test_config)
diff --git a/tests/unit/models/generation/test_vllm_sparse_delta.py b/tests/unit/models/generation/test_vllm_sparse_delta.py
new file mode 100644
index 00000000000..76f0c47814b
--- /dev/null
+++ b/tests/unit/models/generation/test_vllm_sparse_delta.py
@@ -0,0 +1,540 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import math
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from nemo_rl.models.generation.vllm.vllm_sparse_delta import (
+ VllmSparseDeltaApplier,
+)
+from nemo_rl.utils.weight_transfer_sparse_codec import (
+ SparseOperation,
+ encode_sparse_infos,
+)
+from nemo_rl.utils.weight_transfer_sparse_codec import (
+ integer_view as _bits,
+)
+
+_PACKED_LOADER_INFOS = [
+ ("row_slice", (4, 2), [4, 5, 7], [1.0] * 3),
+ ("column_slice", (2, 8), [4, 7, 12, 15], [2.0] * 4),
+ ("split", (6, 2), [2, 5, 8, 11], [3.0, 3.0, 4.0, 4.0]),
+]
+
+
+class _NativeLoaderModel:
+ def __init__(self, **targets: torch.Tensor) -> None:
+ self.targets = targets
+ self.load_calls = 0
+ self.loaded_names: list[str] = []
+
+ def parameters(self):
+ return iter(self.targets.values())
+
+ def buffers(self):
+ return iter(())
+
+ def load_weights(self, weights):
+ self.load_calls += 1
+ loaded = set()
+ for name, source in weights:
+ self.loaded_names.append(name)
+ target = self.targets
+ if name == "weight":
+ target["identity"].copy_(source)
+ elif name == "weight_scale_inv":
+ target["scale"].copy_(source)
+ elif name == "row_slice":
+ target[name].copy_(source[2:4])
+ elif name == "column_slice":
+ target[name].copy_(source[:, 4:8])
+ elif name == "split":
+ target[name][:2].copy_(source[1:3])
+ target[name][2:].copy_(source[4:6])
+ elif name == "exp_transform":
+ target[name].copy_(-torch.exp(source))
+ else:
+ continue
+ loaded.add(name)
+ return loaded
+
+
+def _applier(model: Any) -> VllmSparseDeltaApplier:
+ return VllmSparseDeltaApplier(
+ SimpleNamespace(
+ model=model,
+ vllm_config=SimpleNamespace(),
+ ),
+ torch.device("cpu"),
+ )
+
+
+def _payload(
+ name: str,
+ tensor: torch.Tensor,
+ locations: torch.Tensor | list[int],
+ values: torch.Tensor,
+ operation: SparseOperation = "overwrite",
+) -> Any:
+ return encode_sparse_infos(
+ [(name, tensor, torch.as_tensor(locations), values, operation)]
+ )
+
+
+def _apply_payload(applier: VllmSparseDeltaApplier, payload: Any) -> None:
+ buffer = io.BytesIO()
+ torch.save(payload, buffer)
+ applier.update_weights_from_decoded_sparse_payload(buffer.getvalue())
+
+
+def test_sparse_discovery_reserves_largest_source_without_loading_weights() -> None:
+ target = torch.zeros(8)
+ applier = _applier(_NativeLoaderModel(identity=target))
+
+ applier.discover_native_skips({"weight": ((8,), torch.float32)})
+
+ assert applier._scratch.numel() == target.numel() * target.element_size()
+ assert torch.equal(target, torch.zeros_like(target))
+
+
+def test_sparse_discovery_caches_rank_local_native_loader_skips() -> None:
+ target = torch.zeros(2)
+ model = _NativeLoaderModel(identity=target)
+ applier = _applier(model)
+ info = {
+ "weight": ((2,), torch.float32),
+ "skipped": ((2,), torch.float32),
+ }
+
+ overwrite_names = applier.discover_native_skips(info)
+ payload = _payload("skipped", target, [1], _bits(torch.tensor([3.0])))
+ _apply_payload(applier, payload)
+
+ assert model.loaded_names == ["weight", "skipped"]
+ assert overwrite_names == set()
+ assert torch.equal(target, torch.zeros_like(target))
+
+
+def test_sparse_discovery_classifies_xor_unsafe_native_loads() -> None:
+ model = _NativeLoaderModel(
+ identity=torch.zeros(2),
+ exp_transform=torch.zeros(2),
+ )
+ applier = _applier(model)
+ info = {
+ "weight": ((2,), torch.bfloat16),
+ "exp_transform": ((2,), torch.float32),
+ }
+
+ assert applier.discover_native_skips(info) == {"weight", "exp_transform"}
+ assert torch.equal(model.targets["identity"], torch.zeros(2))
+ assert torch.equal(model.targets["exp_transform"], torch.zeros(2))
+
+
+@pytest.mark.vllm
+def test_backend_applies_decoded_sparse_payload_sources() -> None:
+ from nemo_rl.models.generation.vllm.vllm_backend import (
+ VllmInternalWorkerExtension,
+ )
+
+ ext = VllmInternalWorkerExtension.__new__(VllmInternalWorkerExtension)
+ applier = MagicMock()
+ applier.discover_native_skips.return_value = {"weight"}
+ applier.update_weights_from_decoded_sparse_payload.return_value = {"ok": True}
+ ext._get_sparse_delta_applier = MagicMock(return_value=applier)
+
+ assert ext.prepare_sparse_delta_refit_info({"weight": ((8,), torch.float32)}) == [
+ "weight"
+ ]
+
+ assert ext.update_weights_from_decoded_sparse_payload(b"payload") == {"ok": True}
+ assert ext.update_weights_from_decoded_sparse_payload("first", "second") == {
+ "ok": True
+ }
+ assert [
+ item.args
+ for item in applier.update_weights_from_decoded_sparse_payload.call_args_list
+ ] == [(b"payload",), ("first", "second")]
+ applier.discover_native_skips.assert_called_once_with(
+ {"weight": ((8,), torch.float32)}
+ )
+
+
+def test_sparse_payload_batches_preserve_order(tmp_path) -> None:
+ applier = _applier(_NativeLoaderModel(identity=torch.zeros(1)))
+ payload_paths = [tmp_path / f"payload-{index}.pt" for index in range(3)]
+ payloads = [
+ _payload(
+ f"weight-{index}",
+ torch.empty(1),
+ [0],
+ _bits(torch.tensor([float(index)])),
+ )
+ for index in range(3)
+ ]
+ for path, payload in zip(payload_paths, payloads, strict=True):
+ torch.save(payload, path)
+ decoded_applied: list[Any] = []
+ applier._apply_decoded_items = lambda items: decoded_applied.extend(
+ item for item, _, _ in items
+ )
+ result = applier.update_weights_from_decoded_sparse_payload(
+ *(path.read_bytes() for path in payload_paths)
+ )
+ decoded_result = applier.update_weights_from_decoded_sparse_payload(
+ *(str(path) for path in reversed(payload_paths))
+ )
+
+ assert [item["name"] for item in decoded_applied] == [
+ "weight-0",
+ "weight-1",
+ "weight-2",
+ "weight-2",
+ "weight-1",
+ "weight-0",
+ ]
+ assert result["receiver_deserialize_s"] >= 0.0
+ assert result["receiver_sparse_apply_s"] >= 0.0
+ assert decoded_result["receiver_deserialize_s"] >= 0.0
+
+
+def test_sparse_payload_batch_uses_one_streaming_native_loader_call() -> None:
+ identity = torch.zeros(2)
+ scale = torch.zeros(2)
+ model = _NativeLoaderModel(identity=identity, scale=scale)
+ serialized = []
+ for payload in (
+ _payload("weight", identity, [0], _bits(torch.tensor([2.0]))),
+ _payload("weight_scale_inv", scale, [1], _bits(torch.tensor([3.0]))),
+ ):
+ buffer = io.BytesIO()
+ torch.save(payload, buffer)
+ serialized.append(buffer.getvalue())
+
+ _applier(model).update_weights_from_decoded_sparse_payload(*serialized)
+
+ assert model.load_calls == 1
+ assert torch.equal(identity, torch.tensor([2.0, 0.0]))
+ assert torch.equal(scale, torch.tensor([0.0, 3.0]))
+
+
+def test_compact_sparse_payload_decodes_locations_for_apply(tmp_path) -> None:
+ target = torch.zeros(8)
+ payload = _payload("weight", target, [1, 5], _bits(torch.tensor([2.0, 6.0])))
+ path = tmp_path / "payload.pt"
+ torch.save(payload, path)
+
+ result = _applier(
+ _NativeLoaderModel(identity=target)
+ ).update_weights_from_decoded_sparse_payload(str(path))
+
+ assert torch.equal(target, torch.tensor([0.0, 2.0, 0.0, 0.0, 0.0, 6.0, 0.0, 0.0]))
+ assert result["receiver_sparse_apply_s"] >= 0.0
+
+
+def test_native_loaders_apply_sparse_views_and_transforms() -> None:
+ targets = {
+ "identity": torch.zeros(4),
+ "row_slice": torch.zeros(2, 2),
+ "column_slice": torch.zeros(2, 4),
+ "split": torch.zeros(4, 2),
+ "exp_transform": torch.tensor([-2.0, -4.0]),
+ }
+ infos = [
+ ("weight", (4,), [1, 2], [1.0, 2.0]),
+ *_PACKED_LOADER_INFOS,
+ (
+ "exp_transform",
+ (2,),
+ [0, 1],
+ [math.log(3.0), math.log(2.0)],
+ ),
+ ]
+ payload = encode_sparse_infos(
+ [
+ (
+ name,
+ torch.empty(shape),
+ torch.tensor(locations),
+ _bits(torch.tensor(values, dtype=torch.float32)),
+ "overwrite",
+ )
+ for name, shape, locations, values in infos
+ ]
+ )
+ payload[2][-1]["verification_samples"] = 2
+
+ applier = _applier(_NativeLoaderModel(**targets))
+ _apply_payload(applier, payload)
+ verification = applier.finish_sparse_delta_refit()
+
+ assert torch.equal(targets["identity"], torch.tensor([0.0, 1.0, 2.0, 0.0]))
+ assert targets["row_slice"].view(-1).tolist() == [1.0, 1.0, 0.0, 1.0]
+ assert targets["column_slice"].view(-1).tolist() == [2.0, 0.0, 0.0, 2.0] * 2
+ assert targets["split"].view(-1).tolist() == [
+ 3.0,
+ 0.0,
+ 0.0,
+ 3.0,
+ 4.0,
+ 0.0,
+ 0.0,
+ 4.0,
+ ]
+ assert torch.allclose(targets["exp_transform"], torch.tensor([-3.0, -2.0]))
+ assert verification["verification_candidates"] == 2
+ assert verification["verification_samples"] == 2
+ assert verification["verification_exact_mismatches"] == 0
+
+
+def test_xor_applies_through_packed_native_loaders() -> None:
+ targets = {
+ "row_slice": torch.zeros(2, 2),
+ "column_slice": torch.zeros(2, 4),
+ "split": torch.zeros(4, 2),
+ }
+ payload = encode_sparse_infos(
+ [
+ (
+ name,
+ torch.empty(shape),
+ torch.tensor(locations),
+ _bits(torch.tensor(values)).bitwise_xor(
+ _bits(torch.zeros(len(values)))
+ ),
+ "xor",
+ )
+ for name, shape, locations, values in _PACKED_LOADER_INFOS
+ ]
+ )
+
+ _apply_payload(_applier(_NativeLoaderModel(**targets)), payload)
+
+ assert targets["row_slice"].view(-1).tolist() == [1.0, 1.0, 0.0, 1.0]
+ assert targets["column_slice"].view(-1).tolist() == [2.0, 0.0, 0.0, 2.0] * 2
+ assert targets["split"].view(-1).tolist() == [
+ 3.0,
+ 0.0,
+ 0.0,
+ 3.0,
+ 4.0,
+ 0.0,
+ 0.0,
+ 4.0,
+ ]
+
+
+def test_native_loader_explicit_skip_is_accepted() -> None:
+ model = _NativeLoaderModel(identity=torch.zeros(1))
+ payload = _payload("skipped", torch.empty(1), [0], _bits(torch.tensor([1.0])))
+
+ _apply_payload(_applier(model), payload)
+
+ assert torch.equal(model.targets["identity"], torch.zeros(1))
+
+
+def test_native_loader_claim_without_copy_fails_closed() -> None:
+ model = _NativeLoaderModel(identity=torch.zeros(1))
+ model.load_weights = lambda weights: {name for name, _ in weights}
+ payload = _payload("weight", torch.empty(1), [0], _bits(torch.tensor([1.0])))
+
+ with pytest.raises(RuntimeError, match="without a supported target copy"):
+ _apply_payload(_applier(model), payload)
+
+
+def test_sparse_overwrite_preserves_unselected_transform_inputs() -> None:
+ target = torch.tensor([-2.0, -4.0, -6.0, -8.0])
+ payload = _payload(
+ "exp_transform",
+ target,
+ [1],
+ _bits(torch.tensor([math.log(3.0)])),
+ )
+
+ _apply_payload(_applier(_NativeLoaderModel(exp_transform=target)), payload)
+
+ assert torch.allclose(target, torch.tensor([-2.0, -3.0, -6.0, -8.0]))
+
+
+def test_unknown_sparse_operation_fails_closed() -> None:
+ target = torch.zeros(1)
+ payload = _payload("weight", target, [0], _bits(torch.tensor([1.0])))
+ payload[2][0]["operation"] = "unknown"
+
+ with pytest.raises(ValueError, match="Unsupported sparse-refit operation"):
+ _apply_payload(_applier(_NativeLoaderModel(identity=target)), payload)
+
+
+@pytest.mark.parametrize(
+ ("initial", "verified_value", "exact_mismatches", "mismatches"),
+ [(200.0, 4.0, 0, 0), (2.0, 4.0000005, 1, 0), (2.0, 5.0, 1, 1)],
+)
+def test_sparse_delta_verification_compares_replacement(
+ initial: float,
+ verified_value: float,
+ exact_mismatches: int,
+ mismatches: int,
+) -> None:
+ target = torch.tensor([1.0, initial, 3.0, initial])
+ replacement = torch.tensor([initial + 4.0, initial + 4.0])
+ payload = _payload("weight", target, [1, 3], _bits(replacement))
+ payload[2][0]["verification_samples"] = 2
+ applier = _applier(_NativeLoaderModel(identity=target))
+
+ _apply_payload(applier, payload)
+ target[[1, 3]] = initial + verified_value
+ result = applier.finish_sparse_delta_refit()
+
+ assert torch.equal(
+ target,
+ torch.tensor([1.0, initial + verified_value, 3.0, initial + verified_value]),
+ )
+ assert result["verification_candidates"] == 2
+ assert result["verification_samples"] == 2
+ assert result["verification_exact_mismatches"] == 2 * exact_mismatches
+ assert result["verification_mismatches"] == 2 * mismatches
+
+
+def test_fp8_weight_and_scale_use_exact_bit_overwrite() -> None:
+ target = torch.tensor([0x38, 0x40, 0x48], dtype=torch.uint8).view(
+ torch.float8_e4m3fn
+ )
+ scale = torch.tensor([1.0, 2.0])
+ current = target.clone()
+ current.view(torch.uint8)[1] = 0x41
+ current.view(torch.uint8)[2] = 0x7F
+ current_scale = scale.clone()
+ current_scale[0] = 1.5
+ payload = encode_sparse_infos(
+ [
+ (
+ "weight",
+ current,
+ torch.tensor([1, 2]),
+ current.view(torch.uint8)[1:3],
+ "overwrite",
+ ),
+ (
+ "weight_scale_inv",
+ current_scale,
+ torch.tensor([0]),
+ current_scale.view(torch.int32)[:1],
+ "overwrite",
+ ),
+ ]
+ )
+ payload[2][0]["verification_samples"] = 2
+ payload[2][1]["verification_samples"] = 1
+ applier = _applier(_NativeLoaderModel(identity=target, scale=scale))
+
+ _apply_payload(applier, payload)
+ _apply_payload(applier, payload)
+ result = applier.finish_sparse_delta_refit()
+
+ assert target.view(torch.uint8).tolist() == [0x38, 0x41, 0x7F]
+ assert torch.equal(scale, current_scale)
+ assert result["verification_samples"] == 6
+ assert result["verification_exact_mismatches"] == 0
+ assert result["verification_mismatches"] == 0
+ assert result["verification_abs_sum"] == 0.0
+
+
+def test_xor_applies_exact_bits_and_replay_reverts() -> None:
+ baseline = torch.tensor([1.0, 2.0, 3.0])
+ target = baseline.clone()
+ current = torch.tensor([1.0, 5.0, -0.0])
+ locations = torch.tensor([1, 2])
+ xor_values = _bits(current)[locations].bitwise_xor(_bits(baseline)[locations])
+ payload = encode_sparse_infos([("weight", current, locations, xor_values, "xor")])
+ payload[2][0]["verification_samples"] = 2
+ applier = _applier(_NativeLoaderModel(identity=target))
+
+ _apply_payload(applier, payload)
+ result = applier.finish_sparse_delta_refit()
+
+ assert torch.equal(_bits(target), _bits(current))
+ assert result["verification_exact_mismatches"] == 0
+ assert result["verification_mismatches"] == 0
+
+ _apply_payload(applier, payload)
+ assert torch.equal(_bits(target), _bits(baseline))
+
+
+def test_overwrite_casts_absolute_source_values() -> None:
+ target = torch.zeros(2, dtype=torch.float16)
+ source = torch.tensor([1.25, -2.5], dtype=torch.float32)
+ payload = _payload("weight", source, [0, 1], _bits(source))
+
+ _apply_payload(_applier(_NativeLoaderModel(identity=target)), payload)
+
+ assert torch.equal(target, source.to(torch.float16))
+
+
+@pytest.mark.parametrize(
+ ("name", "source", "targets", "error"),
+ [
+ (
+ "exp_transform",
+ torch.tensor([math.log(2.0)]),
+ {"exp_transform": torch.tensor([-1.0])},
+ "without changing semantics",
+ ),
+ (
+ "weight",
+ torch.tensor([1.0], dtype=torch.float32),
+ {"identity": torch.zeros(1, dtype=torch.float16)},
+ "without changing semantics",
+ ),
+ ],
+)
+def test_xor_rejects_non_bitwise_compatible_targets(
+ name: str,
+ source: torch.Tensor,
+ targets: dict[str, torch.Tensor],
+ error: str,
+) -> None:
+ payload = _payload(name, source, [0], _bits(source), "xor")
+
+ with pytest.raises(RuntimeError, match=error):
+ _apply_payload(_applier(_NativeLoaderModel(**targets)), payload)
+
+
+def test_xor_rejects_overlapping_native_loader_copies() -> None:
+ class RepeatedCopyModel(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.zeros(2))
+
+ def load_weights(self, weights) -> None:
+ for _, source in weights:
+ self.weight.copy_(source)
+ self.weight.copy_(source)
+
+ source = torch.tensor([1.0, 2.0])
+ payload = _payload(
+ "weight",
+ source,
+ [0, 1],
+ _bits(source).bitwise_xor(_bits(torch.zeros_like(source))),
+ "xor",
+ )
+
+ with pytest.raises(RuntimeError, match="without changing semantics"):
+ _apply_payload(_applier(RepeatedCopyModel()), payload)
diff --git a/tests/unit/models/generation/test_vllm_sparse_refit.py b/tests/unit/models/generation/test_vllm_sparse_refit.py
new file mode 100644
index 00000000000..1bee0dfa7ea
--- /dev/null
+++ b/tests/unit/models/generation/test_vllm_sparse_refit.py
@@ -0,0 +1,629 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import io
+import threading
+import time
+from collections.abc import Iterator
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import contextmanager
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, call
+
+import pytest
+import torch
+
+from nemo_rl.models.generation.vllm.vllm_sparse_refit import (
+ VllmSparseRefitReceiver,
+ _stage_sparse_payload,
+)
+from nemo_rl.models.generation.vllm.vllm_worker import VllmGenerationWorkerImpl
+from nemo_rl.models.generation.vllm.vllm_worker_async import (
+ VllmAsyncGenerationWorkerImpl,
+)
+from nemo_rl.utils.weight_transfer_http import (
+ G_VLLM_REFIT_API_KEY_HEADER,
+ G_VLLM_REFIT_FLUSH_PATH,
+ G_VLLM_REFIT_PREPARE_PATH,
+ G_VLLM_REFIT_S3_MANIFEST_PATH,
+ G_VLLM_REFIT_ZMQ_FLUSH_PATH,
+)
+from nemo_rl.utils.weight_transfer_sparse_codec import (
+ encode_sparse_infos,
+ sparse_locations_for_item,
+)
+
+
+@contextmanager
+def _sparse_refit_receiver(
+ *,
+ batch_size: int = 2,
+ futures: list[Future[dict[str, Any]]] | None = None,
+ async_engine: bool = False,
+ config: dict[str, Any] | None = None,
+) -> Iterator[VllmSparseRefitReceiver]:
+ owner = SimpleNamespace(
+ cfg=config or {"vllm_cfg": {"async_engine": async_engine}},
+ llm=MagicMock(),
+ )
+ owner.llm.collective_rpc.return_value = [{"ok": True}]
+ receiver = VllmSparseRefitReceiver(owner)
+ receiver._refit_apply_batch_size = batch_size
+ receiver._refit_apply_futures = list(futures or [])
+ for future in receiver._refit_apply_futures:
+ future.add_done_callback(receiver._notify_refit_apply_waiters)
+ try:
+ yield receiver
+ finally:
+ receiver._refit_apply_executor.shutdown(wait=True)
+ receiver._refit_partition_executor.shutdown(wait=True)
+
+
+def _serialized_sparse_payload() -> bytes:
+ values = torch.tensor([1, 2, 3, 4], dtype=torch.int32)
+ payload = encode_sparse_infos(
+ [
+ (
+ "weight",
+ torch.empty((8,), dtype=torch.float32),
+ torch.tensor([1, 3, 4, 7]),
+ values,
+ "overwrite",
+ )
+ ]
+ )
+ payload[2][0]["verification_samples"] = 2
+ buffer = io.BytesIO()
+ torch.save(payload, buffer)
+ return buffer.getvalue()
+
+
+def _payload_locations(path: str) -> list[int]:
+ packed_locations, _, items = torch.load(path, weights_only=True, mmap=True)
+ return [
+ int(location)
+ for item in items
+ for location in sparse_locations_for_item(item, packed_locations, device="cpu")
+ ]
+
+
+def _stage_payloads(
+ receiver: VllmSparseRefitReceiver,
+ staging_dir: Path,
+ *payloads: bytes,
+):
+ return tuple(
+ receiver._refit_partition_executor.submit(
+ _stage_sparse_payload, payload, str(staging_dir)
+ )
+ for payload in payloads
+ )
+
+
+def test_sparse_refit_queue_batches_payloads_in_fifo_order() -> None:
+ applied: list[tuple[bytes, ...]] = []
+
+ def apply(payloads: tuple[bytes, ...]) -> dict[str, Any]:
+ applied.append(payloads)
+ return {
+ "ok": True,
+ "payloads": len(payloads),
+ "receiver_total_s": float(len(payloads)),
+ }
+
+ with _sparse_refit_receiver(batch_size=3) as receiver:
+ receiver.update_weights_from_serialized_sparse_payloads = apply
+ responses = [
+ receiver._enqueue_sparse_payload_apply(
+ payload, ("transfer", 0, index), str(index)
+ )
+ for index, payload in enumerate((b"0", b"1", b"2", b"3", b"4"))
+ ]
+ response = receiver._flush_queued_sparse_payloads()
+ responses.append(response)
+
+ assert applied == [(b"0", b"1", b"2"), (b"3", b"4")]
+ assert response["payloads"] == 5
+ assert response["batches"] == 2
+ assert sum(result.get("receiver_total_s", 0.0) for result in responses) == 5.0
+ receiver._worker.llm.collective_rpc.assert_called_once_with(
+ "finish_sparse_delta_refit", args=()
+ )
+
+
+def test_sparse_refit_queue_stages_payload_before_batch_is_full(
+ tmp_path: Path,
+) -> None:
+ with _sparse_refit_receiver(batch_size=2) as receiver:
+ receiver._refit_workers_share_node = True
+ receiver._refit_batch_staging_dir = str(tmp_path)
+
+ response = receiver._enqueue_sparse_payload_apply(
+ _serialized_sparse_payload(), ("transfer", 0, 0), "checksum"
+ )
+ pending = receiver._refit_apply_pending_payloads[0]
+ assert isinstance(pending, Future)
+ staged = pending.result(timeout=1.0)
+
+ assert response["ok"]
+ assert Path(staged.path).is_file()
+ assert receiver._refit_apply_futures == []
+ receiver._flush_queued_sparse_payloads()
+
+ assert not list(tmp_path.iterdir())
+ assert receiver._worker.llm.collective_rpc.call_args_list[0].args[0] == (
+ "update_weights_from_decoded_sparse_payload"
+ )
+
+
+def test_sparse_refit_queue_deduplicates_transactional_payloads() -> None:
+ key = ("transfer", 0, 1)
+ with _sparse_refit_receiver() as receiver:
+ receiver.update_weights_from_serialized_sparse_payloads = MagicMock(
+ return_value={"ok": True, "payloads": 1}
+ )
+ assert receiver._enqueue_sparse_payload_apply(b"payload", key, "checksum", 4)[
+ "ok"
+ ]
+ duplicate = receiver._enqueue_sparse_payload_apply(
+ b"payload", key, "checksum", 4
+ )
+ assert duplicate == {"ok": True, "payloads": 0, "duplicate": True}
+ with pytest.raises(ValueError, match="reused with different data"):
+ receiver._enqueue_sparse_payload_apply(b"other", key, "different")
+ response = receiver._flush_queued_sparse_payloads()
+
+ assert response["payloads"] == 1
+ assert response["verification_candidates"] == 4
+ assert receiver._refit_seen_payloads == {}
+
+
+def test_sparse_refit_queue_does_not_deduplicate_failed_enqueue() -> None:
+ failed = Future()
+ failed.set_exception(RuntimeError("prior apply failed"))
+ with _sparse_refit_receiver(futures=[failed]) as receiver:
+ with pytest.raises(RuntimeError, match="prior apply failed"):
+ receiver._enqueue_sparse_payload_apply(
+ b"payload", ("transfer", 0, 1), "checksum"
+ )
+
+ assert receiver._refit_seen_payloads == {}
+ assert receiver._refit_apply_pending_payloads == []
+
+
+def test_sparse_refit_collective_response_merges_verification_metrics() -> None:
+ response = VllmSparseRefitReceiver._refit_collective_response(
+ [
+ {
+ "receiver_total_s": 1.0,
+ "verification_candidates": 4,
+ "verification_samples": 2,
+ "verification_exact_mismatches": 1,
+ "verification_mismatches": 0,
+ "verification_abs_sum": 0.25,
+ "verification_max_abs": 0.25,
+ },
+ {
+ "receiver_total_s": 2.0,
+ "verification_candidates": 4,
+ "verification_samples": 3,
+ "verification_exact_mismatches": 2,
+ "verification_mismatches": 1,
+ "verification_abs_sum": 0.5,
+ "verification_max_abs": 0.4,
+ },
+ ]
+ )
+
+ assert response == {
+ "ok": True,
+ "receiver_total_s": 2.0,
+ "verification_candidates": 4,
+ "verification_samples": 5,
+ "verification_exact_mismatches": 3,
+ "verification_mismatches": 1,
+ "verification_abs_sum": 0.75,
+ "verification_max_abs": 0.4,
+ }
+
+
+def test_sparse_refit_queue_releases_condition_while_backpressured() -> None:
+ first: Future[dict[str, Any]] = Future()
+ second: Future[dict[str, Any]] = Future()
+ started = threading.Event()
+
+ with _sparse_refit_receiver(futures=[first, second]) as receiver:
+ with ThreadPoolExecutor(max_workers=1) as callers:
+ pending_call = callers.submit(
+ lambda: (
+ started.set(),
+ receiver._enqueue_sparse_payload_apply(
+ b"payload", ("transfer", 0, 1), "checksum"
+ ),
+ )[1]
+ )
+ assert started.wait(timeout=1.0)
+ time.sleep(0.05)
+ assert receiver._refit_apply_queue_condition.acquire(timeout=1.0)
+ receiver._refit_apply_queue_condition.release()
+ first.set_result({"ok": True, "payloads": 1})
+ assert pending_call.result(timeout=1.0)["ok"]
+
+
+def test_sparse_refit_batch_stages_compact_payload_before_collective_apply(
+ tmp_path: Path,
+) -> None:
+ with _sparse_refit_receiver() as receiver:
+ staged_locations: list[int] = []
+
+ def collective_rpc(method: str, args: tuple[Any, ...]) -> list[Any]:
+ assert method == "update_weights_from_decoded_sparse_payload"
+ for path in args:
+ staged_locations.extend(_payload_locations(path))
+ return [
+ {"ok": True, "receiver_total_s": 1.0},
+ {"ok": True, "receiver_total_s": 1.0},
+ {"ok": True, "receiver_total_s": 1.0},
+ ]
+
+ receiver._worker.llm = MagicMock(
+ collective_rpc=MagicMock(side_effect=collective_rpc)
+ )
+ response = receiver.update_weights_from_staged_sparse_payloads(
+ _stage_payloads(receiver, tmp_path, _serialized_sparse_payload())
+ )
+ receiver.update_weights_from_staged_sparse_payloads(
+ _stage_payloads(receiver, tmp_path, _serialized_sparse_payload())
+ )
+
+ assert staged_locations == [1, 3, 4, 7] * 2
+ assert not list(tmp_path.iterdir())
+ assert [
+ item.args[0] for item in receiver._worker.llm.collective_rpc.call_args_list
+ ] == [
+ "update_weights_from_decoded_sparse_payload",
+ "update_weights_from_decoded_sparse_payload",
+ ]
+ assert response["payloads"] == 1
+ assert response["receiver_worker_total_s"] == 1.0
+ assert response["receiver_total_s"] >= 0.0
+ assert receiver._refit_verification_candidates == 0
+
+
+def test_sparse_refit_batch_drains_workers_before_error_cleanup(tmp_path: Path) -> None:
+ with _sparse_refit_receiver() as receiver:
+ staged_paths: tuple[str, ...] = ()
+
+ def collective_rpc(method: str, args: tuple[Any, ...]) -> list[Any]:
+ nonlocal staged_paths
+ if method == "update_weights_from_decoded_sparse_payload":
+ staged_paths = args
+ raise RuntimeError("apply failed")
+ assert method == "synchronize_device"
+ assert all(Path(path).is_file() for path in staged_paths)
+ return [True]
+
+ receiver._worker.llm = MagicMock(
+ collective_rpc=MagicMock(side_effect=collective_rpc)
+ )
+ with pytest.raises(RuntimeError, match="apply failed"):
+ receiver.update_weights_from_staged_sparse_payloads(
+ _stage_payloads(receiver, tmp_path, _serialized_sparse_payload())
+ )
+
+ assert [
+ entry.args[0]
+ for entry in receiver._worker.llm.collective_rpc.call_args_list
+ ] == [
+ "update_weights_from_decoded_sparse_payload",
+ "synchronize_device",
+ ]
+ assert not list(tmp_path.iterdir())
+
+
+def test_sparse_refit_batch_uses_one_collective_rpc_across_nodes() -> None:
+ with _sparse_refit_receiver() as receiver:
+ receiver._refit_workers_share_node = False
+ receiver._worker.llm = MagicMock(
+ collective_rpc=MagicMock(
+ return_value=[{"ok": True, "receiver_total_s": 1.0}]
+ )
+ )
+
+ response = receiver.update_weights_from_serialized_sparse_payloads(
+ (_serialized_sparse_payload(),) * 3
+ )
+
+ rpc = receiver._worker.llm.collective_rpc.call_args
+ assert rpc.args[0] == "update_weights_from_decoded_sparse_payload"
+ assert len(rpc.kwargs["args"]) == 3
+ assert all(
+ len(torch.load(io.BytesIO(payload), weights_only=True)[2]) == 1
+ for payload in rpc.kwargs["args"]
+ )
+ assert response == {"ok": True, "receiver_total_s": 1.0, "payloads": 3}
+
+
+@pytest.mark.asyncio
+async def test_async_sparse_refit_batch_bridges_to_async_collective(
+ tmp_path: Path,
+) -> None:
+ with _sparse_refit_receiver(async_engine=True) as receiver:
+ staged_locations: list[int] = []
+
+ class AsyncLlm:
+ async def collective_rpc(
+ self, method: str, args: tuple[Any, ...]
+ ) -> list[Any]:
+ assert method == "update_weights_from_decoded_sparse_payload"
+ for path in args:
+ staged_locations.extend(_payload_locations(path))
+ return [{"ok": True, "receiver_total_s": 1.0}]
+
+ receiver._worker.llm = AsyncLlm()
+ receiver._refit_async_loop = asyncio.get_running_loop()
+ response = await asyncio.to_thread(
+ receiver.update_weights_from_staged_sparse_payloads,
+ _stage_payloads(receiver, tmp_path, _serialized_sparse_payload()),
+ )
+
+ assert staged_locations == [1, 3, 4, 7]
+ assert response["payloads"] == 1
+ assert response["receiver_worker_total_s"] == 1.0
+
+
+@pytest.mark.asyncio
+async def test_sparse_refit_payload_handlers_decode_and_enqueue(monkeypatch) -> None:
+ with _sparse_refit_receiver() as receiver:
+ enqueue = MagicMock(
+ side_effect=[
+ {"ok": True, "payloads": 1},
+ {"ok": True, "payloads": 1},
+ ]
+ )
+ receiver._enqueue_sparse_payload_apply = enqueue
+ monkeypatch.setattr(
+ "nemo_rl.models.generation.vllm.vllm_sparse_refit.download_s3_refit_payload",
+ lambda _manifest: b"s3-payload",
+ )
+ monkeypatch.setattr(
+ "nemo_rl.models.generation.vllm.vllm_sparse_refit.decode_sparse_payload",
+ lambda _body, _checksum: b"zmq-payload",
+ )
+
+ s3_result = await receiver._apply_s3_manifest_payload(
+ {
+ "key": "object-key",
+ "checksum": "checksum",
+ "verification_candidates": 4,
+ }
+ )
+ assert s3_result["ok"]
+
+ zmq_result = receiver._apply_zmq_payload(
+ b"compressed",
+ {
+ "transfer_id": "transfer",
+ "producer_id": 2,
+ "payload_id": 3,
+ "checksum": "checksum",
+ "verification_candidates": 5,
+ },
+ )
+ assert zmq_result["ok"]
+ assert enqueue.call_args_list == [
+ call(b"s3-payload", ("object-key", -1, -1), "checksum", 4),
+ call(b"zmq-payload", ("transfer", 2, 3), "checksum", 5),
+ ]
+
+
+def test_sparse_refit_api_auth_dispatch_and_error_mapping(monkeypatch) -> None:
+ from fastapi import FastAPI
+ from fastapi.testclient import TestClient
+
+ monkeypatch.setenv("NRL_TEST_REFIT_KEY", "secret")
+ config = {
+ "vllm_cfg": {
+ "async_engine": True,
+ "http_refit_api_key_env_var": "NRL_TEST_REFIT_KEY",
+ }
+ }
+ with _sparse_refit_receiver(async_engine=True, config=config) as receiver:
+ receiver._apply_s3_manifest_payload = AsyncMock(
+ return_value={"ok": True, "payloads": 1}
+ )
+ receiver._flush_queued_sparse_payloads = MagicMock(
+ return_value={"ok": True, "payloads": 2}
+ )
+ receiver.flush_zmq_sparse_refit_relay = MagicMock(
+ return_value={"ok": True, "payloads": 3}
+ )
+ receiver._refit_collective_rpc = MagicMock(return_value=[])
+ app = FastAPI()
+ receiver.setup_api_server(app)
+ headers = {G_VLLM_REFIT_API_KEY_HEADER: "secret"}
+
+ with TestClient(app) as client:
+ unauthorized = client.post(G_VLLM_REFIT_S3_MANIFEST_PATH, json={})
+ s3_response = client.post(
+ G_VLLM_REFIT_S3_MANIFEST_PATH,
+ json={"key": "key"},
+ headers=headers,
+ )
+ flush_response = client.post(G_VLLM_REFIT_FLUSH_PATH, headers=headers)
+ zmq_flush_response = client.post(
+ G_VLLM_REFIT_ZMQ_FLUSH_PATH,
+ json={"transfer_id": "transfer"},
+ headers=headers,
+ )
+ prepare_response = client.post(
+ G_VLLM_REFIT_PREPARE_PATH,
+ json={"tensors": {"weight": [[2, 3], "bfloat16"]}},
+ headers=headers,
+ )
+ assert unauthorized.status_code == 403
+ assert s3_response.status_code == 200
+ assert flush_response.status_code == 200
+ assert zmq_flush_response.status_code == 200
+ assert prepare_response.status_code == 200
+ assert receiver._refit_async_loop is not None
+ receiver._apply_s3_manifest_payload.assert_awaited_once_with({"key": "key"})
+ receiver._flush_queued_sparse_payloads.assert_called_once_with()
+ receiver.flush_zmq_sparse_refit_relay.assert_called_once_with("transfer", 0)
+ receiver._refit_collective_rpc.assert_called_once_with(
+ "prepare_sparse_delta_refit_info",
+ ({"weight": ((2, 3), torch.bfloat16)},),
+ )
+
+
+def test_sync_sparse_refit_server_shutdown_cleans_transport_resources(
+ monkeypatch, caplog
+) -> None:
+ import uvicorn
+
+ from nemo_rl.models.generation.vllm import vllm_sparse_refit as refit_module
+
+ configs: list[Any] = []
+ servers: list[Any] = []
+
+ def make_config(app: Any, **kwargs: Any) -> Any:
+ config = SimpleNamespace(app=app, **kwargs)
+ configs.append(config)
+ return config
+
+ class Server:
+ def __init__(self, config: Any) -> None:
+ self.config = config
+ self.should_exit = False
+ self.ran = threading.Event()
+ servers.append(self)
+
+ def run(self) -> None:
+ self.ran.set()
+
+ monkeypatch.setattr(uvicorn, "Config", make_config)
+ monkeypatch.setattr(uvicorn, "Server", Server)
+ monkeypatch.setattr(refit_module, "_get_free_port_local", lambda *_args: 12345)
+ monkeypatch.setattr(refit_module, "_get_node_ip_local", lambda: "10.0.0.1")
+ refit_module._warn_unauthenticated_refit_server.cache_clear()
+ config = {
+ "vllm_cfg": {"async_engine": False, "http_refit_server_port": None},
+ "port_range_low": 10000,
+ "port_range_high": 11000,
+ }
+
+ with _sparse_refit_receiver(config=config) as receiver:
+ receiver._worker.base_url = "http://10.0.0.2:8000/v1"
+ assert receiver.report_refit_server_base_url() == "http://10.0.0.2:8000"
+ receiver.setup_api_server = MagicMock()
+ receiver._setup_vllm_refit_server()
+ assert configs[0].host == "0.0.0.0"
+ assert configs[0].port == 12345
+ assert servers[0].ran.wait(timeout=1.0)
+ assert receiver.report_refit_server_base_url() == "http://10.0.0.1:12345"
+ assert "binding 0.0.0.0 without an API key" in caplog.text
+
+ relay = MagicMock()
+ receiver._zmq_refit_server = (relay, "tcp://relay")
+ receiver._flush_queued_sparse_payloads = MagicMock()
+ receiver._refit_apply_executor = MagicMock()
+ receiver.shutdown()
+
+ relay.close.assert_called_once_with()
+ receiver._flush_queued_sparse_payloads.assert_called_once_with()
+ receiver._refit_apply_executor.shutdown.assert_called_once_with(wait=True)
+ assert servers[0].should_exit is True
+ assert receiver._refit_http_server is None
+
+
+@pytest.mark.asyncio
+async def test_async_sparse_refit_post_init_records_worker_locality() -> None:
+ worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl)
+ worker._sparse_refit_receiver = MagicMock()
+ worker._mtp_load_from_disk = False
+ worker.report_device_id_async = AsyncMock(return_value=["0"])
+ worker.llm = MagicMock()
+ worker.llm.collective_rpc = AsyncMock(return_value=["node-0", "node-0"])
+
+ await worker.post_init_async()
+
+ assert worker.vllm_device_ids == ["0"]
+ worker._sparse_refit_receiver.set_worker_hostnames.assert_called_once_with(
+ ["node-0", "node-0"]
+ )
+ assert worker.llm.collective_rpc.await_args_list == [
+ call("bind_numa", args=()),
+ call("report_node_hostname", args=()),
+ ]
+
+
+def test_async_sparse_refit_exposes_zmq_relay(monkeypatch) -> None:
+ from nemo_rl.models.generation.vllm import vllm_sparse_refit as refit_module
+
+ server = MagicMock()
+ server_type = MagicMock(return_value=server)
+ monkeypatch.setattr(refit_module, "ZmqSparseRefitServer", server_type)
+ monkeypatch.setattr(refit_module, "_get_free_port_local", lambda *_args: 12345)
+ monkeypatch.setattr(refit_module, "_get_node_ip_local", lambda: "10.0.0.1")
+ config = {
+ "vllm_cfg": {
+ "async_engine": True,
+ "zmq_refit_server_port": None,
+ "http_refit_api_key_env_var": None,
+ }
+ }
+
+ with _sparse_refit_receiver(async_engine=True, config=config) as receiver:
+ receiver._worker.base_url = "http://10.0.0.1:8000/v1"
+ assert receiver.start_zmq_sparse_refit_relay() == "tcp://10.0.0.1:12345"
+ server_type.assert_called_once_with(
+ receiver._apply_zmq_payload,
+ bind_address="tcp://0.0.0.0:12345",
+ api_key_env_var=None,
+ timeout_s=600.0,
+ tuning=receiver._refit_config.tuning,
+ )
+ server.start.assert_called_once_with()
+ receiver.configure_zmq_sparse_refit_relay(
+ ["tcp://10.0.0.1:12345", "tcp://10.0.0.2:12345"]
+ )
+ server.configure_tree.assert_called_once_with(
+ ["tcp://10.0.0.1:12345", "tcp://10.0.0.2:12345"],
+ own_address="tcp://10.0.0.1:12345",
+ )
+ server.flush.return_value = {"ok": True, "payloads": 2}
+ assert receiver.flush_zmq_sparse_refit_relay("transfer") == {
+ "ok": True,
+ "payloads": 2,
+ }
+ server.flush.assert_called_once_with("transfer", 0)
+
+ receiver.stop_zmq_sparse_refit_relay()
+ server.close.assert_called_once_with()
+ assert receiver._zmq_refit_server is None
+
+
+def test_vllm_worker_configures_zmq_relay() -> None:
+ worker = VllmGenerationWorkerImpl.__new__(VllmGenerationWorkerImpl)
+ worker._sparse_refit_receiver = MagicMock()
+ addresses = ["tcp://relay-0:19090", "tcp://relay-1:19090"]
+
+ worker.configure_zmq_sparse_refit_relay(addresses)
+
+ worker._sparse_refit_receiver.configure_zmq_sparse_refit_relay.assert_called_once_with(
+ addresses
+ )
diff --git a/tests/unit/models/policy/test_megatron_remote_sparse_refit.py b/tests/unit/models/policy/test_megatron_remote_sparse_refit.py
new file mode 100644
index 00000000000..8b1b697cb9b
--- /dev/null
+++ b/tests/unit/models/policy/test_megatron_remote_sparse_refit.py
@@ -0,0 +1,133 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import torch
+
+from nemo_rl.models.generation.vllm.config import VllmRefitConfig
+from nemo_rl.models.policy.workers.megatron_remote_sparse_refit import (
+ MegatronRemoteSparseRefit,
+)
+
+_REFIT_CONFIG = VllmRefitConfig(
+ delta_compression={
+ "encoding": "overwrite",
+ "sparse_bucket_size_bytes": 1024,
+ },
+ baseline={"in_memory": True},
+)
+
+
+def _worker(weights=()):
+ def export():
+ return iter(weights)
+
+ return SimpleNamespace(_iter_params_with_optional_kv_scales=export)
+
+
+def test_remote_sparse_initializes_canonical_hf_baseline() -> None:
+ weights = [
+ ("embedding.weight", torch.ones(2, 3)),
+ ("linear.weight", torch.ones(4, 3)),
+ ]
+ remote_refit = MegatronRemoteSparseRefit(_worker(weights), _REFIT_CONFIG)
+
+ info = remote_refit.initialize_baseline(
+ shard_rank=0, shard_count=1, transport="zmq"
+ )
+
+ assert info == {
+ name: (tuple(tensor.shape), tensor.dtype) for name, tensor in weights
+ }
+ assert set(vars(remote_refit)) == {"_worker", "_tracker"}
+
+
+def test_remote_sparse_preserves_xor_config() -> None:
+ remote_refit = MegatronRemoteSparseRefit(
+ _worker(),
+ VllmRefitConfig(
+ delta_compression={
+ "encoding": "xor",
+ "sparse_bucket_size_bytes": 1024,
+ },
+ baseline={"in_memory": True},
+ ),
+ )
+
+ assert remote_refit._tracker.encoding == "xor"
+
+
+def test_remote_sparse_streams_one_canonical_path_and_drains_cuda(monkeypatch) -> None:
+ weights = [("model.weight", torch.ones(2))]
+ remote_refit = MegatronRemoteSparseRefit(_worker(weights), _REFIT_CONFIG)
+ expected = {"payloads": 1, "changed_elements": 2, "total_elements": 2}
+ events = []
+
+ def stream(iterator, **kwargs):
+ assert list(iterator) == weights
+ assert kwargs == {
+ "delta_tracker": remote_refit._tracker,
+ "transfer_id": "transfer",
+ "refit_targets": ["tcp://receiver:5555"],
+ "api_key_env_var": None,
+ "timeout_s": 1.0,
+ "shard_rank": 0,
+ "shard_count": 1,
+ }
+ events.append("stream")
+ return expected
+
+ monkeypatch.setattr(
+ "nemo_rl.models.policy.workers.megatron_remote_sparse_refit."
+ "stream_sparse_delta_payloads_via_zmq",
+ stream,
+ )
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "synchronize", lambda: events.append("sync"))
+
+ result = remote_refit.stream(
+ "zmq",
+ ["tcp://receiver:5555"],
+ transfer_id="transfer",
+ api_key_env_var=None,
+ timeout_s=1.0,
+ shard_rank=0,
+ shard_count=1,
+ overwrite_names=["model.weight"],
+ )
+
+ assert result is expected
+ assert remote_refit._tracker.overwrite_names == frozenset({"model.weight"})
+ assert events == ["stream", "sync"]
+
+
+def test_remote_sparse_finishes_single_tracker(monkeypatch) -> None:
+ remote_refit = MegatronRemoteSparseRefit(_worker(), _REFIT_CONFIG)
+ events = []
+ monkeypatch.setattr(
+ remote_refit._tracker,
+ "on_sync_succeeded",
+ lambda: events.append("succeeded"),
+ )
+ monkeypatch.setattr(
+ remote_refit._tracker,
+ "on_sync_failed",
+ lambda: events.append("failed"),
+ )
+
+ remote_refit.finish(True)
+ remote_refit.finish(False)
+
+ assert events == ["succeeded", "failed"]
diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml
index 55cdc1d4247..39aa2459723 100644
--- a/tests/unit/reference_configs/grpo_math_1B.yaml
+++ b/tests/unit/reference_configs/grpo_math_1B.yaml
@@ -332,6 +332,8 @@ policy:
top_k: null
stop_token_ids: null
stop_strings: null
+ refit_transport: null # Set to "vllm_s3_sparse" or "vllm_zmq_sparse" for remote sparse-delta refit.
+ refit_cfg: null # Optional tuning and storage settings for remote sparse-delta refit.
mcore_generation_config:
async_engine: false
max_model_len: ${policy.max_total_sequence_length}
@@ -369,6 +371,9 @@ policy:
num_first_layers_in_bf16: 0
enable_vllm_metrics_logger: true # Set to true to enable vLLM internal metrics logger, turn off for better performance
vllm_metrics_logger_interval: 0.5 # Interval in seconds to collect vLLM logger metrics
+ http_refit_api_key_env_var: null # Optional env var containing the internal refit API key.
+ http_refit_server_port: null # Optional fixed port for Kubernetes targetPorts.
+ zmq_refit_server_port: null # Optional fixed ZeroMQ relay port for Kubernetes targetPorts.
vllm_kwargs: {}
colocated:
# true: generation shares training GPUs
diff --git a/tests/unit/test_recipes_and_test_suites.py b/tests/unit/test_recipes_and_test_suites.py
index 3e7f01365cc..968a90265a1 100644
--- a/tests/unit/test_recipes_and_test_suites.py
+++ b/tests/unit/test_recipes_and_test_suites.py
@@ -255,7 +255,7 @@ def test_all_recipe_yamls_accounted_for_in_test_suites(
)
-def test_nightly_compute_stays_below_3390_hours(nightly_test_suite, tracker):
+def test_nightly_compute_stays_below_3420_hours(nightly_test_suite, tracker):
command = f"DRYRUN=1 HF_HOME=... HF_DATASETS_CACHE=... CONTAINER= ACCOUNT= PARTITION= ./tools/launch {' '.join(nightly_test_suite)}"
print(f"Running command: {command}")
@@ -287,8 +287,8 @@ def test_nightly_compute_stays_below_3390_hours(nightly_test_suite, tracker):
f"Last line of output was not as expected: '{last_line}'"
)
total_gpu_hours = float(last_line.split(":")[-1].strip())
- assert total_gpu_hours <= 3390, (
- f"Total GPU hours exceeded 3390: {last_line}. We should revisit the test suites to reduce the total GPU hours."
+ assert total_gpu_hours <= 3420, (
+ f"Total GPU hours exceeded 3420: {last_line}. We should revisit the test suites to reduce the total GPU hours."
)
tracker.track("total_nightly_gpu_hours", total_gpu_hours)
diff --git a/tests/unit/utils/test_weight_transfer_stream.py b/tests/unit/utils/test_weight_transfer_stream.py
new file mode 100644
index 00000000000..4ab661fc657
--- /dev/null
+++ b/tests/unit/utils/test_weight_transfer_stream.py
@@ -0,0 +1,917 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import io
+import json
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import requests
+import torch
+import zstandard
+
+from nemo_rl.models.generation.vllm.config import VllmRefitConfig
+from nemo_rl.utils import (
+ weight_transfer_http,
+ weight_transfer_stream,
+ weight_transfer_zmq,
+)
+from nemo_rl.utils.weight_transfer_sparse_codec import (
+ DeltaCompressionTracker,
+ encode_sparse_infos,
+ sparse_locations_for_item,
+)
+from nemo_rl.utils.weight_transfer_stream import (
+ SparseRefitTransport,
+ download_s3_refit_payload,
+ sparse_export_chunk_size,
+ sparse_payload_checksum,
+)
+from nemo_rl.utils.weight_transfer_zmq import (
+ ZmqSparseRefitClient,
+ ZmqSparseRefitServer,
+)
+
+
+def _refit_config(
+ *,
+ encoding: str = "overwrite",
+ bucket_bytes: int = 1024,
+ verify_samples: int = 0,
+ s3_bucket: str | None = None,
+ s3_region: str = "us-east-1",
+ s3_prefix: str = "nemo-rl-refit",
+ s3_export_bytes: int = 64 * 1024**2,
+ zmq_export_bytes: int = 256 * 1024**2,
+ s3_encode_workers: int = 8,
+ zmq_encode_workers: int = 8,
+ s3_transfer_workers: int = 32,
+ zmq_transfer_workers: int = 4,
+ zmq_retries: int = 3,
+) -> VllmRefitConfig:
+ return VllmRefitConfig(
+ delta_compression={
+ "encoding": encoding,
+ "sparse_bucket_size_bytes": bucket_bytes,
+ "export_chunk_bytes": {
+ "s3": s3_export_bytes,
+ "zmq": zmq_export_bytes,
+ },
+ },
+ storage={
+ "s3_bucket": s3_bucket,
+ "s3_region": s3_region,
+ "s3_prefix": s3_prefix,
+ },
+ baseline={"in_memory": True},
+ tuning={
+ "encode_workers": {
+ "s3": s3_encode_workers,
+ "zmq": zmq_encode_workers,
+ },
+ "transfer_workers": {
+ "s3": s3_transfer_workers,
+ "zmq": zmq_transfer_workers,
+ },
+ "zmq_retries": zmq_retries,
+ },
+ verify_samples_per_payload=verify_samples,
+ )
+
+
+class _SparsePipelineTracker:
+ sparse_bucket_size_bytes = 1
+ refit_config = _refit_config(
+ bucket_bytes=1,
+ zmq_export_bytes=1,
+ zmq_encode_workers=1,
+ )
+
+ @staticmethod
+ def prepare_sparse_delta_payload(chunk):
+ count = sum(tensor.numel() for _, tensor in chunk)
+ payload = encode_sparse_infos(
+ (
+ (
+ name,
+ tensor,
+ torch.arange(tensor.numel()),
+ tensor.reshape(-1),
+ "overwrite",
+ )
+ for name, tensor in chunk
+ )
+ )
+ return payload, count, count
+
+
+class _BaselineNamesTracker:
+ sparse_bucket_size_bytes = 4
+ refit_config = _refit_config(bucket_bytes=4, zmq_export_bytes=4)
+
+ def __init__(self) -> None:
+ self.names = []
+
+ def snapshot_baseline(self, chunk) -> None:
+ self.names.extend(name for name, _tensor in chunk)
+
+
+def _stream_sparse_test_payloads(tensors, send_payload):
+ cleaned = []
+ transport = SparseRefitTransport(
+ name="zmq",
+ transfer_workers=1,
+ send=lambda body, payload_id, _candidates: send_payload(body, payload_id),
+ cleanup=lambda: cleaned.append(True),
+ )
+ try:
+ return weight_transfer_stream.stream_sparse_delta_payloads(
+ tensors,
+ delta_tracker=_SparsePipelineTracker(),
+ transport=transport,
+ shard_rank=0,
+ shard_count=1,
+ )
+ finally:
+ assert cleaned
+
+
+def _delta_tracker(
+ encoding: str = "overwrite", **config: Any
+) -> DeltaCompressionTracker:
+ return DeltaCompressionTracker(_refit_config(encoding=encoding, **config))
+
+
+def _baseline_names(tensors, *, rank: int):
+ tracker = _BaselineNamesTracker()
+ weight_transfer_stream.init_sparse_delta_baseline_from_iterator(
+ tensors,
+ delta_tracker=tracker,
+ shard_rank=rank,
+ shard_count=2,
+ transport="zmq",
+ )
+ return tracker.names
+
+
+def test_delta_tracker_commits_only_successful_syncs() -> None:
+ tracker = _delta_tracker()
+ tensor = torch.tensor([1.0, 2.0, 3.0])
+ tracker.snapshot_baseline([("weight", tensor)])
+ tensor[1] += 4
+
+ assert tracker.prepare_sparse_delta_payload([("weight", tensor)])[0][2]
+ tracker.on_sync_failed()
+ assert tracker.prepare_sparse_delta_payload([("weight", tensor)])[0][2]
+ tracker.on_sync_succeeded()
+ assert not tracker.prepare_sparse_delta_payload([("weight", tensor)])[0][2]
+
+
+def test_delta_tracker_emits_bounded_verification_budget() -> None:
+ tracker = _delta_tracker(verify_samples=2)
+ tensor = torch.tensor([1.0, 2.0, 3.0, 4.0])
+ tracker.snapshot_baseline([("weight", tensor)])
+ tensor[[1, 3]] += 1
+
+ (_, _, metadata), changed, total = tracker.prepare_sparse_delta_payload(
+ [("weight", tensor)]
+ )
+
+ assert metadata[0]["verification_samples"] == 2
+ assert metadata[0]["operation"] == "overwrite"
+ assert (changed, total) == (2, 4)
+
+
+def test_delta_tracker_commits_exact_source_baseline() -> None:
+ tracker = _delta_tracker()
+ tensor = torch.tensor([1.0])
+ tracker.snapshot_baseline([("weight", tensor)])
+ tensor.add_(0.001)
+
+ (_, value_groups, _), _, _ = tracker.prepare_sparse_delta_payload(
+ [("weight", tensor)]
+ )
+ assert torch.equal(value_groups[0], tensor.view(torch.int32))
+ tracker.on_sync_succeeded()
+ tracker.prepare_sparse_delta_payload([("weight", tensor)])
+
+ assert torch.equal(tracker.baseline["weight"], tensor)
+
+
+def test_delta_tracker_xor_encodes_against_baseline() -> None:
+ tracker = _delta_tracker("xor")
+ tensor = torch.tensor([1.0, 2.0, 3.0])
+ baseline = tensor.clone()
+ tracker.snapshot_baseline([("weight", tensor)])
+ tensor[[0, 2]] = torch.tensor([4.0, 5.0])
+
+ (_, value_groups, metadata), changed, total = tracker.prepare_sparse_delta_payload(
+ [("weight", tensor)]
+ )
+ locations = torch.tensor([0, 2])
+ expected = tensor.view(torch.int32)[locations].bitwise_xor(
+ baseline.view(torch.int32)[locations]
+ )
+
+ assert torch.equal(value_groups[0], expected)
+ assert metadata[0]["operation"] == "xor"
+ assert (changed, total) == (2, 3)
+ tracker.on_sync_succeeded()
+ tracker.prepare_sparse_delta_payload([("weight", tensor)])
+ assert torch.equal(tracker.baseline["weight"], tensor)
+
+
+def test_delta_tracker_uses_overwrite_for_receiver_incompatible_weights() -> None:
+ tracker = _delta_tracker("xor")
+ weight = torch.tensor([1.0])
+ scale = torch.tensor([2.0])
+ tracker.snapshot_baseline([("weight", weight), ("scale", scale)])
+ tracker.overwrite_names = frozenset({"scale"})
+ weight.add_(1)
+ scale.add_(1)
+
+ (_, _, metadata), _, _ = tracker.prepare_sparse_delta_payload(
+ [("weight", weight), ("scale", scale)]
+ )
+
+ assert [item["operation"] for item in metadata] == ["xor", "overwrite"]
+
+
+def test_sparse_index_encoding_preserves_uint64_locations() -> None:
+ locations = torch.tensor([0, 2**32 + 5])
+ packed, _, metadata = encode_sparse_infos(
+ [
+ (
+ "weight",
+ torch.empty(2),
+ locations,
+ torch.ones(2, dtype=torch.int32),
+ "overwrite",
+ )
+ ],
+ )
+
+ decoded = sparse_locations_for_item(metadata[0], packed, device="cpu")
+ assert torch.equal(decoded, locations)
+
+
+def test_sparse_index_encoding_preserves_uint32_locations() -> None:
+ locations = torch.tensor([0, 2**16 + 1])
+ packed, _, metadata = encode_sparse_infos(
+ [
+ (
+ "weight",
+ torch.empty(2),
+ locations,
+ torch.ones(2, dtype=torch.int32),
+ "overwrite",
+ )
+ ],
+ )
+
+ assert metadata[0]["index_encoding"] == "deltas"
+ assert packed.numel() == 2 * 4
+ decoded = sparse_locations_for_item(metadata[0], packed, device="cpu")
+ assert torch.equal(decoded, locations)
+
+
+@pytest.mark.parametrize("encoding", ["xor", "overwrite"])
+def test_delta_tracker_encodes_fp8_weight_and_scale_bits(encoding: str) -> None:
+ tracker = _delta_tracker(encoding, verify_samples=2)
+ weight = torch.tensor([0x38, 0x40, 0x48], dtype=torch.uint8).view(
+ torch.float8_e4m3fn
+ )
+ scale = torch.tensor([1.0, 2.0], dtype=torch.float32)
+ tracker.snapshot_baseline([("weight", weight), ("weight_scale_inv", scale)])
+ weight.view(torch.uint8)[1] = 0x41
+ scale[0] = 1.5
+
+ (locations, value_groups, metadata), changed, total = (
+ tracker.prepare_sparse_delta_payload(
+ [("weight", weight), ("weight_scale_inv", scale)]
+ )
+ )
+
+ assert (changed, total) == (2, 5)
+ assert len(value_groups) == 2
+ assert [item["operation"] for item in metadata] == [encoding, encoding]
+ assert [item["dtype"] for item in metadata] == ["float8_e4m3fn", "float32"]
+ assert [item["verification_samples"] for item in metadata] == [1, 1]
+ assert [
+ sparse_locations_for_item(item, locations, device="cpu").tolist()
+ for item in metadata
+ ] == [
+ [1],
+ [0],
+ ]
+
+ tracker.on_sync_succeeded()
+ assert not tracker.prepare_sparse_delta_payload(
+ [("weight", weight), ("weight_scale_inv", scale)]
+ )[0][2]
+
+
+def test_refit_config_rejects_arithmetic_encoding() -> None:
+ with pytest.raises(ValueError, match="Input should be 'xor' or 'overwrite'"):
+ _delta_tracker("add")
+
+
+def test_s3_download_verifies_checksum(monkeypatch) -> None:
+ compressed = zstandard.ZstdCompressor().compress(b"payload")
+ monkeypatch.setattr(
+ weight_transfer_stream,
+ "_get_manifest_s3_store",
+ lambda *_args: SimpleNamespace(get=lambda _key: bytearray(compressed)),
+ )
+ manifest = {
+ "bucket": "bucket",
+ "region": "region",
+ "key": "key",
+ "checksum": sparse_payload_checksum(compressed),
+ }
+
+ assert download_s3_refit_payload(manifest) == b"payload"
+ manifest["checksum"] = "0" * 32
+ with pytest.raises(ValueError, match="checksum mismatch"):
+ download_s3_refit_payload(manifest)
+
+
+def test_refit_http_session_does_not_retry_application_errors() -> None:
+ retry = weight_transfer_http.refit_http_session().get_adapter("http://").max_retries
+
+ assert 500 not in retry.status_forcelist
+ assert {502, 503, 504} <= set(retry.status_forcelist)
+
+
+def test_refit_http_sessions_share_connection_pool_across_threads() -> None:
+ barrier = threading.Barrier(4)
+
+ def adapter(_):
+ barrier.wait()
+ return weight_transfer_http.refit_http_session().get_adapter("http://")
+
+ with ThreadPoolExecutor(max_workers=4) as executor:
+ adapters = list(executor.map(adapter, range(4)))
+
+ assert all(adapter is adapters[0] for adapter in adapters)
+
+
+def test_refit_http_error_preserves_non_json_status_and_body(monkeypatch) -> None:
+ response = requests.Response()
+ response.status_code = 500
+ response._content = b"gateway failure"
+ session = SimpleNamespace(post=lambda *_args, **_kwargs: response)
+ monkeypatch.setattr(weight_transfer_http, "refit_http_session", lambda: session)
+
+ with pytest.raises(RuntimeError, match="HTTP 500: gateway failure"):
+ weight_transfer_http.post_vllm_refit_endpoints(
+ ["http://receiver/refit"], {}, api_key=None, timeout_s=1.0
+ )
+
+
+def test_sparse_export_finishes_before_blocked_transfers() -> None:
+ exported = threading.Event()
+ release_transfers = threading.Event()
+ result = []
+
+ def tensors():
+ for index in range(4):
+ yield f"weight-{index}", torch.ones(1)
+ exported.set()
+
+ def send_payload(_body, _payload_index):
+ assert release_transfers.wait(timeout=5.0)
+ return {"receiver": {}}
+
+ def run():
+ result.append(_stream_sparse_test_payloads(tensors(), send_payload))
+
+ thread = threading.Thread(target=run)
+ thread.start()
+ try:
+ assert exported.wait(timeout=2.0)
+ finally:
+ release_transfers.set()
+ thread.join(timeout=5.0)
+
+ assert not thread.is_alive()
+ assert result == [{"payloads": 4, "changed_elements": 4, "total_elements": 4}]
+
+
+def test_sparse_export_finishes_before_transfer_error() -> None:
+ exported = []
+
+ def tensors():
+ for index in range(4):
+ exported.append(index)
+ yield f"weight-{index}", torch.ones(1)
+
+ def fail_transfer(_body, _payload_index):
+ raise RuntimeError("transfer failed")
+
+ with pytest.raises(RuntimeError, match="transfer failed"):
+ _stream_sparse_test_payloads(tensors(), fail_transfer)
+
+ assert exported == list(range(4))
+
+
+def test_sparse_transport_cleanup_runs_on_transfer_workers() -> None:
+ class Transport:
+ name = "zmq"
+ transfer_workers = 2
+
+ def __init__(self) -> None:
+ self.barrier = threading.Barrier(2)
+ self.send_threads = set()
+ self.cleanup_threads = set()
+
+ def send(self, _body, _payload_id, _verification_candidates):
+ self.send_threads.add(threading.get_ident())
+ self.barrier.wait(timeout=5.0)
+ return {"receiver": {}}
+
+ def cleanup(self) -> None:
+ self.cleanup_threads.add(threading.get_ident())
+
+ transport = Transport()
+ result = weight_transfer_stream.stream_sparse_delta_payloads(
+ [(f"weight-{index}", torch.ones(1)) for index in range(2)],
+ delta_tracker=_SparsePipelineTracker(),
+ transport=transport,
+ shard_rank=0,
+ shard_count=1,
+ )
+
+ assert result["payloads"] == 2
+ assert transport.send_threads == transport.cleanup_threads
+
+
+def test_sparse_stream_coalesces_export_chunks() -> None:
+ tracker = _delta_tracker(
+ bucket_bytes=8,
+ zmq_export_bytes=4,
+ zmq_encode_workers=2,
+ )
+ tensors = [(f"weight-{index}", torch.zeros(1)) for index in range(4)]
+ tracker.snapshot_baseline(tensors)
+ for _, tensor in tensors:
+ tensor.fill_(1)
+ payloads = {}
+
+ def send(body, payload_index):
+ raw = zstandard.ZstdDecompressor().decompress(body)
+ payloads[payload_index] = torch.load(
+ io.BytesIO(raw), map_location="cpu", weights_only=True
+ )
+ return {"receiver": {}}
+
+ cleaned = []
+ transport = SparseRefitTransport(
+ "zmq",
+ 1,
+ lambda body, payload_id, _candidates: send(body, payload_id),
+ lambda: cleaned.append(True),
+ )
+ result = weight_transfer_stream.stream_sparse_delta_payloads(
+ tensors,
+ delta_tracker=tracker,
+ transport=transport,
+ shard_rank=0,
+ shard_count=1,
+ )
+
+ assert cleaned
+ assert result == {"payloads": 2, "changed_elements": 4, "total_elements": 4}
+ payload_names = [
+ [item["name"] for item in payloads[index][2]] for index in sorted(payloads)
+ ]
+ assert all(len(names) == 2 for names in payload_names)
+ assert sorted(name for names in payload_names for name in names) == [
+ f"weight-{index}" for index in range(4)
+ ]
+ for locations, value_groups, metadata in payloads.values():
+ for item in metadata:
+ assert sparse_locations_for_item(
+ item, locations, device="cpu"
+ ).tolist() == [0]
+ assert value_groups[item["value_group"]][
+ item["value_start"] : item["value_end"]
+ ].tolist() == [1065353216]
+
+
+def test_sparse_export_chunk_defaults_are_transport_specific() -> None:
+ tracker = _delta_tracker(bucket_bytes=1024**3)
+
+ assert sparse_export_chunk_size(tracker, "s3") == 64 * 1024**2
+ assert sparse_export_chunk_size(tracker, "zmq") == 256 * 1024**2
+
+
+def test_sparse_baseline_snapshots_only_owned_export_chunks(capsys) -> None:
+ tensors = [(f"weight-{index}", torch.tensor([float(index)])) for index in range(4)]
+
+ assert _baseline_names(tensors, rank=1) == ["weight-1", "weight-3"]
+ assert "chunks=4" in capsys.readouterr().out
+
+
+def test_sparse_stream_sends_only_owned_export_chunks() -> None:
+ sent = []
+ transport = SparseRefitTransport(
+ "zmq",
+ 1,
+ lambda _body, payload_id, _candidates: (
+ sent.append(payload_id) or {"receiver": {}}
+ ),
+ lambda: None,
+ )
+
+ result = weight_transfer_stream.stream_sparse_delta_payloads(
+ [(f"weight-{index}", torch.ones(1)) for index in range(4)],
+ delta_tracker=_SparsePipelineTracker(),
+ transport=transport,
+ shard_rank=1,
+ shard_count=2,
+ )
+
+ assert result == {"payloads": 2, "changed_elements": 2, "total_elements": 2}
+ assert sorted(sent) == [0, 1]
+
+
+def test_s3_manifest_transport_validates_configuration(monkeypatch) -> None:
+ tracker = SimpleNamespace(refit_config=_refit_config(s3_bucket="bucket"))
+ kwargs = {
+ "iterator": (),
+ "delta_tracker": tracker,
+ "transfer_id": "transfer",
+ "api_key_env_var": None,
+ "timeout_s": 1.0,
+ "shard_rank": 0,
+ "shard_count": 1,
+ }
+ with pytest.raises(ValueError, match="URL is required"):
+ weight_transfer_stream.stream_sparse_delta_payloads_via_s3_manifest(
+ refit_targets=[], **kwargs
+ )
+
+ tracker.refit_config = _refit_config()
+ with pytest.raises(RuntimeError, match="refit_cfg.storage.s3_bucket"):
+ weight_transfer_stream.stream_sparse_delta_payloads_via_s3_manifest(
+ refit_targets=["http://receiver"], **kwargs
+ )
+
+
+def test_s3_manifest_transport_uploads_notifies_and_deletes(monkeypatch) -> None:
+ operations = []
+ posts = []
+
+ class Store:
+ bucket = "bucket"
+ region = "us-west-2"
+
+ def put(self, key, body) -> None:
+ operations.append(("put", key, body))
+
+ def delete(self, key) -> None:
+ operations.append(("delete", key))
+
+ store = Store()
+ monkeypatch.setenv("NRL_TEST_REFIT_KEY", "secret")
+ monkeypatch.setattr(
+ weight_transfer_stream,
+ "_get_manifest_s3_store",
+ lambda *_args: store,
+ )
+
+ def post(endpoints, manifest, **kwargs):
+ posts.append((endpoints, manifest, kwargs))
+ return [
+ {"ok": True, "receiver_total_s": 1.0},
+ {"ok": True, "receiver_total_s": 2.0},
+ ]
+
+ def stream(iterator, **kwargs):
+ assert list(iterator) == [("weight", torch.tensor([1.0]))]
+ transport = kwargs["transport"]
+ assert transport.name == "s3"
+ assert transport.transfer_workers == 3
+ response = transport.send(b"payload", 3, 4)
+ assert response["receiver"] == {"receiver_total_s": 2.0}
+ transport.cleanup()
+ return {"payloads": 1, "changed_elements": 1, "total_elements": 1}
+
+ monkeypatch.setattr(weight_transfer_stream, "post_vllm_refit_endpoints", post)
+ monkeypatch.setattr(weight_transfer_stream, "stream_sparse_delta_payloads", stream)
+
+ result = weight_transfer_stream.stream_sparse_delta_payloads_via_s3_manifest(
+ [("weight", torch.tensor([1.0]))],
+ delta_tracker=SimpleNamespace(
+ refit_config=_refit_config(
+ s3_bucket=store.bucket,
+ s3_region=store.region,
+ s3_prefix="/prefix/",
+ s3_transfer_workers=3,
+ )
+ ),
+ refit_targets=[" http://receiver-a/ ", "http://receiver-b"],
+ transfer_id="transfer",
+ api_key_env_var="NRL_TEST_REFIT_KEY",
+ timeout_s=7.0,
+ shard_rank=2,
+ shard_count=4,
+ )
+
+ key = "prefix/transfer/000002/000003.pt"
+ assert result == {"payloads": 1, "changed_elements": 1, "total_elements": 1}
+ assert operations == [("put", key, b"payload"), ("delete", key)]
+ assert posts == [
+ (
+ [
+ "http://receiver-a/nemo-rl/refit/s3-manifest",
+ "http://receiver-b/nemo-rl/refit/s3-manifest",
+ ],
+ {
+ "bucket": store.bucket,
+ "region": store.region,
+ "key": key,
+ "checksum": sparse_payload_checksum(b"payload"),
+ "verification_candidates": 4,
+ },
+ {"api_key": "secret", "timeout_s": 7.0},
+ )
+ ]
+
+
+def test_zmq_stream_routes_shards_and_closes_clients(monkeypatch) -> None:
+ created = []
+ sent = []
+ closed = []
+ monkeypatch.setenv("NRL_TEST_REFIT_KEY", "secret")
+
+ class Client:
+ def __init__(self, address, **kwargs) -> None:
+ created.append((address, kwargs))
+
+ def send_payload(self, **kwargs):
+ sent.append(kwargs)
+ return {"ok": True, "receiver_total_s": 0.5}
+
+ def close(self) -> None:
+ closed.append(True)
+
+ def stream(_iterator, **kwargs):
+ transport = kwargs["transport"]
+ assert transport.name == "zmq"
+ assert transport.transfer_workers == 2
+ for payload_id in range(2):
+ response = transport.send(
+ f"body-{payload_id}".encode(), payload_id, payload_id + 1
+ )
+ assert response["receiver"]["ok"]
+ transport.cleanup()
+ return {"payloads": 2, "changed_elements": 2, "total_elements": 2}
+
+ monkeypatch.setattr(weight_transfer_zmq, "ZmqSparseRefitClient", Client)
+ monkeypatch.setattr(weight_transfer_zmq, "stream_sparse_delta_payloads", stream)
+
+ with pytest.raises(ValueError, match="address is required"):
+ weight_transfer_zmq.stream_sparse_delta_payloads_via_zmq(
+ (),
+ delta_tracker=SimpleNamespace(),
+ refit_targets=[],
+ transfer_id="transfer",
+ api_key_env_var=None,
+ timeout_s=1.0,
+ shard_rank=0,
+ shard_count=1,
+ )
+
+ kwargs = {
+ "iterator": (),
+ "delta_tracker": SimpleNamespace(
+ refit_config=_refit_config(zmq_transfer_workers=2)
+ ),
+ "refit_targets": ["tcp://receiver-a", " tcp://receiver-b "],
+ "transfer_id": "transfer",
+ "api_key_env_var": "NRL_TEST_REFIT_KEY",
+ "timeout_s": 7.0,
+ "shard_rank": 3,
+ "shard_count": 4,
+ }
+ assert (
+ weight_transfer_zmq.stream_sparse_delta_payloads_via_zmq(**kwargs)["payloads"]
+ == 2
+ )
+ assert (
+ weight_transfer_zmq.stream_sparse_delta_payloads_via_zmq(**kwargs)["payloads"]
+ == 2
+ )
+
+ assert created == 2 * [
+ (
+ "tcp://receiver-b",
+ {
+ "timeout_s": 7.0,
+ "producer_id": 3,
+ "retries": 3,
+ "api_key": "secret",
+ },
+ )
+ ]
+ assert closed == [True, True]
+ assert [item["payload_id"] for item in sent] == [0, 1, 0, 1]
+ assert [item["verification_candidates"] for item in sent] == [1, 2, 1, 2]
+ assert all(
+ item["checksum"] == sparse_payload_checksum(item["body"]) for item in sent
+ )
+
+
+def test_zmq_client_filters_replies_rejects_nack_and_retries(monkeypatch) -> None:
+ class Socket:
+ def __init__(self, replies=(), send_failures=0) -> None:
+ self.replies = list(replies)
+ self.send_failures = send_failures
+ self.sent = []
+
+ def send_multipart(self, frames, **_kwargs) -> None:
+ self.sent.append(frames)
+ if self.send_failures:
+ self.send_failures -= 1
+ raise weight_transfer_zmq.zmq.Again()
+
+ def poll(self, *_args) -> bool:
+ return bool(self.replies)
+
+ def recv_multipart(self):
+ return self.replies.pop(0)
+
+ def client(socket) -> ZmqSparseRefitClient:
+ result = ZmqSparseRefitClient.__new__(ZmqSparseRefitClient)
+ result._address = "tcp://receiver"
+ result._timeout_ms = 1000
+ result._producer_id = 4
+ result._retries = 1
+ result._api_key = "secret"
+ result._socket = socket
+ return result
+
+ success = {
+ "ok": True,
+ "transfer_id": "transfer-a",
+ "producer_id": 4,
+ "payload_id": 7,
+ }
+ socket = Socket(
+ [
+ [b"malformed"],
+ [
+ b"ACK",
+ json.dumps({**success, "transfer_id": "stale"}).encode(),
+ ],
+ [b"ACK", json.dumps(success).encode()],
+ ]
+ )
+ assert _send_zmq_payload(client(socket), 7, b"body") == success
+ assert json.loads(socket.sent[0][1])["api_key"] == "secret"
+
+ denied = Socket([[b"NACK", json.dumps({"ok": False, "error": "denied"}).encode()]])
+ with pytest.raises(RuntimeError, match="denied"):
+ _send_zmq_payload(client(denied), 7, b"body")
+
+ with pytest.raises(TimeoutError, match="payload 7"):
+ _send_zmq_payload(client(Socket(send_failures=2)), 7, b"body")
+
+
+def test_zmq_server_rejects_malformed_messages() -> None:
+ server = ZmqSparseRefitServer.__new__(ZmqSparseRefitServer)
+ server._token = "secret"
+ body = b"body"
+ metadata = {
+ "protocol": "nemo-rl-sparse-zmq-v1",
+ "api_key": "secret",
+ "transfer_id": "transfer",
+ "producer_id": 0,
+ "payload_id": 1,
+ "checksum": sparse_payload_checksum(body),
+ "verification_candidates": 2,
+ }
+
+ def frames(kind: bytes = b"DATA", **updates: object) -> list[bytes]:
+ return [b"id", kind, json.dumps({**metadata, **updates}).encode(), body]
+
+ for message, error, match in (
+ ([], ValueError, "Expected 4"),
+ (frames(b"OTHER"), ValueError, "Unsupported ZeroMQ sparse refit message"),
+ (frames(protocol="other"), ValueError, "protocol"),
+ (frames(api_key="wrong"), PermissionError, "authentication"),
+ (frames(transfer_id=""), ValueError, "identity"),
+ (frames(checksum=""), ValueError, "identity"),
+ (frames(verification_candidates=-1), ValueError, "identity"),
+ ):
+ with pytest.raises(error, match=match):
+ server._parse_data_message(message)
+
+ assert server._parse_data_message(frames())[1] == ("transfer", 0, 1)
+
+
+def _send_zmq_payload(
+ client: ZmqSparseRefitClient,
+ payload_id: int,
+ body: bytes,
+ checksum: str | None = None,
+ transfer_id: str = "transfer-a",
+) -> dict[str, object]:
+ return client.send_payload(
+ transfer_id=transfer_id,
+ payload_id=payload_id,
+ checksum=checksum or sparse_payload_checksum(body),
+ verification_candidates=2,
+ body=body,
+ )
+
+
+def test_zmq_sparse_refit_relay_fans_out_and_rejects_corruption(monkeypatch) -> None:
+ monkeypatch.setenv("NRL_TEST_REFIT_KEY", "secret")
+ received = [[] for _ in range(4)]
+
+ def apply(items):
+ def apply_payload(body, metadata):
+ if metadata["checksum"] != sparse_payload_checksum(body):
+ raise ValueError("checksum mismatch")
+ items.append((dict(metadata), body))
+ return {"ok": True, "receiver_total_s": 0.25}
+
+ return apply_payload
+
+ relays = [
+ ZmqSparseRefitServer(
+ apply(items),
+ bind_address="tcp://127.0.0.1:*",
+ api_key_env_var="NRL_TEST_REFIT_KEY",
+ timeout_s=5.0,
+ tuning=_refit_config().tuning,
+ )
+ for items in received
+ ]
+ addresses = [relay.start() for relay in relays]
+ for relay, address in zip(relays, addresses, strict=True):
+ relay.configure_tree(addresses, own_address=address)
+ unauthenticated_client = ZmqSparseRefitClient(
+ addresses[0],
+ timeout_s=5.0,
+ producer_id=2,
+ retries=3,
+ )
+ client = ZmqSparseRefitClient(
+ addresses[0],
+ timeout_s=5.0,
+ producer_id=3,
+ retries=3,
+ api_key="secret",
+ )
+ body = b"compressed sparse payload"
+ checksum = sparse_payload_checksum(body)
+ try:
+ with pytest.raises(RuntimeError, match="authentication failed"):
+ _send_zmq_payload(unauthenticated_client, 6, body)
+ first = _send_zmq_payload(client, 7, body)
+ duplicate = _send_zmq_payload(client, 7, body)
+ assert first["ok"] and first["staged"]
+ assert duplicate["ok"] and duplicate["staged"]
+ for relay in relays:
+ flushed = relay.flush("transfer-a", expected_payloads=1)
+ assert flushed["payloads"] == 1
+ assert flushed["receiver_total_s"] == 0.25
+ assert [len(items) for items in received] == [1] * 4
+ for items in received:
+ metadata, posted_body = items[0]
+ assert posted_body == body
+ assert metadata["transfer_id"] == "transfer-a"
+ assert metadata["producer_id"] == 3
+ assert metadata["payload_id"] == 7
+ assert metadata["checksum"] == checksum
+ assert metadata["verification_candidates"] == 2
+ assert metadata["api_key"] == "secret"
+
+ with pytest.raises(RuntimeError, match="already flushed"):
+ _send_zmq_payload(client, 8, body)
+ assert _send_zmq_payload(client, 8, body, "0" * 32, "transfer-b")["staged"]
+ with pytest.raises(RuntimeError, match="checksum mismatch"):
+ relays[0].flush("transfer-b")
+ finally:
+ unauthenticated_client.close()
+ client.close()
+ for relay in relays:
+ relay.close()
diff --git a/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py b/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py
new file mode 100644
index 00000000000..0c6f5a6c461
--- /dev/null
+++ b/tests/unit/weight_sync/test_vllm_remote_sparse_weight_synchronizer.py
@@ -0,0 +1,369 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from unittest.mock import ANY, MagicMock, patch
+
+import pytest
+
+from nemo_rl.models.generation.vllm.config import VllmRefitConfig
+from nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer import (
+ VllmRemoteSparseWeightSynchronizer,
+ validate_vllm_remote_sparse_refit,
+)
+
+_MODULE = "nemo_rl.weight_sync.vllm_remote_sparse_weight_synchronizer"
+
+
+@pytest.fixture
+def mock_ray():
+ with patch(f"{_MODULE}.ray") as value:
+ yield value
+
+
+@pytest.fixture
+def post():
+ with patch(f"{_MODULE}.post_vllm_refit_endpoints") as value:
+ yield value
+
+
+def _remote_sparse_sync(
+ mock_ray: MagicMock,
+ transport: str,
+ stream_result: list[dict[str, int]] | RuntimeError,
+) -> tuple[VllmRemoteSparseWeightSynchronizer, MagicMock, MagicMock]:
+ init_refs, stream_refs, commit_refs = [MagicMock()], [MagicMock()], [MagicMock()]
+ policy = MagicMock()
+ policy.worker_group.workers = [object(), object()]
+ policy.worker_group.run_all_workers_multiple_data.side_effect = [
+ init_refs,
+ stream_refs,
+ ]
+ policy.worker_group.run_all_workers_single_data.return_value = commit_refs
+
+ generation = MagicMock()
+ generation.worker_group.run_all_workers_single_data.return_value = [MagicMock()]
+ generation.invalidate_kv_cache.return_value = True
+
+ get_results: list[object] = [["http://receiver"]]
+ if transport == "zmq":
+ get_results.extend([["tcp://relay:19090"], [None]])
+ get_results.extend([[{"weight": ((8,), "float32")}], stream_result])
+ mock_ray.get.side_effect = get_results
+
+ sync = VllmRemoteSparseWeightSynchronizer(policy, generation, transport=transport)
+ with patch(f"{_MODULE}.post_vllm_refit_endpoints"):
+ sync.init_communicator()
+ return sync, policy, generation
+
+
+def _valid_config() -> dict:
+ return {
+ "refit_transport": "vllm_s3_sparse",
+ "refit_cfg": {
+ "delta_compression": {"encoding": "overwrite"},
+ "storage": {"s3_bucket": "bucket"},
+ },
+ "vllm_cfg": {"precision": "bfloat16", "kv_cache_dtype": "auto"},
+ }
+
+
+def test_validate_remote_sparse_refit_accepts_supported_scope():
+ config = _valid_config()
+ assert (
+ validate_vllm_remote_sparse_refit(
+ config, colocated=False, megatron_enabled=True
+ )
+ == "s3"
+ )
+ assert config["refit_cfg"] == VllmRefitConfig(
+ delta_compression={"encoding": "overwrite"},
+ storage={"s3_bucket": "bucket"},
+ )
+
+
+@pytest.mark.parametrize(
+ ("change", "kwargs"),
+ [
+ ({"refit_transport": "unknown"}, {}),
+ ({}, {"colocated": True}),
+ ({}, {"megatron_enabled": False}),
+ ({"refit_cfg": {"storage": {"s3_bucket": None}}}, {}),
+ ({"quant_cfg": "fp8"}, {}),
+ ({"vllm_cfg": {"precision": "fp8", "kv_cache_dtype": "auto"}}, {}),
+ (
+ {"vllm_cfg": {"precision": "bfloat16", "kv_cache_dtype": "fp8_e4m3"}},
+ {},
+ ),
+ ],
+)
+def test_validate_remote_sparse_refit_rejects_unsupported_scope(change, kwargs):
+ config = _valid_config()
+ config.update(change)
+ arguments = {"colocated": False, "megatron_enabled": True}
+ arguments.update(kwargs)
+
+ with pytest.raises(ValueError):
+ validate_vllm_remote_sparse_refit(config, **arguments)
+
+
+class TestVllmRemoteSparseWeightSynchronizer:
+ def test_init_communicator_joins_prelaunched_baseline(self, mock_ray, post):
+ baseline_ref = MagicMock()
+ policy = MagicMock()
+ generation = MagicMock()
+ generation.worker_group.run_all_workers_single_data.return_value = [MagicMock()]
+ mock_ray.get.side_effect = [
+ ["http://receiver"],
+ [{"weight": ((8,), "float32")}],
+ ]
+ sync = VllmRemoteSparseWeightSynchronizer(
+ policy,
+ generation,
+ transport="s3",
+ baseline_init_refs=[baseline_ref],
+ )
+ post.return_value = [{"overwrite_names": ["weight"]}]
+
+ sync.init_communicator()
+
+ policy.worker_group.run_all_workers_multiple_data.assert_not_called()
+ mock_ray.get.assert_any_call([baseline_ref])
+ post.assert_called_once_with(
+ ["http://receiver/nemo-rl/refit/prepare"],
+ {"tensors": {"weight": [[8], "float32"]}},
+ api_key=None,
+ timeout_s=600.0,
+ )
+ assert sync._baseline_init_refs == []
+ assert sync._overwrite_names == ["weight"]
+
+ def test_merge_refit_info_rejects_conflicting_metadata(self):
+ with pytest.raises(ValueError, match="Conflicting sparse refit metadata"):
+ VllmRemoteSparseWeightSynchronizer._merge_refit_info(
+ [
+ {"weight": ((8,), "float32")},
+ {"weight": ((16,), "float32")},
+ ]
+ )
+
+ def test_init_communicator_requires_receiver_endpoints(self, mock_ray):
+ policy = MagicMock()
+ policy.worker_group.workers = [object()]
+ policy.worker_group.run_all_workers_multiple_data.return_value = [MagicMock()]
+ generation = MagicMock()
+ generation.worker_group.workers = [object()]
+ generation.worker_group.run_all_workers_single_data.return_value = [MagicMock()]
+ mock_ray.get.return_value = []
+ sync = VllmRemoteSparseWeightSynchronizer(policy, generation, transport="s3")
+
+ with pytest.raises(ValueError, match="endpoints are missing"):
+ sync.init_communicator()
+
+ def test_shutdown_cancels_pending_work_and_stops_zmq(self, mock_ray):
+ policy = MagicMock()
+ generation = MagicMock()
+ generation.worker_group.workers = [object()]
+ generation.worker_group.run_all_workers_single_data.return_value = [MagicMock()]
+ sync = VllmRemoteSparseWeightSynchronizer(policy, generation, transport="zmq")
+ init_ref, commit_ref = MagicMock(), MagicMock()
+ sync._baseline_init_refs = [init_ref]
+ sync._baseline_commit_refs = [commit_ref]
+ sync._refit_urls = ["http://receiver"]
+ sync._targets = ["tcp://relay"]
+ sync._stale = False
+
+ sync.mark_stale()
+ sync.shutdown()
+
+ assert mock_ray.cancel.call_count == 2
+ mock_ray.cancel.assert_any_call(init_ref, force=False)
+ mock_ray.cancel.assert_any_call(commit_ref, force=False)
+ generation.worker_group.run_all_workers_single_data.assert_called_once_with(
+ "stop_zmq_sparse_refit_relay",
+ run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
+ )
+ assert sync.is_stale
+ assert sync._baseline_init_refs == []
+ assert sync._baseline_commit_refs == []
+ assert sync._refit_urls == []
+ assert sync._targets == []
+ assert sync._overwrite_names == []
+
+ def test_fails_before_transfer_when_kv_cache_invalidation_fails(self, mock_ray):
+ policy = MagicMock()
+ generation = MagicMock()
+ generation.invalidate_kv_cache.return_value = False
+ sync = VllmRemoteSparseWeightSynchronizer(policy, generation, transport="s3")
+
+ with pytest.raises(RuntimeError, match="KV cache invalidation failed"):
+ sync.sync_weights()
+ policy.worker_group.run_all_workers_multiple_data.assert_not_called()
+
+ def test_initializes_streams_commits_and_updates_baseline(
+ self, mock_ray, post, capsys
+ ):
+ sync, policy, generation = _remote_sparse_sync(
+ mock_ray,
+ "zmq",
+ [{"payloads": 3, "changed_elements": 3, "total_elements": 100}],
+ )
+ post.return_value = [
+ {
+ "verification_candidates": 4,
+ "verification_samples": 4,
+ "verification_exact_mismatches": 1,
+ "verification_mismatches": 0,
+ "verification_abs_sum": 1e-9,
+ "verification_max_abs": 1e-9,
+ }
+ ]
+ verification = post.return_value
+ post.side_effect = [
+ [{"ok": True, "payloads": 3, "receiver_relay_fanout_s": 1.0}],
+ verification,
+ ]
+ metrics = sync.sync_weights()
+
+ assert [
+ entry.args[0]
+ for entry in policy.worker_group.run_all_workers_multiple_data.call_args_list
+ ] == ["init_remote_sparse_delta_baseline", "stream_remote_sparse_weights"]
+ assert (
+ policy.worker_group.run_all_workers_multiple_data.call_args_list[1].kwargs[
+ "common_kwargs"
+ ]["overwrite_names"]
+ == []
+ )
+ assert [
+ entry.args[0]
+ for entry in generation.worker_group.run_all_workers_single_data.call_args_list
+ ] == [
+ "report_refit_server_base_url",
+ "start_zmq_sparse_refit_relay",
+ "configure_zmq_sparse_refit_relay",
+ ]
+ generation.worker_group.run_all_workers_single_data.assert_any_call(
+ "start_zmq_sparse_refit_relay",
+ run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
+ )
+ generation.worker_group.run_all_workers_single_data.assert_any_call(
+ "configure_zmq_sparse_refit_relay",
+ run_rank_0_only_axes=["tensor_parallel", "pipeline_parallel"],
+ relay_addresses=["tcp://relay:19090"],
+ )
+ post.assert_any_call(
+ ["http://receiver/nemo-rl/refit/zmq-flush"],
+ {"transfer_id": ANY, "expected_payloads": 3},
+ api_key=None,
+ timeout_s=600.0,
+ )
+ post.assert_any_call(
+ ["http://receiver/nemo-rl/refit/flush"],
+ {},
+ api_key=None,
+ timeout_s=600.0,
+ )
+ policy.worker_group.run_all_workers_single_data.assert_called_once_with(
+ "finish_remote_sparse_delta_sync", succeeded=True
+ )
+ assert (
+ "REFIT_ZMQ_DELTA_CHANGE changed_elements=3 total_elements=100 "
+ "changed_pct=3" in capsys.readouterr().out
+ )
+ assert metrics["delta/changed_pct"] == 3.0
+ assert metrics["delta_verify/candidates"] == 4.0
+ assert metrics["delta_verify/samples"] == 4.0
+ assert metrics["delta_verify/exact_mismatches"] == 1.0
+ assert metrics["delta_verify/mismatches"] == 0.0
+ assert metrics["delta_verify/mean_abs"] == 2.5e-10
+ assert metrics["delta_verify/max_abs"] == 1e-9
+ assert metrics["transfer/payloads"] == 3.0
+ assert metrics["transfer/relay_flush_s"] >= 0.0
+ assert not sync.is_stale
+
+ def test_sample_mismatch_does_not_commit_baseline(self, mock_ray, post):
+ sync, policy, _ = _remote_sparse_sync(
+ mock_ray,
+ "zmq",
+ [{"payloads": 3, "changed_elements": 3, "total_elements": 100}],
+ )
+ post.return_value = [
+ {
+ "verification_samples": 4,
+ "verification_mismatches": 1,
+ "verification_abs_sum": 0.5,
+ "verification_max_abs": 0.5,
+ }
+ ]
+
+ verification = post.return_value
+ post.side_effect = [
+ [{"ok": True, "payloads": 3}],
+ verification,
+ verification,
+ ]
+ with pytest.raises(RuntimeError, match="1 mismatched deltas out of 4"):
+ sync.sync_weights()
+
+ policy.worker_group.run_all_workers_single_data.assert_called_once_with(
+ "finish_remote_sparse_delta_sync", succeeded=False
+ )
+
+ def test_failure_drains_receivers_without_committing_baseline(self, mock_ray, post):
+ sync, policy, _ = _remote_sparse_sync(
+ mock_ray, "s3", RuntimeError("stream failed")
+ )
+
+ with pytest.raises(RuntimeError, match="stream failed"):
+ sync.sync_weights()
+
+ with pytest.raises(RuntimeError, match="poisoned by a prior failed sync"):
+ sync.sync_weights()
+
+ post.assert_called_once_with(
+ ["http://receiver/nemo-rl/refit/flush"],
+ {},
+ api_key=None,
+ timeout_s=60.0,
+ )
+ policy.worker_group.run_all_workers_single_data.assert_called_once_with(
+ "finish_remote_sparse_delta_sync", succeeded=False
+ )
+
+ def test_zmq_relay_failure_drains_before_rejecting_baseline(self, mock_ray, post):
+ sync, policy, _ = _remote_sparse_sync(
+ mock_ray,
+ "zmq",
+ [{"payloads": 3, "changed_elements": 3, "total_elements": 100}],
+ )
+ mock_ray.get.side_effect = [
+ [{"payloads": 3, "changed_elements": 3, "total_elements": 100}],
+ ]
+ post.side_effect = [
+ RuntimeError("fanout failed"),
+ [{"ok": True, "payloads": 3}],
+ [{"ok": True}],
+ ]
+
+ with pytest.raises(RuntimeError, match="fanout failed"):
+ sync.sync_weights()
+
+ assert [call.args[0][0] for call in post.call_args_list] == [
+ "http://receiver/nemo-rl/refit/zmq-flush",
+ "http://receiver/nemo-rl/refit/zmq-flush",
+ "http://receiver/nemo-rl/refit/flush",
+ ]
+ policy.worker_group.run_all_workers_single_data.assert_called_once_with(
+ "finish_remote_sparse_delta_sync", succeeded=False
+ )
diff --git a/tools/refit_bandwidth_calculator.py b/tools/refit_bandwidth_calculator.py
new file mode 100644
index 00000000000..1a2a69f7258
--- /dev/null
+++ b/tools/refit_bandwidth_calculator.py
@@ -0,0 +1,247 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Project current zstd sparse refit against NCCL over candidate Ethernet."""
+
+import argparse
+import json
+import math
+from bisect import bisect_right
+from dataclasses import asdict, dataclass
+from typing import Literal
+
+Transport = Literal["s3", "zmq"]
+
+_REFERENCE_IB_GBPS = 400.0
+_DENSITIES = (3.0, 5.0)
+_SPARSE_ANCHOR_SIZE_GB = 247.2
+_SPARSE_BUCKET_SIZE_BYTES = 512 * 1024**2
+_SPARSE_ANCHOR_LATENCY_S: dict[Transport, tuple[float, float]] = {
+ "s3": (20.233733, 25.790387),
+ "zmq": (24.095243, 33.7759615),
+}
+_SPARSE_ANCHOR_WIRE_GB: dict[Transport, tuple[float, float]] = {
+ "s3": (5.858410822107136, 9.765203805732864),
+ "zmq": (5.608205511, 9.3456392505),
+}
+_NCCL_ANCHORS = (
+ (63.2, 0.84, 1.60),
+ (247.2, 1.46, 1.74),
+ (470.2, 2.31, 2.73),
+ (1342.0, 3.27, 3.46),
+)
+
+
+@dataclass(frozen=True)
+class Estimate:
+ transport: Transport
+ model_size_gb: float
+ changed_pct: float
+ sparse_bucket_size_bytes: int
+ sparse_seconds: float
+ approximate_wire_gb: float
+ nccl_ib_low_s: float
+ nccl_ib_high_s: float
+ candidate_ethernet_gbps: float | None
+ nccl_ethernet_low_s: float | None
+ nccl_ethernet_high_s: float | None
+ break_even_ethernet_low_gbps: float
+ break_even_ethernet_high_gbps: float
+ candidate_winner: str | None
+
+
+def _project_density(anchors: tuple[float, float], changed_pct: float) -> float:
+ exponent = math.log(anchors[1] / anchors[0]) / math.log(5.0 / 3.0)
+ return anchors[0] * (changed_pct / 3.0) ** exponent
+
+
+def predict_sparse_seconds(
+ model_size_gb: float,
+ changed_pct: float,
+ *,
+ transport: Transport,
+) -> float:
+ """Linearly project the latest 120B measurement by model size."""
+ if model_size_gb <= 0 or changed_pct <= 0:
+ raise ValueError("model_size_gb and changed_pct must be positive")
+ anchor = _project_density(_SPARSE_ANCHOR_LATENCY_S[transport], changed_pct)
+ return anchor * model_size_gb / _SPARSE_ANCHOR_SIZE_GB
+
+
+def predict_sparse_wire_gb(
+ model_size_gb: float,
+ changed_pct: float,
+ *,
+ transport: Transport,
+) -> float:
+ """Scale the latest measured zstd wire bytes by model size and density."""
+ if model_size_gb <= 0 or changed_pct <= 0:
+ raise ValueError("model_size_gb and changed_pct must be positive")
+ anchor = _project_density(_SPARSE_ANCHOR_WIRE_GB[transport], changed_pct)
+ return anchor * model_size_gb / _SPARSE_ANCHOR_SIZE_GB
+
+
+def _nccl_reference(model_size_gb: float) -> tuple[float, float]:
+ sizes = tuple(anchor[0] for anchor in _NCCL_ANCHORS)
+ index = min(max(bisect_right(sizes, model_size_gb) - 1, 0), len(sizes) - 2)
+ left, right = _NCCL_ANCHORS[index : index + 2]
+ position = math.log(model_size_gb / left[0]) / math.log(right[0] / left[0])
+ low = left[1] + position * (right[1] - left[1])
+ high = left[2] + position * (right[2] - left[2])
+ return max(0.001, low), max(0.001, high)
+
+
+def estimate(
+ *,
+ model_size_gb: float,
+ changed_pct: float,
+ transport: Transport,
+ candidate_ethernet_gbps: float | None = None,
+) -> Estimate:
+ """Estimate sparse latency and the NCCL-over-Ethernet crossover."""
+ if candidate_ethernet_gbps is not None and candidate_ethernet_gbps <= 0:
+ raise ValueError("candidate_ethernet_gbps must be positive")
+ sparse_seconds = predict_sparse_seconds(
+ model_size_gb,
+ changed_pct,
+ transport=transport,
+ )
+ nccl_low, nccl_high = _nccl_reference(model_size_gb)
+ crossover_low = _REFERENCE_IB_GBPS * nccl_low / sparse_seconds
+ crossover_high = _REFERENCE_IB_GBPS * nccl_high / sparse_seconds
+
+ projected_low = projected_high = None
+ winner = None
+ if candidate_ethernet_gbps is not None:
+ scale = _REFERENCE_IB_GBPS / candidate_ethernet_gbps
+ projected_low, projected_high = nccl_low * scale, nccl_high * scale
+ if sparse_seconds < projected_low and not math.isclose(
+ sparse_seconds, projected_low
+ ):
+ winner = transport
+ elif sparse_seconds > projected_high and not math.isclose(
+ sparse_seconds, projected_high
+ ):
+ winner = "nccl"
+ else:
+ winner = "depends"
+
+ return Estimate(
+ transport,
+ model_size_gb,
+ changed_pct,
+ _SPARSE_BUCKET_SIZE_BYTES,
+ sparse_seconds,
+ predict_sparse_wire_gb(
+ model_size_gb,
+ changed_pct,
+ transport=transport,
+ ),
+ nccl_low,
+ nccl_high,
+ candidate_ethernet_gbps,
+ projected_low,
+ projected_high,
+ crossover_low,
+ crossover_high,
+ winner,
+ )
+
+
+def _positive(value: str) -> float:
+ number = float(value)
+ if number <= 0:
+ raise argparse.ArgumentTypeError("must be positive")
+ return number
+
+
+def _seconds(low: float, high: float) -> str:
+ return f"{low:.3f}-{high:.3f}s"
+
+
+def _print_results(results: list[Estimate]) -> None:
+ first = results[0]
+ print(
+ f"Model: {first.model_size_gb:g} GB indexed BF16; "
+ f"changed: {first.changed_pct:g}%; compression: zstd; "
+ f"bucket: {first.sparse_bucket_size_bytes // 1024**2} MiB"
+ )
+ print(
+ "Measured NCCL on 400 Gbps/rank H100 IB: "
+ f"{_seconds(first.nccl_ib_low_s, first.nccl_ib_high_s)}"
+ )
+ if first.candidate_ethernet_gbps is not None:
+ assert first.nccl_ethernet_low_s is not None
+ assert first.nccl_ethernet_high_s is not None
+ print(
+ f"Projected NCCL on {first.candidate_ethernet_gbps:g} Gbps/rank "
+ f"Ethernet: {_seconds(first.nccl_ethernet_low_s, first.nccl_ethernet_high_s)}"
+ )
+
+ print(
+ "\nPath Sparse Wire Ethernet crossover Winner@candidate"
+ )
+ for result in results:
+ crossover = (
+ f"{result.break_even_ethernet_low_gbps:.2f}-"
+ f"{result.break_even_ethernet_high_gbps:.2f} Gbps/rank"
+ )
+ print(
+ f"{result.transport.upper():<6} {result.sparse_seconds:>9.3f}s "
+ f"{result.approximate_wire_gb:>8.3f} GB {crossover:>26} "
+ f"{result.candidate_winner or '':>18}"
+ )
+ print(
+ "\nBelow the lower crossover sparse wins across the NCCL envelope; "
+ "above the upper crossover NCCL wins."
+ )
+ if not math.isclose(first.model_size_gb, _SPARSE_ANCHOR_SIZE_GB):
+ print(
+ f"Note: sparse time is a linear model-size projection from the "
+ f"{_SPARSE_ANCHOR_SIZE_GB:g} GB measured anchor."
+ )
+ if not _DENSITIES[0] <= first.changed_pct <= _DENSITIES[1]:
+ print("Note: changed density is extrapolated from measured 3% and 5% arms.")
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--model-size-gb", type=_positive, required=True)
+ parser.add_argument(
+ "--changed-pct", "--sparsity-pct", type=_positive, required=True
+ )
+ parser.add_argument("--transport", choices=("all", "s3", "zmq"), default="all")
+ parser.add_argument("--candidate-ethernet-gbps", type=_positive)
+ parser.add_argument("--json", action="store_true")
+ args = parser.parse_args()
+ transports: tuple[Transport, ...] = (
+ ("s3", "zmq") if args.transport == "all" else (args.transport,)
+ )
+ results = [
+ estimate(
+ model_size_gb=args.model_size_gb,
+ changed_pct=args.changed_pct,
+ transport=transport,
+ candidate_ethernet_gbps=args.candidate_ethernet_gbps,
+ )
+ for transport in transports
+ ]
+ if args.json:
+ print(json.dumps([asdict(result) for result in results], indent=2))
+ else:
+ _print_results(results)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/uv.lock b/uv.lock
index 1e5f744404a..0bfea3edf89 100644
--- a/uv.lock
+++ b/uv.lock
@@ -422,6 +422,24 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" },
]
+[[package]]
+name = "awscrt"
+version = "0.35.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/e8/8a/294c2f6cdda8f386057a5f6b349fec9f4838b9c25a98cb67dc503bb80514/awscrt-0.35.0.tar.gz", hash = "sha256:761ae0dda17fd9dfaff4bbb2a376e28e44dfd77dc6410b7bc408297a1fd5600e", size = 37016406, upload-time = "2026-06-25T18:17:26.611Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/d9/5a/44aa794eee204002ae161ddcd1a0d901c6c9ca587f22294e345bb468b3ae/awscrt-0.35.0-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cd20c94a6008164c89eeb420892c3d88b165a79bdd22ef0a4ee383bee8b4cdb", size = 3964249, upload-time = "2026-06-25T18:16:30.389Z" },
+ { url = "https://files.pythonhosted.org/packages/84/03/edba2d4e7bf381eece2ff176836dbdbb5111d5a81f06705415a899848da9/awscrt-0.35.0-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e26cca47c8b84dd968bde8c39385010acd710631e55b686a0e5422573c89a25b", size = 4258920, upload-time = "2026-06-25T18:16:31.611Z" },
+ { url = "https://files.pythonhosted.org/packages/17/56/67c6374c97da326ef4a1af075b6a776cbcdc1bf14ddd259aa9acc613adfa/awscrt-0.35.0-cp311-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:ab050c01fb3a64c4efc7a12baf49321c10e1635903e8ad05d1fe7a88ef5b0f2a", size = 3876609, upload-time = "2026-06-25T18:16:33.032Z" },
+ { url = "https://files.pythonhosted.org/packages/2d/48/40014ff6278699e0065165128742a35b0654cb544e980496bd09b5d9c5db/awscrt-0.35.0-cp311-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:bc604dd77f61c3b6fef06dd73108b390d01699e5535c6d9cf244775cd0bac2f3", size = 4117595, upload-time = "2026-06-25T18:16:34.333Z" },
+ { url = "https://files.pythonhosted.org/packages/bc/42/08e91275771c3c72ab643f87e4e4b93d1fcb09769bc0d4df60834443bb68/awscrt-0.35.0-cp313-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4a7e5f267c494146ca680b1405ad2290b2dcc0821d5f386adcca8aa0bc427c25", size = 3955329, upload-time = "2026-06-25T18:16:39.727Z" },
+ { url = "https://files.pythonhosted.org/packages/a1/b9/db8bae837ff816861c06d423c66bfa3e9ee55e975fb26235134c300e04aa/awscrt-0.35.0-cp313-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9d766a3471f637bf4024fd487a10457e53809e17f3fb2f554bf4b73fad415f58", size = 4252293, upload-time = "2026-06-25T18:16:41.043Z" },
+ { url = "https://files.pythonhosted.org/packages/19/9b/37649041dbd8cc06373b9018d882ed42ca6535416e20358c305669cc720b/awscrt-0.35.0-cp313-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:286995476f2f8fd217fbe4c37c7e5bd69953cff03b6ea8d37946537a43208467", size = 3868495, upload-time = "2026-06-25T18:16:42.345Z" },
+ { url = "https://files.pythonhosted.org/packages/7d/eb/40c92251e4a17c1fdc2364bab2345791ded4f4da043262c6fadc3f51ed2a/awscrt-0.35.0-cp313-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:177f0f9bbddc3227ede427df0bee9353dbb32fbfdf0d424680400e6d76d11ebc", size = 4111934, upload-time = "2026-06-25T18:16:43.789Z" },
+ { url = "https://files.pythonhosted.org/packages/8b/cb/ed8503bcc55c150092278e0019c813cedcb937cad56303e4007e94f43e7a/awscrt-0.35.0-cp313-cp313t-musllinux_1_1_aarch64.whl", hash = "sha256:59976c6063dd79d117d08f93aaddf02448a5b19086b9444c6bdbedf50feff54c", size = 4008528, upload-time = "2026-06-25T18:16:49.675Z" },
+ { url = "https://files.pythonhosted.org/packages/8e/5c/461de896c73b2406bb6014422d53f4e0e3a56cc7b856ae0fff570f3878f7/awscrt-0.35.0-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:0c933fc94db944f4d2504ffd6bd0c2ef2289a475e0f8896cfea6569be2e4796a", size = 4248623, upload-time = "2026-06-25T18:16:51.227Z" },
+]
+
[[package]]
name = "babel"
version = "2.18.0"
@@ -3256,6 +3274,8 @@ docs = [
name = "nemo-rl"
source = { editable = "." }
dependencies = [
+ { name = "awscrt", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
+ { name = "zstandard", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
{ name = "accelerate", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
{ name = "blobfile", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
{ name = "colored", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-fsdp') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-mcore') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-automodel' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-fsdp' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-sglang') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-mcore' and extra == 'extra-7-nemo-rl-vllm') or (sys_platform != 'linux' and extra == 'extra-7-nemo-rl-sglang' and extra == 'extra-7-nemo-rl-vllm')" },
@@ -3424,6 +3444,8 @@ test = [
[package.metadata]
requires-dist = [
+ { name = "awscrt", specifier = ">=0.35.0" },
+ { name = "zstandard" },
{ name = "accelerate", specifier = ">=0.26" },
{ name = "blobfile" },
{ name = "causal-conv1d", marker = "extra == 'automodel'", git = "https://github.com/Dao-AILab/causal-conv1d?rev=4f6ae4e26ae5fe8af9372f8d312ab25cc4595223" },
@@ -7136,6 +7158,20 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/bf/2e/0b49f7e4e53817cfb09a0f6585012b782dfe0b666e8abefcb4fac0570606/z3_solver-4.15.4.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:62c7e9cbdd711932301f29919ad9158de9b2f58b4d281dd259bbcd0a2f408ba1", size = 27226534, upload-time = "2025-10-29T18:11:55.59Z" },
]
+
+[[package]]
+name = "zstandard"
+version = "0.25.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/fd/aa/3e0508d5a5dd96529cdc5a97011299056e14c6505b678fd58938792794b1/zstandard-0.25.0.tar.gz", hash = "sha256:7713e1179d162cf5c7906da876ec2ccb9c3a9dcbdffef0cc7f70c3667a205f0b", size = 711513, upload-time = "2025-09-14T22:15:54.002Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/6d/db/ddb11011826ed7db9d0e485d13df79b58586bfdec56e5c84a928a9a78c1c/zstandard-0.25.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:bfc4e20784722098822e3eee42b8e576b379ed72cca4a7cb856ae733e62192ea", size = 5063001, upload-time = "2025-09-14T22:17:31.044Z" },
+ { url = "https://files.pythonhosted.org/packages/63/4b/e3678b4e776db00f9f7b2fe58e547e8928ef32727d7a1ff01dea010f3f13/zstandard-0.25.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8e735494da3db08694d26480f1493ad2cf86e99bdd53e8e9771b2752a5c0246a", size = 5547173, upload-time = "2025-09-14T22:17:36.084Z" },
+ { url = "https://files.pythonhosted.org/packages/4e/d5/ba05ed95c6b8ec30bd468dfeab20589f2cf709b5c940483e31d991f2ca58/zstandard-0.25.0-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:3a39c94ad7866160a4a46d772e43311a743c316942037671beb264e395bdd611", size = 5046736, upload-time = "2025-09-14T22:17:37.891Z" },
+ { url = "https://files.pythonhosted.org/packages/50/d5/870aa06b3a76c73eced65c044b92286a3c4e00554005ff51962deef28e28/zstandard-0.25.0-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:172de1f06947577d3a3005416977cce6168f2261284c02080e7ad0185faeced3", size = 5576368, upload-time = "2025-09-14T22:17:40.206Z" },
+ { url = "https://files.pythonhosted.org/packages/5d/35/398dc2ffc89d304d59bc12f0fdd931b4ce455bddf7038a0a67733a25f550/zstandard-0.25.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3c83b0188c852a47cd13ef3bf9209fb0a77fa5374958b8c53aaa699398c6bd7b", size = 4954022, upload-time = "2025-09-14T22:17:41.879Z" },
+ { url = "https://files.pythonhosted.org/packages/b2/e5/fbd822d5c6f427cf158316d012c5a12f233473c2f9c5fe5ab1ae5d21f3d8/zstandard-0.25.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:4f187a0bb61b35119d1926aee039524d1f93aaf38a9916b8c4b78ac8514a0aaf", size = 5360113, upload-time = "2025-09-14T22:17:48.893Z" },
+]
[[package]]
name = "zipp"
version = "3.23.1"