| 1 | # NCCL M2N Engine |
| 2 | |
| 3 | The `nccl_m2n` weight transfer engine uses [NCCL M2N](https://github.com/NVIDIA/nccl-extensions) to **reshard** weights between two disjoint meshes of GPUs — the trainer and the inference workers — in a single collective per parameter. Unlike the broadcast-based [NCCL engine](nccl.md), the trainer sends its local shard as-is: no full tensor is ever materialized on the training side, whatever its parallelism layout. |
| 4 | |
| 5 | ## When to Use NCCL M2N |
| 6 | |
| 7 | - The trainer and the inference engine use **different parallelism layouts** (e.g. trainer FSDP/EP, inference TP) and gathering full tensors on the trainer is expensive |
| 8 | - Training and inference are on **separate GPUs**, possibly across nodes |
| 9 | - You want weight sync off the critical path of an async RL loop |
| 10 | |
| 11 | ## Requirements |
| 12 | |
| 13 | This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. It needs NCCL 2.30.5 or newer, and `libnccl_m2n.so` is linked against a specific NCCL build, so vLLM must load the same library: |
| 14 | |
| 15 | ```bash |
| 16 | export VLLM_NCCL_SO_PATH=/path/to/the/same/libnccl.so |
| 17 | ``` |
| 18 | |
| 19 | The backend imports the runtime lazily, so vLLM is unaffected unless you select it. |
| 20 | |
| 21 | ## How It Works |
| 22 | |
| 23 | 1. The trainer ranks and all inference workers join **one** NCCL communicator via `StatelessProcessGroup`. Trainer ranks occupy `[0, T)` and workers `[T, T + N)` — m2n meshes are contiguous rank intervals, which is exactly what vLLM's existing `rank_offset` convention produces. |
| 24 | 2. At the init handshake the trainer ships the full transfer plan: both meshes — one per side, shared by every parameter — then each parameter's name, dtype, shape and placement. The workers resolve the matching **destination placements** and publish them back over the shared group; the meshes themselves are declared, not negotiated. |
| 25 | 3. Each round, both sides issue one `reshard` per parameter in the same order. Every rank in the communicator participates in every reshard, so — as with the NCCL engine — the worker must be inside `update_weights` while the trainer sends. The trainer engine handles that concurrency itself. |
| 26 | |
| 27 | ### Destination resolution |
| 28 | |
| 29 | A parameter is resharded **directly into the live vLLM parameter** when its checkpoint name maps 1:1 onto a model parameter whose shape is either identical to the checkpoint shape (replicated) or differs on exactly one dim by the tensor-parallel factor. That rank then only ever receives its own shard. |
| 30 | |
| 31 | Anything else falls back to receiving the full tensor and letting `load_weights` do the sharding, exactly as the NCCL engine does for every parameter. The fallback covers: |
| 32 | |
| 33 | - fused parameters (`qkv_proj`, `gate_up_proj`, MoE `w13`/`w2`), whose checkpoint names do not name a model parameter |
| 34 | - quantized models, where the weight loader does more than a sharded copy |
| 35 | - pipeline-parallel deployments |
| 36 | |
| 37 | Correctness never depends on a parameter resolving. The engine logs how many parameters took each path at init. |
| 38 | |
| 39 | ## Configuration |
| 40 | |
| 41 | ```python |
| 42 | from vllm import LLM |
| 43 | from vllm.config import WeightTransferConfig |
| 44 | |
| 45 | llm = LLM( |
| 46 | model="my-model", |
| 47 | weight_transfer_config=WeightTransferConfig(backend="nccl_m2n"), |
| 48 | ) |
| 49 | ``` |
| 50 | |
| 51 | ```bash |
| 52 | vllm serve my-model --weight-transfer-config '{"backend": "nccl_m2n"}' |
| 53 | ``` |
| 54 | |
| 55 | ## Trainer Side |
| 56 | |
| 57 | The trainer uses the stateful trainer engine. m2n needs each parameter's source layout, which the base `WeightSource` does not carry, so it takes an `M2NWeightSource`. `DTensorModuleSource` covers the common case: it reads each parameter's device mesh and placements and yields **local shards** — it never calls `full_tensor()`. |
| 58 | |
| 59 | ```python |
| 60 | from vllm.distributed.weight_transfer import ( |
| 61 | RayVLLMWeightSyncClient, |
| 62 | WeightTransferTrainerFactory, |
| 63 | ) |
| 64 | from vllm.distributed.weight_transfer.m2n_source import DTensorModuleSource |
| 65 | from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo |
| 66 | |
| 67 | # Called on every trainer rank; rank 0 is the sender. |
| 68 | engine = WeightTransferTrainerFactory.trainer_init( |
| 69 | init_info=M2NTrainerInitInfo( |
| 70 | master_address=master_address, |
| 71 | master_port=master_port, |
| 72 | world_size=num_trainer_ranks + num_inference_workers, |
| 73 | num_trainer_ranks=num_trainer_ranks, |
| 74 | rank=my_rank, |
| 75 | ), |
| 76 | client=RayVLLMWeightSyncClient(llm_handle), |
| 77 | source=DTensorModuleSource(model, num_trainer_ranks), |
| 78 | ) |
| 79 | |
| 80 | # Each round, on every trainer rank: |
| 81 | engine.send_weights() |
| 82 | ``` |
| 83 | |
| 84 | `send_weights()` drives the whole round — `start_weight_update`, `update_weights`, `finish_weight_update`. Every rank in the communicator must be inside the same reshard, so `update_weights` has to be in flight while the trainer sends; the engine runs it on a helper thread so the trainer loop does not have to open-code that. |
| 85 | |
| 86 | Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass `M2NWeightSource` and supply layouts explicitly. |
| 87 | |
| 88 | See [`examples/rl/rlhf_m2n.py`](../../../examples/rl/rlhf_m2n.py) for a runnable FSDP → TP example. |
| 89 | |
| 90 | ## Limitations |
| 91 | |
| 92 | - Tensor rank 1..3; 4-D parameters are rejected. |
| 93 | - dtypes: int8/uint8, fp8 e4m3/e5m2, fp16, bf16, int32/uint32, fp32, int64/uint64, fp64. No fp4. |
| 94 | - The trainer's device mesh must cover a contiguous rank interval starting at 0, and be 1-D or 2-D. |
| 95 | - `Partial` placements are not supported. |
| 96 | - Models with known model-level checkpoint preprocessing use the full-tensor |
| 97 | fallback for the affected module. |
| 98 | - A reshard cannot be issued from inside a CUDA graph capture (weight updates run outside capture, so this only matters for custom callers). |
| 99 | - The reshard plan is capped per *shard*, not per mesh: at most 16 source shards may feed one destination shard, and one source shard may feed at most 64 destination shards (`MAX_SOURCES` / `MAX_TARGETS` in the m2n build; raising either needs a rebuild). |
| 100 | - A directly-resharded parameter spreads over `tp_size` destination shards, so each one is fed by fewer source shards. A parameter on the fallback path is replicated, i.e. a single destination shard fed by every source shard, so a trainer that **sharded** its parameters is still limited to 16 ranks whenever any parameter falls back — which fused parameters always do. A replicated trainer is unaffected. |
| 101 | - Each communicator holds a cached staging pool (~256 MiB by default: 4 channels x 64 MiB). |