# Weight synchronization

The trainer's weights advance every optimizer step; the inference fleet only
sees new weights on a **refit**. The `WeightSynchronizer` abstraction
(`weight_sync/`) owns that transfer and hides the transport and topology from the
GRPO loop, which never branches on backend type or colocation.

## The interface

`WeightSynchronizer` (`weight_sync/interfaces.py`) is a small ABC the loop drives
explicitly:

| Method | Role |
| --- | --- |
| `init_communicator()` | Called once at setup, after workers exist. Colocated transports prepare refit metadata; the collective transport also initializes its process group. |
| `mark_stale()` | Called after every training step. Flags that the inference weights are behind. |
| `is_stale` | Property — `True` between `mark_stale()` and the next successful `sync_weights()`. |
| `sync_weights()` | Transfers the latest policy weights into the generation backend. |
| `shutdown()` | Releases communication resources. |

The loop's usage is the **staleness handshake**: after an optimizer step it
calls `mark_stale()`; before generating, if `is_stale`, it calls
`sync_weights()` and then `prepare_for_generation()` on the engine.

The interface assumes **global** weight updates — all generation workers update
atomically and share one weight version. In async GRPO, heterogeneous weight
ages are tracked per-sample in the replay buffer (`target_weight_versions`), not
here.

## Transports

`create_weight_synchronizer()` (`weight_sync/factory.py`) selects the transport
from the deployment topology and the generation backend:

| Topology | Backend | Transport | Class |
| --- | --- | --- | --- |
| Non-colocated | vLLM | NCCL collective | `CollectiveWeightSynchronizer` |
| Colocated | vLLM | IPC / ZMQ | `IPCWeightSynchronizer` |
| Colocated | SGLang | HTTP | `HTTPWeightSynchronizer` |

Non-colocated SGLang is not supported (the factory raises `NotImplementedError`).
The non-colocated path additionally requires `train_cluster` and
`inference_world_size`.

### Colocated vs. non-colocated lifecycle

This distinction shapes what `sync_weights()` does:

- **Colocated transports (IPC, HTTP)** own GPU phase transitions internally —
  offload optimizer state before the refit, `prepare_for_generation`, offload
  again after — because the policy and the engine share GPUs and must hand the
  device back and forth.
- **The NCCL collective transport is a pure data mover.** Policy and generation
  run on separate GPU clusters, so there are no phase transitions; it only
  forwards weights (and optional FP8 KV-cache scales) over the collective.

For colocated **SGLang** the engine's GPU memory is released and re-acquired
around each generation phase when `enable_memory_saver` is set: `finish_generation`
posts `release_memory_occupation` so the server stops holding `mem_fraction_static`
of the device during training, and `prepare_for_generation` issues a two-phase
`wake_up` (weights, then KV cache) to resume it. With `enable_memory_saver` off the
server stays resident (the historical always-on behaviour). A colocated generation
worker also reserves only a capped GPU fraction
(`min(0.2, 1/max_colocated_worker_groups)`), while the dedicated, non-colocated
path still claims a full GPU.

`sync_weights()` accepts an optional `kv_scales` dict; only the collective
transport honors it, forwarding to `policy.broadcast_weights_for_collective()`.
IPC and HTTP ignore it.

## The refit seam

Whatever the transport, the trainer must expose its weights in the layout the
inference engine expects — HF parameter names and shapes. The policy worker
implements a **refit seam** that maps its internal representation to that layout:

- The **DTensor** backend gathers sharded `DTensor` parameters and emits the HF
  state dict.
- The **JAX** backend converts NNX arrays (via dlpack on GPU, numpy on CPU) and
  applies a **model-specific** inverse of that model's HF loader. The dense path
  transposes linear kernels back to HF layout; Qwen3-MoE additionally re-expands
  fused expert tensors into the per-expert `experts.{i}.{gate,up,down}_proj`
  layout the inference engine's fused-MoE kernel consumes; Qwen3-Next adds its
  linear-attention parameters; and Gemma 4 maps its multi-parameter router and
  per-layer embedding table, keeping the experts **fused** because that is the
  layout its checkpoint and engine use. Each map is the exact inverse of the
  model's loader, so a load→refit round-trip reproduces the HF state dict.

`prepare_refit_info()` reports the `{hf_name: (shape, dtype)}` map the collective
producer uses to pack the broadcast; the consumer side unpacks it into the
engine's weights. The NCCL broadcast into a live engine is exercised on hardware
(see `handoff/hardware-deferred-validation.md`); the name-map inverse and the
JAX→torch value round-trip are covered by CPU tests.