Delta Weight Transfer#

Delta weight transfer is an optional mode that sends only changed bf16 elements instead of the full model. After a single RL optimizer step, >98% of weights are bit-identical, so the delta is typically 30-40x smaller than the full model.

How It Works#

The trainer config declares supported strategies (["full", "delta"]). RaaS decides which mode to use per pull based on sender capabilities and its own state.

Trainer                SenderAgent (subprocess)           RaaS
  │                         │                              │
  offload()                 │                              │
  (GPU→CPU copy, ~0.5s)    │                              │
  notify_buffer_ready ────► │                              │
                     swap + ack immediately                │
  ◄──── ack ─────────────── │                              │
  save checkpoint           │ _compute_delta() [async]     │
  wait_delta_ready()        │ compare halves (numpy)       │
  ◄──── metrics ────────────│ encode sparse → delta shm    │
                            │ set delta_done_event         │
  notify_version ──────────────────────────────────────► │
                            │         GET /get_capabilities │
                            │ ◄──────── delta_ready? ──────│
                            │         POST /request_transfer│
                            │ ◄──────── mode=delta ────────│
                            │ TCP sendfile(delta) ────────►│
                            │ ZMQ done ───────────────────►│
                            │              mmap patch + load│

Enabling Delta#

Trainer — set the environment variable:

export WEIGHT_TRANSFER_STRATEGIES="full,delta"

RaaS — add to the YAML config under the raas section:

raas:
  delta_full_sync_interval: 10   # full resync every 10 steps

When delta_full_sync_interval is 0 (default), RaaS uses delta whenever available and never forces full. A non-zero value causes a periodic full transfer for resync (e.g. every 10th step).

Mode Decision#

RaaS decides per-pull in _choose_transfer_mode():

Condition

Mode

Reason

First pull (local_v=0)

full

No local weights to patch

delta not in sender strategies

full

Sender only supports full

Delta not ready

full

Async compute not finished

local_v % full_sync_interval == 0

full

Periodic resync

local_v != delta_base_version

full

Version mismatch (>1 step behind)

All checks pass

delta

Apply sparse patch

Sparse Format#

[header 16 bytes][indices][values]

Header:
  [0:8]   num_nonzero     (uint64)
  [8:10]  element_size    (uint16, 2 for bf16)
  [10:12] flags           (uint16, bit 0 = uint64 indices)
  [12:16] reserved

Indices:  num_nonzero x 4 bytes (uint32) or 8 bytes (uint64)
Values:   num_nonzero x element_size bytes (raw bf16)

The indices are element offsets into the flat weight buffer (same layout as the safetensors data section). Values are the new element values (not deltas) — the receiver overwrites at each index.

Delta Computation#

Runs in the sender agent subprocess (CPU, numpy), asynchronously after acking the trainer:

  1. View both buffer halves as uint16 arrays (bf16 representation)

  2. diff_mask = new_arr != old_arr — vectorized comparison

  3. nonzero_indices = np.where(diff_mask) — find changed elements

  4. Encode header + indices + values into the delta shm buffer

Typical performance for a 1.7B model: ~1.8s compute, ~99% sparsity, ~95 MB delta (vs 3.9 GB full = 40x compression).

Delta Application (RaaS)#

RaaS applies the delta via mmap in-place patching:

  1. Open the existing safetensors file in /dev/shm with mmap

  2. Create a numpy view of the weight data section

  3. Vectorized scatter write: weight_2d[indices] = values

  4. mmap.flush() — data is ready for SGLang to load

No read-copy-write cycle — only the changed elements are touched. Typical time: ~1.1s for 1.7B.

Buffer Layout#

/dev/shm/astraflow_buffer_XXXX  (2x model, same as full transfer)
┌─────────────┬─────────────┐
│   Half 0    │   Half 1    │
└─────────────┴─────────────┘

/dev/shm/astraflow_delta_XXXX   (1x model, for sparse delta)
┌─────────────────────────────┐
│  [header][indices][values]  │
└─────────────────────────────┘

The double buffer trick provides both current and previous versions for the delta comparison. No extra “previous copy” buffer is needed.

Guard Barrier#

The delta computation reads the inactive half asynchronously. A guard barrier at the start of offload() prevents all ranks from writing to the inactive half until the previous delta finishes:

# In offload():
self._wait_previous_delta()   # rank 0 waits on delta_done_event
dist.barrier()                 # all ranks sync before writing

Normally instant (<5ms). Only blocks if the training step is faster than delta compute — unlikely for models larger than 1.7B.

Timeline#

Time ──────────────────────────────────────────────────────────────────────►

TRAINER (GPU)
  │ train_step N               │guard│offload│save│wait │notify│ train_step
  │[===========================][---][======]│chkp│delta│async │[==========
  │ forward/backward/optim     │~3ms│~0.5s  │    │ready│      │
  │                            │    │       │    │     │      │

SENDER AGENT (CPU)
  │ · · · · · · · · · · · · · ·│swap│ack│ delta compute      │done│ · · · ·
  │         idle               │idx │   │[==================]│evt │   idle
  │                            │    │   │ ~1.8s (overlapped) │    │

RAAS (GPU)
  │ generating · · · · · · · · · · · · · · │choose│delta│pause│load│resume
  │                                        │mode  │pull │[===]│[==]│
  │                                        │      │~1.1s│3.0s │0.7s│

Measured Performance (1.7B model, 20 steps)#

Metric

Full

Delta

Trainer offload (blocking)

0.5s

0.5s (same)

Delta compute (async)

—

1.8s (overlapped)

Guard barrier

—

<5ms

TCP transfer

1.0s

0.025s

Save/Patch on RaaS

1.4s

1.1s

Pull total

2.4s

1.1s

Sparsity

—

98.4-99.2%

Compression ratio

1x

21-43x

Wandb Metrics#

Metric

Source

Description

weight_transfer/offload_guard_time

trainer

Guard barrier wait

weight_transfer/offload_copy_time

trainer

GPU→CPU copy

weight_transfer/offload_total_time

trainer

End-to-end offload

weight_transfer/delta_sparsity

sender agent

Fraction unchanged

weight_transfer/delta_size_mb

sender agent

Sparse delta size

weight_transfer/delta_compute_time

sender agent

Async compute time

weight_transfer/use_full

RaaS via AstraFlow

0=delta, 1=full