[RFC] Delta weight sync in SkyRL

[SumanthRH]

Summary

This RFC proposes delta-based weight sync for disaggregated RL with SkyRL. The trainer computes byte-exact XOR deltas against the previously published BF16 checkpoint-format weights, compresses the deltas, publishes a manifest, and asks inference workers to fetch/process the next version before pausing generation. During the weight update, inference workers load the prepared weights from disk.

Motivation

RL training on disaggregated clusters allows users to better utilize available resources in disparate regions , different cloud providers or different hardware. A key bottleneck for disaggregated RL is weight transfer between inference and trainer nodes. During RL training, most of the weights are typically not updated in a given training step. A recent study found that only about 1-3% of the weights change after an optimization step, observed across model families in both sync and async RL setups (link). Disaggregated RL is thus enabled with delta weight updates: transmit only the compressed deltas every weight sync.

Goals

1. Support RL training in fully disaggregated clusters: Support RL with training and inference running on separate clusters. ¹
2. Support weight sync with delta weights: Support weight syncing in SkyRL with delta weights for checkpoint-to-storage: Weight transfer via checkpoint-to-disk or cloud storage with training and inference workers in separate clusters.
3. Support weight transfer for different quantization schemes: Support weight transfer between training and inference where inference engine can use different quantization schemes or kernels (specifically, support weight transfer in checkpoint format in vLLM)
   1. One interesting use-case: Train on NVIDIA in BF16 but do inference on AMD in FP8.

Non-Goals

1. Support sparse weight transfer via NCCL: Sparse weight transfer can also be used for RL training in the same ray cluster to optimize payload during weight sync. However, this is not a focus of this doc. There is additional latency in delta computation that leads to worse performance. Further, there are additional complications around vLLM weight loading here (to support goal 3.) that would complicate the weight transfer design.

Functional Requirements

1. Compute exact deltas of weights on the trainer without any precision related rounding errors.
2. Transfer deltas from the trainer in BF16 checkpoint-format² weights with some form of lossless compression to save on network bandwidth
3. Compute the updated weights on the inference engine in an exact way from the synced deltas and from the previous BF16 weights. Verify with reference post-delta checksum on the trainer. A downstream effect we would like to avoid is any potential drift from repeated weight sync of deltas.
4. Scale delta-based weight sync to large model sizes (350B+).

Design

Before going into the design, it is helpful to categorize the steps on the inference side as two operations:

1. Fetch + Process: Fetch the delta weights from storage and process them to compute the new updated weights
2. Load: Load the new weights into the inference engines. This can involve additional postprocessing ex: quantization

Overall trainer flow:

1. Store a CPU based snapshot of the previous version of the weights in uint8 on rank 0
2. During a weight sync, iterate over the parameters chunk by chunk and all gather the weights for a chunk
3. Copy a chunk from GPU to CPU
4. Convert to bytes and compute XOR deltas for all the weights
5. Compress the deltas using lossless compression and store the deltas on CPU
6. Update the CPU based snapshot to use the new uint8 bytes for the current weight chunk
7. After computing deltas for all the weights, publish to cloud storage

Overall inference flow:

1. Split up Fetch + Process and Load into two separate stages
2. Run Fetch + Process before pausing generation (with a /fetch_weights endpoint).
3. Pause generation and run Load (with /update_weights)

The sequence of operations in a weight update are: /fetch_weights -> /pause -> /start_weight_update -> /update_weights -> /finish_weight_update -> /resume

Other Design Options Considered

Option 1: CPU Snapshot on Trainer, CPU Snapshot on Inference, Deltas applied to the live parameters

Note that we need a CPU snapshot of BF16 dtype weights on the inference side due to requirement 3. Applying the diff to the live parameters could be done in the following way:

1. Store a snapshot of the BF16 weights on CPU for each worker. Each worker only needs to think about its shard, so the snapshot contains sharded weights per parameter and it is in worker format (see : Delta weight updates in SkyRL)
2. On receiving weights from the trainer, stage the sharded snapshot on CPU into the layerwise reload buffers on GPU
3. Process the received deltas via model.load_weights as a dummy full parameter tensor (NaNs in elements that are not present)
4. In the final copy operation, we copy the transformed weights into the buffers. However, instead of a direct copy, we do a modified version where we apply the non-NaN parameters to the layerwise buffer as a delta, and then copy into parameter memory
5. Update the sharded snapshot on CPU

This approach can also support in-cluster, NCCL based delta weight updates in a similar way - deltas can be transmitted via NCCL and the same approach for applying deltas holds.

This approach has the following cons:

1. Layerwise reload complexity : Need to deal with internals of layerwise reloading and delta some internal vllm loading functions
2. Fully synchronous fetch and load
3. Large CPU memory requirements on Inference: (number of parameters x num_engines_per_node x 2) bytes
4. Large CPU memory requirements on the Trainer: (number of parameters x 2)

Option 2: CPU Snapshot on Trainer, Disk based Snapshot on Inference, Deltas applied on the fly

At a high level, the flow on the inference side is:

• Inference ranks will store the current version of the weights in BF16 model loading format (“checkpoint format”) on local disk.
• During a weight update, one of the inference ranks will first download the delta to the local disk if not already done.
• Each rank will iterate over the saved weights on disk parameter by parameter, load the corresponding delta and compute the updated weights.
• One of the inference ranks in each node will trigger a background update and save the updated weights in a new location on disk
• The updated weights are fed to model.load_weights as an iterator.
• The updated weights on disk will be used as the reference weights for the next update.

This approach has the following cons:

1. Fully synchronous fetch + load
2. Large CPU memory requirements on the Trainer: (number of parameters x 2)
3. Slower disk based process on the inference side: Even with some pipelining, the bottleneck is synchronous writes into disk during each weight update

The chosen design has the same CPU memory requirements on the Trainer as the above options, but it enables much smaller pause times with decoupled fetch and load.

Detailed Design

For the detailed design, please refer to the design doc

Detailed Flow

Trainer Flow

For each sync:

1. All-gather weights on GPU and convert to Huggingface format on rank 0
2. Compare each tensor against the trainer's last published snapshot on CPU
3. Build XOR delta bytes for changed tensors.
4. Compress each delta with zstd.
5. Upload compressed deltas as safetensor files with deltas keyed by parameter name, along with manifest.json
6. Advance the trainer snapshot after the version has been written successfully.
7. Ask inference workers to fetch and apply that exact version.

Note: the current implementation only performs the delta computation and upload on rank 0, but the design allows for multiple senders.

Receiver Flow

/fetch_weights

1. Inference ranks receive a weight update request from the trainer/ controller with the delta manifest.
2. Inference ranks will inspect the current local checkpoint directory and try to grab the writer lock. The rank that grabs the lock first will download the new delta into the local directory and write the new version of the weights. Other ranks wait until the lock is released.
3. The writer rank will check the target_version by the controller as well as the version currently available in local_checkpoint_dir by reading state.json.
4. The writer rank marks write_in_progress as true in state.json
5. Until the current version becomes the target_version, perform the following:
   1. Download the delta file for version +1
   2. For each delta tensor present in the delta file, load the corresponding tensor from weights/ on disk using the tensor metadata
   3. Decompress the delta and apply the delta to the tensor to get the new weights
   4. Save the new tensors to disk
6. Writer rank releases the writer lock and marks write_in_progress as false.

/update_weights³

1. Each rank inspects the delta manifest received from the trainer
2. Ranks inspect the local_checkpoint_dir to verify that the target_version of the weights are present on disk. If not present, ranks error out.
3. If present, the ranks will proceed to iteratively load the weights from disk to GPU memory

Implementation

Component	Category	Responsibility
DeltaWeightSyncConfig	Config	User-facing configuration for checkpoint-delta weight sync
DeltaTransferStrategy	Weight Sync	Creates the sender, receiver init payload and vLLM transfer engine
DeltaWeightTransferSender	Trainer	Owns the trainer-side weight send logic
DeltaPublisher	Trainer	Maintains the CPU byte snapshot, computes XOR deltas, compresses payloads and publishes manifests
DeltaManifest / DeltaTensorRecord	Data Model	Version and per-tensor metadata for a published delta
RemoteInferenceClient.fetch_weights	Control Plane	Runs the pre-pause receiver-side fetch+process phase on every inference worker
DeltaWeightTransferEngine	Inference	vLLM transfer engine for delta weight sync
LocalCheckpointStore	Inference	Maintains the local checkpoint and applies deltas into it

For more details on the implementation, please see the design doc

Prior Work

The current design is inspired by the decoupled fetch and load design implemented in Slime, and is inspired by blogs from Fireworks and Cognition on delta weight updates which detail the same ideas. The SkyRL design differs in the trainer integration, the manifest schema, the local checkpoint store behaviour and in integrating with vLLM.

Further reading

For more details on limitations and performance with the above approach, refer to the doc.

Delta weight sync roadmap

We will put up a roadmap of features for delta weight sync in the RFC after initial support is merged.

Footnotes

1. We previously had this with the GLOO-based weight sync (over TCP/IP) in SkyRL (back in the old days), but it is no longer supported ↩

2. See Appendix: Delta weight updates in SkyRL for terminology on weight loading ↩

3. Weight updates have three API calls: /start_weight_update, /update_weights and /finish_weight_update. We only mention the second here since it contains the core weight update operations ↩