Weight Synchronization

Weight Synchronization

Overview

The WeightSyncCoordinator class provides centralized management of weight transfers from training GPUs to inference replicas via NCCL. It serves as the single source of truth for sender/receiver IP mapping, GPU IDs, and NCCL ports, enabling both training and inference clients to coordinate weight synchronization without duplicating topology state.

Configuration

WeightSyncConfig

Weight synchronization is configured using WeightSyncConfig:

from arctic_platform.rl.config import WeightSyncConfig

config = WeightSyncConfig(
    training_sharding="dp",           # "dp" or "fsdp"
    training_gpus=8,                  # Number of training GPUs
    inference_replicas=4,             # Number of inference replicas
    inference_tp=2,                   # Tensor parallelism degree per replica
    base_port=29500,                  # Base NCCL port
    bucket_size=256 * 1024 * 1024,   # Bucket size in bytes (256 MB default)
)

Key configuration fields:

Core Concepts

Topology and Scheduling

The coordinator uses TransferSchedule to statically assign training TP workers to sender GPUs. The schedule determines:

• Which training GPU ranks actively send weights
• Which inference replicas receive from each sender
• NCCL group membership and world sizes

The number of active senders is the minimum of training_gpus and inference_replicas * inference_tp.

NCCL Sender Management

Senders are created lazily on first use via get_or_create_sender(). Each sender encapsulates:

• NCCL group initialization
• Connection to target inference receivers
• Data transfer via bucketed sends

API Reference

Creating a Coordinator

from arctic_platform.rl.weight_sync import WeightSyncCoordinator

coordinator = WeightSyncCoordinator(config)

Topology Queries

sender_ranks

active_ranks: list[int] = coordinator.sender_ranks

Returns the list of training GPU ranks that are actively sending weights.

num_groups

n_groups: int = coordinator.num_groups

Returns the total number of NCCL groups.

Weight Transfer

sync_weights()

result = coordinator.sync_weights(
    rank=0,
    weights=[("layer.0.weight", tensor1), ("layer.0.bias", tensor2), ...],
    client=rl_client,
    master_addr="192.168.1.10",
    device=torch.device("cuda", 0),
    direct=False,
)

Sends weights from a training rank to its assigned inference targets.

Parameters:

• rank (int): Training GPU rank initiating the transfer
• weights (Iterable[tuple[str, torch.Tensor]]): Iterable of (name, tensor) pairs
• client (ArcticRLClient): RL client for HTTP coordination
• master_addr (str, optional): Master node IP address. Required on the first call for each rank; may be omitted on subsequent calls
• device (torch.device, optional): GPU device. If omitted, defaults to cuda:{rank}. Required on the first call for each rank
• direct (bool, optional): If True, use direct transfer mode. Defaults to False

Returns:

A dictionary with transfer results (e.g., timing, bytes transferred).

Thread Safety:

This method is fully thread-safe. All active ranks should call concurrently. Synchronization logic:

1. The first rank to arrive locks and submits an HTTP request to the server via client.update_weights(), which dispatches all receivers
2. All ranks then push their data through NCCL independently
3. The last rank to finish resets shared state for the next round

Example (multi-rank):

import threading

def rank_worker(rank):
    coordinator.sync_weights(
        rank=rank,
        weights=training_model.named_parameters(),
        client=rl_client,
        master_addr="192.168.1.10",  # First call only
        device=torch.device("cuda", rank),  # First call only
    )

# All active ranks call concurrently
threads = [
    threading.Thread(target=rank_worker, args=(rank,))
    for rank in coordinator.sender_ranks
]
for t in threads:
    t.start()
for t in threads:
    t.join()

Inference Server Setup

ensure_server_ready()

coordinator.ensure_server_ready(
    client=rl_client,
    master_addr="192.168.1.10",
)

Idempotent call to prepare the inference server once. Safe to call from every training rank — only the first call performs the HTTP request; subsequent calls are no-ops.

Parameters:

• client (ArcticRLClient): RL client for HTTP communication
• master_addr (str): Master node IP address

prepare_inference_server()

result = coordinator.prepare_inference_server(
    client=rl_client,
    master_addr="192.168.1.10",
)

Explicitly tells the inference server to create NCCL receiver engines via an HTTP request with engine_only=True. This performs the NCCL rendezvous without actually receiving data.

Parameters:

• client (ArcticRLClient): RL client for HTTP communication
• master_addr (str): Master node IP address

Returns:

A dictionary with server response.

Cleanup

destroy()

coordinator.destroy()

Destroys all NCCL sender connections, clears sender state, resets server-ready flag, and shuts down the internal thread pool.

Call this when the coordinator is no longer needed (e.g., at application shutdown) to free NCCL resources.

Typical Workflow

from arctic_platform.rl.weight_sync import WeightSyncCoordinator
from arctic_platform.rl.config import WeightSyncConfig
from arctic_platform.rl.client import ArcticRLClient

# 1. Create configuration
config = WeightSyncConfig(
    training_gpus=8,
    inference_replicas=4,
    inference_tp=2,
)

# 2. Initialize coordinator
coordinator = WeightSyncCoordinator(config)

# 3. Prepare inference server (once, before any weight syncs)
rl_client = ArcticRLClient(...)
coordinator.ensure_server_ready(
    client=rl_client,
    master_addr="192.168.1.10",
)

# 4. Synchronize weights from training ranks (concurrent calls)
result = coordinator.sync_weights(
    rank=0,
    weights=model.named_parameters(),
    client=rl_client,
    master_addr="192.168.1.10",
    device=torch.device("cuda", 0),
)

# 5. Clean up when done
coordinator.destroy()

Internal Architecture

TransferSchedule

The coordinator reuses TransferSchedule from arctic_inference.server.weight_sync.schedule to:

• Map training TP workers to sender GPU ranks
• Assign each sender to a set of inference replicas
• Configure NCCL group topology (world size, ranks per group)

WeightSender

The WeightSender class (from arctic_inference.server.weight_sync.sender) encapsulates:

• NCCL group initialization via master_addr and base_port + group_id * inference_tp
• Bucketed tensor transfer using NCCL send/receive primitives
• Support for both standard and direct-mode transfers

Synchronization Primitives

The coordinator uses:

• threading.Lock (_sync_lock): Protects shared state during concurrent rank calls
• ThreadPoolExecutor: Executes HTTP requests asynchronously without blocking NCCL sends
• Future: Tracks completion of server-side receiver setup

Error Handling

• Missing master_addr on first call: If sync_weights() is called for a new rank without master_addr, a ValueError is raised
• Sender initialization errors: Exceptions during WeightSender creation propagate to the caller
• Cleanup errors: Exceptions during sender.destroy() are logged as warnings and do not halt the cleanup process