Awex is a high-performance RL training-inference weight synchronization framework, designed to enable second-level parameter updates from training to inference in RL workflows. It minimizes iteration latency, ensuring rollout phases consistently use the latest model.
The Awex weight exchange framework consists primarily of three components:
- WeightWriter: Runs within each training process, responsible for metadata collection and reporting of weight shards for the current training process, weight convert, resharding transfer plan construction, weight transmission, and other functions;
- WeightReader: Runs on the control process of each inference instance, which starts a WorkerWeightsReader on each GPU managed by the inference instance, corresponding to the WeightWriter of the training process. Responsible for metadata collection and reporting of weight shards for each inference process, weight convert, resharding transfer plan construction, weight reception, and other functions;
- MetaServer: Job-level global server for service discovery and weight metadata exchange between training and inference engines, as well as event notification functions in co-located scenarios;
The core functional modules of weight exchange consist mainly of 6 parts:
- Unified training-inference weight convert: Responsible for converting weights from training and inference engines with different parallelism strategies and tensor layouts into a unified format for subsequent weight metadata calculation and weight transmission;
- Global weight metadata calculation and exchange: After converting training and inference weights into a unified format, collects all weight shard metadata from each worker and reports to Meta Server for subsequent weight transmission plan construction;
- P2P weight transmission execution plan: Training and inference engines obtain global weight shard metadata from all workers, then separately construct peer-to-peer deterministic transfer plan for sending and receiving;
- NCCL weight transmission: Uses NCCL's send/recv API for peer-to-peer weight transmission based on the constructed transmission plan;
- Mooncake weight transmission: Executes shard transfers using one-sided reads over RDMA or CPU-buffer TCP;
- AState weight transmission: Uses AState distributed tables to exchange weight shards through step-scoped put/get operations;
Due to different computational workloads, training and inference engines generally adopt different parallelism strategies. Megatron training engine uses 5D parallelism strategy, DeepSpeed/FSDP uses Zero + DP data parallelism, while SGLang and VLLM inference engines mostly use DP + TP + EP. Additionally, different engines perform fusion, transposition, and quantization optimizations on weights after loading to adapt to high-performance operators.
To eliminate differences between different engines for subsequent weight exchange, Awex constructs a unified weight convert layer that performs the following converts:
- Weight splitting: Splits merged weights (such as FFN's gate/up) into independent weights, supporting cross-TP Resharding;
- Weight name unification: Converts all internal weights from all engines to the same namespace, establishing weight mapping relationships between training and inference engines;
- Attention weight ReGroup: On the training engine side, regroups and aligns QKV weights along the inference engine's TP/DPAttention parallelism strategy, avoiding shard explosion from fine-grained splitting;
- Quantization, precision, and format conversion: Automatically converts weights on the training side according to the precision and format of the inference side, and this low-precision conversion can also reduce the amount of transmitted data.
The entire weight convert adaptation layer is implemented as a pluggable structure that can be fully customized at the engine layer, model weight convert, and sharding layer to meet the customization needs of complex scenarios.
Each training and inference process needs to be aware of the weight metadata of all training and inference processes globally for constructing subsequent weight transfer plan. Awex also performs consistency validation of weight metadata between training and inference at this step. The main workflow is as follows:
- Each process in the training engine performs weight convert and obtains metadata for the converted shards
- Through all_gather_object, each rank obtains global training shard metadata
- Rank0 on the training side serializes global metadata and reports it to Meta Server
- Inference instance 0 on the inference side performs similar work; other inference instances have identical metadata and don't need additional computation
- All training and inference processes obtain global metadata from MetaServer
- Training and inference engines each perform shard-level metadata consistency and compatibility validation
After obtaining global weight metadata, Awex constructs a deterministic point-to-point transmission plan within each training and inference process.
Core Strategy (NCCL mode):
- For each replica of the same tensor shard, assign training shards to inference shards through Round Robin to ensure uniform pulling;
- For overlapping shard intervals, if perfectly aligned, directly map; otherwise, use two sends to different shards;
- Pre-filter shards related to the current process to avoid constructing a global plan (shards can reach tens of millions for trillion-parameter models);
- Ensure strict order consistency of NCCL send/recv;
Mooncake reuses this shard plan and checks the source/target slice signatures before copying received data into the inference weights.
The comm_backend configuration selects the transport:
| Backend | Runtime and requirements | Update mode |
|---|---|---|
file |
Engine checkpoint save/load hooks and a shared model path | File reload; colocated use also requires memory hooks |
nccl |
CUDA workers and NCCL | Separate workers, or colocated CUDA IPC plus inference-side NCCL |
hccl |
Experimental NPU runtime with torch-npu/MindSpeed | NCCL-style process-group path; CPU IPC for colocated use |
astate |
Separately installed AState service/runtime | Distributed table put/get with completion per step |
mooncake |
Mooncake Transfer Engine; CPU or CUDA workers | Separate or colocated workers; RDMA or CPU-buffer TCP |
Every non-file backend requires a shared meta_server_addr, even for one
inference instance. All workers must use the same AWEX and per-framework version
set. Rank counts come from runtime metadata: SGLang DP Attention partitions TP
ranks, while vLLM internal DP adds ranks to the instance.
NCCL transmission mode primarily uses NCCL's send/recv interface for weight transmission. There are some implementation differences in Awex for separated and co-located modes, which we will detail here.
NCCL Separated Weight Transmission
In separated transmission mode, Awex first constructs a joint training-inference NCCL Process Group for subsequent weight transmission from training to inference. Next, based on the NCCL send and transfer plan created during initialization, P2P peer-to-peer ordered weight sending and receiving is performed on each process of the training and inference engines. The overall workflow is as follows:
CUDA IPC Co-located Zero-Copy Weight Mapping
In colocated mode, training and inference processes share devices. AWEX avoids placing both processes on the same device in one NCCL communicator.
In this case, Awex uses CUDA IPC to zero-copy map the training process's GPU memory to the inference process, establishes a global communication group for all inference processes, then uses this communication group for NCCL send/recv to complete weight exchange from training to inference engines:
In implementation, we have also made some performance optimizations:
- Problem: Each CUDA IPC Handle's Open/Close has significant overhead; MOE and other models may have thousands to tens of thousands of weight tensors per card requiring IPC serialization;
- Solution: Before IPC serialization, merge tensors by shape and dtype, reducing the count to dozens, greatly reducing CUDA IPC overhead;
(Note: CUDA IPC does not support CUDA virtual memory. Future plans include allocating additional physical memory space for weight merging and transmission when enabling virtual GPU memory in the training engine)
For the default colocated NCCL transport, set AWEX_CHUNK_MB on every reader to
bound transient chunk allocations; zero selects the unchunked path. This is not a
limit on total model or converted-weight memory. Measure peak memory and update
latency when choosing a value. BoundedMemoryNcclColocateStreamBatchTransport
is an extension implementation with expert packing, not the default factory
selection; a custom reader must select it consistently on all ranks.
Set comm_backend="mooncake" on both roles. Mooncake transfers data without a
joint training-inference process group; AWEX's metadata server exchanges buffer
addresses and completion notifications. Set AWEX_MOONCAKE_PROTOCOL=rdma on RDMA
nodes or tcp when RDMA is unavailable. TCP stages CUDA tensors through CPU
memory. Protocol and step order must agree on every participating worker.
The sender owns registered snapshots until every reader has acknowledged them. Colocated updates retain these snapshots on CPU while training weights are offloaded, so the engine must implement memory release/resume hooks. The receiver stages one slice at a time to support strided target views.
Topology and transfer plans are fixed for an initialized job. Changing inference rank counts requires reinitialization; dynamic scaling is not implemented by this backend. After a failed exchange, restart workers with a fresh metadata server: published buffers cannot be released while remote accesses may still be active.
See the installation and protocol settings. The native TCP test covers transport correctness; model/framework and RDMA deployments require their own integration validation. Performance claims must identify the backend, protocol, versions, hardware, tensor sizes and timing scope.
Set comm_backend="astate" on both training and inference roles and install the
AState runtime separately. Both roles still use the shared AWEX metadata server
to exchange model and rank metadata.
AWEX creates the weights_exchange remote table with each role's parallelism
configuration. The training writer publishes converted weight shards through
multi_put(step_id, ...); inference workers load shards into their weight tensors
through multi_get(step_id, ...). Shard keys identify the parameter name, global
shape and shard offset, allowing AState to handle the data movement. Each role
calls complete(step_id) after finishing that update.